Compare commits

...
Author SHA1 Message Date
easyai 3056cf8fca feat(gateway): 补齐桌面端高级媒体直连接口
ci / verify (pull_request) Successful in 15m34s
新增图片矢量化、视频超分、每日用量、计价与任务隔离能力,并通过环境变量解析平台凭据。

已通过 Go 全量门禁、迁移检查、镜像构建以及 Vectorizer 五格式和 Topaz 3 秒视频真实 DEV 验收。
2026-07-22 14:02:53 +08:00
easyai cbebfd7baa chore(ci): 推进生产迁移基线至 v0.4.4 (#23)
ci / verify (push) Successful in 11m18s
生产发布、健康检查与隔离恢复演练均已通过;仅推进生产迁移基线。
2026-07-22 11:19:59 +08:00
easyai ee0256adb4 chore(ci): 推进生产迁移基线至 v0.4.4
ci / verify (pull_request) Successful in 14m30s
2026-07-22 11:04:42 +08:00
easyai fce76a30ba fix(keling): 统一 Omni 模型别名与生产绑定 (#22)
ci / verify (push) Successful in 12m40s
release-ci / verify-tag (push) Successful in 14m2s
线上真实任务验证发现官方兼容接口的模型别名与生产绑定不一致;统一使用原生模型名匹配候选,同时保留通用接口旧别名兼容。
2026-07-22 10:23:18 +08:00
easyai d37fb3ea60 fix(keling): 统一 Omni 模型别名与生产绑定
ci / verify (pull_request) Successful in 21m44s
2026-07-22 10:00:30 +08:00
easyai 9872068596 Merge pull request #21: 统一公开 API 为 /api/v1 前缀
ci / verify (push) Successful in 12m50s
release-ci / verify-tag (push) Successful in 16m11s
合并公开 API V1 前缀统一、兼容路由、OpenAPI、部署健康检查与验证测试。
2026-07-22 09:10:09 +08:00
easyai a23b28c27d test(api): 等待异步任务完成后校验结果
ci / verify (pull_request) Successful in 12m38s
真实队列环境中 X-Async 请求会先返回 queued。集成测试改为轮询任务完成并重新读取详情,避免把正确的异步受理状态误判为失败。
2026-07-22 08:56:30 +08:00
easyai eb37b568ae fix(deploy): 使用 /api/v1 执行发布健康检查
ci / verify (pull_request) Failing after 3m41s
将 Compose 发布脚本的 API、就绪和 Web 反代探针切换到统一公开前缀,避免发布完成后仍只验证旧兼容入口。\n\n验证:bash -n;ShellCheck;ci-build-images-test.sh;pipeline-test.sh。
2026-07-22 08:50:55 +08:00
easyai a8d1c550ef chore(deps): 升级 fast-uri 修复高危漏洞
ci / verify (pull_request) Failing after 3m45s
将构建链路中的 fast-uri 锁定到 3.1.4,修复 host confusion 高危公告,不改变应用运行时依赖。\n\n验证:pnpm install --frozen-lockfile;pnpm audit --audit-level high;pnpm lint;pnpm test;pnpm build。
2026-07-22 08:48:40 +08:00
easyai 5432760cf7 fix(identity): 避免未配置 SSF 时触发空指针
SecurityEventReceiver 不再把空 ConnectionManager 包装为非空接口;未启用 SSF 时公开端点明确返回 404。补充运行时和 HTTP 回归测试。\n\n验证:go test ./internal/identityruntime ./internal/httpapi。
2026-07-22 08:48:36 +08:00
easyai 7c5a999e32 feat(api): 统一公开接口为 /api/v1 前缀
将通用生成、Gemini、可灵、火山、健康检查与 OpenAPI 的推荐入口统一到 /api/v1,并保留历史路径作为兼容别名。同步更新代理配置、接入文档、接口清单和前缀回归测试。\n\n验证:go vet ./...;go test ./...;pnpm openapi;pnpm lint;pnpm test;pnpm build;公开 OpenAPI 71 个方法与接口清单机器比对一致。
2026-07-22 08:48:32 +08:00
easyai f7a5f2e808 Merge pull request #20: 前移生产迁移基线至 v0.4.2
ci / verify (push) Successful in 9m43s
2026-07-22 05:42:59 +08:00
easyai 152f9d1206 chore(deploy): 前移生产迁移基线至 v0.4.2
ci / verify (pull_request) Successful in 9m42s
2026-07-22 05:32:39 +08:00
easyai d95cecd0eb Merge pull request #19: 修复 Gemini 兼容平台版本路径重复拼接
ci / verify (push) Successful in 9m45s
release-ci / verify-tag (push) Successful in 9m54s
修复 Gemini 兼容平台已包含 /v1 时重复追加 /v1beta 的问题。

验证:精确 ci / verify (pull_request) 已通过。
2026-07-22 05:04:12 +08:00
easyai 5954afef55 fix(gemini): 保留兼容平台的 API 版本路径
ci / verify (pull_request) Successful in 9m42s
2026-07-22 04:53:42 +08:00
easyai 929a7a172c Merge pull request #18: 修复 Gemini 图像默认参数与 server-main 脚本兼容
ci / verify (push) Successful in 9m44s
release-ci / verify-tag (push) Successful in 9m46s
合并已通过 ci / verify (pull_request) 的图像参数与脚本兼容修复。
2026-07-22 04:17:44 +08:00
easyai 16149260c7 fix(universal): 映射脚本上游模型名
ci / verify (pull_request) Successful in 9m51s
让 server-main 参数预处理兼容上下文使用平台绑定的真实模型名,避免把用户侧别名提交给上游。
2026-07-22 03:56:00 +08:00
easyai c72c43aaaa fix(universal): 兼容 server-main 脚本语义
ci / verify (pull_request) Successful in 9m52s
补充历史脚本依赖的参数预处理上下文,并识别嵌套同步结果,使迁移后的 Universal 图片通道可按原配置执行。
2026-07-22 03:50:59 +08:00
easyai 84aea01b5b build(web): 修复高危间接依赖漏洞
ci / verify (pull_request) Successful in 9m48s
将 fast-uri 约束到已修复的 3.1.3,确保高危依赖审计通过。
2026-07-22 03:39:48 +08:00
easyai 293ef24bb7 build(api): 升级存在漏洞的 Go 依赖 2026-07-22 03:37:34 +08:00
easyai 0f1765b226 fix(gemini): 初始化图像生成配置对象 2026-07-22 03:37:28 +08:00
easyai 1fa58ba901 Merge pull request #17: 前移生产迁移基线至 v0.4.0
ci / verify (push) Successful in 10m15s
生产版本 v0.4.0 已完成发布、健康检查和恢复演练,前移生产迁移安全基线。
2026-07-22 02:59:54 +08:00
easyai 55595570c2 chore(deploy): 前移生产迁移基线至 v0.4.0
ci / verify (pull_request) Successful in 10m18s
2026-07-22 02:48:38 +08:00
easyai fba9759bc7 Merge pull request 'fix(kling): 补齐全模态基础视频能力' (#16) from codex/omni-default-video-capabilities into main
ci / verify (push) Successful in 12m8s
release-ci / verify-tag (push) Successful in 9m56s
2026-07-22 02:00:02 +08:00
easyai 3d3460ce63 test(kling): 兼容全模态模型展示别名
ci / verify (pull_request) Successful in 11m22s
集成测试通过现有可灵别名归一规则校验解析模型,兼容共享数据库返回 kling-o1 等目录展示名,同时保留任务成功与推断模型类型断言。
2026-07-22 01:48:10 +08:00
easyai 000ee1bbfd fix(kling): 补齐全模态基础视频能力
ci / verify (pull_request) Failing after 4m10s
让 omni_video 模型默认同时注册 video_generate 与 image_to_video,兼容通用视频接口在不显式传 modelType 时的能力推断。

新增迁移同步可灵 O1 与 3.0 Omni 的基础目录、平台映射、能力详情和默认快照;集成测试覆盖两个模型的无 modelType 文生视频与图生视频。
2026-07-22 01:41:23 +08:00
wangbo d0cfd0a385 Merge pull request 'feat: 合并网关能力与兼容性优化' (#15) from codex/gateway-feature-batch-20260721 into main
ci / verify (push) Successful in 10m38s
2026-07-22 01:17:45 +08:00
wangbo 0818f55235 fix(billing): 保留规则集未覆盖的基础资源价格
ci / verify (pull_request) Successful in 12m21s
2026-07-22 00:57:50 +08:00
wangbo fe83da56d2 test(e2e): 增加视频生产验收脚本与记录
ci / verify (pull_request) Failing after 3m43s
2026-07-22 00:25:22 +08:00
wangbo 002422b753 fix(catalog): 同步 GLM-5.2 文本能力配置 2026-07-22 00:25:22 +08:00
wangbo 24b778b3ba docs(api): 同步模型权限与 Kling 接口契约 2026-07-22 00:25:22 +08:00
wangbo d7951cfdd2 feat(volces): 接入任务兼容查询与上游取消 2026-07-22 00:25:22 +08:00
wangbo ddd68cfebd feat(seedance): 接入真人资产管理与引用 2026-07-22 00:25:22 +08:00
wangbo 5a71643099 feat(kling): 接入 Omni 视频兼容接口 2026-07-22 00:25:05 +08:00
wangbo b04a7d9d3d feat(web): 支持文档页真实提交请求 2026-07-22 00:24:36 +08:00
wangbo 6b675c406e fix(billing): 修正模型计费配置继承优先级 2026-07-22 00:24:36 +08:00
wangbo 56d4a3a6b7 fix(access): 修正 API Key 可分配模型范围 2026-07-22 00:24:36 +08:00
wangbo 276c0612d8 refactor(task): 复用网关任务创建流程 2026-07-22 00:24:36 +08:00
wangbo d818e7947a 调整 Gemini 模型接口路由 2026-07-22 00:24:36 +08:00
easyai d5c2c58c67 Merge pull request 'feat(kling): 接入 O1 与 3.0 Omni 兼容接口' (#14) from codex/kling-o1-v3-compat into main
ci / verify (push) Successful in 35m39s
release-ci / verify-tag (push) Successful in 13m50s
2026-07-22 00:09:49 +08:00
easyai 9d4501bc42 feat(kling): 接入O1与3.0 Omni兼容接口
ci / verify (pull_request) Successful in 15m32s
新增中国区可灵 V1 AK/SK Omni 协议与 API 2.0 兼容路径,补齐任务隔离、外部任务幂等、参数校验和 OpenAPI 文档。\n\n验证:O1 与 3.0 Omni 真实 V1 任务成功;Go、前端、依赖审计、迁移及 CI 脚本门禁通过。
2026-07-21 23:47:13 +08:00
chengcheng e280c0875c Merge pull request 'fix: 修正视频线性计费与在线测试模型筛选' (#13) from codex/fix-linear-video-billing into main
ci / verify (push) Successful in 15m9s
release-ci / verify-tag (push) Successful in 16m46s
2026-07-21 14:15:19 +08:00
chengcheng 142dcc7932 chore(git): 合并视频时长线性计费修复
ci / verify (pull_request) Successful in 13m41s
2026-07-21 13:34:52 +08:00
chengcheng e3dfe8162b fix(billing): 视频时长按实际秒数线性计费
将五秒基础价按 duration / 5 比例结算,保留 provider 返回的小数时长,避免六秒视频被按两个完整单位收费。

影响:所有使用 5s 视频基础价的模型,quantity 与 durationUnitCount 允许小数。新增迁移同步现存规则的旧 ceil 公式元数据。

验证:go vet ./...;pnpm lint;pnpm test;pnpm build;./tests/ci/migrations-test.sh
2026-07-21 13:32:23 +08:00
chengcheng 69b0c107d3 fix(web): 修正在线测试模型类型筛选
移除可能将 image_to_video 误判为图像模型的子串回退,在没有匹配模型时显示空状态。

新增图像生成、图像编辑和视频模式回归测试。验证通过:前端 94 项测试、类型检查、lint 和生产构建。
2026-07-21 13:12:40 +08:00
chengcheng e533ec2367 Merge pull request 'fix(auth): 修复生产登录卡死与有界超时' (#11) from codex/fix-api-key-pool-deadlock into main
ci / verify (push) Successful in 11m46s
2026-07-21 12:09:39 +08:00
chengcheng bfa17a3aba fix(auth): 为登录链路增加有界超时
ci / verify (pull_request) Successful in 11m22s
为 PostgreSQL 连接、就绪检查和本地登录设置分层超时,数据库不可用时返回稳定 503 错误码并记录无凭据的连接池统计。

前端登录在 10 秒后取消请求并兼容调用方 AbortSignal,Nginx 登录精确路由限制上游为 15 秒,同时更新 OpenAPI 和回归测试。
2026-07-21 11:47:09 +08:00
chengcheng 86c374b5c2 fix(auth): 消除 API Key 校验连接池死锁
先完整收集同前缀候选项并关闭查询结果,再执行 bcrypt 比对和 last_used_at 更新,避免小连接池下查询与更新相互等待。

新增 Rows 关闭顺序、前缀碰撞、MaxConns=1 和 8 并发真实 PostgreSQL 回归测试。
2026-07-21 11:47:09 +08:00
141 changed files with 15892 additions and 4960 deletions
+7
View File
@@ -49,6 +49,13 @@ AI_GATEWAY_API_RUNTIME_IMAGE=alpine:3.22
AI_GATEWAY_NODE_BUILD_IMAGE=node:22-alpine
AI_GATEWAY_WEB_RUNTIME_IMAGE=nginx:1.27-alpine
# Opt-in, billable China Kling V1 integration tests. Keep real AK/SK only in
# .env.local (gitignored); never commit them.
KELING_LIVE_TEST=0
KELING_TEST_BASE_URL=https://api-beijing.klingai.com/v1
KELING_TEST_ACCESS_KEY=
KELING_TEST_SECRET_KEY=
# Used when the gateway delegates OpenAPI sk-* validation, user/group sync, file upload, and settlement callbacks.
SERVER_MAIN_BASE_URL=http://localhost:3000
SERVER_MAIN_INTERNAL_TOKEN=change-me
+1 -1
View File
@@ -36,7 +36,7 @@ RUN --mount=type=cache,target=/go/pkg/mod \
FROM ${API_RUNTIME_IMAGE} AS api
RUN apk add --no-cache ca-certificates tzdata wget && \
RUN apk add --no-cache ca-certificates tzdata wget ffmpeg && \
adduser -D -H -u 10001 appuser
WORKDIR /app
+7 -3
View File
@@ -70,8 +70,10 @@ scripts/deploy-compose.sh
部署成功后默认访问地址:
- Web: `http://127.0.0.1:5178`
- API: `http://127.0.0.1:8088/healthz`
- Web 反代 API: `http://127.0.0.1:5178/gateway-api/healthz`
- API: `http://127.0.0.1:8088/api/v1/healthz`
- Web 反代公开 API: `http://127.0.0.1:5178/api/v1/healthz`
公开接口统一使用 `/api/v1` 前缀,完整分组清单见 [公开 API V1 清单](docs/public-api-v1.md)。
常用覆盖项:
@@ -99,7 +101,7 @@ scripts/deploy-compose.sh clean
docker login --username=<your-aliyun-account> registry.cn-shanghai.aliyuncs.com
```
Web 容器的 Nginx 配置通过 bind mount 挂载自仓库文件 [docker/nginx.conf](docker/nginx.conf),可直接修改该文件调整静态资源 `/gateway-api` 反向代理配置。修改后执行以下命令使配置生效:
Web 容器的 Nginx 配置通过 bind mount 挂载自仓库文件 [docker/nginx.conf](docker/nginx.conf),可直接修改该文件调整静态资源、规范 `/api/v1` 公开入口和旧 `/gateway-api` 兼容反向代理。修改后执行以下命令使配置生效:
```bash
docker compose -f docker-compose.yml restart web
@@ -125,6 +127,8 @@ AI_GATEWAY_COMPOSE_DATABASE_URL='postgresql://easyai:easyai2025@postgres:5432/ea
pnpm openapi
```
中国区可灵 O1 / 3.0 Omni 的 V1 AK/SK 与 API 2.0 兼容接入方式见 [可灵兼容接口说明](docs/kling-compatible-api.md)。
默认 EasyAI 部署里,`easyai-pgvector` 在容器网络内的连接串是:
+1811 -2631
View File
File diff suppressed because it is too large Load Diff
+1195 -1735
View File
File diff suppressed because it is too large Load Diff
+2 -2
View File
@@ -36,7 +36,7 @@ require (
github.com/tidwall/pretty v1.2.1 // indirect
github.com/tidwall/sjson v1.2.5 // indirect
go.uber.org/goleak v1.3.0 // indirect
golang.org/x/sync v0.20.0 // indirect
golang.org/x/text v0.37.0 // indirect
golang.org/x/sync v0.21.0 // indirect
golang.org/x/text v0.39.0 // indirect
gopkg.in/yaml.v3 v3.0.1 // indirect
)
+4 -4
View File
@@ -73,10 +73,10 @@ 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/oauth2 v0.36.0 h1:peZ/1z27fi9hUOFCAZaHyrpWG5lwe0RJEEEeH0ThlIs=
golang.org/x/oauth2 v0.36.0/go.mod h1:YDBUJMTkDnJS+A4BP4eZBjCqtokkg1hODuPjwiGPO7Q=
golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4=
golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
golang.org/x/text v0.37.0 h1:Cqjiwd9eSg8e0QAkyCaQTNHFIIzWtidPahFWR83rTrc=
golang.org/x/text v0.37.0/go.mod h1:a5sjxXGs9hsn/AJVwuElvCAo9v8QYLzvavO5z2PiM38=
golang.org/x/sync v0.21.0 h1:HLII4xRRTtCRkxYp4HNFF0Js/Og6q2i++KXbg0gHCwM=
golang.org/x/sync v0.21.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
golang.org/x/text v0.39.0 h1:UbZz4pLOvn600D6Oh6GGEI6VAmndrEBLv8/6BEXzyus=
golang.org/x/text v0.39.0/go.mod h1:3UwRclnC2g0TU9x8PZiyfOajCd1zaUNHF9cvqcQZ+ZM=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
@@ -0,0 +1,282 @@
package clients
import (
"context"
"encoding/json"
"io"
"net"
"net/http"
"net/http/httptest"
"strings"
"sync"
"testing"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestVectorizerClientMultipartAndPrivateResumeState(t *testing.T) {
var receivedPath string
var receivedFields map[string]string
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
receivedPath = r.URL.Path
if username, password, ok := r.BasicAuth(); !ok || username != "id" || password != "secret" {
t.Fatalf("unexpected vectorizer auth")
}
if err := r.ParseMultipartForm(1 << 20); err != nil {
t.Fatalf("parse multipart: %v", err)
}
receivedFields = map[string]string{}
for key, values := range r.MultipartForm.Value {
receivedFields[key] = values[0]
}
w.Header().Set("Content-Type", "image/svg+xml")
w.Header().Set("X-Image-Token", "private-image-token")
w.Header().Set("X-Receipt", "private-receipt")
_, _ = io.WriteString(w, `<svg xmlns="http://www.w3.org/2000/svg"/>`)
}))
defer server.Close()
privateState := map[string]any{}
response, err := (VectorizerClient{LookupIP: func(context.Context, string) ([]net.IPAddr, error) {
return []net.IPAddr{{IP: net.ParseIP("93.184.216.34")}}, nil
}}).Run(context.Background(), Request{
Kind: "images.vectorize", Model: "easy-image-vectorizer-1",
Body: map[string]any{"source": map[string]any{"url": "https://example.com/input.png"}, "format": "svg", "maxColors": 16, "cleanupLevel": "strong"},
Candidate: storeCandidate(server.URL, "vectorizer", map[string]any{"accessKey": "id", "secretKey": "secret"}),
OnRemoteTaskSubmitted: func(_ string, payload map[string]any) error { privateState = payload; return nil },
})
if err != nil {
t.Fatalf("run vectorizer: %v", err)
}
if receivedPath != "/vectorize" || receivedFields["image.url"] == "" || receivedFields["output.file_format"] != "svg" || receivedFields["processing.max_colors"] != "16" {
t.Fatalf("unexpected multipart request: path=%s fields=%+v", receivedPath, receivedFields)
}
if privateState["imageToken"] != "private-image-token" || privateState["receipt"] != "private-receipt" {
t.Fatalf("private resume state missing: %+v", privateState)
}
serialized := strings.ToLower(toJSONForTest(response.Result))
if strings.Contains(serialized, "private-image-token") || strings.Contains(serialized, "private-receipt") {
t.Fatalf("private vectorizer state leaked into response: %s", serialized)
}
}
func TestVectorizerClientRejectsPrivateSourceBeforeProviderCall(t *testing.T) {
providerCalls := 0
server := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { providerCalls++ }))
defer server.Close()
_, err := (VectorizerClient{LookupIP: func(context.Context, string) ([]net.IPAddr, error) {
return []net.IPAddr{{IP: net.ParseIP("127.0.0.1")}}, nil
}}).Run(context.Background(), Request{
Body: map[string]any{"source": map[string]any{"url": "https://assets.example.test/input.png"}},
Candidate: storeCandidate(server.URL, "vectorizer", map[string]any{"apiKey": "key"}),
})
if err == nil || !strings.Contains(err.Error(), "blocked network") {
t.Fatalf("expected blocked source error, got %v", err)
}
if providerCalls != 0 {
t.Fatalf("provider received %d calls for blocked source", providerCalls)
}
}
func TestVectorizerClientReusesImageTokenForExtraFormat(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/download" {
t.Fatalf("unexpected path %s", r.URL.Path)
}
if err := r.ParseMultipartForm(1 << 20); err != nil {
t.Fatal(err)
}
if r.FormValue("image.token") != "token" || r.FormValue("receipt") != "receipt" || r.FormValue("output.file_format") != "pdf" {
t.Fatalf("unexpected reuse form: %+v", r.MultipartForm.Value)
}
w.Header().Set("Content-Type", "application/pdf")
_, _ = w.Write([]byte("%PDF-1.7\n"))
}))
defer server.Close()
response, err := (VectorizerClient{}).Run(context.Background(), Request{
Body: map[string]any{"_vectorizer_image_token": "token", "_vectorizer_receipt": "receipt", "format": "pdf"},
Candidate: storeCandidate(server.URL, "vectorizer", map[string]any{"apiKey": "key"}),
})
if err != nil {
t.Fatal(err)
}
item := response.Result["data"].([]any)[0].(map[string]any)
if item["type"] != "file" || item["mime_type"] != "application/pdf" {
t.Fatalf("unexpected pdf result: %+v", item)
}
}
func TestTopazClientUploadsPartsPollsAndValidatesOutput(t *testing.T) {
var mu sync.Mutex
called := map[string]int{}
var server *httptest.Server
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
mu.Lock()
called[r.Method+" "+r.URL.Path]++
mu.Unlock()
switch {
case r.Method == http.MethodGet && r.URL.Path == "/source.mp4":
_, _ = w.Write([]byte("small-video-source"))
case r.Method == http.MethodPost && r.URL.Path == "/video/":
if r.Header.Get("X-API-Key") != "topaz-key" {
t.Fatalf("missing Topaz API key")
}
_, _ = io.WriteString(w, `{"requestId":"job-1"}`)
case r.Method == http.MethodPatch && r.URL.Path == "/video/job-1/accept":
_, _ = io.WriteString(w, `{"urls":["`+server.URL+`/upload/1","`+server.URL+`/upload/2"]}`)
case r.Method == http.MethodPut && strings.HasPrefix(r.URL.Path, "/upload/"):
w.Header().Set("ETag", `"etag-`+strings.TrimPrefix(r.URL.Path, "/upload/")+`"`)
w.WriteHeader(http.StatusOK)
case r.Method == http.MethodPatch && r.URL.Path == "/video/job-1/complete-upload/":
w.WriteHeader(http.StatusOK)
_, _ = io.WriteString(w, `{}`)
case r.Method == http.MethodGet && r.URL.Path == "/video/job-1/status":
_, _ = io.WriteString(w, `{"status":"completed","download":{"url":"`+server.URL+`/output.mp4"}}`)
default:
http.NotFound(w, r)
}
}))
defer server.Close()
phases := []string{}
client := TopazClient{Probe: func(_ context.Context, target string) (TopazVideoMetadata, error) {
if strings.HasPrefix(target, "http") {
return TopazVideoMetadata{Width: 640, Height: 360, Duration: 3, FrameRate: 24, HasAudio: true}, nil
}
return TopazVideoMetadata{Width: 320, Height: 180, Duration: 3, FrameRate: 24, FrameCount: 72, HasAudio: true}, nil
}}
response, err := client.Run(context.Background(), Request{
Model: "easy-proteus-standard-4",
Body: map[string]any{"video_url": server.URL + "/source.mp4", "output_width": 640, "output_height": 360, "preserve_audio": true},
Candidate: storeCandidateWithConfig(server.URL, "topaz", map[string]any{"apiKey": "topaz-key"}, map[string]any{"allowPrivateSourceDownloads": true, "pollIntervalMs": 1, "pollTimeoutMs": 100}),
OnRemoteTaskSubmitted: func(_ string, payload map[string]any) error {
phases = append(phases, payload["phase"].(string))
return nil
},
OnRemoteTaskPolled: func(_ string, payload map[string]any) error {
phases = append(phases, payload["phase"].(string))
return nil
},
})
if err != nil {
t.Fatalf("run Topaz: %v", err)
}
if response.RequestID != "job-1" || !containsStringForTest(phases, "created") || !containsStringForTest(phases, "uploaded") {
t.Fatalf("unexpected Topaz response/phases: response=%+v phases=%+v", response, phases)
}
if called["PUT /upload/1"] != 1 || called["PUT /upload/2"] != 1 || called["PATCH /video/job-1/complete-upload/"] != 1 {
t.Fatalf("upload lifecycle incomplete: %+v", called)
}
}
func TestTopazClientResumesUploadedTaskWithoutDownloadingSource(t *testing.T) {
var sourceCalls int
var server *httptest.Server
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path == "/source.mp4" {
sourceCalls++
}
if r.URL.Path == "/video/job-resume/status" {
_, _ = io.WriteString(w, `{"status":"completed","download":{"url":"`+server.URL+`/result.mp4"}}`)
return
}
http.NotFound(w, r)
}))
defer server.Close()
_, err := (TopazClient{Probe: func(context.Context, string) (TopazVideoMetadata, error) {
return TopazVideoMetadata{Width: 640, Height: 360, Duration: 3, FrameRate: 24}, nil
}}).Run(context.Background(), Request{
RemoteTaskID: "job-resume", RemoteTaskPayload: map[string]any{"phase": "uploaded", "targetResolution": map[string]any{"width": 640, "height": 360}},
Body: map[string]any{"video_url": server.URL + "/source.mp4"},
Candidate: storeCandidateWithConfig(server.URL, "topaz", map[string]any{"apiKey": "topaz-key"}, map[string]any{"pollIntervalMs": 1, "pollTimeoutMs": 100}),
})
if err != nil {
t.Fatal(err)
}
if sourceCalls != 0 {
t.Fatalf("resume downloaded source %d times", sourceCalls)
}
}
func TestTopazClientReportsProviderFailureAndPollingTimeout(t *testing.T) {
tests := []struct {
name string
statusBody string
timeoutMS int
want string
}{
{name: "provider failure", statusBody: `{"status":"failed","message":"upstream rejected"}`, timeoutMS: 100, want: "upstream rejected"},
{name: "poll timeout", statusBody: `{"status":"processing"}`, timeoutMS: 3, want: "polling timed out"},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/video/job/status" {
http.NotFound(w, r)
return
}
_, _ = io.WriteString(w, test.statusBody)
}))
defer server.Close()
_, err := (TopazClient{}).Run(context.Background(), Request{
RemoteTaskID: "job", RemoteTaskPayload: map[string]any{"phase": "uploaded", "targetResolution": map[string]any{"width": 640, "height": 360}},
Body: map[string]any{"video_url": "https://assets.example.test/source.mp4"},
Candidate: storeCandidateWithConfig(server.URL, "topaz", map[string]any{"apiKey": "topaz-key"}, map[string]any{
"pollIntervalMs": 1, "pollTimeoutMs": test.timeoutMS,
}),
})
if err == nil || !strings.Contains(err.Error(), test.want) {
t.Fatalf("expected %q error, got %v", test.want, err)
}
})
}
}
func TestTopazClientRejectsOversizedInputBeforeProbe(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path == "/source.mp4" {
_, _ = io.WriteString(w, "too-large")
return
}
http.NotFound(w, r)
}))
defer server.Close()
probeCalls := 0
_, err := (TopazClient{Probe: func(context.Context, string) (TopazVideoMetadata, error) {
probeCalls++
return TopazVideoMetadata{}, nil
}}).Run(context.Background(), Request{
Body: map[string]any{"video_url": server.URL + "/source.mp4"},
Candidate: storeCandidateWithConfig(server.URL, "topaz", map[string]any{"apiKey": "topaz-key"}, map[string]any{
"allowPrivateSourceDownloads": true, "maxInputBytes": 3,
}),
})
if err == nil || !strings.Contains(err.Error(), "size limit") {
t.Fatalf("expected input size limit error, got %v", err)
}
if probeCalls != 0 {
t.Fatalf("oversized input was probed %d times", probeCalls)
}
}
func storeCandidate(baseURL, specType string, credentials map[string]any) store.RuntimeModelCandidate {
return storeCandidateWithConfig(baseURL, specType, credentials, nil)
}
func storeCandidateWithConfig(baseURL, specType string, credentials, config map[string]any) store.RuntimeModelCandidate {
return store.RuntimeModelCandidate{BaseURL: baseURL, SpecType: specType, Provider: specType, Credentials: credentials, PlatformConfig: config}
}
func toJSONForTest(value any) string {
raw, _ := json.Marshal(value)
return string(raw)
}
func containsStringForTest(values []string, expected string) bool {
for _, value := range values {
if value == expected {
return true
}
}
return false
}
+164 -4
View File
@@ -1306,6 +1306,25 @@ func TestGeminiClientImageGenerateBuildsNativeImageBody(t *testing.T) {
}
}
func TestGeminiGenerationConfigInitializesMissingImageConfig(t *testing.T) {
config := geminiGenerationConfig(map[string]any{
"aspect_ratio": "16:9",
"resolution": "4K",
}, true)
imageConfig, ok := config["imageConfig"].(map[string]any)
if !ok {
t.Fatalf("imageConfig should be initialized, got %+v", config)
}
if imageConfig["aspectRatio"] != "16:9" || imageConfig["imageSize"] != "4K" {
t.Fatalf("unexpected imageConfig: %+v", imageConfig)
}
modalities, ok := config["responseModalities"].([]any)
if !ok || len(modalities) != 1 || modalities[0] != "IMAGE" {
t.Fatalf("image response modality should be initialized, got %+v", config)
}
}
func TestGeminiClientImageEditPreservesNativeContentsAndFileData(t *testing.T) {
var captured map[string]any
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -1539,10 +1558,39 @@ func TestGeminiClientChatConvertsFunctionCallResponse(t *testing.T) {
}
func TestGeminiURLAcceptsVersionedBaseURL(t *testing.T) {
got := geminiURL("https://generativelanguage.googleapis.com/v1beta", "gemini-2.5-flash", "test-key")
want := "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent?key=test-key"
if got != want {
t.Fatalf("unexpected gemini url: %s", got)
tests := []struct {
name string
baseURL string
want string
}{
{
name: "default version",
baseURL: "https://generativelanguage.googleapis.com",
want: "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent?key=test-key",
},
{
name: "google beta version",
baseURL: "https://generativelanguage.googleapis.com/v1beta",
want: "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent?key=test-key",
},
{
name: "compatible v1 endpoint",
baseURL: "https://cloud.dataeyes.ai/v1",
want: "https://cloud.dataeyes.ai/v1/models/gemini-2.5-flash:generateContent?key=test-key",
},
{
name: "openai suffix after version",
baseURL: "https://generativelanguage.googleapis.com/v1beta/openai",
want: "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent?key=test-key",
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
got := geminiURL(test.baseURL, "gemini-2.5-flash", "test-key")
if got != test.want {
t.Fatalf("unexpected gemini url: %s", got)
}
})
}
}
@@ -1934,6 +1982,77 @@ func TestVolcesClientVideoSubmitsAndPollsTask(t *testing.T) {
}
}
func TestVolcesClientVideoRetriesTransientPollAndKeepsOfficialResult(t *testing.T) {
polls := 0
persisted := make([]string, 0)
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.Method + " " + r.URL.Path {
case "POST /contents/generations/tasks":
_ = json.NewEncoder(w).Encode(map[string]any{"id": "cgt-retry"})
case "GET /contents/generations/tasks/cgt-retry":
polls++
if polls == 1 {
w.WriteHeader(http.StatusServiceUnavailable)
_, _ = w.Write([]byte(`{"error":{"message":"try later"}}`))
return
}
_ = json.NewEncoder(w).Encode(map[string]any{
"id": "cgt-retry", "model": "doubao-seedance-2-0-mini-260615", "status": "succeeded",
"created_at": 123, "updated_at": 124, "content": map[string]any{"video_url": "https://example.com/retry.mp4"},
"usage": map[string]any{"total_tokens": 8}, "seed": 7,
})
default:
t.Fatalf("unexpected request %s %s", r.Method, r.URL.Path)
}
}))
defer server.Close()
response, err := (VolcesClient{HTTPClient: server.Client()}).Run(context.Background(), Request{
Kind: "videos.generations", Model: "seedance", Body: map[string]any{"model": "seedance", "prompt": "retry"},
Candidate: store.RuntimeModelCandidate{BaseURL: server.URL, ProviderModelName: "doubao-seedance-2-0-mini-260615", Credentials: map[string]any{"apiKey": "key"}, PlatformConfig: map[string]any{"volcesPollIntervalMs": 100, "volcesPollRetryMaxMs": 100, "volcesPollTimeoutSeconds": 2}},
OnRemoteTaskPolled: func(remoteTaskID string, payload map[string]any) error {
persisted = append(persisted, remoteTaskID+":"+stringFromAny(payload["status"]))
return nil
},
})
if err != nil {
t.Fatalf("run retrying Volces video: %v", err)
}
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)
}
}
func TestVolcesClientDeletesOfficialVideoTask(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodDelete || r.URL.Path != "/contents/generations/tasks/cgt-delete" {
t.Fatalf("unexpected delete request %s %s", r.Method, r.URL.Path)
}
if r.Header.Get("Authorization") != "Bearer delete-key" {
t.Fatalf("unexpected delete authorization: %q", r.Header.Get("Authorization"))
}
_ = json.NewEncoder(w).Encode(map[string]any{"id": "cgt-delete", "status": "cancelled"})
}))
defer server.Close()
result, _, err := (VolcesClient{HTTPClient: server.Client()}).DeleteVideoTask(context.Background(), Request{
Candidate: store.RuntimeModelCandidate{BaseURL: server.URL, Credentials: map[string]any{"apiKey": "delete-key"}},
RemoteTaskID: "cgt-delete",
})
if err != nil || result["status"] != "cancelled" {
t.Fatalf("unexpected delete response result=%+v err=%v", result, err)
}
}
func TestVolcesCancelledTaskUsesDedicatedCancellationCode(t *testing.T) {
if got := volcesTaskErrorCode(map[string]any{"status": "cancelled"}); got != "volces_task_cancelled" {
t.Fatalf("cancelled task error code = %q", got)
}
}
func TestVolcesClientVideoRejectsDuplicateFirstFrameBeforeSubmit(t *testing.T) {
var submitted bool
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -2612,6 +2731,47 @@ func TestKelingOmniPayloadConvertsGatewayContent(t *testing.T) {
}
}
func TestKelingOmniPayloadPreservesLegacyV1Options(t *testing.T) {
payload, cleanupIDs, err := (KelingClient{}).kelingOmniPayload(context.Background(), Request{
Kind: "videos.generations",
ModelType: "omni_video",
Body: map[string]any{
"multi_shot": true,
"shot_type": "customize",
"multi_prompt": []any{
map[string]any{"index": 1, "prompt": "镜头一", "duration": 7},
map[string]any{"index": 2, "prompt": "镜头二", "duration": 8},
},
"resolution": "1080p",
"callback_url": "https://example.com/callback",
"external_task_id": "client-task-1",
"watermark_info": map[string]any{"enabled": true},
"voice_list": []any{map[string]any{"voice_id": "voice-1"}},
},
Candidate: store.RuntimeModelCandidate{
Provider: "keling",
ProviderModelName: "kling-v3-omni",
Capabilities: map[string]any{"omni_video": map[string]any{}},
},
}, "token")
if err != nil {
t.Fatalf("build keling legacy V1 payload: %v", err)
}
if len(cleanupIDs) != 0 {
t.Fatalf("unexpected cleanup ids: %+v", cleanupIDs)
}
if payload["multi_shot"] != true || payload["shot_type"] != "customize" || payload["duration"] != "15" {
t.Fatalf("unexpected multi-shot payload: %+v", payload)
}
if payload["callback_url"] != "https://example.com/callback" || payload["external_task_id"] != "client-task-1" {
t.Fatalf("legacy task options were not preserved: %+v", payload)
}
watermark := mapFromAny(payload["watermark_info"])
if watermark["enabled"] != true || len(mapListFromAny(payload["voice_list"])) != 1 {
t.Fatalf("watermark or voice options were not preserved: %+v", payload)
}
}
func TestKelingClientVideoResumePollsWithoutSubmitting(t *testing.T) {
var submitCalled bool
var pollPath string
+14 -3
View File
@@ -61,11 +61,13 @@ func geminiURL(baseURL string, model string, apiKey string) string {
base = "https://generativelanguage.googleapis.com"
}
base = strings.TrimSuffix(base, "/openai")
if strings.HasSuffix(base, "/v1beta") {
base = strings.TrimSuffix(base, "/v1beta")
if !strings.HasSuffix(base, "/v1") &&
!strings.HasSuffix(base, "/v1beta") &&
!strings.HasSuffix(base, "/v1alpha") {
base += "/v1beta"
}
escapedModel := url.PathEscape(model)
return fmt.Sprintf("%s/v1beta/models/%s:generateContent?key=%s", base, escapedModel, url.QueryEscape(apiKey))
return fmt.Sprintf("%s/models/%s:generateContent?key=%s", base, escapedModel, url.QueryEscape(apiKey))
}
func geminiBody(request Request) map[string]any {
@@ -195,14 +197,23 @@ func geminiApplyRequestOptions(body map[string]any, request Request, imageRespon
func geminiGenerationConfig(body map[string]any, imageResponse bool) map[string]any {
source := mapFromAny(firstPresent(body["generationConfig"], body["generation_config"]))
out := cloneMapAny(source)
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"])))
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 imageConfig == nil {
imageConfig = map[string]any{}
}
imageConfig["imageSize"] = imageSize
out["imageConfig"] = imageConfig
delete(out, "image_config")
+145 -28
View File
@@ -338,8 +338,11 @@ func kelingVideoPayload(ctx context.Context, request Request) (map[string]any, s
if value, ok := body["cfg_scale"]; ok && numericValue(value, 0) > 0 {
payload["cfg_scale"] = value
}
if boolValue(body, "audio") || boolValue(body, "output_audio") {
payload["sound"] = "on"
if sound, ok := kelingSoundSetting(body); ok {
if sound == "on" && !kelingSupportsGeneratedSound(request.Candidate) {
return nil, "", &ClientError{Code: "invalid_parameter", Message: "kling-video-o1 does not support generated audio", StatusCode: http.StatusBadRequest, Retryable: false}
}
payload["sound"] = sound
}
if mode := kelingModeByResolution(firstNonEmptyStringValue(body, "resolution", "size")); mode != "" {
payload["mode"] = mode
@@ -420,15 +423,21 @@ func (c KelingClient) kelingOmniPayload(ctx context.Context, request Request, to
}
uploadedElementIDs = append(uploadedElementIDs, createdIDs...)
shots := kelingShotPrompts(content)
hasMultiPrompt := len(shots) > 0
rawMultiPrompt := mapListFromAny(body["multi_prompt"])
hasMultiPrompt := len(shots) > 0 || len(rawMultiPrompt) > 0
multiShot := boolValue(body, "multi_shot") || hasMultiPrompt
hasVideo := len(videos) > 0
hasVideoEdit := kelingHasBaseVideo(videos)
hasFirstFrame := kelingHasFirstFrame(images)
watermarkEnabled := boolValue(body, "watermark")
if watermarkInfo := mapFromAny(body["watermark_info"]); watermarkInfo != nil {
watermarkEnabled = boolValue(watermarkInfo, "enabled")
}
payload := map[string]any{
"model_name": upstreamModelName(request.Candidate),
"model_name": kelingOmniUpstreamModelName(request.Candidate),
"mode": kelingModeByResolution(firstNonEmptyStringValue(body, "resolution", "size")),
"watermark_info": map[string]any{"enabled": false},
"watermark_info": map[string]any{"enabled": watermarkEnabled},
"negative_prompt": strings.TrimSpace(stringFromAny(body["negative_prompt"])),
}
if !hasMultiPrompt {
@@ -449,29 +458,67 @@ func (c KelingClient) kelingOmniPayload(ctx context.Context, request Request, to
if len(elements) > 0 {
payload["element_list"] = elements
}
if (boolValue(body, "audio") || boolValue(body, "output_audio")) && !hasVideo {
payload["sound"] = "on"
if voices := mapListFromAny(body["voice_list"]); len(voices) > 0 {
payload["voice_list"] = voices
}
if hasMultiPrompt {
payload["multi_shot"] = true
payload["shot_type"] = "customize"
total := 0.0
multiPrompt := make([]any, 0, len(shots))
for index, shot := range shots {
duration := shot.duration
if duration <= 0 {
duration = 5
}
total += duration
multiPrompt = append(multiPrompt, map[string]any{
"index": index + 1,
"prompt": shot.text,
"duration": fmtDuration(duration, 5),
})
if sound, ok := kelingSoundSetting(body); ok {
if sound == "on" && !kelingSupportsGeneratedSound(request.Candidate) {
return nil, nil, &ClientError{Code: "invalid_parameter", Message: "kling-video-o1 does not support generated audio", StatusCode: http.StatusBadRequest, Retryable: false}
}
delete(payload, "prompt")
payload["multi_prompt"] = multiPrompt
payload["duration"] = fmtDuration(total, 0)
if !hasVideo {
payload["sound"] = sound
}
}
if multiShot {
payload["multi_shot"] = true
shotType := strings.TrimSpace(firstNonEmptyStringValue(body, "shot_type", "shotType"))
if shotType == "" {
if hasMultiPrompt {
shotType = "customize"
} else {
shotType = "intelligence"
}
}
payload["shot_type"] = shotType
if shotType == "customize" {
total := 0.0
multiPrompt := make([]any, 0, len(rawMultiPrompt)+len(shots))
if len(rawMultiPrompt) > 0 {
for index, item := range rawMultiPrompt {
duration := numericValue(item["duration"], 0)
total += duration
multiPrompt = append(multiPrompt, map[string]any{
"index": int(math.Round(numericValue(item["index"], float64(index+1)))),
"prompt": strings.TrimSpace(stringFromAny(item["prompt"])),
"duration": fmtDuration(duration, 0),
})
}
} else {
for index, shot := range shots {
duration := shot.duration
if duration <= 0 {
duration = 5
}
total += duration
multiPrompt = append(multiPrompt, map[string]any{
"index": index + 1,
"prompt": shot.text,
"duration": fmtDuration(duration, 5),
})
}
}
delete(payload, "prompt")
payload["multi_prompt"] = multiPrompt
if total > 0 {
payload["duration"] = fmtDuration(total, 0)
}
}
}
if callbackURL := strings.TrimSpace(firstNonEmptyStringValue(body, "callback_url", "callbackUrl")); callbackURL != "" {
payload["callback_url"] = callbackURL
}
if externalTaskID := strings.TrimSpace(firstNonEmptyStringValue(body, "external_task_id", "externalTaskId")); externalTaskID != "" {
payload["external_task_id"] = externalTaskID
}
deleteEmptyStringFields(payload)
if hasVideoEdit {
@@ -689,6 +736,18 @@ func kelingIsOmniRequest(request Request) bool {
request.Candidate.Capabilities["omni"] != nil
}
func kelingOmniUpstreamModelName(candidate store.RuntimeModelCandidate) string {
model := strings.TrimSpace(upstreamModelName(candidate))
switch strings.ToLower(model) {
case "kling-o1":
return "kling-video-o1"
case "kling-3.0-omni":
return "kling-v3-omni"
default:
return model
}
}
func kelingIs30TurboRequest(request Request) bool {
switch strings.ToLower(strings.TrimSpace(upstreamModelName(request.Candidate))) {
case "kling-3.0-turbo", "kling-v3-turbo", "kling-3-0-turbo":
@@ -1034,6 +1093,54 @@ func kelingModeByResolution(resolution string) string {
}
}
func kelingSoundSetting(body map[string]any) (string, bool) {
if sound := strings.ToLower(strings.TrimSpace(stringFromAny(body["sound"]))); sound == "on" || sound == "off" {
return sound, true
}
for _, key := range []string{"audio", "output_audio", "generate_audio"} {
if enabled, ok := kelingBoolFieldValue(body, key); ok {
if enabled {
return "on", true
}
return "off", true
}
}
return "", false
}
func kelingSupportsGeneratedSound(candidate store.RuntimeModelCandidate) bool {
switch strings.ToLower(strings.TrimSpace(upstreamModelName(candidate))) {
case "kling-o1", "kling-video-o1":
return false
default:
return true
}
}
func kelingWatermarkEnabled(body map[string]any) bool {
if enabled, ok := kelingBoolFieldValue(body, "watermark"); ok {
return enabled
}
info := mapFromAny(body["watermark_info"])
if info == nil {
return false
}
enabled, _ := kelingBoolFieldValue(info, "enabled")
return enabled
}
func kelingBoolFieldValue(body map[string]any, key string) (bool, bool) {
if body == nil {
return false, false
}
value, ok := body[key]
if !ok {
return false, false
}
typed, ok := value.(bool)
return typed, ok
}
func kelingCameraControl(body map[string]any) map[string]any {
cameraControl := strings.TrimSpace(stringFromAny(body["camera_control"]))
if cameraControl == "" {
@@ -1140,20 +1247,30 @@ func kelingVideoSuccessResult(request Request, upstreamTaskID string, raw map[st
continue
}
item := map[string]any{"url": url, "video_url": url, "type": "video"}
if duration := intFromAny(video["duration"]); duration > 0 {
if id := strings.TrimSpace(stringFromAny(video["id"])); id != "" {
item["id"] = id
}
if duration := firstPresent(video["duration"]); duration != nil && strings.TrimSpace(stringFromAny(duration)) != "" {
item["duration"] = duration
}
if watermarkURL := strings.TrimSpace(stringFromAny(video["watermark_url"])); watermarkURL != "" {
item["watermark_url"] = watermarkURL
}
items = append(items, item)
}
created := intFromAny(data["created_at"])
if created == 0 {
created = int(nowUnix())
}
modelName := upstreamModelName(request.Candidate)
if kelingIsOmniRequest(request) {
modelName = kelingOmniUpstreamModelName(request.Candidate)
}
return map[string]any{
"id": upstreamTaskID,
"object": "video.generation",
"created": created,
"model": upstreamModelName(request.Candidate),
"model": modelName,
"status": "succeeded",
"upstream_task_id": upstreamTaskID,
"data": items,
@@ -0,0 +1,166 @@
package clients
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestKelingOmniPayloadPreservesCompatibleSettings(t *testing.T) {
payload, cleanupIDs, err := (KelingClient{}).kelingOmniPayload(context.Background(), Request{
Kind: "videos.generations",
ModelType: "omni_video",
Body: map[string]any{
"prompt": "A product reveal",
"duration": 3,
"resolution": "720p",
"aspect_ratio": "16:9",
"audio": false,
"watermark_info": map[string]any{"enabled": true},
"external_task_id": "external-1",
},
Candidate: store.RuntimeModelCandidate{
Provider: "keling",
ProviderModelName: "kling-video-o1",
Capabilities: map[string]any{"omni_video": map[string]any{}},
},
}, "token")
if err != nil {
t.Fatalf("build compatible Omni payload: %v", err)
}
if len(cleanupIDs) != 0 ||
payload["model_name"] != "kling-video-o1" ||
payload["mode"] != "std" ||
payload["sound"] != "off" ||
payload["duration"] != "3" ||
payload["aspect_ratio"] != "16:9" ||
payload["external_task_id"] != "external-1" {
t.Fatalf("unexpected compatible Omni payload: %+v", payload)
}
watermark, _ := payload["watermark_info"].(map[string]any)
if watermark["enabled"] != true {
t.Fatalf("watermark setting was not preserved: %+v", payload)
}
}
func TestKelingOmniUpstreamModelNameSeparatesGatewayAliases(t *testing.T) {
tests := map[string]string{
"kling-o1": "kling-video-o1",
"kling-video-o1": "kling-video-o1",
"kling-3.0-omni": "kling-v3-omni",
"kling-v3-omni": "kling-v3-omni",
}
for configured, want := range tests {
got := kelingOmniUpstreamModelName(store.RuntimeModelCandidate{ProviderModelName: configured})
if got != want {
t.Fatalf("configured=%s got=%s want=%s", configured, got, want)
}
}
}
func TestKelingOmniRejectsGeneratedAudioForO1(t *testing.T) {
_, _, err := (KelingClient{}).kelingOmniPayload(context.Background(), Request{
Body: map[string]any{
"prompt": "A beach",
"duration": 5,
"resolution": "1080p",
"aspect_ratio": "9:16",
"audio": true,
},
Candidate: store.RuntimeModelCandidate{ProviderModelName: "kling-video-o1"},
}, "token")
if err == nil || ErrorCode(err) != "invalid_parameter" {
t.Fatalf("expected generated-audio rejection for O1, got %v", err)
}
}
func TestKelingOmniResumeReturnsUpstreamFailureCodeAndModel(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodGet || r.URL.Path != "/videos/omni-video/remote-failed" {
t.Fatalf("unexpected request %s %s", r.Method, r.URL.Path)
}
if r.Header.Get("Authorization") != "Bearer upstream-key" {
t.Fatalf("unexpected Authorization header: %q", r.Header.Get("Authorization"))
}
_ = json.NewEncoder(w).Encode(map[string]any{
"code": 0,
"request_id": "failure-request",
"data": map[string]any{
"task_id": "remote-failed",
"task_status": "failed",
"task_status_msg": "content policy rejection",
},
})
}))
defer server.Close()
_, err := (KelingClient{HTTPClient: server.Client()}).Run(context.Background(), Request{
Kind: "videos.generations",
ModelType: "omni_video",
RemoteTaskID: "remote-failed",
RemoteTaskPayload: map[string]any{"endpoint": "/videos/omni-video"},
Candidate: store.RuntimeModelCandidate{
BaseURL: server.URL,
Provider: "keling",
ProviderModelName: "kling-v3-omni",
Credentials: map[string]any{"apiKey": "upstream-key"},
PlatformConfig: map[string]any{
"kelingPollIntervalMs": 10,
"kelingPollTimeoutSeconds": 1,
},
},
})
if err == nil || ErrorCode(err) != "keling_task_failed" || !strings.Contains(err.Error(), "content policy rejection") {
t.Fatalf("expected preserved Keling task failure, got code=%q err=%v", ErrorCode(err), err)
}
}
func TestKelingOmniPayloadPreservesIntelligentMultiShot(t *testing.T) {
payload, _, err := (KelingClient{}).kelingOmniPayload(context.Background(), Request{
Kind: "videos.generations",
ModelType: "omni_video",
Body: map[string]any{
"prompt": "Create three coherent shots",
"duration": 5,
"resolution": "1080p",
"aspect_ratio": "9:16",
"multi_shot": true,
"shot_type": "intelligence",
},
Candidate: store.RuntimeModelCandidate{
ProviderModelName: "kling-v3-omni",
Capabilities: map[string]any{"omni_video": map[string]any{}},
},
}, "token")
if err != nil {
t.Fatalf("build intelligent multi-shot payload: %v", err)
}
if payload["multi_shot"] != true || payload["shot_type"] != "intelligence" || payload["prompt"] != "Create three coherent shots" {
t.Fatalf("unexpected intelligent multi-shot payload: %+v", payload)
}
}
func TestKelingVideoSuccessResultPreservesOfficialVideoMetadata(t *testing.T) {
result := kelingVideoSuccessResult(Request{Candidate: store.RuntimeModelCandidate{ProviderModelName: "kling-video-o1"}}, "remote-1", map[string]any{
"data": map[string]any{
"task_result": map[string]any{
"videos": []any{map[string]any{
"id": "video-1",
"url": "https://example.com/video.mp4",
"watermark_url": "https://example.com/watermarked.mp4",
"duration": "3",
}},
},
},
})
data, _ := result["data"].([]any)
video, _ := data[0].(map[string]any)
if video["id"] != "video-1" || video["watermark_url"] != "https://example.com/watermarked.mp4" || video["duration"] != "3" {
t.Fatalf("official video metadata was lost: %+v", video)
}
}
@@ -0,0 +1,83 @@
package clients
import (
"context"
"os"
"strings"
"testing"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
// TestKelingLegacyV1LiveOmni is opt-in because it creates billable upstream
// video tasks. Credentials must only be supplied through local environment
// variables; the test never prints them.
func TestKelingLegacyV1LiveOmni(t *testing.T) {
if strings.TrimSpace(os.Getenv("KELING_LIVE_TEST")) != "1" {
t.Skip("set KELING_LIVE_TEST=1 to run billable Kling V1 integration tests")
}
baseURL := strings.TrimRight(strings.TrimSpace(os.Getenv("KELING_TEST_BASE_URL")), "/")
accessKey := strings.TrimSpace(os.Getenv("KELING_TEST_ACCESS_KEY"))
secretKey := strings.TrimSpace(os.Getenv("KELING_TEST_SECRET_KEY"))
if baseURL == "" || accessKey == "" || secretKey == "" {
t.Fatal("KELING_TEST_BASE_URL, KELING_TEST_ACCESS_KEY, and KELING_TEST_SECRET_KEY are required")
}
models := []string{"kling-video-o1", "kling-v3-omni"}
if selected := strings.TrimSpace(os.Getenv("KELING_LIVE_TEST_MODELS")); selected != "" {
models = strings.Split(selected, ",")
}
for _, model := range models {
model = strings.TrimSpace(model)
t.Run(model, func(t *testing.T) {
duration := 3
if model == "kling-video-o1" {
duration = 5
}
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
defer cancel()
response, err := (KelingClient{}).Run(ctx, Request{
Kind: "videos.generations",
ModelType: "omni_video",
Model: model,
Body: map[string]any{
"prompt": "清晨的湖面上,一只白色纸鹤缓慢飞过,镜头平稳推进",
"duration": duration,
"aspect_ratio": "16:9",
"resolution": "720p",
"sound": "off",
},
Candidate: store.RuntimeModelCandidate{
BaseURL: baseURL,
Provider: "keling",
AuthType: "AccessKey-SecretKey",
ProviderModelName: model,
Credentials: map[string]any{
"accessKey": accessKey,
"secretKey": secretKey,
},
Capabilities: map[string]any{"omni_video": map[string]any{}},
PlatformConfig: map[string]any{
"kelingPollIntervalMs": 5000,
"kelingPollTimeoutSeconds": 840,
},
},
})
if err != nil {
t.Fatalf("Kling V1 %s live task failed: %v", model, err)
}
items, _ := response.Result["data"].([]any)
if len(items) == 0 {
t.Fatalf("Kling V1 %s returned no video", model)
}
item, _ := items[0].(map[string]any)
if strings.TrimSpace(stringFromAny(item["url"])) == "" {
t.Fatalf("Kling V1 %s returned an empty video URL", model)
}
if strings.TrimSpace(response.RequestID) == "" {
t.Fatalf("Kling V1 %s returned an empty request id", model)
}
})
}
}
+617
View File
@@ -0,0 +1,617 @@
package clients
import (
"bytes"
"context"
"crypto/md5"
"encoding/hex"
"encoding/json"
"errors"
"fmt"
"io"
"math"
"net"
"net/http"
"net/url"
"os"
"os/exec"
"path/filepath"
"strconv"
"strings"
"time"
)
const (
topazDefaultPollInterval = 15 * time.Second
topazDefaultPollTimeout = 60 * time.Minute
topazDefaultMaxInput = int64(2 << 30)
)
type TopazClient struct {
HTTPClient *http.Client
Probe func(context.Context, string) (TopazVideoMetadata, error)
}
type TopazVideoMetadata struct {
Width int
Height int
Duration float64
FrameRate float64
FrameCount int
HasAudio bool
}
type topazSource struct {
Path string
Container string
Size int64
MD5 string
Resolution map[string]any
Duration float64
FrameRate float64
FrameCount int
}
func (c TopazClient) Run(ctx context.Context, request Request) (Response, error) {
startedAt := time.Now()
apiKey := credential(request.Candidate.Credentials, "apiKey", "api_key", "key", "token")
if apiKey == "" {
return Response{}, &ClientError{Code: "missing_credentials", Message: "Topaz API key is required", Retryable: false}
}
requestID := strings.TrimSpace(request.RemoteTaskID)
payload := cloneMapAny(request.RemoteTaskPayload)
var source *topazSource
var target map[string]int
if requestID == "" || strings.TrimSpace(firstNonEmptyString(payload["phase"])) != "uploaded" {
videoURL := strings.TrimSpace(firstNonEmptyString(request.Body["video_url"], request.Body["videoUrl"]))
if videoURL == "" {
return Response{}, &ClientError{Code: "invalid_parameter", Message: "video_url is required", Param: "video_url", StatusCode: http.StatusBadRequest}
}
prepared, err := c.prepareSource(ctx, request, videoURL)
if err != nil {
return Response{}, err
}
source = &prepared
defer os.Remove(prepared.Path)
target = topazTargetResolution(request.Body, prepared.Resolution)
if requestID == "" {
created, err := c.createRequest(ctx, request, apiKey, prepared, target)
if err != nil {
return Response{}, annotateResponseError(err, "", startedAt, time.Now())
}
requestID = strings.TrimSpace(firstNonEmptyString(created["requestId"], created["request_id"], created["id"]))
if requestID == "" {
return Response{}, &ClientError{Code: "invalid_response", Message: "Topaz create response is missing requestId", Retryable: false}
}
payload = map[string]any{"phase": "created", "targetResolution": target}
if request.OnRemoteTaskSubmitted != nil {
if err := request.OnRemoteTaskSubmitted(requestID, payload); err != nil {
return Response{}, err
}
}
}
if err := c.uploadSource(ctx, request, apiKey, requestID, prepared); err != nil {
return Response{}, annotateResponseError(err, requestID, startedAt, time.Now())
}
payload = map[string]any{"phase": "uploaded", "targetResolution": target}
if request.OnRemoteTaskPolled != nil {
if err := request.OnRemoteTaskPolled(requestID, payload); err != nil {
return Response{}, err
}
}
}
if target == nil {
target = topazTargetFromPayload(payload, request.Body)
}
completed, err := c.poll(ctx, request, apiKey, requestID, target)
if err != nil {
return Response{}, annotateResponseError(err, requestID, startedAt, time.Now())
}
outputURL := topazDownloadURL(completed)
if outputURL == "" {
return Response{}, &ClientError{Code: "invalid_response", Message: "Topaz completed without a download URL", RequestID: requestID, Retryable: false}
}
metadata, err := c.probe(ctx, outputURL)
if err != nil {
return Response{}, &ClientError{Code: "invalid_response", Message: "cannot validate Topaz output metadata: " + err.Error(), RequestID: requestID, Retryable: true}
}
if target["width"] > 0 && target["height"] > 0 && (metadata.Width+4 < target["width"] || metadata.Height+4 < target["height"]) {
return Response{}, &ClientError{Code: "invalid_response", Message: fmt.Sprintf("Topaz output resolution %dx%d does not reach target %dx%d", metadata.Width, metadata.Height, target["width"], target["height"]), RequestID: requestID, Retryable: false}
}
finishedAt := time.Now()
resultItem := map[string]any{
"type": "video",
"url": outputURL,
"video_url": outputURL,
"width": metadata.Width,
"height": metadata.Height,
"target_resolution": fmt.Sprintf("%dx%d", target["width"], target["height"]),
"duration": metadata.Duration,
"target_frame_rate": metadata.FrameRate,
"slow_motion_rate": numericValue(request.Body["slow_motion_rate"], 1),
}
if source != nil {
resultItem["source_resolution"] = fmt.Sprintf("%vx%v", source.Resolution["width"], source.Resolution["height"])
resultItem["source_frame_rate"] = source.FrameRate
resultItem["duration"] = source.Duration
}
return Response{
Result: map[string]any{
"status": "success",
"model": request.Model,
"task_id": requestID,
"upstream_task_id": requestID,
"data": []any{resultItem},
},
RequestID: requestID,
Progress: append(providerProgress(request), Progress{Phase: "polling", Progress: 0.9, Message: "Topaz upscale completed", Payload: map[string]any{"upstreamTaskId": requestID}}),
ResponseStartedAt: startedAt,
ResponseFinishedAt: finishedAt,
ResponseDurationMS: responseDurationMS(startedAt, finishedAt),
}, nil
}
func (c TopazClient) prepareSource(ctx context.Context, request Request, rawURL string) (topazSource, error) {
parsed, err := url.Parse(rawURL)
if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") {
return topazSource{}, &ClientError{Code: "invalid_parameter", Message: "video_url must be an http(s) URL", Param: "video_url", StatusCode: http.StatusBadRequest}
}
req, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil)
if err != nil {
return topazSource{}, err
}
downloadClient := topazSourceHTTPClient(httpClient(request.HTTPClient, c.HTTPClient), boolishDefault(request.Candidate.PlatformConfig["allowPrivateSourceDownloads"], false))
resp, err := downloadClient.Do(req)
if err != nil {
return topazSource{}, &ClientError{Code: "network", Message: err.Error(), Retryable: true}
}
defer resp.Body.Close()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
return topazSource{}, &ClientError{Code: "invalid_parameter", Message: "video_url download failed: " + resp.Status, Param: "video_url", StatusCode: http.StatusBadRequest}
}
limit := int64(numericValue(firstPresent(request.Candidate.PlatformConfig["maxInputBytes"], request.Candidate.PlatformConfig["max_input_bytes"]), float64(topazDefaultMaxInput)))
if limit <= 0 {
limit = topazDefaultMaxInput
}
tmp, err := os.CreateTemp("", "easyai-topaz-*"+topazContainerExtension(parsed.Path))
if err != nil {
return topazSource{}, err
}
path := tmp.Name()
cleanup := func() {
_ = tmp.Close()
_ = os.Remove(path)
}
hash := md5.New()
written, err := io.Copy(io.MultiWriter(tmp, hash), io.LimitReader(resp.Body, limit+1))
closeErr := tmp.Close()
if err != nil || closeErr != nil || written <= 0 || written > limit {
cleanup()
if written > limit {
return topazSource{}, &ClientError{Code: "invalid_parameter", Message: "video source exceeds configured size limit", Param: "video_url", StatusCode: http.StatusBadRequest}
}
return topazSource{}, firstError(err, closeErr)
}
metadata, err := c.probe(ctx, path)
if err != nil || metadata.Width <= 0 || metadata.Height <= 0 {
cleanup()
return topazSource{}, &ClientError{Code: "invalid_parameter", Message: "cannot probe source video metadata", Param: "video_url", StatusCode: http.StatusBadRequest}
}
frameRate := metadata.FrameRate
if frameRate <= 0 {
frameRate = numericValue(request.Candidate.PlatformConfig["frameRate"], 30)
}
duration := math.Max(1, metadata.Duration)
frameCount := metadata.FrameCount
if frameCount <= 0 {
frameCount = int(math.Round(duration * frameRate))
}
return topazSource{
Path: path, Container: strings.TrimPrefix(strings.ToLower(filepath.Ext(path)), "."), Size: written,
MD5: hex.EncodeToString(hash.Sum(nil)), Resolution: map[string]any{"width": metadata.Width, "height": metadata.Height},
Duration: duration, FrameRate: frameRate, FrameCount: frameCount,
}, nil
}
func (c TopazClient) createRequest(ctx context.Context, request Request, apiKey string, source topazSource, target map[string]int) (map[string]any, error) {
model := topazModelName(firstNonEmptyString(request.Candidate.ProviderModelName, request.Model))
scale := math.Max(float64(target["width"])/numericValue(source.Resolution["width"], 1), float64(target["height"])/numericValue(source.Resolution["height"], 1))
filter := map[string]any{"model": model}
if scale > 1 {
key := "scale"
if model == "slf-2" || model == "slhq-1" || model == "slm-1" || model == "slp-2.5" {
key = "upscaling_factor"
}
filter[key] = math.Round(scale*1000) / 1000
}
preserveAudio := boolishDefault(request.Body["preserve_audio"], true)
filters := []any{filter}
targetFrameRate := numericValue(firstPresent(request.Body["target_frame_rate"], request.Body["output_frame_rate"]), source.FrameRate)
slowMotionRate := numericValue(request.Body["slow_motion_rate"], 1)
if targetFrameRate <= 0 {
targetFrameRate = source.FrameRate
}
if slowMotionRate <= 0 {
slowMotionRate = 1
}
if targetFrameRate != source.FrameRate || slowMotionRate > 1 {
interpolationModel := strings.TrimSpace(firstNonEmptyString(request.Body["frame_interpolation_model"]))
if interpolationModel == "" {
return nil, &ClientError{Code: "invalid_parameter", Message: "frame_interpolation_model is required when output frame rate or slow motion changes", Param: "frame_interpolation_model", StatusCode: http.StatusBadRequest}
}
interpolation := map[string]any{"model": topazModelName(interpolationModel)}
if targetFrameRate > 0 {
interpolation["fps"] = targetFrameRate
}
if slowMotionRate > 1 {
interpolation["slowmo"] = slowMotionRate
}
filters = append(filters, interpolation)
}
body := map[string]any{
"source": map[string]any{"container": source.Container, "size": source.Size, "duration": source.Duration, "frameCount": source.FrameCount, "frameRate": source.FrameRate, "resolution": source.Resolution},
"filters": filters,
"output": map[string]any{"resolution": target, "frameRate": targetFrameRate, "audioCodec": "AAC", "audioTransfer": map[bool]string{true: "Copy", false: "None"}[preserveAudio], "dynamicCompressionLevel": "High", "videoEncoder": "H265", "videoProfile": "Main", "container": "mp4"},
}
return c.topazJSON(ctx, request, apiKey, http.MethodPost, "/video/", body)
}
func (c TopazClient) uploadSource(ctx context.Context, request Request, apiKey, requestID string, source topazSource) error {
accepted, err := c.topazJSON(ctx, request, apiKey, http.MethodPatch, "/video/"+url.PathEscape(requestID)+"/accept", nil)
if err != nil {
return err
}
urls := stringList(accepted["urls"])
if len(urls) == 0 {
return &ClientError{Code: "invalid_response", Message: "Topaz accept response has no upload URLs", RequestID: requestID}
}
file, err := os.Open(source.Path)
if err != nil {
return err
}
defer file.Close()
chunkSize := (source.Size + int64(len(urls)) - 1) / int64(len(urls))
results := make([]any, 0, len(urls))
for index, uploadURL := range urls {
remaining := source.Size - int64(index)*chunkSize
if remaining <= 0 {
break
}
length := minInt64(chunkSize, remaining)
section := io.NewSectionReader(file, int64(index)*chunkSize, length)
req, err := http.NewRequestWithContext(ctx, http.MethodPut, uploadURL, section)
if err != nil {
return err
}
req.ContentLength = length
req.Header.Set("Content-Type", topazContainerContentType(source.Container))
resp, err := httpClient(request.HTTPClient, c.HTTPClient).Do(req)
if err != nil {
return &ClientError{Code: "network", Message: err.Error(), Retryable: true}
}
_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 1<<20))
resp.Body.Close()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
return &ClientError{Code: "provider_failed", Message: "Topaz part upload failed: " + resp.Status, StatusCode: resp.StatusCode, Retryable: resp.StatusCode >= 500}
}
etag := strings.Trim(resp.Header.Get("ETag"), `"`)
if etag == "" {
return &ClientError{Code: "invalid_response", Message: "Topaz part upload is missing ETag", Retryable: false}
}
results = append(results, map[string]any{"partNum": index + 1, "eTag": etag})
}
_, err = c.topazJSON(ctx, request, apiKey, http.MethodPatch, "/video/"+url.PathEscape(requestID)+"/complete-upload/", map[string]any{"md5Hash": source.MD5, "uploadResults": results})
return err
}
func (c TopazClient) poll(ctx context.Context, request Request, apiKey, requestID string, target map[string]int) (map[string]any, error) {
interval := universalDurationConfig(request.Candidate.PlatformConfig, topazDefaultPollInterval, "pollIntervalMs", "poll_interval_ms")
timeout := universalDurationConfig(request.Candidate.PlatformConfig, topazDefaultPollTimeout, "pollTimeoutMs", "poll_timeout_ms", "timeoutMs")
deadline := time.NewTimer(timeout)
defer deadline.Stop()
for {
status, err := c.topazJSON(ctx, request, apiKey, http.MethodGet, "/video/"+url.PathEscape(requestID)+"/status", nil)
if err != nil {
return nil, err
}
state := strings.ToLower(strings.TrimSpace(firstNonEmptyString(status["status"], status["state"])))
if request.OnRemoteTaskPolled != nil {
if err := request.OnRemoteTaskPolled(requestID, map[string]any{"phase": "uploaded", "status": state, "targetResolution": target}); err != nil {
return nil, err
}
}
switch state {
case "complete", "completed", "success", "succeeded":
return status, nil
case "failed", "failure", "cancelled", "canceled", "error":
return nil, &ClientError{Code: "provider_failed", Message: firstNonEmptyString(status["message"], "Topaz task failed"), RequestID: requestID, Retryable: false}
}
select {
case <-ctx.Done():
return nil, &ClientError{Code: "cancelled", Message: ctx.Err().Error(), RequestID: requestID, Retryable: true}
case <-deadline.C:
return nil, &ClientError{Code: "timeout", Message: "Topaz task polling timed out", RequestID: requestID, Retryable: true}
case <-time.After(interval):
}
}
}
func (c TopazClient) topazJSON(ctx context.Context, request Request, apiKey, method, path string, body map[string]any) (map[string]any, error) {
var reader io.Reader
if body != nil {
raw, _ := json.Marshal(body)
reader = bytes.NewReader(raw)
}
req, err := http.NewRequestWithContext(ctx, method, providerURL(request.Candidate.BaseURL, path), reader)
if err != nil {
return nil, err
}
req.Header.Set("X-API-Key", apiKey)
req.Header.Set("Accept", "application/json")
if body != nil {
req.Header.Set("Content-Type", "application/json")
}
resp, err := httpClient(request.HTTPClient, c.HTTPClient).Do(req)
if err != nil {
return nil, &ClientError{Code: "network", Message: err.Error(), Retryable: true}
}
result, decodeErr := decodeHTTPResponse(resp)
if decodeErr != nil {
return nil, decodeErr
}
return result, nil
}
func (c TopazClient) probe(ctx context.Context, target string) (TopazVideoMetadata, error) {
if c.Probe != nil {
return c.Probe(ctx, target)
}
return probeTopazVideo(ctx, target)
}
func probeTopazVideo(ctx context.Context, target string) (TopazVideoMetadata, error) {
if _, err := exec.LookPath("ffprobe"); err != nil {
return TopazVideoMetadata{}, err
}
probeCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
defer cancel()
output, err := exec.CommandContext(probeCtx, "ffprobe", "-v", "error", "-show_entries", "format=duration:stream=codec_type,width,height,avg_frame_rate,nb_frames", "-of", "json", target).Output()
if err != nil {
return TopazVideoMetadata{}, err
}
var decoded struct {
Format struct {
Duration string `json:"duration"`
} `json:"format"`
Streams []struct {
CodecType string `json:"codec_type"`
Width int `json:"width"`
Height int `json:"height"`
AverageRate string `json:"avg_frame_rate"`
FrameCount string `json:"nb_frames"`
} `json:"streams"`
}
if err := json.Unmarshal(output, &decoded); err != nil {
return TopazVideoMetadata{}, err
}
metadata := TopazVideoMetadata{}
metadata.Duration, _ = strconv.ParseFloat(decoded.Format.Duration, 64)
for _, stream := range decoded.Streams {
if stream.CodecType == "audio" {
metadata.HasAudio = true
}
if stream.CodecType != "video" || stream.Width <= 0 || stream.Height <= 0 {
continue
}
metadata.Width, metadata.Height = stream.Width, stream.Height
metadata.FrameRate = parseTopazRate(stream.AverageRate)
metadata.FrameCount, _ = strconv.Atoi(stream.FrameCount)
}
if metadata.Width <= 0 || metadata.Height <= 0 {
return TopazVideoMetadata{}, errors.New("video stream metadata is missing")
}
return metadata, nil
}
func topazTargetResolution(body map[string]any, source map[string]any) map[string]int {
width := int(numericValue(body["output_width"], 0))
height := int(numericValue(body["output_height"], 0))
if width > 0 && height > 0 {
return map[string]int{"width": width, "height": height}
}
value := strings.ToLower(strings.TrimSpace(firstNonEmptyString(body["target_resolution"], body["output_resolution"], "1080p")))
if parsedWidth, parsedHeight, ok := parseTopazSize(value); ok {
return map[string]int{"width": parsedWidth, "height": parsedHeight}
}
base := map[string][2]int{"480p": {854, 480}, "720p": {1280, 720}, "1080p": {1920, 1080}, "1440p": {2560, 1440}, "2k": {2560, 1440}, "2160p": {3840, 2160}, "4k": {3840, 2160}}
value = strings.TrimSuffix(value, "_upscale")
target, ok := base[value]
if !ok {
target = base["1080p"]
}
sourceWidth := numericValue(source["width"], 1)
sourceHeight := numericValue(source["height"], 1)
if sourceWidth == sourceHeight {
side := minInt(target[0], target[1])
return map[string]int{"width": side, "height": side}
}
if sourceWidth > sourceHeight {
return map[string]int{"width": int(math.Round(sourceWidth / sourceHeight * float64(target[1]))), "height": target[1]}
}
return map[string]int{"width": target[1], "height": int(math.Round(sourceHeight / sourceWidth * float64(target[1])))}
}
func topazTargetFromPayload(payload map[string]any, body map[string]any) map[string]int {
if raw, ok := payload["targetResolution"].(map[string]any); ok {
return map[string]int{"width": int(numericValue(raw["width"], 0)), "height": int(numericValue(raw["height"], 0))}
}
return topazTargetResolution(body, map[string]any{"width": 16, "height": 9})
}
func topazDownloadURL(result map[string]any) string {
if download, ok := result["download"].(map[string]any); ok {
return firstNonEmptyString(download["url"], download["downloadUrl"])
}
return firstNonEmptyString(result["download_url"], result["downloadUrl"], result["url"])
}
func topazModelName(model string) string {
switch strings.TrimSpace(model) {
case "easy-proteus-standard-4":
return "prob-4"
case "easy-starlight-fast-2":
return "slf-2"
case "easy-starlight-hq-1":
return "slhq-1"
case "easy-starlight-mini-1":
return "slm-1"
default:
return strings.TrimSpace(model)
}
}
func topazContainerExtension(path string) string {
ext := strings.ToLower(filepath.Ext(strings.Split(path, "?")[0]))
switch ext {
case ".mov", ".mkv", ".webm", ".mp4":
return ext
default:
return ".mp4"
}
}
func topazContainerContentType(container string) string {
switch strings.ToLower(container) {
case "mov":
return "video/quicktime"
case "mkv":
return "video/x-matroska"
case "webm":
return "video/webm"
default:
return "video/mp4"
}
}
func parseTopazRate(value string) float64 {
parts := strings.Split(value, "/")
if len(parts) == 2 {
numerator, _ := strconv.ParseFloat(parts[0], 64)
denominator, _ := strconv.ParseFloat(parts[1], 64)
if denominator > 0 {
return numerator / denominator
}
}
parsed, _ := strconv.ParseFloat(value, 64)
return parsed
}
func parseTopazSize(value string) (int, int, bool) {
parts := strings.Split(strings.ReplaceAll(value, " ", ""), "x")
if len(parts) != 2 {
return 0, 0, false
}
width, errWidth := strconv.Atoi(parts[0])
height, errHeight := strconv.Atoi(parts[1])
return width, height, errWidth == nil && errHeight == nil && width > 0 && height > 0
}
func stringList(value any) []string {
items, ok := value.([]any)
if !ok {
if values, ok := value.([]string); ok {
return values
}
return nil
}
out := make([]string, 0, len(items))
for _, item := range items {
if text := strings.TrimSpace(fmt.Sprint(item)); text != "" {
out = append(out, text)
}
}
return out
}
func boolishDefault(value any, fallback bool) bool {
if value == nil {
return fallback
}
switch typed := value.(type) {
case bool:
return typed
case string:
parsed, err := strconv.ParseBool(strings.TrimSpace(typed))
if err == nil {
return parsed
}
}
return fallback
}
func firstError(values ...error) error {
for _, value := range values {
if value != nil {
return value
}
}
return nil
}
func minInt64(left, right int64) int64 {
if left < right {
return left
}
return right
}
func minInt(left, right int) int {
if left < right {
return left
}
return right
}
func topazSourceHTTPClient(base *http.Client, allowPrivate bool) *http.Client {
client := *base
client.CheckRedirect = func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }
if client.Timeout <= 0 {
client.Timeout = 10 * time.Minute
}
if allowPrivate {
return &client
}
transport := http.DefaultTransport.(*http.Transport).Clone()
transport.Proxy = nil
transport.DialContext = func(ctx context.Context, network, address string) (net.Conn, error) {
host, port, err := net.SplitHostPort(address)
if err != nil {
return nil, err
}
addresses, err := net.DefaultResolver.LookupIPAddr(ctx, host)
if err != nil || len(addresses) == 0 {
return nil, errors.New("video source DNS resolution failed")
}
for _, address := range addresses {
if topazBlockedAddress(address.IP) {
return nil, errors.New("video source resolved to a blocked network")
}
}
dialer := &net.Dialer{Timeout: 10 * time.Second}
var attempts []error
for _, resolved := range addresses {
connection, dialErr := dialer.DialContext(ctx, network, net.JoinHostPort(resolved.IP.String(), port))
if dialErr == nil {
return connection, nil
}
attempts = append(attempts, dialErr)
}
return nil, errors.Join(attempts...)
}
client.Transport = transport
return &client
}
func topazBlockedAddress(ip net.IP) bool {
return ip == nil || ip.IsLoopback() || ip.IsPrivate() || ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() || ip.IsUnspecified() || ip.IsMulticast()
}
+1
View File
@@ -20,6 +20,7 @@ type Request struct {
RemoteTaskID string
RemoteTaskPayload map[string]any
OnRemoteTaskSubmitted func(remoteTaskID string, payload map[string]any) error
OnRemoteTaskPolled func(remoteTaskID string, payload map[string]any) error
Stream bool
StreamDelta StreamDelta
UpstreamProtocol string
+25
View File
@@ -44,6 +44,8 @@ func (c UniversalClient) Run(ctx context.Context, request Request) (Response, er
if err != nil {
return Response{}, annotateResponseError(err, submitRequestID, startedAt, time.Now())
}
submitResult = universalEffectiveResult(submitResult)
submitRequestID = firstNonEmptyString(submitRequestID, requestIDFromResult(submitResult))
if isUniversalSuccess(submitResult) && submitResult["data"] != nil {
return Response{
Result: normalizeUniversalResult(request, submitResult, ""),
@@ -157,6 +159,7 @@ func (c UniversalClient) universalPollUntilDone(ctx context.Context, executor *s
if err != nil {
return nil, "", annotateResponseError(err, firstNonEmptyString(pollRequestID, requestID, upstreamTaskID), pollStarted, pollFinished)
}
result = universalEffectiveResult(result)
lastResult = result
requestID = firstNonEmptyString(pollRequestID, requestID, requestIDFromResult(result), upstreamTaskID)
if isUniversalSuccess(result) {
@@ -239,6 +242,14 @@ func universalScriptContext(request Request, modelType string, payload map[strin
return selectedBase + "/" + strings.TrimLeft(path, "/")
}
context["creatRequestURL"] = context["createRequestURL"]
processedParams := cloneBody(request.Body)
processedParams["model"] = upstreamModelName(request.Candidate)
context["processedParams"] = processedParams
context["preProcessParams"] = func(params map[string]any, _ ...string) map[string]any {
processed := cloneMapAny(params)
processed["model"] = upstreamModelName(request.Candidate)
return processed
}
context["resolveGetTaskURL"] = func(taskID string) string {
return resolveUniversalTaskURL(request.Candidate.PlatformConfig, taskID)
}
@@ -389,6 +400,20 @@ func universalStatus(result map[string]any) string {
return strings.ToLower(strings.TrimSpace(firstNonEmptyString(result["status"], result["state"], result["task_status"])))
}
func universalEffectiveResult(result map[string]any) map[string]any {
nested, ok := result["result"].(map[string]any)
if !ok || nested == nil || isUniversalFailure(result) {
return result
}
out := cloneMapAny(nested)
for _, key := range []string{"status", "request_id", "requestId", "upstream_task_id", "task_id", "taskId", "id"} {
if out[key] == nil && result[key] != nil {
out[key] = result[key]
}
}
return out
}
func universalTaskID(result map[string]any) string {
return firstNonEmptyString(result["upstream_task_id"], result["task_id"], result["taskId"], result["id"])
}
@@ -92,6 +92,47 @@ func TestUniversalClientDefaultSubmitAndPoll(t *testing.T) {
}
}
func TestUniversalClientSupportsServerMainScriptContextAndNestedResult(t *testing.T) {
request := Request{
Kind: "images.generations",
ModelType: "image_generate",
Model: "custom-image",
Body: map[string]any{"model": "custom-image", "prompt": "hello"},
Candidate: testUniversalCandidate(map[string]any{
"customGetParamsScript": map[string]any{
"image_generate": `async function getParams(params, context) {
const processed = await context.preProcessParams(params, context.type);
return { prompt: processed.prompt + "-" + processed.model + "-" + context.processedParams.model };
}`,
},
"customSubmitScript": map[string]any{
"image_generate": `async function submitTask(payload) {
return {
status: "success",
result: {
status: "success",
data: [{ url: "https://cdn.example/" + payload.prompt + ".png" }]
}
};
}`,
},
}),
}
response, err := (UniversalClient{}).Run(context.Background(), request)
if err != nil {
t.Fatalf("run failed: %v", err)
}
data, ok := response.Result["data"].([]any)
if !ok || len(data) != 1 {
t.Fatalf("unexpected nested result: %#v", response.Result)
}
image, ok := data[0].(map[string]any)
if !ok || image["url"] != "https://cdn.example/hello-provider-model-provider-model.png" {
t.Fatalf("unexpected image result: %#v", response.Result)
}
}
func TestUniversalClientResumeSkipsSubmit(t *testing.T) {
request := Request{
Kind: "videos.generations",
+271
View File
@@ -0,0 +1,271 @@
package clients
import (
"bytes"
"context"
"crypto/sha256"
"encoding/base64"
"encoding/hex"
"fmt"
"io"
"mime/multipart"
"net"
"net/http"
"net/textproto"
"net/url"
"strings"
"time"
)
const vectorizerMaxResponseBytes = 128 << 20
// VectorizerClient implements the Vectorizer.AI binary vectorization API.
// Gateway async semantics are supplied by the outer River-backed task runner.
type VectorizerClient struct {
HTTPClient *http.Client
LookupIP func(context.Context, string) ([]net.IPAddr, error)
}
func (c VectorizerClient) Run(ctx context.Context, request Request) (Response, error) {
startedAt := time.Now()
format := strings.ToLower(strings.TrimSpace(firstNonEmptyString(request.Body["format"], request.Body["output_format"])))
if format == "" {
format = "svg"
}
if !vectorizerFormatAllowed(format) {
return Response{}, &ClientError{Code: "invalid_parameter", Message: "vectorizer format must be svg, eps, pdf, dxf, or png", Param: "format", StatusCode: http.StatusBadRequest}
}
imageToken := strings.TrimSpace(firstNonEmptyString(request.Body["_vectorizer_image_token"]))
receipt := strings.TrimSpace(firstNonEmptyString(request.Body["_vectorizer_receipt"]))
endpoint := "vectorize"
fields := map[string]string{"output.file_format": format}
if imageToken != "" {
endpoint = "download"
fields["image.token"] = imageToken
if receipt != "" {
fields["receipt"] = receipt
}
} else {
imageURL := vectorizerImageURL(request.Body)
if imageURL == "" {
return Response{}, &ClientError{Code: "invalid_parameter", Message: "vectorizer source image URL is required", Param: "source.url", StatusCode: http.StatusBadRequest}
}
if err := c.validateSourceURL(ctx, imageURL, boolishDefault(request.Candidate.PlatformConfig["allowPrivateSourceDownloads"], false)); err != nil {
return Response{}, err
}
fields["image.url"] = imageURL
fields["mode"] = firstNonEmptyString(request.Candidate.PlatformConfig["mode"], "production")
fields["policy.retention_days"] = firstNonEmptyString(request.Candidate.PlatformConfig["retentionDays"], request.Candidate.PlatformConfig["retention_days"], "7")
fields["processing.shapes.min_area_px"] = vectorizerCleanupMinArea(request.Body)
if maxColors := vectorizerMaxColors(request.Body); maxColors > 0 {
fields["processing.max_colors"] = fmt.Sprint(maxColors)
}
}
appendVectorizerOutputFields(fields, format)
payload, contentType, err := vectorizerMultipartBody(fields)
if err != nil {
return Response{}, err
}
url := providerURL(request.Candidate.BaseURL, endpoint)
req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, payload)
if err != nil {
return Response{}, err
}
req.Header.Set("Content-Type", contentType)
applyVectorizerAuth(req, request.Candidate.Credentials)
client := httpClient(request.HTTPClient, c.HTTPClient)
resp, err := client.Do(req)
if err != nil {
return Response{}, &ClientError{Code: "network", Message: err.Error(), Retryable: true}
}
defer resp.Body.Close()
requestID := requestIDFromHTTPResponse(resp)
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
result, decodeErr := decodeHTTPResponse(resp)
if decodeErr != nil {
return Response{}, annotateResponseError(decodeErr, requestID, startedAt, time.Now())
}
return Response{}, &ClientError{Code: "provider_failed", Message: firstNonEmptyString(result["message"], result["error"], resp.Status), RequestID: requestID, StatusCode: resp.StatusCode, Retryable: resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode >= 500}
}
raw, err := io.ReadAll(io.LimitReader(resp.Body, vectorizerMaxResponseBytes+1))
if err != nil {
return Response{}, &ClientError{Code: "invalid_response", Message: err.Error(), RequestID: requestID, Retryable: true}
}
if len(raw) == 0 || len(raw) > vectorizerMaxResponseBytes {
return Response{}, &ClientError{Code: "invalid_response", Message: "vectorizer response is empty or too large", RequestID: requestID, Retryable: false}
}
returnedToken := strings.TrimSpace(resp.Header.Get("X-Image-Token"))
returnedReceipt := strings.TrimSpace(resp.Header.Get("X-Receipt"))
if request.OnRemoteTaskSubmitted != nil && returnedToken != "" {
digest := sha256.Sum256([]byte(returnedToken))
if err := request.OnRemoteTaskSubmitted("vectorizer-"+hex.EncodeToString(digest[:8]), map[string]any{
"imageToken": returnedToken,
"receipt": returnedReceipt,
}); err != nil {
return Response{}, err
}
}
finishedAt := time.Now()
mimeType := strings.TrimSpace(strings.Split(resp.Header.Get("Content-Type"), ";")[0])
if mimeType == "" || mimeType == "application/octet-stream" {
mimeType = vectorizerContentType(format)
}
return Response{
Result: map[string]any{
"status": "success",
"model": request.Model,
"data": []any{map[string]any{
"type": vectorizerOutputKind(format),
"b64_json": base64.StdEncoding.EncodeToString(raw),
"mime_type": mimeType,
"format": format,
}},
"vectorizer": map[string]any{
"format": format,
"creditsCharged": numericHeader(resp.Header.Get("X-Credits-Charged")),
"creditsCalculated": numericHeader(resp.Header.Get("X-Credits-Calculated")),
},
},
RequestID: requestID,
Progress: append(providerProgress(request), Progress{Phase: "uploading", Progress: 0.9, Message: "vectorizer result received"}),
ResponseStartedAt: startedAt,
ResponseFinishedAt: finishedAt,
ResponseDurationMS: responseDurationMS(startedAt, finishedAt),
}, nil
}
func (c VectorizerClient) validateSourceURL(ctx context.Context, rawURL string, allowPrivate bool) error {
parsed, err := url.Parse(rawURL)
if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") || strings.TrimSpace(parsed.Hostname()) == "" || parsed.User != nil {
return &ClientError{Code: "invalid_parameter", Message: "vectorizer source URL must be a public http(s) URL without userinfo", Param: "source.url", StatusCode: http.StatusBadRequest}
}
if allowPrivate {
return nil
}
lookup := c.LookupIP
if lookup == nil {
lookup = net.DefaultResolver.LookupIPAddr
}
addresses, err := lookup(ctx, parsed.Hostname())
if err != nil || len(addresses) == 0 {
return &ClientError{Code: "invalid_parameter", Message: "vectorizer source DNS resolution failed", Param: "source.url", StatusCode: http.StatusBadRequest}
}
for _, address := range addresses {
if topazBlockedAddress(address.IP) {
return &ClientError{Code: "invalid_parameter", Message: "vectorizer source resolved to a blocked network", Param: "source.url", StatusCode: http.StatusBadRequest}
}
}
return nil
}
func vectorizerImageURL(body map[string]any) string {
if source, ok := body["source"].(map[string]any); ok {
return firstNonEmptyString(source["url"], source["image_url"], source["imageUrl"])
}
return firstNonEmptyString(body["image_url"], body["imageUrl"])
}
func vectorizerMultipartBody(fields map[string]string) (*bytes.Buffer, string, error) {
var payload bytes.Buffer
writer := multipart.NewWriter(&payload)
for key, value := range fields {
header := make(textproto.MIMEHeader)
header.Set("Content-Disposition", fmt.Sprintf(`form-data; name="%s"`, strings.ReplaceAll(key, `"`, `\"`)))
part, err := writer.CreatePart(header)
if err != nil {
return nil, "", err
}
if _, err := io.WriteString(part, value); err != nil {
return nil, "", err
}
}
if err := writer.Close(); err != nil {
return nil, "", err
}
return &payload, writer.FormDataContentType(), nil
}
func applyVectorizerAuth(req *http.Request, credentials map[string]any) {
username := credential(credentials, "username", "accessKey", "access_key", "apiId", "api_id", "id")
password := credential(credentials, "password", "secretKey", "secret_key", "apiSecret", "api_secret", "secret")
if username != "" || password != "" {
req.SetBasicAuth(username, password)
return
}
if apiKey := credential(credentials, "apiKey", "api_key", "token"); apiKey != "" {
req.Header.Set("Authorization", "Bearer "+apiKey)
}
}
func vectorizerFormatAllowed(format string) bool {
switch format {
case "svg", "eps", "pdf", "dxf", "png":
return true
default:
return false
}
}
func vectorizerCleanupMinArea(body map[string]any) string {
cleanup := strings.ToLower(strings.TrimSpace(firstNonEmptyString(body["cleanupLevel"], body["cleanup_level"])))
switch cleanup {
case "low":
return "0"
case "strong":
return "1"
default:
return "0.125"
}
}
func vectorizerMaxColors(body map[string]any) int {
for _, key := range []string{"maxColors", "max_colors"} {
if value := int(numericValue(body[key], 0)); value == 0 || value == 2 || value == 4 || value == 8 || value == 16 || value == 32 {
return value
}
}
return 0
}
func appendVectorizerOutputFields(fields map[string]string, format string) {
if format == "svg" {
fields["output.svg.version"] = "svg_1_1"
fields["output.svg.fixed_size"] = "false"
fields["output.svg.adobe_compatibility_mode"] = "true"
}
if format == "dxf" {
fields["output.dxf.compatibility_level"] = "lines_and_arcs"
}
}
func vectorizerContentType(format string) string {
switch format {
case "svg":
return "image/svg+xml"
case "eps":
return "application/postscript"
case "pdf":
return "application/pdf"
case "dxf":
return "application/dxf"
default:
return "image/png"
}
}
func vectorizerOutputKind(format string) string {
if format == "svg" || format == "png" {
return "image"
}
return "file"
}
func numericHeader(value string) any {
value = strings.TrimSpace(value)
if value == "" {
return nil
}
return numericValue(value, 0)
}
+132 -54
View File
@@ -100,66 +100,105 @@ func (c VolcesClient) runVideo(ctx context.Context, request Request, apiKey stri
timeout := volcesPollTimeout(request)
deadline := time.NewTimer(timeout)
defer deadline.Stop()
ticker := time.NewTicker(interval)
defer ticker.Stop()
nextPoll := time.NewTimer(0)
defer nextPoll.Stop()
var lastResult map[string]any
lastRequestID := firstNonEmpty(submitRequestID, upstreamTaskID)
transientFailures := 0
for {
select {
case <-ctx.Done():
return Response{}, &ClientError{Code: "cancelled", Message: ctx.Err().Error(), RequestID: submitRequestID, Retryable: true}
default:
}
pollStartedAt := time.Now()
pollResult, pollRequestID, err := c.getJSON(ctx, request, request.Candidate.BaseURL, taskPath+"/"+upstreamTaskID, apiKey)
pollFinishedAt := time.Now()
requestID := firstNonEmpty(pollRequestID, submitRequestID, upstreamTaskID)
if err != nil {
return Response{}, annotateResponseError(err, requestID, pollStartedAt, pollFinishedAt)
}
lastResult = pollResult
switch volcesTaskStatus(pollResult) {
case "succeeded":
result := volcesVideoSuccessResult(request, upstreamTaskID, pollResult)
return Response{
Result: result,
RequestID: requestID,
Usage: volcesVideoUsage(pollResult),
Progress: volcesVideoProgress(request, upstreamTaskID),
ResponseStartedAt: submitStartedAt,
ResponseFinishedAt: pollFinishedAt,
ResponseDurationMS: responseDurationMS(submitStartedAt, pollFinishedAt),
}, nil
case "failed", "cancelled":
return Response{}, &ClientError{
Code: volcesTaskErrorCode(pollResult),
Message: volcesTaskErrorMessage(pollResult),
RequestID: requestID,
ResponseStartedAt: submitStartedAt,
ResponseFinishedAt: pollFinishedAt,
ResponseDurationMS: responseDurationMS(submitStartedAt, pollFinishedAt),
Retryable: false,
}
}
select {
case <-ctx.Done():
return Response{}, &ClientError{Code: "cancelled", Message: ctx.Err().Error(), RequestID: requestID, Retryable: true}
return Response{}, &ClientError{Code: "cancelled", Message: ctx.Err().Error(), RequestID: lastRequestID, Retryable: true}
case <-deadline.C:
return Response{}, &ClientError{
Code: "timeout",
Message: fmt.Sprintf("volces video task %s did not finish before timeout; last status: %s", upstreamTaskID, volcesTaskStatus(lastResult)),
RequestID: requestID,
RequestID: lastRequestID,
Retryable: true,
}
case <-ticker.C:
case <-nextPoll.C:
pollStartedAt := time.Now()
pollResult, pollRequestID, err := c.getJSON(ctx, request, request.Candidate.BaseURL, taskPath+"/"+upstreamTaskID, apiKey)
pollFinishedAt := time.Now()
requestID := firstNonEmpty(pollRequestID, submitRequestID, upstreamTaskID)
lastRequestID = requestID
if err != nil {
err = annotateResponseError(err, requestID, pollStartedAt, pollFinishedAt)
if !IsRetryable(err) {
return Response{}, err
}
transientFailures++
resetVolcesPollTimer(nextPoll, volcesRetryPollInterval(request, interval, transientFailures))
continue
}
transientFailures = 0
lastResult = pollResult
if request.OnRemoteTaskPolled != nil {
if err := request.OnRemoteTaskPolled(upstreamTaskID, pollResult); err != nil {
return Response{}, err
}
}
switch volcesTaskStatus(pollResult) {
case "succeeded":
result := volcesVideoSuccessResult(request, upstreamTaskID, pollResult)
return Response{
Result: result,
RequestID: requestID,
Usage: volcesVideoUsage(pollResult),
Progress: volcesVideoProgress(request, upstreamTaskID),
ResponseStartedAt: submitStartedAt,
ResponseFinishedAt: pollFinishedAt,
ResponseDurationMS: responseDurationMS(submitStartedAt, pollFinishedAt),
}, nil
case "failed", "cancelled":
return Response{}, &ClientError{
Code: volcesTaskErrorCode(pollResult),
Message: volcesTaskErrorMessage(pollResult),
RequestID: requestID,
ResponseStartedAt: submitStartedAt,
ResponseFinishedAt: pollFinishedAt,
ResponseDurationMS: responseDurationMS(submitStartedAt, pollFinishedAt),
Retryable: false,
}
}
resetVolcesPollTimer(nextPoll, interval)
}
}
}
// DeleteVideoTask calls the official contents-generations cancellation endpoint.
// It is intentionally separate from Run so task cancellation can use the same
// provider credentials that submitted the remote task.
func (c VolcesClient) DeleteVideoTask(ctx context.Context, request Request) (map[string]any, string, error) {
apiKey := credential(request.Candidate.Credentials, "apiKey", "api_key", "key", "token")
if apiKey == "" {
return nil, "", &ClientError{Code: "missing_credentials", Message: "volces api key is required", Retryable: false}
}
remoteTaskID := strings.TrimSpace(request.RemoteTaskID)
if remoteTaskID == "" {
return nil, "", &ClientError{Code: "invalid_parameter", Message: "volces remote task id is required", Retryable: false}
}
taskPath := volcesVideoTaskPath(request) + "/" + remoteTaskID
req, err := http.NewRequestWithContext(ctx, http.MethodDelete, joinURL(request.Candidate.BaseURL, taskPath), nil)
if err != nil {
return nil, "", err
}
req.Header.Set("Authorization", "Bearer "+apiKey)
response, err := httpClient(request.HTTPClient, c.HTTPClient).Do(req)
if err != nil {
return nil, "", &ClientError{Code: "network", Message: err.Error(), Retryable: true}
}
requestID := requestIDFromHTTPResponse(response)
result, err := decodeHTTPResponse(response)
if err != nil {
return result, requestID, annotateResponseError(err, requestID, time.Now(), time.Now())
}
result, envelopeRequestID, err := normalizeVolcesCompatibleResult(result)
return result, firstNonEmpty(requestID, envelopeRequestID), err
}
func volcesVideoTaskPath(request Request) string {
path := firstNonEmptyStringValue(
request.Candidate.PlatformConfig,
@@ -997,6 +1036,9 @@ func volcesTaskErrorCode(result map[string]any) string {
return code
}
status := volcesTaskStatus(result)
if status == "cancelled" {
return "volces_task_cancelled"
}
if status != "" {
return status
}
@@ -1015,6 +1057,10 @@ 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{}
}
content, _ := raw["content"].(map[string]any)
videoURL := strings.TrimSpace(stringFromAny(content["video_url"]))
created := intFromAny(raw["created_at"])
@@ -1025,16 +1071,17 @@ func volcesVideoSuccessResult(request Request, upstreamTaskID string, raw map[st
if videoURL != "" {
data = append(data, map[string]any{"url": videoURL, "type": "video"})
}
return map[string]any{
"id": upstreamTaskID,
"object": "video.generation",
"created": created,
"model": upstreamModelName(request.Candidate),
"status": "succeeded",
"upstream_task_id": upstreamTaskID,
"data": data,
"raw": raw,
result["id"] = firstNonEmpty(stringFromAny(raw["id"]), upstreamTaskID)
if strings.TrimSpace(stringFromAny(result["model"])) == "" {
result["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
}
func volcesVideoUsage(raw map[string]any) Usage {
@@ -1074,6 +1121,37 @@ func volcesPollTimeout(request Request) time.Duration {
return time.Duration(seconds) * time.Second
}
func volcesRetryPollInterval(request Request, normal time.Duration, failures int) time.Duration {
if failures < 1 {
return normal
}
max := time.Duration(numericValue(firstPresent(request.Candidate.PlatformConfig["volcesPollRetryMaxMs"], request.Body["pollRetryMaxMs"], request.Body["poll_retry_max_ms"]), 30000)) * time.Millisecond
if max < normal {
max = normal
}
delay := normal
for attempt := 1; attempt < failures && delay < max; attempt++ {
delay *= 2
}
if delay > max {
return max
}
return delay
}
func resetVolcesPollTimer(timer *time.Timer, delay time.Duration) {
if delay < 100*time.Millisecond {
delay = 100 * time.Millisecond
}
if !timer.Stop() {
select {
case <-timer.C:
default:
}
}
timer.Reset(delay)
}
func firstNonEmpty(values ...string) string {
for _, value := range values {
if strings.TrimSpace(value) != "" {
+194
View File
@@ -0,0 +1,194 @@
package clients
import (
"bytes"
"context"
"crypto/hmac"
"crypto/sha256"
"encoding/hex"
"encoding/json"
"fmt"
"net/http"
"net/url"
"sort"
"strings"
"time"
)
const (
volcesAssetDefaultEndpoint = "https://ark.cn-beijing.volcengineapi.com"
volcesAssetRegion = "cn-beijing"
volcesAssetService = "ark"
volcesAssetVersion = "2024-01-01"
)
type VolcesAssetClient struct {
HTTPClient *http.Client
Now func() time.Time
}
type VolcesAssetCredentials struct {
AccessKey string
SecretKey string
Endpoint string
}
type VolcesAssetResult struct {
ID string `json:"Id"`
Name string `json:"Name,omitempty"`
URL string `json:"URL,omitempty"`
AssetType string `json:"AssetType,omitempty"`
GroupID string `json:"GroupId,omitempty"`
Status string `json:"Status,omitempty"`
Error map[string]any `json:"Error,omitempty"`
ProjectName string `json:"ProjectName,omitempty"`
CreateTime string `json:"CreateTime,omitempty"`
UpdateTime string `json:"UpdateTime,omitempty"`
}
func (c VolcesAssetClient) CreateAsset(ctx context.Context, credentials VolcesAssetCredentials, body map[string]any) (VolcesAssetResult, string, error) {
var result struct {
ID string `json:"Id"`
}
requestID, err := c.call(ctx, credentials, "CreateAsset", body, &result)
return VolcesAssetResult{ID: result.ID}, requestID, err
}
func (c VolcesAssetClient) GetAsset(ctx context.Context, credentials VolcesAssetCredentials, body map[string]any) (VolcesAssetResult, string, error) {
var result VolcesAssetResult
requestID, err := c.call(ctx, credentials, "GetAsset", body, &result)
return result, requestID, err
}
func (c VolcesAssetClient) call(ctx context.Context, credentials VolcesAssetCredentials, action string, body map[string]any, target any) (string, error) {
accessKey := strings.TrimSpace(credentials.AccessKey)
secretKey := strings.TrimSpace(credentials.SecretKey)
if accessKey == "" || secretKey == "" {
return "", &ClientError{Code: "missing_credentials", Message: "volces portrait asset accessKey and secretKey are required", Retryable: false}
}
endpoint := strings.TrimRight(strings.TrimSpace(credentials.Endpoint), "/")
if endpoint == "" {
endpoint = volcesAssetDefaultEndpoint
}
baseURL, err := url.Parse(endpoint)
if err != nil || baseURL.Scheme == "" || baseURL.Host == "" {
return "", &ClientError{Code: "invalid_configuration", Message: "invalid volces portrait asset endpoint", Retryable: false}
}
bodyJSON, err := json.Marshal(body)
if err != nil {
return "", fmt.Errorf("marshal volces asset request: %w", err)
}
now := time.Now().UTC()
if c.Now != nil {
now = c.Now().UTC()
}
xDate := now.Format("20060102T150405Z")
contentSHA := sha256HexBytes(bodyJSON)
requestURL := *baseURL
requestURL.Path = "/"
requestURL.RawPath = ""
requestURL.RawQuery = canonicalVolcesAssetQuery(map[string]string{"Action": action, "Version": volcesAssetVersion})
headers := map[string]string{
"content-type": "application/json",
"host": baseURL.Host,
"x-content-sha256": contentSHA,
"x-date": xDate,
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, requestURL.String(), bytes.NewReader(bodyJSON))
if err != nil {
return "", err
}
req.Host = baseURL.Host
req.Header.Set("Content-Type", headers["content-type"])
req.Header.Set("X-Content-Sha256", headers["x-content-sha256"])
req.Header.Set("X-Date", headers["x-date"])
req.Header.Set("Authorization", volcesAssetAuthorization(accessKey, secretKey, http.MethodPost, "/", requestURL.RawQuery, headers, contentSHA, xDate))
response, err := httpClient(nil, c.HTTPClient).Do(req)
if err != nil {
return "", &ClientError{Code: "network", Message: err.Error(), Retryable: true}
}
defer response.Body.Close()
var envelope struct {
ResponseMetadata struct {
RequestID string `json:"RequestId"`
Error struct {
Code string `json:"Code"`
Message string `json:"Message"`
} `json:"Error"`
} `json:"ResponseMetadata"`
Result json.RawMessage `json:"Result"`
}
if err := json.NewDecoder(response.Body).Decode(&envelope); err != nil {
return requestIDFromHTTPResponse(response), &ClientError{Code: "invalid_response", Message: "decode volces portrait asset response: " + err.Error(), Retryable: HTTPRetryable(response.StatusCode), StatusCode: response.StatusCode}
}
requestID := firstNonEmpty(requestIDFromHTTPResponse(response), envelope.ResponseMetadata.RequestID)
if envelope.ResponseMetadata.Error.Code != "" || response.StatusCode >= http.StatusBadRequest {
message := strings.TrimSpace(envelope.ResponseMetadata.Error.Message)
if message == "" {
message = strings.TrimSpace(envelope.ResponseMetadata.Error.Code)
}
if message == "" {
message = fmt.Sprintf("volces %s failed with status %d", action, response.StatusCode)
}
return requestID, &ClientError{Code: firstNonEmpty(envelope.ResponseMetadata.Error.Code, "volces_asset_error"), Message: message, RequestID: requestID, StatusCode: response.StatusCode, Retryable: HTTPRetryable(response.StatusCode)}
}
if len(envelope.Result) == 0 || string(envelope.Result) == "null" {
return requestID, &ClientError{Code: "invalid_response", Message: "volces " + action + " returned empty result", RequestID: requestID, Retryable: false}
}
if err := json.Unmarshal(envelope.Result, target); err != nil {
return requestID, &ClientError{Code: "invalid_response", Message: "decode volces " + action + " result: " + err.Error(), RequestID: requestID, Retryable: false}
}
return requestID, nil
}
func volcesAssetAuthorization(accessKey string, secretKey string, method string, path string, canonicalQuery string, headers map[string]string, bodySHA string, xDate string) string {
signedHeaders := []string{"content-type", "host", "x-content-sha256", "x-date"}
canonicalHeaderLines := make([]string, 0, len(signedHeaders))
for _, key := range signedHeaders {
canonicalHeaderLines = append(canonicalHeaderLines, key+":"+strings.TrimSpace(headers[key]))
}
canonicalRequest := strings.Join([]string{
strings.ToUpper(method), path, canonicalQuery,
strings.Join(canonicalHeaderLines, "\n") + "\n",
strings.Join(signedHeaders, ";"), bodySHA,
}, "\n")
date := xDate
if len(date) >= 8 {
date = date[:8]
}
scope := strings.Join([]string{date, volcesAssetRegion, volcesAssetService, "request"}, "/")
stringToSign := strings.Join([]string{"HMAC-SHA256", xDate, scope, sha256HexString(canonicalRequest)}, "\n")
kDate := hmacSHA256([]byte(secretKey), date)
kRegion := hmacSHA256(kDate, volcesAssetRegion)
kService := hmacSHA256(kRegion, volcesAssetService)
kSigning := hmacSHA256(kService, "request")
signature := hex.EncodeToString(hmacSHA256(kSigning, stringToSign))
return "HMAC-SHA256 Credential=" + accessKey + "/" + scope + ", SignedHeaders=" + strings.Join(signedHeaders, ";") + ", Signature=" + signature
}
func canonicalVolcesAssetQuery(values map[string]string) string {
keys := make([]string, 0, len(values))
for key := range values {
keys = append(keys, key)
}
sort.Strings(keys)
parts := make([]string, 0, len(keys))
for _, key := range keys {
parts = append(parts, url.QueryEscape(key)+"="+url.QueryEscape(values[key]))
}
return strings.ReplaceAll(strings.Join(parts, "&"), "+", "%20")
}
func sha256HexBytes(value []byte) string {
digest := sha256.Sum256(value)
return hex.EncodeToString(digest[:])
}
func sha256HexString(value string) string { return sha256HexBytes([]byte(value)) }
func hmacSHA256(key []byte, value string) []byte {
mac := hmac.New(sha256.New, key)
_, _ = mac.Write([]byte(value))
return mac.Sum(nil)
}
@@ -0,0 +1,66 @@
package clients
import (
"context"
"crypto/sha256"
"encoding/hex"
"encoding/json"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
)
func TestVolcesAssetClientSignsCreateAndReadsAsset(t *testing.T) {
var calls []string
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
calls = append(calls, r.URL.Query().Get("Action"))
if r.Method != http.MethodPost || r.URL.Path != "/" || r.URL.Query().Get("Version") != "2024-01-01" {
t.Fatalf("unexpected asset request %s %s?%s", r.Method, r.URL.Path, r.URL.RawQuery)
}
if r.Header.Get("X-Date") != "20260718T010203Z" {
t.Fatalf("unexpected x-date: %q", r.Header.Get("X-Date"))
}
if !strings.HasPrefix(r.Header.Get("Authorization"), "HMAC-SHA256 Credential=asset-ak/") {
t.Fatalf("missing Volces authorization: %q", r.Header.Get("Authorization"))
}
var body map[string]any
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
t.Fatalf("decode body: %v", err)
}
raw, _ := json.Marshal(body)
digest := sha256.Sum256(raw)
if got := r.Header.Get("X-Content-Sha256"); got != hex.EncodeToString(digest[:]) {
t.Fatalf("content hash mismatch got=%q", got)
}
switch r.URL.Query().Get("Action") {
case "CreateAsset":
if body["GroupId"] != "group-1" || body["AssetType"] != "Image" {
t.Fatalf("unexpected CreateAsset body: %+v", body)
}
_ = json.NewEncoder(w).Encode(map[string]any{"ResponseMetadata": map[string]any{"RequestId": "create-rid"}, "Result": map[string]any{"Id": "asset-1"}})
case "GetAsset":
_ = json.NewEncoder(w).Encode(map[string]any{"ResponseMetadata": map[string]any{"RequestId": "get-rid"}, "Result": map[string]any{"Id": "asset-1", "Status": "Active", "AssetType": "Image"}})
default:
t.Fatalf("unexpected Action: %q", r.URL.Query().Get("Action"))
}
}))
defer server.Close()
client := VolcesAssetClient{HTTPClient: server.Client(), Now: func() time.Time {
return time.Date(2026, 7, 18, 1, 2, 3, 0, time.UTC)
}}
credentials := VolcesAssetCredentials{AccessKey: "asset-ak", SecretKey: "asset-sk", Endpoint: server.URL}
created, requestID, err := client.CreateAsset(context.Background(), credentials, map[string]any{"GroupId": "group-1", "URL": "https://example.com/person.png", "AssetType": "Image", "ProjectName": "default"})
if err != nil || created.ID != "asset-1" || requestID != "create-rid" {
t.Fatalf("unexpected CreateAsset result=%+v requestID=%s err=%v", created, requestID, err)
}
asset, requestID, err := client.GetAsset(context.Background(), credentials, map[string]any{"Id": "asset-1", "ProjectName": "default"})
if err != nil || asset.Status != "Active" || requestID != "get-rid" {
t.Fatalf("unexpected GetAsset result=%+v requestID=%s err=%v", asset, requestID, err)
}
if strings.Join(calls, ",") != "CreateAsset,GetAsset" {
t.Fatalf("unexpected actions: %+v", calls)
}
}
@@ -58,6 +58,32 @@ func (s *Server) listAPIKeyAccessRules(w http.ResponseWriter, r *http.Request) {
writeJSON(w, http.StatusOK, map[string]any{"items": items})
}
// listAPIKeyAssignableModels godoc
// @Summary 列出 API Key 可分配模型
// @Description 按当前用户自身的用户、租户和用户组权限返回可分配给 API Key 的启用模型,不受任何 API Key 权限规则影响。
// @Tags api-keys
// @Produce json
// @Security BearerAuth
// @Success 200 {object} PlatformModelListResponse
// @Failure 401 {object} ErrorEnvelope
// @Failure 503 {object} ErrorEnvelope
// @Failure 500 {object} ErrorEnvelope
// @Router /api/v1/api-keys/assignable-models [get]
func (s *Server) listAPIKeyAssignableModels(w http.ResponseWriter, r *http.Request) {
user, _ := auth.UserFromContext(r.Context())
models, err := s.store.ListAPIKeyAssignablePlatformModels(r.Context(), user)
if err != nil {
if errors.Is(err, store.ErrLocalUserRequired) {
writeLocalUserRequired(w)
return
}
s.logger.Error("list api key assignable models failed", "error", err)
writeError(w, http.StatusInternalServerError, "list api key assignable models failed")
return
}
writeJSON(w, http.StatusOK, map[string]any{"items": s.platformModelResponses(r.Context(), models)})
}
// createAccessRule godoc
// @Summary 创建访问规则
// @Description 管理端创建一条访问控制规则。
@@ -0,0 +1,62 @@
package httpapi
import (
"testing"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestAdvancedMediaScopeAliases(t *testing.T) {
tests := []struct {
kind string
scope string
}{
{kind: "images.vectorize", scope: "image_vectorize"},
{kind: "images.vectorize", scope: "image"},
{kind: "images.vectorize", scope: "vectorize"},
{kind: "videos.upscales", scope: "video_enhance"},
{kind: "videos.upscales", scope: "video"},
{kind: "videos.upscales", scope: "video_upscale"},
}
for _, test := range tests {
user := &auth.User{APIKeyID: "key", APIKeyScopes: []string{test.scope}}
if !apiKeyScopeAllowed(user, test.kind) {
t.Fatalf("scope %q should allow %q", test.scope, test.kind)
}
}
if apiKeyScopeAllowed(&auth.User{APIKeyID: "key", APIKeyScopes: []string{"chat"}}, "videos.upscales") {
t.Fatal("chat scope must not allow video upscale")
}
}
func TestFillDailyTokenUsageDaysKeepsGapsAndStreaks(t *testing.T) {
location, err := time.LoadLocation("Asia/Shanghai")
if err != nil {
t.Fatal(err)
}
from := time.Date(2026, 7, 20, 0, 0, 0, 0, location)
to := time.Date(2026, 7, 24, 0, 0, 0, 0, location)
items, summary := fillDailyTokenUsageDays(map[string]store.DailyTokenUsage{
"2026-07-20": {Date: "2026-07-20", TotalTokens: 10, TaskCount: 1},
"2026-07-22": {Date: "2026-07-22", TotalTokens: 20, TaskCount: 1},
"2026-07-23": {Date: "2026-07-23", TotalTokens: 30, TaskCount: 2},
"2026-07-24": {Date: "2026-07-24", TotalTokens: 40, TaskCount: 1},
}, from, to)
if len(items) != 5 || items[1].Date != "2026-07-21" || items[1].TaskCount != 0 {
t.Fatalf("daily usage should contain continuous zero days: %+v", items)
}
if summary.CumulativeTokens != 100 || summary.PeakDailyTokens != 40 || summary.CurrentStreakDays != 3 || summary.LongestStreakDays != 3 {
t.Fatalf("unexpected daily usage summary: %+v", summary)
}
}
func TestAdvancedMediaDefaultModels(t *testing.T) {
if got := canonicalTaskModelName("images.vectorize", ""); got != "easy-image-vectorizer-1" {
t.Fatalf("vectorizer default model=%q", got)
}
if got := canonicalTaskModelName("videos.upscales", ""); got != "easy-proteus-standard-4" {
t.Fatalf("Topaz default model=%q", got)
}
}
@@ -0,0 +1,45 @@
package httpapi
import "net/http"
// createImageVectorizeTask godoc
// @Summary 图片矢量化
// @Description 将 URL 位图转换为 SVG、EPS、PDF、DXF 或 PNG;设置 X-Async=true 时返回统一异步任务结构。可使用同一所有者历史任务的 vectorizerTaskId 复用上游 image token。
// @Tags images
// @Accept json
// @Produce json
// @Security BearerAuth
// @Param X-Async header bool false "true 时异步创建任务并返回 202"
// @Param input body ImageVectorizeRequest true "图片矢量化请求"
// @Success 200 {object} CompatibleResponse
// @Success 202 {object} TaskAcceptedResponse
// @Failure 400 {object} ErrorEnvelope
// @Failure 401 {object} ErrorEnvelope
// @Failure 403 {object} ErrorEnvelope
// @Failure 404 {object} ErrorEnvelope
// @Failure 502 {object} ErrorEnvelope
// @Router /api/v1/images/vectorize [post]
func (s *Server) createImageVectorizeTask() http.Handler {
return s.createTask("images.vectorize", true)
}
// createVideoUpscaleTask godoc
// @Summary 视频超分
// @Description 原生执行 Topaz 视频增强流程;设置 X-Async=true 时返回统一异步任务结构,结果在任务完成前持久化到 Gateway 文件存储。
// @Tags videos
// @Accept json
// @Produce json
// @Security BearerAuth
// @Param X-Async header bool false "true 时异步创建任务并返回 202"
// @Param input body VideoUpscaleRequest true "视频超分请求"
// @Success 200 {object} CompatibleResponse
// @Success 202 {object} TaskAcceptedResponse
// @Failure 400 {object} ErrorEnvelope
// @Failure 401 {object} ErrorEnvelope
// @Failure 403 {object} ErrorEnvelope
// @Failure 404 {object} ErrorEnvelope
// @Failure 502 {object} ErrorEnvelope
// @Router /api/v1/videos/upscales [post]
func (s *Server) createVideoUpscaleTask() http.Handler {
return s.createTask("videos.upscales", true)
}
@@ -9,8 +9,8 @@ import (
const (
opsManagementSkillDownloadPath = "/api/v1/public/skills/ai-gateway-ops-management/download"
apiDocsJSONPath = "/api-docs-json"
apiDocsYAMLPath = "/api-docs-yaml"
apiDocsJSONPath = "/api/v1/openapi.json"
apiDocsYAMLPath = "/api/v1/openapi.yaml"
)
// getOpsManagementSkillMetadata godoc
@@ -64,7 +64,7 @@ func (s *Server) downloadOpsManagementSkill(w http.ResponseWriter, _ *http.Reque
// @Tags agent-resources
// @Produce json
// @Success 200 {object} map[string]interface{}
// @Router /api-docs-json [get]
// @Router /api/v1/openapi.json [get]
func (s *Server) apiDocsJSON(w http.ResponseWriter, _ *http.Request) {
w.Header().Set("Content-Type", "application/json; charset=utf-8")
w.WriteHeader(http.StatusOK)
@@ -77,7 +77,7 @@ func (s *Server) apiDocsJSON(w http.ResponseWriter, _ *http.Request) {
// @Tags agent-resources
// @Produce application/yaml
// @Success 200 {string} string
// @Router /api-docs-yaml [get]
// @Router /api/v1/openapi.yaml [get]
func (s *Server) apiDocsYAML(w http.ResponseWriter, _ *http.Request) {
w.Header().Set("Content-Type", "application/yaml; charset=utf-8")
w.WriteHeader(http.StatusOK)
@@ -32,7 +32,7 @@ func TestGetOpsManagementSkillMetadata(t *testing.T) {
if len(metadata.Modules) != 1 || metadata.Modules[0] != "model-runtime" {
t.Fatalf("unexpected metadata modules: %+v", metadata.Modules)
}
if metadata.APIDocsJSONPath != "/api-docs-json" || metadata.APIDocsYAMLPath != "/api-docs-yaml" {
if metadata.APIDocsJSONPath != "/api/v1/openapi.json" || metadata.APIDocsYAMLPath != "/api/v1/openapi.yaml" {
t.Fatalf("unexpected API docs paths: %+v", metadata)
}
}
@@ -74,7 +74,7 @@ func TestEmbeddedAPIDocs(t *testing.T) {
server := &Server{}
jsonResponse := httptest.NewRecorder()
server.apiDocsJSON(jsonResponse, httptest.NewRequest(http.MethodGet, "/api-docs-json", nil))
server.apiDocsJSON(jsonResponse, httptest.NewRequest(http.MethodGet, "/api/v1/openapi.json", nil))
if jsonResponse.Code != http.StatusOK {
t.Fatalf("expected JSON docs status 200, got %d", jsonResponse.Code)
}
@@ -88,7 +88,7 @@ func TestEmbeddedAPIDocs(t *testing.T) {
t.Fatalf("decode embedded Swagger JSON: %v", err)
}
for _, path := range []string{
"/api-docs-json",
"/api/v1/openapi.json",
"/api/v1/public/skills/ai-gateway-ops-management/download",
"/api/admin/catalog/providers",
"/api/admin/catalog/base-models",
@@ -102,7 +102,7 @@ func TestEmbeddedAPIDocs(t *testing.T) {
}
yamlResponse := httptest.NewRecorder()
server.apiDocsYAML(yamlResponse, httptest.NewRequest(http.MethodGet, "/api-docs-yaml", nil))
server.apiDocsYAML(yamlResponse, httptest.NewRequest(http.MethodGet, "/api/v1/openapi.yaml", nil))
if yamlResponse.Code != http.StatusOK {
t.Fatalf("expected YAML docs status 200, got %d", yamlResponse.Code)
}
@@ -0,0 +1,150 @@
package httpapi
import (
"context"
"io"
"log/slog"
"net/http"
"net/http/httptest"
"os"
"strconv"
"strings"
"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 TestAPIKeyAssignableModelsIgnoreAPIKeyRules(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 API key assignable-model 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()
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",
CORSAllowedOrigin: "*",
}, db, slog.New(slog.NewTextHandler(io.Discard, nil))))
defer server.Close()
suffixText := strconv.FormatInt(time.Now().UnixNano(), 10)
username := "api_key_assignable_" + suffixText
password := "password123"
var registerResponse struct {
AccessToken string `json:"accessToken"`
}
doJSON(t, server.URL, http.MethodPost, "/api/v1/auth/register", "", map[string]any{
"username": username,
"email": username + "@example.com",
"password": password,
}, http.StatusCreated, &registerResponse)
testPool, err := pgxpool.New(ctx, databaseURL)
if err != nil {
t.Fatalf("connect test pool: %v", err)
}
defer testPool.Close()
if _, err := testPool.Exec(ctx, `UPDATE gateway_users SET roles = '["admin"]'::jsonb WHERE username = $1`, username); err != nil {
t.Fatalf("promote test user: %v", err)
}
var loginResponse struct {
AccessToken string `json:"accessToken"`
}
doJSON(t, server.URL, http.MethodPost, "/api/v1/auth/login", "", map[string]any{
"account": username,
"password": password,
}, http.StatusOK, &loginResponse)
createAPIKey := func(name string) string {
t.Helper()
var response struct {
APIKey struct {
ID string `json:"id"`
} `json:"apiKey"`
}
doJSON(t, server.URL, http.MethodPost, "/api/v1/api-keys", loginResponse.AccessToken, map[string]any{
"name": name,
}, http.StatusCreated, &response)
return response.APIKey.ID
}
firstAPIKeyID := createAPIKey("first assignable key")
secondAPIKeyID := createAPIKey("second assignable key")
var platform struct {
ID string `json:"id"`
}
doJSON(t, server.URL, http.MethodPost, "/api/admin/platforms", loginResponse.AccessToken, map[string]any{
"provider": "openai",
"platformKey": "api-key-assignable-" + suffixText,
"name": "API Key Assignable Test",
"baseUrl": "https://api.openai.com/v1",
"authType": "bearer",
"credentials": map[string]any{"mode": "simulation"},
"config": map[string]any{"testMode": true},
}, http.StatusCreated, &platform)
var model struct {
ID string `json:"id"`
}
modelName := "api-key-assignable-model-" + suffixText
doJSON(t, server.URL, http.MethodPost, "/api/admin/platforms/"+platform.ID+"/models", loginResponse.AccessToken, map[string]any{
"canonicalModelKey": "openai:gpt-4o-mini",
"modelName": modelName,
"modelAlias": modelName,
"modelType": []string{"text_generate"},
"displayName": "API Key Assignable Model",
}, http.StatusCreated, &model)
assertAssignable := func() {
t.Helper()
var response struct {
Items []struct {
ID string `json:"id"`
ModelName string `json:"modelName"`
} `json:"items"`
}
doJSON(t, server.URL, http.MethodGet, "/api/v1/api-keys/assignable-models", loginResponse.AccessToken, nil, http.StatusOK, &response)
if !modelListContains(response.Items, model.ID) {
t.Fatalf("user-owned model should remain assignable regardless of API key rules: %+v", response.Items)
}
}
assignModel := func(apiKeyID string) {
t.Helper()
doJSON(t, server.URL, http.MethodPost, "/api/v1/api-keys/access-rules/batch", loginResponse.AccessToken, map[string]any{
"subjectType": "api_key",
"subjectId": apiKeyID,
"effect": "allow",
"upsertResources": []map[string]any{{
"resourceType": "platform_model",
"resourceId": model.ID,
"priority": 100,
"minPermissionLevel": 0,
"status": "active",
}},
"deleteResources": []map[string]any{},
}, http.StatusOK, nil)
}
assertAssignable()
assignModel(firstAPIKeyID)
assertAssignable()
assignModel(secondAPIKeyID)
assertAssignable()
}
@@ -0,0 +1,111 @@
package httpapi
import (
"bytes"
"context"
"encoding/json"
"io"
"log/slog"
"net/http"
"net/http/httptest"
"net/url"
"os"
"strings"
"testing"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestReadyReturnsPostgresUnavailableWithinTwoSeconds(t *testing.T) {
db := newExhaustedPostgresStore(t)
server := &Server{store: db, logger: slog.New(slog.NewJSONHandler(io.Discard, nil))}
requestContext, cancel := context.WithTimeout(context.Background(), 4*time.Second)
defer cancel()
request := httptest.NewRequest(http.MethodGet, "/readyz", nil).WithContext(requestContext)
recorder := httptest.NewRecorder()
startedAt := time.Now()
server.ready(recorder, request)
elapsed := time.Since(startedAt)
assertUnavailableResponse(t, recorder, "POSTGRES_UNAVAILABLE", "postgres unavailable")
if elapsed > 3*time.Second {
t.Fatalf("readiness timeout took %s, want no more than 3s", elapsed)
}
}
func TestLoginReturnsAuthStoreUnavailableWithinFiveSeconds(t *testing.T) {
db := newExhaustedPostgresStore(t)
var logs bytes.Buffer
server := &Server{store: db, logger: slog.New(slog.NewJSONHandler(&logs, nil))}
requestContext, cancel := context.WithTimeout(context.Background(), 7*time.Second)
defer cancel()
request := httptest.NewRequest(http.MethodPost, "/api/v1/auth/login", strings.NewReader(`{"account":"timeout-test-account","password":"timeout-test-password"}`)).WithContext(requestContext)
recorder := httptest.NewRecorder()
startedAt := time.Now()
server.login(recorder, request)
elapsed := time.Since(startedAt)
assertUnavailableResponse(t, recorder, "AUTH_STORE_UNAVAILABLE", "authentication service temporarily unavailable")
if elapsed > 6*time.Second {
t.Fatalf("login timeout took %s, want no more than 6s", elapsed)
}
logOutput := logs.String()
for _, field := range []string{"postgres_pool_max_connections", "postgres_pool_acquired_connections", "postgres_pool_idle_connections", "postgres_pool_empty_acquire_count", "postgres_pool_canceled_acquire_count"} {
if !strings.Contains(logOutput, field) {
t.Fatalf("login failure log did not include %q: %s", field, logOutput)
}
}
if strings.Contains(logOutput, "timeout-test-account") || strings.Contains(logOutput, "timeout-test-password") {
t.Fatalf("login failure log exposed credentials: %s", logOutput)
}
}
func newExhaustedPostgresStore(t *testing.T) *store.Store {
t.Helper()
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
if databaseURL == "" {
t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run PostgreSQL availability timeout tests")
}
parsed, err := url.Parse(databaseURL)
if err != nil {
t.Fatalf("parse test database URL: %v", err)
}
query := parsed.Query()
query.Set("pool_max_conns", "1")
query.Set("pool_min_conns", "0")
parsed.RawQuery = query.Encode()
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
db, err := store.Connect(ctx, parsed.String())
if err != nil {
t.Fatalf("connect timeout test store: %v", err)
}
t.Cleanup(db.Close)
connection, err := db.Pool().Acquire(ctx)
if err != nil {
t.Fatalf("exhaust timeout test pool: %v", err)
}
t.Cleanup(connection.Release)
return db
}
func assertUnavailableResponse(t *testing.T, recorder *httptest.ResponseRecorder, expectedCode, expectedMessage string) {
t.Helper()
if recorder.Code != http.StatusServiceUnavailable {
t.Fatalf("status = %d, want 503; body=%s", recorder.Code, recorder.Body.String())
}
var envelope ErrorEnvelope
if err := json.Unmarshal(recorder.Body.Bytes(), &envelope); err != nil {
t.Fatalf("decode unavailable response: %v", err)
}
if envelope.Error.Code != expectedCode {
t.Fatalf("error code = %q, want %q; body=%s", envelope.Error.Code, expectedCode, recorder.Body.String())
}
if envelope.Error.Message != expectedMessage {
t.Fatalf("error message = %q, want %q; body=%s", envelope.Error.Message, expectedMessage, recorder.Body.String())
}
}
@@ -58,17 +58,18 @@ func TestPlanTaskResponseTreatsAPIV1EmbeddingAndRerankAsSynchronousCompatibleRes
}
}
func TestPlanTaskResponseKeepsAsyncTaskModeForOtherAPIV1Tasks(t *testing.T) {
req := httptest.NewRequest(http.MethodPost, "/api/v1/images/generations", nil)
req.Header.Set("X-Async", "true")
plan := planTaskResponse("images.generations", false, map[string]any{"stream": true}, req)
if !plan.asyncMode {
t.Fatal("non-chat /api/v1 task endpoints should keep X-Async task mode")
func TestPlanTaskResponseUsesCompatibleAPIV1MediaResponsesAndKeepsAsyncOptIn(t *testing.T) {
defaultRequest := httptest.NewRequest(http.MethodPost, "/api/v1/images/generations", nil)
defaultPlan := planTaskResponse("images.generations", true, map[string]any{}, defaultRequest)
if defaultPlan.asyncMode || !defaultPlan.compatibleMode {
t.Fatalf("canonical /api/v1 media endpoints should default to synchronous compatible responses, got %+v", defaultPlan)
}
if plan.compatibleMode {
t.Fatal("non-compatible /api/v1 task endpoints should not return OpenAI-compatible payloads")
asyncRequest := httptest.NewRequest(http.MethodPost, "/api/v1/images/generations", nil)
asyncRequest.Header.Set("X-Async", "true")
asyncPlan := planTaskResponse("images.generations", true, map[string]any{}, asyncRequest)
if !asyncPlan.asyncMode || !asyncPlan.compatibleMode {
t.Fatalf("canonical /api/v1 media endpoints should keep compatible mode when X-Async is enabled, got %+v", asyncPlan)
}
}
@@ -504,7 +504,7 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
Result map[string]any `json:"result"`
} `json:"task"`
}
doJSON(t, server.URL, http.MethodPost, "/api/v1/images/generations", apiKeyResponse.Secret, map[string]any{
doJSONWithHeaders(t, server.URL, http.MethodPost, "/api/v1/images/generations", apiKeyResponse.Secret, map[string]any{
"model": defaultImageModel,
"runMode": "simulation",
"prompt": "a tiny gateway console",
@@ -512,7 +512,9 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
"quality": "medium",
"simulation": true,
"simulationDurationMs": 5,
}, http.StatusAccepted, &imageResponse)
}, map[string]string{"X-Async": "true"}, http.StatusAccepted, &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)
}
@@ -524,7 +526,7 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
Result map[string]any `json:"result"`
} `json:"task"`
}
doJSON(t, server.URL, http.MethodPost, "/api/v1/images/edits", apiKeyResponse.Secret, map[string]any{
doJSONWithHeaders(t, server.URL, http.MethodPost, "/api/v1/images/edits", apiKeyResponse.Secret, map[string]any{
"model": defaultImageModel,
"runMode": "simulation",
"prompt": "replace background with clean studio light",
@@ -532,7 +534,9 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
"mask": "https://example.com/mask.png",
"simulation": true,
"simulationDurationMs": 5,
}, http.StatusAccepted, &imageEditResponse)
}, map[string]string{"X-Async": "true"}, http.StatusAccepted, &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"] == "" {
t.Fatalf("unexpected image edit task: %+v", imageEditResponse.Task)
}
@@ -1196,17 +1200,20 @@ WHERE reference_type = 'gateway_task'
}, http.StatusCreated, &videoRoutePlatformModel)
var textToVideoTask struct {
Task struct {
ID string `json:"id"`
Status string `json:"status"`
ModelType string `json:"modelType"`
} `json:"task"`
}
doJSON(t, server.URL, http.MethodPost, "/api/v1/videos/generations", apiKeyResponse.Secret, map[string]any{
doJSONWithHeaders(t, server.URL, http.MethodPost, "/api/v1/videos/generations", apiKeyResponse.Secret, map[string]any{
"model": videoRouteModel,
"runMode": "simulation",
"simulation": true,
"simulationDurationMs": 5,
"prompt": "text to video route",
}, http.StatusAccepted, &textToVideoTask)
}, map[string]string{"X-Async": "true"}, http.StatusAccepted, &textToVideoTask)
waitForTaskStatus(t, server.URL, apiKeyResponse.Secret, textToVideoTask.Task.ID, []string{"succeeded"}, 10*time.Second)
doJSON(t, server.URL, http.MethodGet, "/api/v1/tasks/"+textToVideoTask.Task.ID, apiKeyResponse.Secret, nil, http.StatusOK, &textToVideoTask.Task)
if textToVideoTask.Task.Status != "succeeded" || textToVideoTask.Task.ModelType != "video_generate" {
t.Fatalf("text-to-video request should use video_generate model_type: %+v", textToVideoTask.Task)
}
@@ -1218,14 +1225,16 @@ WHERE reference_type = 'gateway_task'
Metrics map[string]any `json:"metrics"`
} `json:"task"`
}
doJSON(t, server.URL, http.MethodPost, "/api/v1/videos/generations", apiKeyResponse.Secret, map[string]any{
doJSONWithHeaders(t, server.URL, http.MethodPost, "/api/v1/videos/generations", apiKeyResponse.Secret, map[string]any{
"model": videoRouteModel,
"runMode": "simulation",
"simulation": true,
"simulationDurationMs": 5,
"prompt": "image to video route",
"image": "https://example.com/source.png",
}, http.StatusAccepted, &imageToVideoTask)
}, map[string]string{"X-Async": "true"}, http.StatusAccepted, &imageToVideoTask)
waitForTaskStatus(t, server.URL, apiKeyResponse.Secret, imageToVideoTask.Task.ID, []string{"succeeded"}, 10*time.Second)
doJSON(t, server.URL, http.MethodGet, "/api/v1/tasks/"+imageToVideoTask.Task.ID, apiKeyResponse.Secret, nil, http.StatusOK, &imageToVideoTask.Task)
if imageToVideoTask.Task.Status != "succeeded" || imageToVideoTask.Task.ModelType != "image_to_video" {
t.Fatalf("image-to-video request should use image_to_video model_type: %+v", imageToVideoTask.Task)
}
@@ -0,0 +1,101 @@
package httpapi
import (
"net/http"
"strings"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
type dailyTokenUsageSummary struct {
CumulativeTokens int64 `json:"cumulativeTokens"`
PeakDailyTokens int64 `json:"peakDailyTokens"`
CurrentStreakDays int `json:"currentStreakDays"`
LongestStreakDays int `json:"longestStreakDays"`
}
// dailyTokenUsage godoc
// @Summary 查询每日 Token 与资源点用量
// @Description 按 IANA 时区返回连续自然日用量;API Key 仅统计当前 Key,JWT 统计当前用户。
// @Tags workspace
// @Produce json
// @Security BearerAuth
// @Param from query string true "开始日期 YYYY-MM-DD"
// @Param to query string true "结束日期 YYYY-MM-DD"
// @Param timezone query string false "IANA 时区" default(UTC)
// @Success 200 {object} DailyTokenUsageResponse
// @Failure 400 {object} ErrorEnvelope
// @Failure 401 {object} ErrorEnvelope
// @Failure 500 {object} ErrorEnvelope
// @Router /api/workspace/token-usage/daily [get]
func (s *Server) dailyTokenUsage(w http.ResponseWriter, r *http.Request) {
user, ok := auth.UserFromContext(r.Context())
if !ok {
writeError(w, http.StatusUnauthorized, "unauthorized")
return
}
timezone := strings.TrimSpace(r.URL.Query().Get("timezone"))
if timezone == "" {
timezone = "UTC"
}
location, err := time.LoadLocation(timezone)
if err != nil {
writeError(w, http.StatusBadRequest, "invalid IANA timezone")
return
}
from, err := time.ParseInLocation("2006-01-02", strings.TrimSpace(r.URL.Query().Get("from")), location)
if err != nil {
writeError(w, http.StatusBadRequest, "invalid from date")
return
}
to, err := time.ParseInLocation("2006-01-02", strings.TrimSpace(r.URL.Query().Get("to")), location)
if err != nil || to.Before(from) || to.Sub(from) > 399*24*time.Hour {
writeError(w, http.StatusBadRequest, "invalid to date or date range exceeds 400 days")
return
}
stored, err := s.store.ListDailyTokenUsage(r.Context(), user, from, to, timezone)
if err != nil {
s.logger.Error("list daily token usage failed", "error", err)
writeError(w, http.StatusInternalServerError, "list daily token usage failed")
return
}
byDate := make(map[string]store.DailyTokenUsage, len(stored))
for _, item := range stored {
byDate[item.Date] = item
}
items, summary := fillDailyTokenUsageDays(byDate, from, to)
writeJSON(w, http.StatusOK, map[string]any{
"items": items, "tokenDays": items, "summary": summary,
"range": map[string]string{"from": from.Format("2006-01-02"), "to": to.Format("2006-01-02"), "timezone": timezone},
})
}
func fillDailyTokenUsageDays(byDate map[string]store.DailyTokenUsage, from, to time.Time) ([]store.DailyTokenUsage, dailyTokenUsageSummary) {
items := make([]store.DailyTokenUsage, 0, 32)
summary := dailyTokenUsageSummary{}
currentRun := 0
for day := from; !day.After(to); day = day.AddDate(0, 0, 1) {
key := day.Format("2006-01-02")
item, found := byDate[key]
if !found {
item.Date = key
}
items = append(items, item)
summary.CumulativeTokens += item.TotalTokens
if item.TotalTokens > summary.PeakDailyTokens {
summary.PeakDailyTokens = item.TotalTokens
}
if item.TaskCount > 0 {
currentRun++
if currentRun > summary.LongestStreakDays {
summary.LongestStreakDays = currentRun
}
} else {
currentRun = 0
}
}
summary.CurrentStreakDays = currentRun
return items, summary
}
@@ -26,7 +26,6 @@ const maxGatewayUploadBytes = 256 << 20
// @Failure 502 {object} ErrorEnvelope
// @Failure 503 {object} ErrorEnvelope
// @Router /api/v1/files/upload [post]
// @Router /v1/files/upload [post]
func (s *Server) uploadFile(w http.ResponseWriter, r *http.Request) {
r.Body = http.MaxBytesReader(w, r.Body, maxGatewayUploadBytes)
if err := r.ParseMultipartForm(32 << 20); err != nil {
@@ -0,0 +1,68 @@
package httpapi
import (
"context"
"fmt"
"net/http"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
type gatewayTaskCreationStage string
const (
gatewayTaskCreationPrepare gatewayTaskCreationStage = "prepare"
gatewayTaskCreationStore gatewayTaskCreationStage = "store"
)
type gatewayTaskCreationError struct {
Stage gatewayTaskCreationStage
Err error
}
func (e *gatewayTaskCreationError) Error() string {
if e == nil || e.Err == nil {
return "gateway task creation failed"
}
return e.Err.Error()
}
func (e *gatewayTaskCreationError) Unwrap() error {
if e == nil {
return nil
}
return e.Err
}
func (s *Server) prepareAndCreateGatewayTask(
ctx context.Context,
r *http.Request,
user *auth.User,
kind string,
model string,
body map[string]any,
async bool,
) (store.GatewayTask, error) {
prepared, err := s.prepareTaskRequest(ctx, r, user, body)
if err != nil {
return store.GatewayTask{}, &gatewayTaskCreationError{Stage: gatewayTaskCreationPrepare, Err: err}
}
task, err := s.store.CreateTask(ctx, store.CreateTaskInput{
Kind: kind,
Model: model,
RunMode: runModeFromRequest(prepared.Body),
Async: async,
Request: prepared.Body,
ConversationID: prepared.ConversationID,
NewMessageCount: prepared.NewMessageCount,
MessageRefs: prepared.MessageRefs,
}, user)
if err != nil {
return store.GatewayTask{}, &gatewayTaskCreationError{
Stage: gatewayTaskCreationStore,
Err: fmt.Errorf("create task: %w", err),
}
}
return task, nil
}
+48 -2
View File
@@ -40,9 +40,9 @@ type geminiUploadSession struct {
}
var geminiGenerateContentRoutePrefixes = []string{
"/api/v1/models/",
"/v1beta/models/",
"/v1/models/",
"/models/",
}
func (s *Server) registerGeminiGenerateContentRoutes(mux *http.ServeMux) {
@@ -75,6 +75,21 @@ func geminiGenerateContentModelFromPath(prefix string, requestPath string) (stri
return model, true
}
// geminiGenerateContent godoc
// @Summary Gemini generateContent 兼容接口
// @Description 使用统一 /api/v1 前缀接收 Gemini generateContent 请求;旧 /v1 和 /v1beta 路径保留兼容。
// @Tags gemini-compatible
// @Accept json
// @Produce json
// @Security BearerAuth
// @Param model path string true "模型名称"
// @Param input body map[string]interface{} true "Gemini generateContent 请求"
// @Success 200 {object} map[string]interface{}
// @Failure 400 {object} ErrorEnvelope
// @Failure 401 {object} ErrorEnvelope
// @Failure 403 {object} ErrorEnvelope
// @Failure 502 {object} ErrorEnvelope
// @Router /api/v1/models/{model}:generateContent [post]
func (s *Server) geminiGenerateContent(w http.ResponseWriter, r *http.Request) {
user, ok := auth.UserFromContext(r.Context())
if !ok {
@@ -395,6 +410,18 @@ func geminiUsageMetadataFromOutput(output map[string]any) map[string]any {
return meta
}
// geminiFilesUpload godoc
// @Summary Gemini Files 上传接口
// @Description 使用统一 /api/v1 前缀启动或直接完成 Gemini Files 上传。
// @Tags gemini-compatible
// @Accept octet-stream
// @Produce json
// @Security BearerAuth
// @Param version path string true "Gemini 版本(v1 或 v1beta"
// @Success 200 {object} map[string]interface{}
// @Failure 400 {object} ErrorEnvelope
// @Failure 401 {object} ErrorEnvelope
// @Router /api/v1/gemini/upload/{version}/files [post]
func (s *Server) geminiFilesUpload(w http.ResponseWriter, r *http.Request) {
if err := validateGeminiFilesVersion(r.PathValue("version")); err != nil {
writeError(w, http.StatusBadRequest, err.Error(), clients.ErrorCode(err))
@@ -412,6 +439,18 @@ func (s *Server) geminiFilesUpload(w http.ResponseWriter, r *http.Request) {
})
}
// geminiFilesUploadFinalize godoc
// @Summary 完成 Gemini Files 分段上传
// @Tags gemini-compatible
// @Accept octet-stream
// @Produce json
// @Security BearerAuth
// @Param version path string true "Gemini 版本(v1 或 v1beta"
// @Param uploadID path string true "上传会话 ID"
// @Success 200 {object} map[string]interface{}
// @Failure 400 {object} ErrorEnvelope
// @Failure 401 {object} ErrorEnvelope
// @Router /api/v1/gemini/upload/{version}/files/{uploadID} [post]
func (s *Server) geminiFilesUploadFinalize(w http.ResponseWriter, r *http.Request) {
if err := validateGeminiFilesVersion(r.PathValue("version")); err != nil {
writeError(w, http.StatusBadRequest, err.Error(), clients.ErrorCode(err))
@@ -440,11 +479,18 @@ func (s *Server) startGeminiFilesUpload(w http.ResponseWriter, r *http.Request)
CreatedAt: time.Now(),
}
s.geminiUploadSessions.Store(uploadID, session)
w.Header().Set("X-Goog-Upload-URL", absoluteRequestURL(r, "/upload/"+session.Version+"/files/"+uploadID))
w.Header().Set("X-Goog-Upload-URL", absoluteRequestURL(r, geminiUploadPath(r, session.Version, uploadID)))
w.Header().Set("X-Goog-Upload-Status", "active")
writeJSON(w, http.StatusOK, map[string]any{})
}
func geminiUploadPath(r *http.Request, version string, uploadID string) string {
if strings.HasPrefix(r.URL.Path, "/api/v1/gemini/upload/") {
return "/api/v1/gemini/upload/" + version + "/files/" + uploadID
}
return "/upload/" + version + "/files/" + uploadID
}
func (s *Server) finalizeGeminiFilesUpload(w http.ResponseWriter, r *http.Request, uploadID string, session geminiUploadSession) {
if uploadID == "" {
uploadID = newGeminiUploadID()
@@ -1,6 +1,14 @@
package httpapi
import "testing"
import (
"io"
"log/slog"
"net/http"
"net/http/httptest"
"testing"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
)
func TestGeminiGenerateContentModelFromPath(t *testing.T) {
tests := []struct {
@@ -25,9 +33,9 @@ func TestGeminiGenerateContentModelFromPath(t *testing.T) {
wantOK: true,
},
{
name: "bare model path",
prefix: "/models/",
requestPath: "/models/gemini-image:generateContent",
name: "gateway api v1 model",
prefix: "/api/v1/models/",
requestPath: "/api/v1/models/gemini-image:generateContent",
wantModel: "gemini-image",
wantOK: true,
},
@@ -61,6 +69,51 @@ func TestGeminiGenerateContentModelFromPath(t *testing.T) {
}
}
func TestRegisterGeminiGenerateContentRoutes(t *testing.T) {
server := &Server{
auth: auth.New("test-secret", "", ""),
logger: slog.New(slog.NewTextHandler(io.Discard, nil)),
}
mux := http.NewServeMux()
mux.HandleFunc("GET /api/v1/models", func(w http.ResponseWriter, _ *http.Request) {
w.WriteHeader(http.StatusNoContent)
})
server.registerGeminiGenerateContentRoutes(mux)
tests := []struct {
method string
path string
status int
}{
{method: http.MethodGet, path: "/api/v1/models", status: http.StatusNoContent},
{method: http.MethodPost, path: "/api/v1/models/gemini-image:generateContent", status: http.StatusUnauthorized},
{method: http.MethodPost, path: "/v1/models/gemini-image:generateContent", status: http.StatusUnauthorized},
{method: http.MethodPost, path: "/v1beta/models/gemini-image:generateContent", status: http.StatusUnauthorized},
{method: http.MethodPost, path: "/models/gemini-image:generateContent", status: http.StatusNotFound},
}
for _, tt := range tests {
t.Run(tt.method+" "+tt.path, func(t *testing.T) {
response := httptest.NewRecorder()
mux.ServeHTTP(response, httptest.NewRequest(tt.method, tt.path, nil))
if response.Code != tt.status {
t.Fatalf("status = %d, want %d; body=%s", response.Code, tt.status, response.Body.String())
}
})
}
}
func TestGeminiUploadPathKeepsCanonicalV1Prefix(t *testing.T) {
canonical := httptest.NewRequest(http.MethodPost, "/api/v1/gemini/upload/v1beta/files", nil)
if got := geminiUploadPath(canonical, "v1beta", "upload-1"); got != "/api/v1/gemini/upload/v1beta/files/upload-1" {
t.Fatalf("canonical upload path = %q", got)
}
legacy := httptest.NewRequest(http.MethodPost, "/upload/v1beta/files", nil)
if got := geminiUploadPath(legacy, "v1beta", "upload-1"); got != "/upload/v1beta/files/upload-1" {
t.Fatalf("legacy upload path = %q", got)
}
}
func TestGeminiImageTaskBodyMapsTextOnlyToImageGenerate(t *testing.T) {
mapping, err := geminiImageTaskBody("gemini-image", map[string]any{
"contents": []any{
+102 -33
View File
@@ -17,13 +17,21 @@ import (
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
const (
postgresReadinessTimeout = 2 * time.Second
localLoginStoreTimeout = 5 * time.Second
errorCodePostgresDown = "POSTGRES_UNAVAILABLE"
errorCodeAuthStoreDown = "AUTH_STORE_UNAVAILABLE"
authStoreUnavailableError = "authentication service temporarily unavailable"
)
// health godoc
// @Summary 健康检查
// @Description 返回服务进程、运行环境和身份模式,供负载均衡或人工排障使用。
// @Tags system
// @Produce json
// @Success 200 {object} HealthResponse
// @Router /healthz [get]
// @Router /api/v1/healthz [get]
func (s *Server) health(w http.ResponseWriter, r *http.Request) {
writeJSON(w, http.StatusOK, map[string]any{
"ok": true,
@@ -40,10 +48,13 @@ func (s *Server) health(w http.ResponseWriter, r *http.Request) {
// @Produce json
// @Success 200 {object} ReadyResponse
// @Failure 503 {object} ErrorEnvelope
// @Router /readyz [get]
// @Router /api/v1/readyz [get]
func (s *Server) ready(w http.ResponseWriter, r *http.Request) {
if err := s.store.Ping(r.Context()); err != nil {
writeError(w, http.StatusServiceUnavailable, "postgres unavailable")
ctx, cancel := context.WithTimeout(r.Context(), postgresReadinessTimeout)
defer cancel()
if err := s.store.Ping(ctx); err != nil {
s.logPostgresUnavailable("postgres readiness check failed")
writeError(w, http.StatusServiceUnavailable, "postgres unavailable", errorCodePostgresDown)
return
}
writeJSON(w, http.StatusOK, map[string]any{"ok": true})
@@ -121,6 +132,7 @@ func (s *Server) register(w http.ResponseWriter, r *http.Request) {
// @Failure 401 {object} ErrorEnvelope
// @Failure 403 {object} ErrorEnvelope
// @Failure 500 {object} ErrorEnvelope
// @Failure 503 {object} ErrorEnvelope
// @Router /api/v1/auth/login [post]
func (s *Server) login(w http.ResponseWriter, r *http.Request) {
var input store.LocalLoginInput
@@ -128,12 +140,19 @@ func (s *Server) login(w http.ResponseWriter, r *http.Request) {
writeError(w, http.StatusBadRequest, "invalid json body")
return
}
user, err := s.store.AuthenticateLocalUser(r.Context(), input)
ctx, cancel := context.WithTimeout(r.Context(), localLoginStoreTimeout)
defer cancel()
user, err := s.store.AuthenticateLocalUser(ctx, input)
if err != nil {
if errors.Is(err, store.ErrInvalidCredentials) {
writeError(w, http.StatusUnauthorized, "invalid account or password")
return
}
if store.IsPostgresUnavailable(err) {
s.logPostgresUnavailable("login authentication store unavailable")
writeError(w, http.StatusServiceUnavailable, authStoreUnavailableError, errorCodeAuthStoreDown)
return
}
s.logger.Error("login local user failed", "error", err)
writeError(w, http.StatusInternalServerError, "login failed")
return
@@ -145,6 +164,22 @@ func (s *Server) login(w http.ResponseWriter, r *http.Request) {
s.writeAuthResponse(w, http.StatusOK, user)
}
func (s *Server) logPostgresUnavailable(message string) {
if s.logger == nil || s.store == nil || s.store.Pool() == nil {
return
}
statistics := s.store.Pool().Stat()
s.logger.Error(message,
"error_category", "postgres_unavailable",
"postgres_pool_max_connections", statistics.MaxConns(),
"postgres_pool_total_connections", statistics.TotalConns(),
"postgres_pool_acquired_connections", statistics.AcquiredConns(),
"postgres_pool_idle_connections", statistics.IdleConns(),
"postgres_pool_empty_acquire_count", statistics.EmptyAcquireCount(),
"postgres_pool_canceled_acquire_count", statistics.CanceledAcquireCount(),
)
}
func (s *Server) localIdentityEnabled() bool {
mode := strings.ToLower(strings.TrimSpace(s.cfg.IdentityMode))
return mode == "" || mode == "standalone" || mode == "hybrid"
@@ -986,7 +1021,7 @@ func (s *Server) listModelRateLimitStatuses(w http.ResponseWriter, r *http.Reque
// createTask godoc
// @Summary 创建或执行 AI 任务
// @Description 网关任务接口按 model 选择平台模型;除 /api/v1/chat/completions 以外的 /api/v1 任务路径返回任务受理结果,OpenAI-compatible 路径同步返回兼容响应或 SSE 流
// @Description 统一公开入口按 model 选择平台模型并默认同步返回兼容响应;设置 X-Async=true 时异步创建任务并返回 202
// @Tags tasks
// @Accept json
// @Produce json
@@ -1012,22 +1047,6 @@ func (s *Server) listModelRateLimitStatuses(w http.ResponseWriter, r *http.Reque
// @Router /api/v1/music/generations [post]
// @Router /api/v1/speech/generations [post]
// @Router /api/v1/voice_clone [post]
// @Router /embeddings [post]
// @Router /v1/embeddings [post]
// @Router /reranks [post]
// @Router /v1/reranks [post]
// @Router /images/generations [post]
// @Router /v1/images/generations [post]
// @Router /images/edits [post]
// @Router /v1/images/edits [post]
// @Router /song/generations [post]
// @Router /v1/song/generations [post]
// @Router /music/generations [post]
// @Router /v1/music/generations [post]
// @Router /speech/generations [post]
// @Router /v1/speech/generations [post]
// @Router /voice_clone [post]
// @Router /v1/voice_clone [post]
func (s *Server) createTask(kind string, compatible bool) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
user, ok := auth.UserFromContext(r.Context())
@@ -1051,11 +1070,15 @@ func (s *Server) createTask(kind string, compatible bool) http.Handler {
return
}
}
model := requestModelName(body)
requestedModel := requestModelName(body)
model := canonicalTaskModelName(kind, requestedModel)
if model == "" {
writeError(w, http.StatusBadRequest, "model is required")
return
}
if model != requestedModel {
body["model"] = model
}
if !apiKeyScopeAllowed(user, kind) {
writeError(w, http.StatusForbidden, "api key scope does not allow this capability")
return
@@ -1179,8 +1202,6 @@ func (s *Server) createAPIV1ChatCompletions() http.Handler {
// @Failure 403 {object} ErrorEnvelope
// @Failure 429 {object} ErrorEnvelope
// @Failure 502 {object} ErrorEnvelope
// @Router /chat/completions [post]
// @Router /v1/chat/completions [post]
func openAIChatCompletionsDoc() {}
// openAIResponsesDoc godoc
@@ -1199,8 +1220,6 @@ func openAIChatCompletionsDoc() {}
// @Failure 402 {object} ErrorEnvelope
// @Failure 403 {object} ErrorEnvelope
// @Failure 503 {object} ErrorEnvelope "response_chain_unavailable"
// @Router /responses [post]
// @Router /v1/responses [post]
// @Router /api/v1/responses [post]
func openAIResponsesDoc() {}
@@ -1376,6 +1395,12 @@ func apiKeyScopeAllowed(user *auth.User, kind string) bool {
if required == "voice_clone" && (scope == "audio" || scope == "text_to_speech" || scope == "speech" || scope == "tts") {
return true
}
if required == "image_vectorize" && (scope == "image" || scope == "vectorize") {
return true
}
if required == "video_enhance" && (scope == "video" || scope == "video_upscale" || scope == "upscale") {
return true
}
}
return false
}
@@ -1387,6 +1412,25 @@ func requestModelName(body map[string]any) string {
return modelNameFromValue(body["model"])
}
func canonicalTaskModelName(kind string, model string) string {
model = strings.TrimSpace(model)
if model == "" {
switch kind {
case "images.vectorize":
return "easy-image-vectorizer-1"
case "videos.upscales":
return "easy-proteus-standard-4"
}
}
if kind != "videos.generations" {
return model
}
if canonical, ok := canonicalKlingOmniModel(model); ok {
return canonical
}
return model
}
func modelNameFromValue(value any) string {
switch typed := value.(type) {
case string:
@@ -1411,8 +1455,12 @@ func scopeForTaskKind(kind string) string {
return "rerank"
case "images.generations", "images.edits":
return "image"
case "images.vectorize":
return "image_vectorize"
case "videos.generations":
return "video"
case "videos.upscales":
return "video_enhance"
case "song.generations", "music.generations":
return "music"
case "speech.generations":
@@ -1625,7 +1673,6 @@ func matchedRateLimitRule(policy map[string]any, metric string) map[string]any {
// @Failure 500 {object} ErrorEnvelope
// @Router /api/workspace/tasks [get]
// @Router /api/v1/tasks [get]
// @Router /tasks [get]
func (s *Server) listTasks(w http.ResponseWriter, r *http.Request) {
user, ok := auth.UserFromContext(r.Context())
if !ok {
@@ -1734,10 +1781,18 @@ func boolValue(body map[string]any, key string) bool {
// @Failure 500 {object} ErrorEnvelope
// @Router /api/workspace/tasks/{taskID} [get]
// @Router /api/v1/tasks/{taskID} [get]
// @Router /tasks/{taskID} [get]
func (s *Server) getTask(w http.ResponseWriter, r *http.Request) {
user, ok := auth.UserFromContext(r.Context())
if !ok {
writeError(w, http.StatusUnauthorized, "unauthorized")
return
}
task, err := s.store.GetTask(r.Context(), r.PathValue("taskID"))
if err == nil {
if !runner.TaskAccessibleToUser(task, user) {
writeError(w, http.StatusNotFound, "task not found")
return
}
cancelState := runner.DescribeTaskCancellation(task)
task.Cancellable = &cancelState.Cancellable
task.Submitted = &cancelState.Submitted
@@ -1767,8 +1822,6 @@ func (s *Server) getTask(w http.ResponseWriter, r *http.Request) {
// @Failure 500 {object} ErrorEnvelope
// @Router /api/workspace/tasks/{taskID}/cancel [post]
// @Router /api/v1/tasks/{taskID}/cancel [post]
// @Router /v1/tasks/{taskID}/cancel [post]
// @Router /tasks/{taskID}/cancel [post]
func (s *Server) cancelTask(w http.ResponseWriter, r *http.Request) {
user, ok := auth.UserFromContext(r.Context())
if !ok {
@@ -1805,8 +1858,12 @@ func (s *Server) cancelTask(w http.ResponseWriter, r *http.Request) {
// @Failure 500 {object} ErrorEnvelope
// @Router /api/workspace/tasks/{taskID}/param-preprocessing [get]
// @Router /api/v1/tasks/{taskID}/param-preprocessing [get]
// @Router /tasks/{taskID}/param-preprocessing [get]
func (s *Server) taskParamPreprocessing(w http.ResponseWriter, r *http.Request) {
user, ok := auth.UserFromContext(r.Context())
if !ok {
writeError(w, http.StatusUnauthorized, "unauthorized")
return
}
task, err := s.store.GetTask(r.Context(), r.PathValue("taskID"))
if err != nil {
if store.IsNotFound(err) {
@@ -1817,6 +1874,10 @@ func (s *Server) taskParamPreprocessing(w http.ResponseWriter, r *http.Request)
writeError(w, http.StatusInternalServerError, "get task failed")
return
}
if !runner.TaskAccessibleToUser(task, user) {
writeError(w, http.StatusNotFound, "task not found")
return
}
logs, err := s.store.ListTaskParamPreprocessingLogs(r.Context(), task.ID)
if err != nil {
s.logger.Error("list task parameter preprocessing logs failed", "taskID", task.ID, "error", err)
@@ -1839,8 +1900,12 @@ func (s *Server) taskParamPreprocessing(w http.ResponseWriter, r *http.Request)
// @Failure 500 {object} ErrorEnvelope
// @Router /api/workspace/tasks/{taskID}/events [get]
// @Router /api/v1/tasks/{taskID}/events [get]
// @Router /tasks/{taskID}/events [get]
func (s *Server) taskEvents(w http.ResponseWriter, r *http.Request) {
user, ok := auth.UserFromContext(r.Context())
if !ok {
writeError(w, http.StatusUnauthorized, "unauthorized")
return
}
task, err := s.store.GetTask(r.Context(), r.PathValue("taskID"))
if err != nil {
if store.IsNotFound(err) {
@@ -1850,6 +1915,10 @@ func (s *Server) taskEvents(w http.ResponseWriter, r *http.Request) {
writeError(w, http.StatusInternalServerError, "get task failed")
return
}
if !runner.TaskAccessibleToUser(task, user) {
writeError(w, http.StatusNotFound, "task not found")
return
}
w.Header().Set("Content-Type", "text/event-stream")
w.Header().Set("Cache-Control", "no-cache")
@@ -0,0 +1,885 @@
package httpapi
import (
"context"
"crypto/rand"
"encoding/hex"
"encoding/json"
"errors"
"fmt"
"math"
"net/http"
"strconv"
"strings"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
const kelingOmniCompatibilityMarker = "keling_omni_v1"
type kelingCompatRequestIDKey struct{}
type KelingOmniVideoRequest struct {
ModelName string `json:"model_name" example:"kling-v3-omni"`
Prompt string `json:"prompt" example:"A quiet street in the rain with natural ambient sound"`
MultiShot bool `json:"multi_shot" example:"false"`
ShotType string `json:"shot_type,omitempty" example:"customize"`
MultiPrompt []KelingOmniMultiPrompt `json:"multi_prompt,omitempty"`
ImageList []KelingOmniImageInput `json:"image_list,omitempty"`
ElementList []KelingOmniElementInput `json:"element_list,omitempty"`
VideoList []KelingOmniVideoInput `json:"video_list,omitempty"`
Sound string `json:"sound" enums:"on,off" example:"on"`
Mode string `json:"mode" enums:"std,pro,4k" example:"pro"`
AspectRatio string `json:"aspect_ratio" enums:"16:9,9:16,1:1" example:"9:16"`
Duration any `json:"duration" swaggertype:"string" example:"5"`
WatermarkInfo KelingOmniWatermarkInfo `json:"watermark_info,omitempty"`
CallbackURL string `json:"callback_url,omitempty"`
ExternalTask string `json:"external_task_id,omitempty"`
}
type KelingOmniMultiPrompt struct {
Index int `json:"index" example:"1"`
Prompt string `json:"prompt" example:"A wide establishing shot"`
Duration any `json:"duration" swaggertype:"string" example:"3"`
}
type KelingOmniImageInput struct {
ImageURL string `json:"image_url"`
Type string `json:"type,omitempty" enums:"first_frame,end_frame"`
}
type KelingOmniElementInput struct {
ElementID any `json:"element_id"`
}
type KelingOmniVideoInput struct {
VideoURL string `json:"video_url"`
ReferType string `json:"refer_type,omitempty" enums:"base,feature"`
KeepOriginalSound string `json:"keep_original_sound,omitempty" enums:"yes,no"`
}
type KelingOmniWatermarkInfo struct {
Enabled bool `json:"enabled" example:"false"`
}
type KelingCompatibleEnvelope struct {
Code int `json:"code" example:"0"`
Message string `json:"message" example:"SUCCEED"`
RequestID string `json:"request_id"`
Data any `json:"data,omitempty"`
}
type kelingCompatError struct {
HTTPStatus int
Code int
Message string
RequestID string
}
func (e *kelingCompatError) Error() string {
if e == nil {
return "keling compatibility error"
}
return e.Message
}
func newKelingCompatError(status int, code int, message string) *kelingCompatError {
return &kelingCompatError{HTTPStatus: status, Code: code, Message: message}
}
func (s *Server) requireKelingAPIKey(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
requestID := newKelingCompatRequestID(r)
r = r.WithContext(context.WithValue(r.Context(), kelingCompatRequestIDKey{}, requestID))
user, err := s.auth.Authenticate(r)
if err != nil {
code := 1002
message := "Authorization is invalid"
if strings.TrimSpace(r.Header.Get("Authorization")) == "" {
code = 1001
message = "Authorization is required"
}
writeKelingCompatError(w, requestID, newKelingCompatError(http.StatusUnauthorized, code, message))
return
}
if user == nil || strings.TrimSpace(user.APIKeyID) == "" {
writeKelingCompatError(w, requestID, newKelingCompatError(http.StatusUnauthorized, 1002, "a Gateway API Key is required"))
return
}
if !apiKeyScopeAllowed(user, "videos.generations") {
writeKelingCompatError(w, requestID, newKelingCompatError(http.StatusForbidden, 1103, "API Key scope does not allow video generation"))
return
}
next.ServeHTTP(w, r.WithContext(auth.WithUser(r.Context(), user)))
})
}
// createKelingOmniVideo godoc
// @Summary 创建 Kling Omni 视频任务
// @Description 兼容 Kling 旧版 /v1/videos/omni-videoBearer token 必须为 Gateway API Key。任务固定异步执行,返回的 task_id 为网关任务 ID。
// @Tags kling-compatible
// @Accept json
// @Produce json
// @Security BearerAuth
// @Param input body KelingOmniVideoRequest true "Kling Omni 官方兼容请求"
// @Success 200 {object} KelingCompatibleEnvelope
// @Failure 400 {object} KelingCompatibleEnvelope
// @Failure 401 {object} KelingCompatibleEnvelope
// @Failure 403 {object} KelingCompatibleEnvelope
// @Failure 429 {object} KelingCompatibleEnvelope
// @Failure 500 {object} KelingCompatibleEnvelope
// @Failure 503 {object} KelingCompatibleEnvelope
// @Router /api/v1/videos/omni-video [post]
func (s *Server) createKelingOmniVideo(w http.ResponseWriter, r *http.Request) {
requestID := kelingCompatRequestID(r)
user, ok := auth.UserFromContext(r.Context())
if !ok || user == nil {
writeKelingCompatError(w, requestID, newKelingCompatError(http.StatusUnauthorized, 1002, "Authorization is invalid"))
return
}
body, err := s.decodeTaskRequestBody(r.Context(), w, r, "videos.generations")
if err != nil {
writeKelingCompatError(w, requestID, newKelingCompatError(http.StatusBadRequest, 1201, err.Error()))
return
}
normalized, compatErr := normalizeKelingOmniRequest(body)
if compatErr != nil {
writeKelingCompatError(w, requestID, compatErr)
return
}
model := strings.TrimSpace(stringFromKelingCompat(normalized["model"]))
if normalized["resolution"] == "2160p" {
candidates, candidateErr := s.store.ListModelCandidates(r.Context(), model, "omni_video", user)
if candidateErr != nil {
writeKelingCompatError(w, requestID, kelingCompatGatewayError(candidateErr))
return
}
if !kelingCompatCandidatesSupport4K(candidates) {
writeKelingCompatError(w, requestID, newKelingCompatError(http.StatusBadRequest, 1201, "mode=4k is not enabled by the selected model capabilities"))
return
}
}
task, createErr := s.prepareAndCreateGatewayTask(
r.Context(),
r,
user,
"videos.generations",
model,
normalized,
true,
)
if createErr != nil {
var staged *gatewayTaskCreationError
if errors.As(createErr, &staged) && staged.Stage == gatewayTaskCreationPrepare {
writeKelingCompatError(w, requestID, kelingCompatGatewayError(staged.Err))
return
}
s.logger.Error("create Kling-compatible task failed", "error", createErr)
writeKelingCompatError(w, requestID, newKelingCompatError(http.StatusInternalServerError, 5000, "create task failed"))
return
}
if err := s.runner.EnqueueAsyncTask(r.Context(), task); err != nil {
s.logger.Error("enqueue Kling-compatible task failed", "taskId", task.ID, "error", err)
writeKelingCompatError(w, requestID, newKelingCompatError(http.StatusServiceUnavailable, 5001, "video task queue is unavailable"))
return
}
writeJSON(w, http.StatusOK, KelingCompatibleEnvelope{
Code: 0,
Message: "SUCCEED",
RequestID: requestID,
Data: kelingCompatTaskData(task),
})
}
// getKelingOmniVideo godoc
// @Summary 查询 Kling Omni 视频任务
// @Description 按创建接口返回的网关 task_id 查询任务;仅允许创建任务的 Gateway 用户访问。
// @Tags kling-compatible
// @Produce json
// @Security BearerAuth
// @Param taskID path string true "网关任务 ID"
// @Success 200 {object} KelingCompatibleEnvelope
// @Failure 401 {object} KelingCompatibleEnvelope
// @Failure 403 {object} KelingCompatibleEnvelope
// @Failure 404 {object} KelingCompatibleEnvelope
// @Failure 500 {object} KelingCompatibleEnvelope
// @Router /api/v1/videos/omni-video/{taskID} [get]
func (s *Server) getKelingOmniVideo(w http.ResponseWriter, r *http.Request) {
requestID := kelingCompatRequestID(r)
user, ok := auth.UserFromContext(r.Context())
if !ok || user == nil {
writeKelingCompatError(w, requestID, newKelingCompatError(http.StatusUnauthorized, 1002, "Authorization is invalid"))
return
}
task, err := s.store.GetTask(r.Context(), strings.TrimSpace(r.PathValue("taskID")))
if err != nil {
if store.IsNotFound(err) {
writeKelingCompatError(w, requestID, newKelingCompatError(http.StatusNotFound, 1203, "task not found"))
return
}
s.logger.Error("get Kling-compatible task failed", "error", err)
writeKelingCompatError(w, requestID, newKelingCompatError(http.StatusInternalServerError, 5000, "query task failed"))
return
}
if !kelingCompatTaskOwnedBy(task, user) || !isKelingCompatTask(task) {
writeKelingCompatError(w, requestID, newKelingCompatError(http.StatusNotFound, 1203, "task not found"))
return
}
writeJSON(w, http.StatusOK, KelingCompatibleEnvelope{
Code: 0,
Message: "SUCCEED",
RequestID: requestID,
Data: kelingCompatTaskData(task),
})
}
func normalizeKelingOmniRequest(input map[string]any) (map[string]any, *kelingCompatError) {
if input == nil {
input = map[string]any{}
}
if callbackURL := strings.TrimSpace(stringFromKelingCompat(input["callback_url"])); callbackURL != "" {
return nil, newKelingCompatError(http.StatusBadRequest, 1201, "callback_url is not supported by this Gateway endpoint")
}
requestedModel := strings.TrimSpace(stringFromKelingCompat(input["model_name"]))
if requestedModel == "" {
requestedModel = "kling-video-o1"
}
model, maxDuration, ok := kelingCompatModel(requestedModel)
if !ok {
return nil, newKelingCompatError(http.StatusNotFound, 1203, "unsupported model_name: "+requestedModel)
}
mode := strings.ToLower(strings.TrimSpace(stringFromKelingCompat(input["mode"])))
if mode == "" {
mode = "pro"
}
resolutionByMode := map[string]string{"std": "720p", "pro": "1080p", "4k": "2160p"}
resolution := resolutionByMode[mode]
if resolution == "" {
return nil, newKelingCompatError(http.StatusBadRequest, 1201, "mode must be std, pro, or 4k")
}
sound := strings.ToLower(strings.TrimSpace(stringFromKelingCompat(input["sound"])))
if sound == "" {
sound = "off"
}
if sound != "on" && sound != "off" {
return nil, newKelingCompatError(http.StatusBadRequest, 1201, "sound must be on or off")
}
if model == klingO1Model && sound == "on" {
return nil, newKelingCompatError(http.StatusBadRequest, 1201, "kling-video-o1 does not support generated audio; sound must be off")
}
content := make([]any, 0)
prompt := strings.TrimSpace(stringFromKelingCompat(input["prompt"]))
images, hasFirstFrame, imageErr := normalizeKelingImageList(input["image_list"])
if imageErr != nil {
return nil, imageErr
}
content = append(content, images...)
elements, elementErr := normalizeKelingElementList(input["element_list"])
if elementErr != nil {
return nil, elementErr
}
content = append(content, elements...)
videos, hasBaseVideo, hasVideo, videoErr := normalizeKelingVideoList(input["video_list"])
if videoErr != nil {
return nil, videoErr
}
content = append(content, videos...)
if hasVideo && sound == "on" {
return nil, newKelingCompatError(http.StatusBadRequest, 1201, "sound must be off when video_list is provided")
}
multiShot, multiShotPresent, boolErr := kelingCompatOptionalBool(input, "multi_shot")
if boolErr != nil {
return nil, boolErr
}
shotType := strings.ToLower(strings.TrimSpace(stringFromKelingCompat(input["shot_type"])))
multiPrompts, shotDuration, multiPromptErr := normalizeKelingMultiPrompts(input["multi_prompt"])
if multiPromptErr != nil {
return nil, multiPromptErr
}
if !multiShotPresent {
multiShot = false
}
if multiShot {
if shotType != "customize" && shotType != "intelligence" {
return nil, newKelingCompatError(http.StatusBadRequest, 1201, "shot_type must be customize or intelligence when multi_shot is true")
}
if shotType == "customize" && len(multiPrompts) == 0 {
return nil, newKelingCompatError(http.StatusBadRequest, 1201, "multi_prompt is required for customized multi-shot generation")
}
if shotType == "intelligence" && len(multiPrompts) > 0 {
return nil, newKelingCompatError(http.StatusBadRequest, 1201, "multi_prompt is only supported when shot_type is customize")
}
} else if len(multiPrompts) > 0 || shotType != "" {
return nil, newKelingCompatError(http.StatusBadRequest, 1201, "shot_type and multi_prompt require multi_shot=true")
}
if (len(multiPrompts) == 0 || shotType == "intelligence") && prompt == "" {
return nil, newKelingCompatError(http.StatusBadRequest, 1201, "prompt is required")
}
if prompt != "" {
content = append([]any{map[string]any{"type": "text", "text": prompt}}, content...)
}
content = append(content, multiPrompts...)
duration, durationProvided, durationErr := kelingCompatOptionalInt(input, "duration")
if durationErr != nil {
return nil, durationErr
}
if hasBaseVideo {
if durationProvided {
return nil, newKelingCompatError(http.StatusBadRequest, 1201, "duration is not supported for base video editing")
}
} else {
if len(multiPrompts) > 0 {
if durationProvided && duration != shotDuration {
return nil, newKelingCompatError(http.StatusBadRequest, 1201, "duration must equal the sum of multi_prompt durations")
}
duration = shotDuration
} else if !durationProvided {
duration = 5
}
if duration < 3 || duration > maxDuration {
return nil, newKelingCompatError(http.StatusBadRequest, 1201, fmt.Sprintf("duration for %s must be an integer between 3 and %d seconds", requestedModel, maxDuration))
}
if model == klingO1Model && (len(images) == 0 || hasFirstFrame) && duration != 5 && duration != 10 {
return nil, newKelingCompatError(http.StatusBadRequest, 1201, "kling-video-o1 text-to-video and first-frame generation only support 5 or 10 seconds")
}
}
aspectRatio := strings.TrimSpace(stringFromKelingCompat(input["aspect_ratio"]))
if aspectRatio != "" && aspectRatio != "16:9" && aspectRatio != "9:16" && aspectRatio != "1:1" {
return nil, newKelingCompatError(http.StatusBadRequest, 1201, "aspect_ratio must be 16:9, 9:16, or 1:1")
}
if hasFirstFrame || hasBaseVideo {
if aspectRatio != "" {
return nil, newKelingCompatError(http.StatusBadRequest, 1201, "aspect_ratio is not supported with a first frame or base video")
}
} else if aspectRatio == "" {
return nil, newKelingCompatError(http.StatusBadRequest, 1201, "aspect_ratio is required when no first frame or base video is provided")
}
watermarkEnabled, watermarkErr := kelingCompatWatermarkEnabled(input["watermark_info"])
if watermarkErr != nil {
return nil, watermarkErr
}
externalTaskID := strings.TrimSpace(stringFromKelingCompat(input["external_task_id"]))
normalized := map[string]any{
"model": model,
"model_name": requestedModel,
"modelType": "omni_video",
"runMode": "real",
"content": content,
"resolution": resolution,
"mode": mode,
"sound": sound,
"audio": sound == "on",
"multi_shot": multiShot,
"watermark": watermarkEnabled,
"watermark_info": map[string]any{"enabled": watermarkEnabled},
"external_task_id": externalTaskID,
"_gateway_compatibility": kelingOmniCompatibilityMarker,
}
if prompt != "" {
normalized["prompt"] = prompt
}
if shotType != "" {
normalized["shot_type"] = shotType
}
if !hasBaseVideo {
normalized["duration"] = duration
}
if aspectRatio != "" {
normalized["aspect_ratio"] = aspectRatio
}
return normalized, nil
}
func normalizeKelingImageList(value any) ([]any, bool, *kelingCompatError) {
items, err := kelingCompatObjectList(value, "image_list")
if err != nil {
return nil, false, err
}
out := make([]any, 0, len(items))
hasFirstFrame := false
hasEndFrame := false
for index, item := range items {
url := strings.TrimSpace(stringFromKelingCompat(item["image_url"]))
if url == "" {
return nil, false, newKelingCompatError(http.StatusBadRequest, 1201, fmt.Sprintf("image_list[%d].image_url is required", index))
}
frameType := strings.TrimSpace(stringFromKelingCompat(item["type"]))
role := "reference_image"
switch frameType {
case "":
case "first_frame":
role = "first_frame"
hasFirstFrame = true
case "end_frame":
role = "last_frame"
hasEndFrame = true
default:
return nil, false, newKelingCompatError(http.StatusBadRequest, 1201, fmt.Sprintf("image_list[%d].type must be first_frame or end_frame", index))
}
out = append(out, map[string]any{
"type": "image_url",
"role": role,
"image_url": map[string]any{"url": url},
})
}
if hasEndFrame && !hasFirstFrame {
return nil, false, newKelingCompatError(http.StatusBadRequest, 1201, "end_frame requires first_frame")
}
return out, hasFirstFrame, nil
}
func normalizeKelingElementList(value any) ([]any, *kelingCompatError) {
items, err := kelingCompatObjectList(value, "element_list")
if err != nil {
return nil, err
}
out := make([]any, 0, len(items))
for index, item := range items {
id := item["element_id"]
if strings.TrimSpace(stringFromKelingCompat(id)) == "" {
return nil, newKelingCompatError(http.StatusBadRequest, 1201, fmt.Sprintf("element_list[%d].element_id is required", index))
}
out = append(out, map[string]any{"type": "element", "element": map[string]any{"element_id": id}})
}
return out, nil
}
func normalizeKelingVideoList(value any) ([]any, bool, bool, *kelingCompatError) {
items, err := kelingCompatObjectList(value, "video_list")
if err != nil {
return nil, false, false, err
}
if len(items) > 1 {
return nil, false, false, newKelingCompatError(http.StatusBadRequest, 1201, "video_list supports at most one video")
}
out := make([]any, 0, len(items))
hasBase := false
for index, item := range items {
url := strings.TrimSpace(stringFromKelingCompat(item["video_url"]))
if url == "" {
return nil, false, false, newKelingCompatError(http.StatusBadRequest, 1201, fmt.Sprintf("video_list[%d].video_url is required", index))
}
referType := strings.ToLower(strings.TrimSpace(stringFromKelingCompat(item["refer_type"])))
if referType == "" {
referType = "base"
}
if referType != "base" && referType != "feature" {
return nil, false, false, newKelingCompatError(http.StatusBadRequest, 1201, fmt.Sprintf("video_list[%d].refer_type must be base or feature", index))
}
keepSound := strings.ToLower(strings.TrimSpace(stringFromKelingCompat(item["keep_original_sound"])))
if keepSound != "" && keepSound != "yes" && keepSound != "no" {
return nil, false, false, newKelingCompatError(http.StatusBadRequest, 1201, fmt.Sprintf("video_list[%d].keep_original_sound must be yes or no", index))
}
nested := map[string]any{"url": url, "refer_type": referType}
if keepSound != "" {
nested["keep_original_sound"] = keepSound
}
role := "video_feature"
if referType == "base" {
role = "video_base"
hasBase = true
}
out = append(out, map[string]any{"type": "video_url", "role": role, "video_url": nested})
}
return out, hasBase, len(items) > 0, nil
}
func normalizeKelingMultiPrompts(value any) ([]any, int, *kelingCompatError) {
items, err := kelingCompatObjectList(value, "multi_prompt")
if err != nil {
return nil, 0, err
}
if len(items) > 6 {
return nil, 0, newKelingCompatError(http.StatusBadRequest, 1201, "multi_prompt supports at most six shots")
}
out := make([]any, 0, len(items))
seen := map[int]bool{}
total := 0
for index, item := range items {
shotIndex, ok := kelingCompatInt(item["index"])
if !ok || shotIndex < 1 || shotIndex > 6 || seen[shotIndex] {
return nil, 0, newKelingCompatError(http.StatusBadRequest, 1201, fmt.Sprintf("multi_prompt[%d].index must be a unique integer from 1 to 6", index))
}
seen[shotIndex] = true
prompt := strings.TrimSpace(stringFromKelingCompat(item["prompt"]))
if prompt == "" {
return nil, 0, newKelingCompatError(http.StatusBadRequest, 1201, fmt.Sprintf("multi_prompt[%d].prompt is required", index))
}
duration, ok := kelingCompatInt(item["duration"])
if !ok || duration < 1 {
return nil, 0, newKelingCompatError(http.StatusBadRequest, 1201, fmt.Sprintf("multi_prompt[%d].duration must be an integer of at least 1 second", index))
}
total += duration
out = append(out, map[string]any{
"type": "text",
"role": "shot_prompt",
"shot_index": shotIndex,
"text": prompt,
"duration": duration,
})
}
return out, total, nil
}
func kelingCompatModel(value string) (string, int, bool) {
model, ok := canonicalKlingOmniModel(value)
if !ok {
return "", 0, false
}
if model == klingO1Model {
return model, 10, true
}
return model, 15, true
}
func kelingCompatObjectList(value any, field string) ([]map[string]any, *kelingCompatError) {
if value == nil {
return nil, nil
}
raw, ok := value.([]any)
if !ok {
return nil, newKelingCompatError(http.StatusBadRequest, 1201, field+" must be an array")
}
out := make([]map[string]any, 0, len(raw))
for index, item := range raw {
object, ok := item.(map[string]any)
if !ok {
return nil, newKelingCompatError(http.StatusBadRequest, 1201, fmt.Sprintf("%s[%d] must be an object", field, index))
}
out = append(out, object)
}
return out, nil
}
func kelingCompatOptionalInt(body map[string]any, key string) (int, bool, *kelingCompatError) {
value, present := body[key]
if !present || value == nil || strings.TrimSpace(stringFromKelingCompat(value)) == "" {
return 0, false, nil
}
parsed, ok := kelingCompatInt(value)
if !ok {
return 0, true, newKelingCompatError(http.StatusBadRequest, 1201, key+" must be an integer")
}
return parsed, true, nil
}
func kelingCompatOptionalBool(body map[string]any, key string) (bool, bool, *kelingCompatError) {
value, present := body[key]
if !present || value == nil {
return false, false, nil
}
parsed, ok := value.(bool)
if !ok {
return false, true, newKelingCompatError(http.StatusBadRequest, 1201, key+" must be a boolean")
}
return parsed, true, nil
}
func kelingCompatWatermarkEnabled(value any) (bool, *kelingCompatError) {
if value == nil {
return false, nil
}
object, ok := value.(map[string]any)
if !ok {
return false, newKelingCompatError(http.StatusBadRequest, 1201, "watermark_info must be an object")
}
enabled, present := object["enabled"]
if !present {
return false, nil
}
result, ok := enabled.(bool)
if !ok {
return false, newKelingCompatError(http.StatusBadRequest, 1201, "watermark_info.enabled must be a boolean")
}
return result, nil
}
func kelingCompatInt(value any) (int, bool) {
switch typed := value.(type) {
case int:
return typed, true
case int64:
return int(typed), true
case float64:
if math.Abs(typed-math.Round(typed)) > 1e-9 {
return 0, false
}
return int(math.Round(typed)), true
case json.Number:
parsed, err := strconv.Atoi(typed.String())
return parsed, err == nil
case string:
parsed, err := strconv.Atoi(strings.TrimSpace(typed))
return parsed, err == nil
default:
return 0, false
}
}
func stringFromKelingCompat(value any) string {
switch typed := value.(type) {
case string:
return typed
case json.Number:
return typed.String()
case float64:
if math.Abs(typed-math.Round(typed)) < 1e-9 {
return strconv.FormatInt(int64(math.Round(typed)), 10)
}
return strconv.FormatFloat(typed, 'f', -1, 64)
case int:
return strconv.Itoa(typed)
case int64:
return strconv.FormatInt(typed, 10)
default:
return ""
}
}
func kelingCompatTaskData(task store.GatewayTask) map[string]any {
data := map[string]any{
"task_id": task.ID,
"task_status": kelingCompatTaskStatus(task.Status),
"task_info": map[string]any{
"external_task_id": strings.TrimSpace(stringFromKelingCompat(task.Request["external_task_id"])),
},
"created_at": task.CreatedAt.UnixMilli(),
"updated_at": task.UpdatedAt.UnixMilli(),
"watermark_info": map[string]any{
"enabled": kelingCompatTaskWatermark(task.Request),
},
}
if message := kelingCompatTaskMessage(task); message != "" {
data["task_status_msg"] = message
}
if kelingCompatTaskStatus(task.Status) == "failed" {
data["task_status_code"] = kelingCompatBusinessCode(task.ErrorCode, kelingCompatTaskMessage(task))
}
if videos := kelingCompatTaskVideos(task.Result); len(videos) > 0 {
data["task_result"] = map[string]any{"videos": videos}
}
if task.FinalChargeAmount > 0 {
data["final_unit_deduction"] = strconv.FormatFloat(task.FinalChargeAmount, 'f', -1, 64)
}
return data
}
func kelingCompatTaskStatus(status string) string {
switch strings.ToLower(strings.TrimSpace(status)) {
case "succeeded", "success", "completed":
return "succeed"
case "failed", "cancelled", "canceled":
return "failed"
case "running", "processing":
return "processing"
default:
return "submitted"
}
}
func kelingCompatTaskVideos(result map[string]any) []any {
raw, _ := result["data"].([]any)
out := make([]any, 0, len(raw))
for _, itemValue := range raw {
item, _ := itemValue.(map[string]any)
if item == nil {
continue
}
url := strings.TrimSpace(stringFromKelingCompat(firstKelingCompatValue(item["url"], item["video_url"])))
if url == "" {
continue
}
video := map[string]any{"url": url}
if id := strings.TrimSpace(stringFromKelingCompat(item["id"])); id != "" {
video["id"] = id
}
if watermarkURL := strings.TrimSpace(stringFromKelingCompat(item["watermark_url"])); watermarkURL != "" {
video["watermark_url"] = watermarkURL
}
if duration := strings.TrimSpace(stringFromKelingCompat(item["duration"])); duration != "" {
video["duration"] = duration
}
out = append(out, video)
}
return out
}
func firstKelingCompatValue(values ...any) any {
for _, value := range values {
if strings.TrimSpace(stringFromKelingCompat(value)) != "" {
return value
}
}
return nil
}
func kelingCompatTaskMessage(task store.GatewayTask) string {
return strings.TrimSpace(firstNonEmpty(task.ErrorMessage, task.Error, task.Message))
}
func kelingCompatTaskWatermark(request map[string]any) bool {
if value, ok := request["watermark"].(bool); ok {
return value
}
if info, ok := request["watermark_info"].(map[string]any); ok {
value, _ := info["enabled"].(bool)
return value
}
return false
}
func kelingCompatCandidatesSupport4K(candidates []store.RuntimeModelCandidate) bool {
for _, candidate := range candidates {
if strings.ToLower(strings.TrimSpace(candidate.Provider)) != "keling" {
continue
}
capability, _ := candidate.Capabilities["omni_video"].(map[string]any)
if capability == nil {
capability, _ = candidate.Capabilities["omni"].(map[string]any)
}
for _, resolution := range kelingCompatStringList(capability["output_resolutions"]) {
switch strings.ToLower(strings.TrimSpace(resolution)) {
case "2160p", "4k":
return true
}
}
}
return false
}
func kelingCompatStringList(value any) []string {
switch typed := value.(type) {
case []string:
return typed
case []any:
result := make([]string, 0, len(typed))
for _, item := range typed {
if text := strings.TrimSpace(stringFromKelingCompat(item)); text != "" {
result = append(result, text)
}
}
return result
default:
return nil
}
}
func kelingCompatGatewayError(err error) *kelingCompatError {
if err == nil {
return newKelingCompatError(http.StatusInternalServerError, 5000, "unknown gateway error")
}
codeText := clients.ErrorCode(err)
businessCode := kelingCompatBusinessCode(codeText, err.Error())
status := http.StatusInternalServerError
switch businessCode {
case 1101:
status = http.StatusPaymentRequired
case 1103:
status = http.StatusForbidden
case 1201:
status = http.StatusBadRequest
case 1203:
status = http.StatusNotFound
case 1302, 1303:
status = http.StatusTooManyRequests
case 5001:
status = http.StatusBadGateway
}
return newKelingCompatError(status, businessCode, err.Error())
}
func kelingCompatBusinessCode(errorCode string, message string) int {
combined := strings.ToLower(strings.TrimSpace(errorCode + " " + message))
containsAny := func(values ...string) bool {
for _, value := range values {
if strings.Contains(combined, value) {
return true
}
}
return false
}
switch {
case containsAny("insufficient_balance", "insufficient balance", "balance_not_enough", "wallet balance", "余额不足", "欠费", "quota exceeded"):
return 1101
case containsAny("permission_denied", "permission denied", "forbidden", "access denied", "scope does not allow"):
return 1103
case containsAny("concurrent", "concurrency"):
return 1303
case containsAny("rate_limit", "rate limit", "too many requests", "rpm", "tpm"):
return 1302
case containsAny("no_model_candidate", "no model candidate", "model_not_found", "unsupported model", "resource not found"):
return 1203
case containsAny("invalid_parameter", "invalid parameter", "bad_request", "parameter_preprocessing", "duration", "aspect_ratio"):
return 1201
case containsAny("upload_", "request_asset_", "network", "timeout", "upstream", "service unavailable", "bad gateway"):
return 5001
default:
return 5000
}
}
func isKelingCompatTask(task store.GatewayTask) bool {
return task.Kind == "videos.generations" && strings.TrimSpace(stringFromKelingCompat(task.Request["_gateway_compatibility"])) == kelingOmniCompatibilityMarker
}
func kelingCompatTaskOwnedBy(task store.GatewayTask, user *auth.User) bool {
if user == nil {
return false
}
taskOwner := strings.TrimSpace(firstNonEmpty(task.GatewayUserID, task.UserID))
requestOwner := strings.TrimSpace(firstNonEmpty(user.GatewayUserID, user.ID))
return taskOwner != "" && requestOwner != "" && taskOwner == requestOwner
}
func newKelingCompatRequestID(r *http.Request) string {
if r != nil {
if value := strings.TrimSpace(firstNonEmpty(r.Header.Get("X-Request-ID"), r.Header.Get("X-Request-Id"))); value != "" {
return value
}
}
random := make([]byte, 16)
if _, err := rand.Read(random); err == nil {
return hex.EncodeToString(random)
}
return strconv.FormatInt(time.Now().UnixNano(), 36)
}
func kelingCompatRequestID(r *http.Request) string {
if r != nil {
if value, ok := r.Context().Value(kelingCompatRequestIDKey{}).(string); ok && strings.TrimSpace(value) != "" {
return strings.TrimSpace(value)
}
}
return newKelingCompatRequestID(r)
}
func writeKelingCompatError(w http.ResponseWriter, requestID string, err *kelingCompatError) {
if err == nil {
err = newKelingCompatError(http.StatusInternalServerError, 5000, "internal error")
}
if err.RequestID != "" {
requestID = err.RequestID
}
if requestID == "" {
requestID = newKelingCompatRequestID(nil)
}
status := err.HTTPStatus
if status == 0 {
status = http.StatusInternalServerError
}
writeJSON(w, status, KelingCompatibleEnvelope{
Code: err.Code,
Message: err.Message,
RequestID: requestID,
})
}
@@ -0,0 +1,294 @@
package httpapi
import (
"context"
"encoding/json"
"errors"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestNormalizeKelingOmniRequestMapsOfficialFields(t *testing.T) {
normalized, err := normalizeKelingOmniRequest(map[string]any{
"model_name": "kling-v3-omni",
"prompt": "A rainy street with natural ambience",
"mode": "pro",
"sound": "on",
"aspect_ratio": "9:16",
"duration": "5",
"external_task_id": "external-1",
"watermark_info": map[string]any{"enabled": true},
})
if err != nil {
t.Fatalf("normalize Kling request: %v", err)
}
if normalized["model"] != "kling-v3-omni" ||
normalized["modelType"] != "omni_video" ||
normalized["resolution"] != "1080p" ||
normalized["aspect_ratio"] != "9:16" ||
normalized["duration"] != 5 ||
normalized["audio"] != true ||
normalized["sound"] != "on" ||
normalized["watermark"] != true ||
normalized["external_task_id"] != "external-1" ||
normalized["_gateway_compatibility"] != kelingOmniCompatibilityMarker {
t.Fatalf("unexpected normalized request: %+v", normalized)
}
content, _ := normalized["content"].([]any)
if len(content) != 1 {
t.Fatalf("unexpected content: %+v", normalized["content"])
}
text, _ := content[0].(map[string]any)
if text["type"] != "text" || text["text"] != "A rainy street with natural ambience" {
t.Fatalf("unexpected text content: %+v", text)
}
}
func TestCanonicalTaskModelNameNormalizesKelingOmniAliases(t *testing.T) {
tests := map[string]string{
"kling-o1": "kling-video-o1",
"kling-video-o1": "kling-video-o1",
"kling-3.0-omni": "kling-v3-omni",
"kling-3-omni": "kling-v3-omni",
"kling-v3-omni": "kling-v3-omni",
}
for input, expected := range tests {
if got := canonicalTaskModelName("videos.generations", input); got != expected {
t.Fatalf("canonicalTaskModelName(%q) = %q, want %q", input, got, expected)
}
}
if got := canonicalTaskModelName("chat.completions", "kling-o1"); got != "kling-o1" {
t.Fatalf("non-video model must not be rewritten, got %q", got)
}
}
func TestNormalizeKelingOmniRequestBuildsMultiShotMedia(t *testing.T) {
normalized, err := normalizeKelingOmniRequest(map[string]any{
"model_name": "kling-3.0-omni",
"multi_shot": true,
"shot_type": "customize",
"aspect_ratio": "16:9",
"duration": 5,
"multi_prompt": []any{
map[string]any{"index": 1, "prompt": "First shot", "duration": "2"},
map[string]any{"index": 2, "prompt": "Second shot", "duration": 3},
},
"image_list": []any{
map[string]any{"image_url": "https://example.com/reference.png"},
},
"element_list": []any{
map[string]any{"element_id": float64(123)},
},
})
if err != nil {
t.Fatalf("normalize multi-shot request: %v", err)
}
if normalized["model"] != "kling-v3-omni" || normalized["duration"] != 5 || normalized["multi_shot"] != true || normalized["shot_type"] != "customize" {
t.Fatalf("unexpected multi-shot fields: %+v", normalized)
}
content, _ := normalized["content"].([]any)
if len(content) != 4 {
t.Fatalf("expected image, element and two shot prompts, got %+v", content)
}
}
func TestNormalizeKelingOmniRequestRejectsUnsupportedCombinations(t *testing.T) {
tests := []struct {
name string
body map[string]any
}{
{
name: "callback",
body: map[string]any{"callback_url": "https://example.com/callback"},
},
{
name: "unknown model",
body: map[string]any{"model_name": "kling-unknown", "prompt": "x", "aspect_ratio": "16:9"},
},
{
name: "o1 duration",
body: map[string]any{"model_name": "kling-video-o1", "prompt": "x", "aspect_ratio": "16:9", "duration": 11},
},
{
name: "o1 text to video three seconds",
body: map[string]any{"model_name": "kling-video-o1", "prompt": "x", "aspect_ratio": "16:9", "duration": 3},
},
{
name: "o1 generated audio",
body: map[string]any{"model_name": "kling-video-o1", "prompt": "x", "sound": "on", "aspect_ratio": "16:9", "duration": 5},
},
{
name: "video sound",
body: map[string]any{
"model_name": "kling-v3-omni",
"prompt": "edit",
"sound": "on",
"video_list": []any{map[string]any{"video_url": "https://example.com/base.mp4", "refer_type": "base"}},
},
},
{
name: "first frame ratio",
body: map[string]any{
"prompt": "animate",
"aspect_ratio": "16:9",
"image_list": []any{map[string]any{"image_url": "https://example.com/first.png", "type": "first_frame"}},
},
},
}
for _, item := range tests {
t.Run(item.name, func(t *testing.T) {
_, err := normalizeKelingOmniRequest(item.body)
if err == nil || err.Code != 1201 && err.Code != 1203 {
t.Fatalf("expected official compatibility error, got %v", err)
}
})
}
}
func TestNormalizeKelingO1AllowsThreeSecondsWithReferenceImage(t *testing.T) {
normalized, err := normalizeKelingOmniRequest(map[string]any{
"model_name": "kling-video-o1",
"prompt": "Use the landscape as a visual reference",
"aspect_ratio": "16:9",
"duration": 3,
"image_list": []any{
map[string]any{"image_url": "https://placehold.co/1024x1024/png"},
},
})
if err != nil {
t.Fatalf("reference-image O1 request should allow three seconds: %v", err)
}
if normalized["duration"] != 3 {
t.Fatalf("unexpected duration: %+v", normalized)
}
}
func TestKelingCompatTaskDataAndOwnership(t *testing.T) {
createdAt := time.Unix(100, 0)
task := store.GatewayTask{
ID: "task-1",
Kind: "videos.generations",
GatewayUserID: "user-1",
Status: "succeeded",
FinalChargeAmount: 2.5,
Request: map[string]any{
"_gateway_compatibility": kelingOmniCompatibilityMarker,
"external_task_id": "external-1",
"watermark": true,
},
Result: map[string]any{"data": []any{map[string]any{
"id": "video-1",
"url": "https://example.com/video.mp4",
"watermark_url": "https://example.com/watermarked.mp4",
"duration": "5",
}}},
CreatedAt: createdAt,
UpdatedAt: createdAt.Add(time.Second),
}
if !isKelingCompatTask(task) || !kelingCompatTaskOwnedBy(task, &auth.User{GatewayUserID: "user-1"}) {
t.Fatalf("expected task ownership and compatibility marker")
}
if kelingCompatTaskOwnedBy(task, &auth.User{GatewayUserID: "user-2"}) {
t.Fatalf("cross-user task access must be rejected")
}
data := kelingCompatTaskData(task)
if data["task_status"] != "succeed" || data["final_unit_deduction"] != "2.5" {
t.Fatalf("unexpected task data: %+v", data)
}
result, _ := data["task_result"].(map[string]any)
videos, _ := result["videos"].([]any)
video, _ := videos[0].(map[string]any)
if video["id"] != "video-1" || video["watermark_url"] != "https://example.com/watermarked.mp4" || video["duration"] != "5" {
t.Fatalf("unexpected compatible video: %+v", video)
}
}
func TestRequireKelingAPIKeyWritesOfficialAuthEnvelope(t *testing.T) {
server := &Server{auth: auth.New("secret", "", "")}
server.auth.LocalAPIKeyVerifier = func(_ context.Context, key string) (*auth.User, error) {
if key != "sk-gw-valid" {
return nil, auth.ErrUnauthorized
}
return &auth.User{ID: "user-1", GatewayUserID: "user-1", APIKeyID: "key-1", APIKeyScopes: []string{"video"}}, nil
}
handler := server.requireKelingAPIKey(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if _, ok := auth.UserFromContext(r.Context()); !ok {
t.Fatal("authenticated user is missing")
}
w.WriteHeader(http.StatusNoContent)
}))
missing := httptest.NewRecorder()
handler.ServeHTTP(missing, httptest.NewRequest(http.MethodPost, "/v1/videos/omni-video", nil))
if missing.Code != http.StatusUnauthorized {
t.Fatalf("missing auth status=%d body=%s", missing.Code, missing.Body.String())
}
var missingBody KelingCompatibleEnvelope
if err := json.Unmarshal(missing.Body.Bytes(), &missingBody); err != nil || missingBody.Code != 1001 || missingBody.RequestID == "" {
t.Fatalf("unexpected missing auth envelope: %+v err=%v", missingBody, err)
}
invalid := httptest.NewRecorder()
invalidRequest := httptest.NewRequest(http.MethodPost, "/v1/videos/omni-video", nil)
invalidRequest.Header.Set("Authorization", "Bearer sk-gw-invalid")
handler.ServeHTTP(invalid, invalidRequest)
var invalidBody KelingCompatibleEnvelope
if err := json.Unmarshal(invalid.Body.Bytes(), &invalidBody); err != nil || invalid.Code != http.StatusUnauthorized || invalidBody.Code != 1002 {
t.Fatalf("unexpected invalid auth envelope: status=%d body=%+v err=%v", invalid.Code, invalidBody, err)
}
valid := httptest.NewRecorder()
validRequest := httptest.NewRequest(http.MethodPost, "/v1/videos/omni-video", nil)
validRequest.Header.Set("Authorization", "Bearer sk-gw-valid")
handler.ServeHTTP(valid, validRequest)
if valid.Code != http.StatusNoContent {
t.Fatalf("valid API Key status=%d body=%s", valid.Code, valid.Body.String())
}
}
func TestKelingCompatErrorImplementsError(t *testing.T) {
err := newKelingCompatError(http.StatusBadRequest, 1201, "invalid")
if !errors.Is(err, err) || err.Error() != "invalid" {
t.Fatalf("unexpected error behavior: %v", err)
}
}
func TestKelingCompatBusinessCodeMapping(t *testing.T) {
tests := []struct {
code string
message string
want int
}{
{code: "insufficient_balance", want: 1101},
{code: "permission_denied", want: 1103},
{code: "invalid_parameter", want: 1201},
{code: "no_model_candidate", want: 1203},
{code: "rate_limit_exceeded", want: 1302},
{code: "concurrency_limit", want: 1303},
{code: "network", want: 5001},
{code: "unknown", want: 5000},
}
for _, item := range tests {
if got := kelingCompatBusinessCode(item.code, item.message); got != item.want {
t.Fatalf("code=%q message=%q got=%d want=%d", item.code, item.message, got, item.want)
}
}
}
func TestKelingCompatCandidatesSupport4K(t *testing.T) {
candidates := []store.RuntimeModelCandidate{
{Provider: "keling", Capabilities: map[string]any{"omni_video": map[string]any{"output_resolutions": []any{"720p", "1080p"}}}},
{Provider: "keling", Capabilities: map[string]any{"omni_video": map[string]any{"output_resolutions": []any{"2160p"}}}},
}
if !kelingCompatCandidatesSupport4K(candidates) {
t.Fatal("expected 2160p Keling capability to enable mode=4k")
}
if kelingCompatCandidatesSupport4K(candidates[:1]) {
t.Fatal("mode=4k must stay disabled without an explicit capability")
}
}
@@ -0,0 +1,293 @@
package httpapi
import (
"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"
)
func TestKelingOmniCompatibleHTTPFlow(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 Kling-compatible HTTP integration flow")
}
var upstreamTaskSequence atomic.Int64
var upstreamPayloadMu sync.Mutex
var upstreamPayloads []map[string]any
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Header.Get("Authorization") != "Bearer upstream-keling-key" {
t.Fatalf("unexpected upstream Authorization: %q", r.Header.Get("Authorization"))
}
switch {
case r.Method == http.MethodPost && r.URL.Path == "/videos/omni-video":
var payload map[string]any
if err := json.NewDecoder(r.Body).Decode(&payload); err != nil {
t.Fatalf("decode upstream request: %v", err)
}
upstreamPayloadMu.Lock()
upstreamPayloads = append(upstreamPayloads, payload)
upstreamPayloadMu.Unlock()
id := "upstream-" + strconv.FormatInt(upstreamTaskSequence.Add(1), 10)
_ = json.NewEncoder(w).Encode(map[string]any{
"code": 0,
"request_id": "submit-" + id,
"data": map[string]any{"task_id": id, "task_status": "submitted"},
})
case r.Method == http.MethodGet && strings.HasPrefix(r.URL.Path, "/videos/omni-video/upstream-"):
id := strings.TrimPrefix(r.URL.Path, "/videos/omni-video/")
_ = json.NewEncoder(w).Encode(map[string]any{
"code": 0,
"request_id": "poll-" + id,
"data": map[string]any{
"task_id": id,
"task_status": "succeed",
"created_at": time.Now().UnixMilli(),
"task_result": map[string]any{"videos": []any{map[string]any{
"id": "video-" + id,
"url": "https://example.com/" + id + ".mp4",
"watermark_url": "https://example.com/" + id + "-watermark.mp4",
"duration": "3",
}}},
},
})
default:
t.Fatalf("unexpected upstream request %s %s", r.Method, r.URL.Path)
}
}))
defer upstream.Close()
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)
platform, err := db.CreatePlatform(ctx, store.CreatePlatformInput{
Provider: "keling",
PlatformKey: "keling-compatible-test-" + suffix,
Name: "Kling Compatible Test",
BaseURL: upstream.URL,
AuthType: "APIKey",
Credentials: map[string]any{"apiKey": "upstream-keling-key"},
Config: map[string]any{
"kelingPollIntervalMs": 100,
"kelingPollTimeoutSeconds": 5,
},
Priority: 1,
Status: "enabled",
})
if err != nil {
t.Fatalf("create test platform: %v", err)
}
_, err = db.CreatePlatformModel(ctx, store.CreatePlatformModelInput{
PlatformID: platform.ID,
CanonicalModelKey: "keling:kling-video-o1",
ModelName: "kling-video-o1",
ProviderModelName: "kling-video-o1",
ModelAlias: "",
ModelType: store.StringList{"omni_video", "video_generate"},
DisplayName: "Kling O1 Compatible Test",
Capabilities: map[string]any{
"omni_video": map[string]any{
"supported_modes": []any{"text_to_video", "image_reference"},
"output_resolutions": []any{"720p", "1080p"},
"aspect_ratio_allowed": []any{"16:9", "9:16", "1:1"},
"duration_options": []any{3, 4, 5, 6, 7, 8, 9, 10},
"output_audio": false,
"max_images": 7,
},
"video_generate": map[string]any{
"supported_modes": []any{"text_to_video"},
"output_resolutions": []any{"720p", "1080p"},
"aspect_ratio_allowed": []any{"16:9", "9:16", "1:1"},
"duration_options": []any{3, 4, 5, 6, 7, 8, 9, 10},
"output_audio": true,
},
},
Enabled: true,
})
if err != nil {
t.Fatalf("create test platform model: %v", err)
}
serverCtx, cancelServer := context.WithCancel(ctx)
defer cancelServer()
gateway := httptest.NewServer(NewServerWithContext(serverCtx, config.Config{
AppEnv: "test",
HTTPAddr: ":0",
DatabaseURL: databaseURL,
IdentityMode: "hybrid",
JWTSecret: "test-secret",
CORSAllowedOrigin: "*",
}, db, slog.New(slog.NewTextHandler(io.Discard, nil))))
defer gateway.Close()
firstUserToken, firstAPIKey := createKelingCompatIntegrationUser(t, ctx, db, gateway.URL, "first", suffix, true)
_ = firstUserToken
_, secondAPIKey := createKelingCompatIntegrationUser(t, ctx, db, gateway.URL, "second", suffix, false)
var created KelingCompatibleEnvelope
doJSON(t, gateway.URL, http.MethodPost, "/api/v1/videos/omni-video", firstAPIKey, map[string]any{
"model_name": "kling-video-o1",
"prompt": "A clean product reveal",
"mode": "std",
"aspect_ratio": "16:9",
"duration": "3",
"sound": "off",
"image_list": []any{map[string]any{"image_url": "https://example.com/reference.png"}},
"external_task_id": "compat-http-1",
}, http.StatusOK, &created)
createdData, _ := created.Data.(map[string]any)
if created.Code != 0 || created.RequestID == "" || strings.TrimSpace(stringFromKelingCompat(createdData["task_id"])) == "" || createdData["task_status"] != "submitted" {
t.Fatalf("unexpected compatible create response: %+v", created)
}
taskID := stringFromKelingCompat(createdData["task_id"])
var hidden KelingCompatibleEnvelope
doJSON(t, gateway.URL, http.MethodGet, "/api/v1/videos/omni-video/"+taskID, secondAPIKey, nil, http.StatusNotFound, &hidden)
if hidden.Code != 1203 {
t.Fatalf("cross-user task must be hidden: %+v", hidden)
}
completed := waitForKelingCompatTask(t, gateway.URL, firstAPIKey, taskID, 5*time.Second)
if completed.Code != 0 {
t.Fatalf("compatible task failed: %+v", completed)
}
completedData, _ := completed.Data.(map[string]any)
if completedData["task_status"] != "succeed" {
t.Fatalf("compatible task did not succeed: %+v", completedData)
}
taskResult, _ := completedData["task_result"].(map[string]any)
videos, _ := taskResult["videos"].([]any)
video, _ := videos[0].(map[string]any)
if video["id"] == "" || video["watermark_url"] == "" || video["duration"] != "3" {
t.Fatalf("compatible result lost video metadata: %+v", video)
}
var standard struct {
TaskID string `json:"taskId"`
}
doJSONWithHeaders(t, gateway.URL, http.MethodPost, "/api/v1/videos/generations", firstAPIKey, map[string]any{
"model": "kling-o1",
"prompt": "A second product reveal",
"resolution": "720p",
"aspect_ratio": "16:9",
"duration": 3,
"audio": false,
}, map[string]string{"X-Async": "true"}, http.StatusAccepted, &standard)
if standard.TaskID == "" {
t.Fatal("standard video generation did not return taskId")
}
waitForTaskStatus(t, gateway.URL, firstAPIKey, standard.TaskID, []string{"succeeded"}, 5*time.Second)
var unsupportedAudio struct {
TaskID string `json:"taskId"`
}
doJSONWithHeaders(t, gateway.URL, http.MethodPost, "/api/v1/videos/generations", firstAPIKey, map[string]any{
"model": "kling-o1",
"prompt": "An O1 request that must not silently ignore audio",
"resolution": "1080p",
"aspect_ratio": "9:16",
"duration": 5,
"audio": true,
}, map[string]string{"X-Async": "true"}, http.StatusAccepted, &unsupportedAudio)
if unsupportedAudio.TaskID == "" {
t.Fatal("unsupported O1 audio request did not return taskId")
}
waitForTaskStatus(t, gateway.URL, firstAPIKey, unsupportedAudio.TaskID, []string{"failed"}, 5*time.Second)
var failedAudioTask store.GatewayTask
doJSON(t, gateway.URL, http.MethodGet, "/api/v1/tasks/"+unsupportedAudio.TaskID, firstAPIKey, nil, http.StatusOK, &failedAudioTask)
if failedAudioTask.ErrorCode != "invalid_parameter" || !strings.Contains(failedAudioTask.ErrorMessage, "does not support generated audio") {
t.Fatalf("O1 audio request must fail visibly before upstream submission: %+v", failedAudioTask)
}
upstreamPayloadMu.Lock()
defer upstreamPayloadMu.Unlock()
if len(upstreamPayloads) != 2 {
t.Fatalf("expected two upstream submissions, got %d", len(upstreamPayloads))
}
compatiblePayload := upstreamPayloads[0]
if compatiblePayload["model_name"] != "kling-video-o1" || compatiblePayload["mode"] != "std" || compatiblePayload["sound"] != "off" || compatiblePayload["duration"] != "3" || compatiblePayload["aspect_ratio"] != "16:9" || compatiblePayload["external_task_id"] != "compat-http-1" {
t.Fatalf("unexpected compatible upstream payload: %+v", compatiblePayload)
}
}
func createKelingCompatIntegrationUser(t *testing.T, ctx context.Context, db *store.Store, baseURL string, prefix string, suffix string, fund bool) (string, string) {
t.Helper()
username := fmt.Sprintf("kling_compat_%s_%s", prefix, suffix)
password := "password123"
var registered struct {
AccessToken string `json:"accessToken"`
}
doJSON(t, baseURL, http.MethodPost, "/api/v1/auth/register", "", map[string]any{
"username": username,
"email": username + "@example.com",
"password": password,
}, http.StatusCreated, &registered)
var apiKey struct {
Secret string `json:"secret"`
}
doJSON(t, baseURL, http.MethodPost, "/api/v1/api-keys", registered.AccessToken, map[string]any{
"name": "Kling compatible integration key",
"scopes": []string{"video"},
}, http.StatusCreated, &apiKey)
if fund {
if _, err := db.Pool().Exec(ctx, `UPDATE gateway_users SET roles = '["admin"]'::jsonb WHERE username = $1`, username); err != nil {
t.Fatalf("promote integration user: %v", err)
}
var loggedIn struct {
AccessToken string `json:"accessToken"`
}
doJSON(t, baseURL, http.MethodPost, "/api/v1/auth/login", "", map[string]any{
"account": username,
"password": password,
}, http.StatusOK, &loggedIn)
var gatewayUserID string
if err := db.Pool().QueryRow(ctx, `SELECT id::text FROM gateway_users WHERE username = $1`, username).Scan(&gatewayUserID); err != nil {
t.Fatalf("read integration user id: %v", err)
}
doJSON(t, baseURL, http.MethodPatch, "/api/admin/users/"+gatewayUserID+"/wallet", loggedIn.AccessToken, map[string]any{
"currency": "resource",
"balance": 1000,
"reason": "seed Kling compatible integration wallet",
}, http.StatusOK, nil)
}
return registered.AccessToken, apiKey.Secret
}
func waitForKelingCompatTask(t *testing.T, baseURL string, apiKey string, taskID string, timeout time.Duration) KelingCompatibleEnvelope {
t.Helper()
deadline := time.Now().Add(timeout)
for time.Now().Before(deadline) {
var response KelingCompatibleEnvelope
doJSON(t, baseURL, http.MethodGet, "/api/v1/videos/omni-video/"+taskID, apiKey, nil, http.StatusOK, &response)
data, _ := response.Data.(map[string]any)
switch data["task_status"] {
case "succeed":
return response
case "failed":
t.Fatalf("Kling-compatible task failed: %+v", response)
}
time.Sleep(50 * time.Millisecond)
}
t.Fatalf("timed out waiting for Kling-compatible task %s", taskID)
return KelingCompatibleEnvelope{}
}
@@ -134,16 +134,22 @@ WHERE username = $1`, username); err != nil {
ResponseDurationMS int64 `json:"responseDurationMs"`
} `json:"task"`
}
doJSON(
doJSONWithHeaders(
t,
server.URL,
http.MethodPost,
"/api/v1/videos/generations",
apiKeyResponse.Secret,
request,
map[string]string{"X-Async": "true"},
http.StatusAccepted,
&response,
)
if response.Task.ID == "" {
t.Fatal("async Kling simulation response did not return a task id")
}
waitForTaskStatus(t, server.URL, apiKeyResponse.Secret, response.Task.ID, []string{"succeeded"}, 10*time.Second)
doJSON(t, server.URL, http.MethodGet, "/api/v1/tasks/"+response.Task.ID, apiKeyResponse.Secret, nil, http.StatusOK, &response.Task)
task := response.Task
if task.ID == "" ||
+921
View File
@@ -0,0 +1,921 @@
package httpapi
import (
"encoding/base64"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"strconv"
"strings"
"time"
"unicode/utf8"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
const (
klingCompatProvider = "kling"
klingO1Model = "kling-video-o1"
klingV3OmniModel = "kling-v3-omni"
)
func (s *Server) registerKlingCompatibilityRoutes(mux *http.ServeMux) {
handler := func(next http.HandlerFunc) http.Handler {
return s.requireUser(auth.PermissionBasic, http.HandlerFunc(next))
}
// /api/v1 is the canonical public prefix. The historical /kling paths
// remain registered below so existing clients can migrate without downtime.
mux.Handle("POST /api/v1/kling/v1/videos/omni-video", handler(s.klingV1CreateOmniVideo))
mux.Handle("GET /api/v1/kling/v1/videos/omni-video", handler(s.klingV1ListOmniVideos))
mux.Handle("GET /api/v1/kling/v1/videos/omni-video/{taskID}", handler(s.klingV1GetOmniVideo))
mux.Handle("POST /api/v1/kling/v2/omni-video/{model}", handler(s.klingV2CreateOmniVideo))
mux.Handle("GET /api/v1/kling/v2/tasks", handler(s.klingV2GetTasks))
mux.Handle("POST /api/v1/kling/v2/tasks", handler(s.klingV2ListTasks))
mux.Handle("POST /kling/v1/videos/omni-video", handler(s.klingV1CreateOmniVideo))
mux.Handle("GET /kling/v1/videos/omni-video", handler(s.klingV1ListOmniVideos))
mux.Handle("GET /kling/v1/videos/omni-video/{taskID}", handler(s.klingV1GetOmniVideo))
// Kling API 2.0 uses model-specific paths and a shared /tasks resource.
mux.Handle("POST /kling/omni-video/{model}", handler(s.klingV2CreateOmniVideo))
mux.Handle("GET /kling/tasks", handler(s.klingV2GetTasks))
mux.Handle("POST /kling/tasks", handler(s.klingV2ListTasks))
// Versioned aliases help clients that keep the protocol version in their base path.
mux.Handle("POST /kling/v2/omni-video/{model}", handler(s.klingV2CreateOmniVideo))
mux.Handle("GET /kling/v2/tasks", handler(s.klingV2GetTasks))
mux.Handle("POST /kling/v2/tasks", handler(s.klingV2ListTasks))
}
// klingV1CreateOmniVideo godoc
// @Summary 创建可灵 V1 Omni 视频任务
// @Description 兼容中国区可灵 V1 /v1/videos/omni-video;用户使用网关 API Key,网关在服务端使用 AK/SK 调用上游。
// @Tags kling-compatible
// @Accept json
// @Produce json
// @Security BearerAuth
// @Param input body map[string]interface{} true "可灵 V1 Omni 请求"
// @Success 200 {object} map[string]interface{}
// @Failure 400 {object} map[string]interface{}
// @Failure 401 {object} ErrorEnvelope
// @Router /api/v1/kling/v1/videos/omni-video [post]
func (s *Server) klingV1CreateOmniVideo(w http.ResponseWriter, r *http.Request) {
var native map[string]any
if err := decodeKlingJSON(r, &native); err != nil {
writeKlingCompatError(w, http.StatusBadRequest, "invalid json body", "invalid_json_body")
return
}
model := strings.TrimSpace(stringFromRequestAny(native["model_name"]))
if model == "" {
model = klingO1Model
}
s.createKlingCompatTask(w, r, "v1", model, native)
}
// klingV2CreateOmniVideo godoc
// @Summary 创建可灵 API 2.0 Omni 视频任务
// @Description 兼容可灵 API 2.0 的模型路径;调用方使用网关 API Key,网关转换并使用中国区 V1 AK/SK 上游。
// @Tags kling-compatible
// @Accept json
// @Produce json
// @Security BearerAuth
// @Param model path string true "模型路径(kling-o1 或 kling-v3-omni"
// @Param input body map[string]interface{} true "可灵 API 2.0 请求"
// @Success 200 {object} map[string]interface{}
// @Failure 400 {object} map[string]interface{}
// @Failure 401 {object} ErrorEnvelope
// @Router /api/v1/kling/v2/omni-video/{model} [post]
func (s *Server) klingV2CreateOmniVideo(w http.ResponseWriter, r *http.Request) {
model, ok := klingV2ProviderModel(r.PathValue("model"))
if !ok {
writeKlingCompatError(w, http.StatusNotFound, "unsupported Kling Omni model", "model_not_found")
return
}
var native map[string]any
if err := decodeKlingJSON(r, &native); err != nil {
writeKlingCompatError(w, http.StatusBadRequest, "invalid json body", "invalid_json_body")
return
}
s.createKlingCompatTask(w, r, "v2", model, native)
}
func (s *Server) createKlingCompatTask(w http.ResponseWriter, r *http.Request, version string, model string, native map[string]any) {
user, ok := auth.UserFromContext(r.Context())
if !ok {
writeKlingCompatError(w, http.StatusUnauthorized, "unauthorized", "unauthorized")
return
}
if !apiKeyScopeAllowed(user, "videos.generations") {
writeKlingCompatError(w, http.StatusForbidden, "api key scope does not allow video generation", "permission_denied")
return
}
body, externalTaskID, err := klingCompatTaskBody(version, model, native)
if err != nil {
writeKlingCompatError(w, http.StatusBadRequest, err.Error(), clients.ErrorCode(err))
return
}
prepared, err := s.prepareTaskRequest(r.Context(), r, user, body)
if err != nil {
writeKlingCompatError(w, http.StatusBadRequest, err.Error(), clients.ErrorCode(err))
return
}
createInput := store.CreateTaskInput{
Kind: "videos.generations",
Model: model,
ExternalTaskID: externalTaskID,
RunMode: runModeFromRequest(prepared.Body),
Async: true,
Request: prepared.Body,
}
if idempotencyKey, hasKey, keyErr := optionalTaskIdempotencyKey(r); keyErr != nil {
writeKlingCompatError(w, http.StatusBadRequest, "invalid Idempotency-Key", "invalid_idempotency_key")
return
} else if hasKey {
createInput.IdempotencyKeyHash = taskIdempotencyKeyHash(idempotencyKey)
createInput.IdempotencyRequestHash = taskIdempotencyRequestHash(createInput.Kind, true, false, prepared.Body)
}
created, err := s.store.CreateTaskIdempotent(r.Context(), createInput, user)
if err != nil {
switch {
case errors.Is(err, store.ErrIdempotencyKeyReused):
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")
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")
}
return
}
task := created.Task
if !created.Replayed {
if err := s.runner.EnqueueAsyncTask(r.Context(), task); err != nil {
writeKlingCompatError(w, http.StatusInternalServerError, err.Error(), "enqueue_failed")
return
}
} else {
w.Header().Set("Idempotent-Replayed", "true")
}
if version == "v2" {
writeJSON(w, http.StatusOK, klingV2Envelope(task))
return
}
writeJSON(w, http.StatusOK, klingV1Envelope(task))
}
func klingCompatTaskBody(version string, model string, native map[string]any) (map[string]any, string, error) {
if model != klingO1Model && model != klingV3OmniModel {
return nil, "", &clients.ClientError{Code: "invalid_parameter", Message: "model_name must be kling-video-o1 or kling-v3-omni", StatusCode: 400, Retryable: false}
}
if native == nil {
native = map[string]any{}
}
body := cloneMap(native)
if version == "v2" {
body = klingV2ToLegacyBody(native)
}
body["model"] = model
body["modelType"] = "omni_video"
body["_compat_provider"] = klingCompatProvider
body["_kling_compat_version"] = version
body["content"] = klingLegacyContent(body)
mode := strings.TrimSpace(stringFromRequestAny(body["mode"]))
if mode == "" && version == "v1" {
// The legacy Omni API defaults to professional (1080p) mode.
mode = "pro"
body["mode"] = mode
}
if mode != "" {
resolution, ok := klingResolutionFromMode(mode)
if !ok {
return nil, "", &clients.ClientError{Code: "invalid_parameter", Message: "mode must be std, pro, or 4k", StatusCode: 400, Retryable: false}
}
body["resolution"] = resolution
}
if strings.EqualFold(strings.TrimSpace(stringFromRequestAny(body["sound"])), "on") {
body["audio"] = true
}
externalTaskID := strings.TrimSpace(stringFromRequestAny(body["external_task_id"]))
if err := validateKlingCompatBody(model, body); err != nil {
return nil, "", err
}
return body, externalTaskID, nil
}
func klingV2ToLegacyBody(native map[string]any) map[string]any {
body := map[string]any{}
for _, key := range []string{"runMode", "simulation", "simulationDurationMs", "simulationProfile"} {
if value, ok := native[key]; ok {
body[key] = value
}
}
settings, _ := native["settings"].(map[string]any)
options, _ := native["options"].(map[string]any)
if settings == nil {
settings = map[string]any{}
}
if options == nil {
options = map[string]any{}
}
if resolution := strings.TrimSpace(stringFromRequestAny(settings["resolution"])); resolution != "" {
switch strings.ToLower(resolution) {
case "720p":
body["mode"] = "std"
case "1080p":
body["mode"] = "pro"
case "4k", "2160p":
body["mode"] = "4k"
default:
body["mode"] = resolution
}
}
for _, key := range []string{"aspect_ratio", "duration", "multi_shot", "shot_type", "multi_prompt"} {
if value, ok := settings[key]; ok {
body[key] = value
}
}
audio := strings.ToLower(strings.TrimSpace(stringFromRequestAny(settings["audio"])))
if audio == "native" || audio == "on" {
body["sound"] = "on"
} else {
body["sound"] = "off"
}
for _, key := range []string{"callback_url", "external_task_id", "watermark_info"} {
if value, ok := options[key]; ok {
body[key] = value
}
}
contents, _ := native["contents"].([]any)
imageList := make([]any, 0)
videoList := make([]any, 0)
elementList := make([]any, 0)
for _, raw := range contents {
item, _ := raw.(map[string]any)
kind := strings.ToLower(strings.TrimSpace(stringFromRequestAny(item["type"])))
switch kind {
case "prompt":
body["prompt"] = stringFromRequestAny(item["text"])
case "first_frame", "last_frame", "refer_image", "reference_image":
image := map[string]any{"image_url": firstNonEmptyRequestString(item, "url", "image_url")}
if kind == "first_frame" {
image["type"] = "first_frame"
} else if kind == "last_frame" {
image["type"] = "end_frame"
}
imageList = append(imageList, image)
case "feature_video", "base_video", "refer_video", "reference_video":
referType := "feature"
if kind == "base_video" {
referType = "base"
}
video := map[string]any{
"video_url": firstNonEmptyRequestString(item, "url", "video_url"),
"refer_type": referType,
"keep_original_sound": firstNonEmptyRequestString(item, "keep_original_sound", "keepOriginalSound"),
}
if audio == "original" && video["keep_original_sound"] == "" {
video["keep_original_sound"] = "yes"
}
videoList = append(videoList, video)
case "element":
elementList = append(elementList, map[string]any{"element_id": firstPresentRequest(item["element_id"], item["id"])})
}
}
if len(imageList) > 0 {
body["image_list"] = imageList
}
if len(videoList) > 0 {
body["video_list"] = videoList
}
if len(elementList) > 0 {
body["element_list"] = elementList
}
return body
}
func klingLegacyContent(body map[string]any) []any {
content := make([]any, 0)
if prompt := strings.TrimSpace(stringFromRequestAny(body["prompt"])); prompt != "" {
content = append(content, map[string]any{"type": "text", "text": prompt})
}
for _, raw := range mapListFromRequest(body["image_list"]) {
role := "reference_image"
switch strings.TrimSpace(stringFromRequestAny(raw["type"])) {
case "first_frame":
role = "first_frame"
case "end_frame", "last_frame":
role = "last_frame"
}
content = append(content, map[string]any{
"type": "image_url", "role": role,
"image_url": map[string]any{"url": firstNonEmptyRequestString(raw, "image_url", "url")},
})
}
for _, raw := range mapListFromRequest(body["video_list"]) {
referType := firstNonEmptyRequestString(raw, "refer_type", "referType")
role := "video_feature"
if referType == "base" {
role = "video_base"
}
content = append(content, map[string]any{
"type": "video_url", "role": role,
"video_url": map[string]any{
"url": firstNonEmptyRequestString(raw, "video_url", "url"),
"refer_type": referType,
"keep_original_sound": firstNonEmptyRequestString(raw, "keep_original_sound", "keepOriginalSound"),
},
})
}
for _, raw := range mapListFromRequest(body["element_list"]) {
content = append(content, map[string]any{
"type": "element",
"element": map[string]any{"element_id": firstPresentRequest(raw["element_id"], raw["id"])},
})
}
return content
}
func validateKlingCompatBody(model string, body map[string]any) error {
maxDuration := 10
if model == klingV3OmniModel {
maxDuration = 15
}
if duration, ok := klingCompatInt(body["duration"]); ok && (duration < 3 || duration > maxDuration) {
return &clients.ClientError{Code: "invalid_parameter", Message: fmt.Sprintf("duration must be between 3 and %d seconds for %s", maxDuration, model), StatusCode: 400, Retryable: false}
}
if model == klingO1Model && boolFromRequestAny(body["multi_shot"]) {
return &clients.ClientError{Code: "invalid_parameter", Message: "kling-video-o1 does not support multi_shot", StatusCode: 400, Retryable: false}
}
if model == klingO1Model && strings.EqualFold(strings.TrimSpace(stringFromRequestAny(body["mode"])), "4k") {
return &clients.ClientError{Code: "invalid_parameter", Message: "kling-video-o1 does not support 4k mode", StatusCode: 400, Retryable: false}
}
if ratio := strings.TrimSpace(stringFromRequestAny(body["aspect_ratio"])); ratio != "" && ratio != "16:9" && ratio != "9:16" && ratio != "1:1" {
return &clients.ClientError{Code: "invalid_parameter", Message: "aspect_ratio must be 16:9, 9:16, or 1:1", StatusCode: 400, Retryable: false}
}
if sound := strings.ToLower(strings.TrimSpace(stringFromRequestAny(body["sound"]))); sound != "" && sound != "on" && sound != "off" {
return &clients.ClientError{Code: "invalid_parameter", Message: "sound must be on or off", StatusCode: 400, Retryable: false}
}
prompt := strings.TrimSpace(stringFromRequestAny(body["prompt"]))
if utf8.RuneCountInString(prompt) > 2500 {
return &clients.ClientError{Code: "invalid_parameter", Message: "prompt must not exceed 2500 characters", StatusCode: 400, Retryable: false}
}
images := mapListFromRequest(body["image_list"])
videos := mapListFromRequest(body["video_list"])
elements := mapListFromRequest(body["element_list"])
for _, image := range images {
if firstNonEmptyRequestString(image, "image_url", "url") == "" {
return &clients.ClientError{Code: "invalid_parameter", Message: "every image_list item requires image_url", StatusCode: 400, Retryable: false}
}
}
for _, video := range videos {
if firstNonEmptyRequestString(video, "video_url", "url") == "" {
return &clients.ClientError{Code: "invalid_parameter", Message: "every video_list item requires video_url", StatusCode: 400, Retryable: false}
}
if referType := strings.TrimSpace(firstNonEmptyRequestString(video, "refer_type", "referType")); referType != "" && referType != "base" && referType != "feature" {
return &clients.ClientError{Code: "invalid_parameter", Message: "video refer_type must be base or feature", StatusCode: 400, Retryable: false}
}
}
for _, element := range elements {
if klingStringAny(firstPresentRequest(element["element_id"], element["id"])) == "" {
return &clients.ClientError{Code: "invalid_parameter", Message: "every element_list item requires element_id", StatusCode: 400, Retryable: false}
}
}
if model == klingO1Model && len(images) == 0 && len(videos) == 0 && len(elements) == 0 {
if duration, ok := klingCompatInt(body["duration"]); ok && duration != 5 && duration != 10 {
return &clients.ClientError{Code: "invalid_parameter", Message: "kling-video-o1 text-only generation supports duration 5 or 10 seconds", StatusCode: 400, Retryable: false}
}
}
if len(videos) > 1 {
return &clients.ClientError{Code: "invalid_parameter", Message: "video_list supports at most one video", StatusCode: 400, Retryable: false}
}
if len(images)+len(elements) > 7 {
return &clients.ClientError{Code: "invalid_parameter", Message: "image_list and element_list support at most seven combined references", StatusCode: 400, Retryable: false}
}
if len(videos) > 0 && len(images)+len(elements) > 4 {
return &clients.ClientError{Code: "invalid_parameter", Message: "requests with video input support at most four image and element references", StatusCode: 400, Retryable: false}
}
if len(videos) > 0 && strings.EqualFold(strings.TrimSpace(stringFromRequestAny(body["sound"])), "on") {
return &clients.ClientError{Code: "invalid_parameter", Message: "sound must be off when video_list is provided", StatusCode: 400, Retryable: false}
}
if model == klingV3OmniModel && len(videos) > 0 {
if duration, ok := klingCompatInt(body["duration"]); ok && duration > 10 {
return &clients.ClientError{Code: "invalid_parameter", Message: "kling-v3-omni video-reference generation supports at most 10 seconds", StatusCode: 400, Retryable: false}
}
}
multiShot := boolFromRequestAny(body["multi_shot"])
shotType := strings.ToLower(strings.TrimSpace(stringFromRequestAny(body["shot_type"])))
if shotType != "" && shotType != "customize" && shotType != "intelligence" {
return &clients.ClientError{Code: "invalid_parameter", Message: "shot_type must be customize or intelligence", StatusCode: 400, Retryable: false}
}
if multiShot && shotType == "customize" {
multiPrompt := mapListFromRequest(body["multi_prompt"])
if len(multiPrompt) == 0 || len(multiPrompt) > 6 {
return &clients.ClientError{Code: "invalid_parameter", Message: "customize multi-shot requires between one and six multi_prompt items", StatusCode: 400, Retryable: false}
}
totalDuration := 0
for _, shot := range multiPrompt {
shotPrompt := strings.TrimSpace(stringFromRequestAny(shot["prompt"]))
shotDuration, ok := klingCompatInt(shot["duration"])
if shotPrompt == "" || utf8.RuneCountInString(shotPrompt) > 2500 || !ok || shotDuration <= 0 {
return &clients.ClientError{Code: "invalid_parameter", Message: "every multi_prompt item requires prompt and a positive integer duration", StatusCode: 400, Retryable: false}
}
totalDuration += shotDuration
}
if totalDuration < 3 || totalDuration > maxDuration {
return &clients.ClientError{Code: "invalid_parameter", Message: fmt.Sprintf("multi_prompt duration must total between 3 and %d seconds", maxDuration), StatusCode: 400, Retryable: false}
}
}
if (!multiShot || shotType == "intelligence" || shotType == "") && prompt == "" {
return &clients.ClientError{Code: "invalid_parameter", Message: "prompt is required for single-shot and intelligence multi-shot generation", StatusCode: 400, Retryable: false}
}
return nil
}
// klingV1GetOmniVideo godoc
// @Summary 查询可灵 V1 Omni 视频任务
// @Tags kling-compatible
// @Produce json
// @Security BearerAuth
// @Param taskID path string true "任务 ID"
// @Success 200 {object} map[string]interface{}
// @Failure 404 {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())
task, err := s.store.GetCompatTask(r.Context(), user, klingCompatProvider, "v1", r.PathValue("taskID"))
if err != nil {
if store.IsNotFound(err) {
writeKlingCompatError(w, http.StatusNotFound, "task not found", "task_not_found")
return
}
writeKlingCompatError(w, http.StatusInternalServerError, "get task failed", "task_query_failed")
return
}
writeJSON(w, http.StatusOK, klingV1Envelope(task))
}
// klingV1ListOmniVideos godoc
// @Summary 分页查询可灵 V1 Omni 视频任务
// @Tags kling-compatible
// @Produce json
// @Security BearerAuth
// @Param pageNum query int false "页码" default(1)
// @Param pageSize query int false "每页数量" default(30)
// @Success 200 {object} map[string]interface{}
// @Router /api/v1/kling/v1/videos/omni-video [get]
func (s *Server) klingV1ListOmniVideos(w http.ResponseWriter, r *http.Request) {
user, _ := auth.UserFromContext(r.Context())
page, err := positiveQueryInt(r.URL.Query().Get("pageNum"), 1)
if err != nil || page > 1000 {
writeKlingCompatError(w, http.StatusBadRequest, "invalid pageNum", "invalid_parameter")
return
}
pageSize, err := positiveQueryInt(r.URL.Query().Get("pageSize"), 30)
if err != nil || pageSize > 500 {
writeKlingCompatError(w, http.StatusBadRequest, "invalid pageSize", "invalid_parameter")
return
}
result, err := s.store.ListCompatTasks(r.Context(), user, store.CompatTaskListFilter{Provider: klingCompatProvider, Version: "v1", Page: page, PageSize: pageSize})
if err != nil {
writeKlingCompatError(w, http.StatusInternalServerError, "list tasks failed", "task_query_failed")
return
}
data := make([]any, 0, len(result.Items))
for _, task := range result.Items {
data = append(data, klingV1TaskData(task))
}
writeJSON(w, http.StatusOK, map[string]any{"code": 0, "message": "success", "request_id": klingRequestID(result.Items), "data": data})
}
// klingV2GetTasks godoc
// @Summary 按 ID 查询可灵 API 2.0 任务
// @Tags kling-compatible
// @Produce json
// @Security BearerAuth
// @Param task_ids query string false "逗号分隔的任务 ID"
// @Param external_task_ids query string false "逗号分隔的外部任务 ID"
// @Success 200 {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())
taskIDs := splitKlingIDs(r.URL.Query().Get("task_ids"))
externalIDs := splitKlingIDs(r.URL.Query().Get("external_task_ids"))
if (len(taskIDs) == 0) == (len(externalIDs) == 0) {
writeKlingCompatError(w, http.StatusBadRequest, "choose exactly one of task_ids or external_task_ids", "invalid_parameter")
return
}
identifiers := taskIDs
if len(externalIDs) > 0 {
identifiers = externalIDs
}
data := make([]any, 0, len(identifiers))
for _, identifier := range identifiers {
task, err := s.store.GetCompatTask(r.Context(), user, klingCompatProvider, "v2", identifier)
if err != nil {
if store.IsNotFound(err) {
continue
}
writeKlingCompatError(w, http.StatusInternalServerError, "get task failed", "task_query_failed")
return
}
data = append(data, klingV2TaskData(task))
}
writeJSON(w, http.StatusOK, map[string]any{"code": 0, "message": "success", "request_id": klingRequestIDFromAny(data), "data": data})
}
// klingV2ListTasks godoc
// @Summary 分页查询可灵 API 2.0 任务
// @Tags kling-compatible
// @Accept json
// @Produce json
// @Security BearerAuth
// @Param input body map[string]interface{} true "游标、数量、时间范围和筛选条件"
// @Success 200 {object} map[string]interface{}
// @Router /api/v1/kling/v2/tasks [post]
func (s *Server) klingV2ListTasks(w http.ResponseWriter, r *http.Request) {
user, _ := auth.UserFromContext(r.Context())
var body map[string]any
if err := decodeKlingJSON(r, &body); err != nil {
writeKlingCompatError(w, http.StatusBadRequest, "invalid json body", "invalid_json_body")
return
}
page, err := klingCursorPage(stringFromRequestAny(body["cursor"]))
if err != nil {
writeKlingCompatError(w, http.StatusBadRequest, "invalid cursor", "invalid_parameter")
return
}
limit, ok := klingCompatInt(body["limit"])
if !ok || limit <= 0 {
limit = 100
}
if limit > 500 {
writeKlingCompatError(w, http.StatusBadRequest, "limit must not exceed 500", "invalid_parameter")
return
}
createdFrom, err := klingMillisTime(body["start_time"])
if err != nil {
writeKlingCompatError(w, http.StatusBadRequest, "invalid start_time", "invalid_parameter")
return
}
createdTo, err := klingMillisTime(body["end_time"])
if err != nil {
writeKlingCompatError(w, http.StatusBadRequest, "invalid end_time", "invalid_parameter")
return
}
statuses := klingInternalStatuses(body["filters"])
result, err := s.store.ListCompatTasks(r.Context(), user, store.CompatTaskListFilter{
Provider: klingCompatProvider, Version: "v2", Statuses: statuses,
CreatedFrom: createdFrom, CreatedTo: createdTo, Page: page, PageSize: limit,
})
if err != nil {
writeKlingCompatError(w, http.StatusInternalServerError, "list tasks failed", "task_query_failed")
return
}
items := make([]any, 0, len(result.Items))
for _, task := range result.Items {
items = append(items, klingV2TaskData(task))
}
hasMore := page*limit < result.Total
nextCursor := ""
if hasMore {
nextCursor = base64.RawURLEncoding.EncodeToString([]byte(strconv.Itoa(page + 1)))
}
writeJSON(w, http.StatusOK, map[string]any{
"code": 0, "message": "success", "request_id": klingRequestID(result.Items),
"data": map[string]any{"result": items, "count": len(items), "next_cursor": nextCursor, "has_more": hasMore},
})
}
func klingV1Envelope(task store.GatewayTask) map[string]any {
return map[string]any{"code": 0, "message": "success", "request_id": firstNonEmpty(task.RequestID, task.ID), "data": klingV1TaskData(task)}
}
func klingV1TaskData(task store.GatewayTask) map[string]any {
data := map[string]any{
"task_id": task.ID, "task_status": klingV1Status(task.Status),
"task_info": map[string]any{"external_task_id": task.ExternalTaskID},
"created_at": task.CreatedAt.UnixMilli(), "updated_at": task.UpdatedAt.UnixMilli(),
}
if task.ErrorMessage != "" || task.Error != "" {
data["task_status_msg"] = firstNonEmpty(task.ErrorMessage, task.Error)
}
if watermarkInfo, ok := task.Request["watermark_info"].(map[string]any); ok {
data["watermark_info"] = watermarkInfo
}
videos := klingTaskVideos(task)
if len(videos) > 0 {
data["task_result"] = map[string]any{"videos": videos}
}
return data
}
func klingV2Envelope(task store.GatewayTask) map[string]any {
return map[string]any{"code": 0, "message": "success", "request_id": firstNonEmpty(task.RequestID, task.ID), "data": klingV2TaskData(task)}
}
func klingV2TaskData(task store.GatewayTask) map[string]any {
data := map[string]any{
"id": task.ID, "status": klingV2Status(task.Status),
"create_time": task.CreatedAt.UnixMilli(), "update_time": task.UpdatedAt.UnixMilli(),
"external_id": task.ExternalTaskID,
}
if message := firstNonEmpty(task.ErrorMessage, task.Error); message != "" {
data["message"] = message
}
if outputs := klingV2Outputs(task); len(outputs) > 0 {
data["outputs"] = outputs
}
return data
}
func klingTaskVideos(task store.GatewayTask) []any {
items, _ := task.Result["data"].([]any)
videos := make([]any, 0, len(items))
for index, raw := range items {
item, _ := raw.(map[string]any)
url := firstNonEmptyRequestString(item, "url", "video_url")
if url == "" {
continue
}
video := map[string]any{"id": firstNonEmptyRequestString(item, "id")}
if video["id"] == "" {
video["id"] = fmt.Sprintf("%s-%d", task.ID, index+1)
}
video["url"] = url
if watermarkURL := firstNonEmptyRequestString(item, "watermark_url"); watermarkURL != "" {
video["watermark_url"] = watermarkURL
}
if duration := klingStringAny(item["duration"]); duration != "" {
video["duration"] = duration
}
videos = append(videos, video)
}
return videos
}
func klingV2Outputs(task store.GatewayTask) []any {
videos := klingTaskVideos(task)
outputs := make([]any, 0, len(videos))
for _, raw := range videos {
video, _ := raw.(map[string]any)
output := cloneMap(video)
output["type"] = "video"
outputs = append(outputs, output)
}
return outputs
}
func klingV1Status(status string) string {
switch status {
case "succeeded":
return "succeed"
case "failed", "cancelled":
return "failed"
case "running":
return "processing"
default:
return "submitted"
}
}
func klingV2Status(status string) string {
switch status {
case "succeeded":
return "succeeded"
case "failed", "cancelled":
return "failed"
case "running":
return "processing"
default:
return "submitted"
}
}
func klingV2ProviderModel(pathModel string) (string, bool) {
return canonicalKlingOmniModel(pathModel)
}
func canonicalKlingOmniModel(value string) (string, bool) {
switch strings.ToLower(strings.TrimSpace(value)) {
case "kling-o1", "kling-video-o1":
return klingO1Model, true
case "kling-v3-omni", "kling-3.0-omni", "kling-3-omni":
return klingV3OmniModel, true
default:
return "", false
}
}
func klingResolutionFromMode(mode string) (string, bool) {
switch strings.ToLower(strings.TrimSpace(mode)) {
case "std":
return "720p", true
case "pro":
return "1080p", true
case "4k":
return "2160p", true
default:
return "", false
}
}
func decodeKlingJSON(r *http.Request, target any) error {
decoder := json.NewDecoder(r.Body)
decoder.UseNumber()
if err := decoder.Decode(target); err != nil {
return err
}
if err := decoder.Decode(&struct{}{}); !errors.Is(err, io.EOF) {
if err != nil {
return err
}
return errors.New("multiple json values")
}
return nil
}
func writeKlingCompatError(w http.ResponseWriter, status int, message string, code string) {
if strings.TrimSpace(code) == "" {
code = "invalid_request"
}
writeJSON(w, status, map[string]any{
"code": klingCompatErrorCode(status),
"message": message,
"request_id": "",
"error": code,
})
}
func klingCompatErrorCode(status int) int {
switch status {
case http.StatusBadRequest:
return 1001
case http.StatusUnauthorized:
return 1100
case http.StatusForbidden:
return 1302
case http.StatusNotFound:
return 1201
case http.StatusConflict:
return 1200
case http.StatusTooManyRequests:
return 1400
default:
return 5000
}
}
func mapListFromRequest(value any) []map[string]any {
items, _ := value.([]any)
if len(items) == 0 {
if typed, ok := value.([]map[string]any); ok {
return typed
}
}
out := make([]map[string]any, 0, len(items))
for _, item := range items {
if mapped, ok := item.(map[string]any); ok {
out = append(out, mapped)
}
}
return out
}
func firstPresentRequest(values ...any) any {
for _, value := range values {
if value == nil {
continue
}
if text, ok := value.(string); ok {
if strings.TrimSpace(text) != "" {
return value
}
continue
}
return value
}
return nil
}
func boolFromRequestAny(value any) bool {
switch typed := value.(type) {
case bool:
return typed
case string:
parsed, _ := strconv.ParseBool(strings.TrimSpace(typed))
return parsed
default:
return false
}
}
func klingCompatInt(value any) (int, bool) {
text := klingStringAny(value)
if text == "" {
return 0, false
}
number, err := strconv.Atoi(text)
return number, err == nil
}
func splitKlingIDs(raw string) []string {
parts := strings.Split(raw, ",")
out := make([]string, 0, len(parts))
for _, part := range parts {
if part = strings.TrimSpace(part); part != "" {
out = append(out, part)
}
}
return out
}
func klingCursorPage(cursor string) (int, error) {
cursor = strings.TrimSpace(cursor)
if cursor == "" {
return 1, nil
}
decoded, err := base64.RawURLEncoding.DecodeString(cursor)
if err != nil {
return 0, err
}
page, err := strconv.Atoi(string(decoded))
if err != nil || page <= 0 {
return 0, errors.New("invalid cursor")
}
return page, nil
}
func klingMillisTime(value any) (*time.Time, error) {
text := klingStringAny(value)
if text == "" {
return nil, nil
}
millis, err := strconv.ParseInt(text, 10, 64)
if err != nil || millis < 0 {
return nil, errors.New("invalid millisecond timestamp")
}
parsed := time.UnixMilli(millis)
return &parsed, nil
}
func klingInternalStatuses(filters any) []string {
statuses := make([]string, 0)
for _, filter := range mapListFromRequest(filters) {
if stringFromRequestAny(filter["key"]) != "status" {
continue
}
values, _ := filter["values"].([]any)
for _, value := range values {
switch strings.ToLower(strings.TrimSpace(stringFromRequestAny(value))) {
case "submitted":
statuses = append(statuses, "queued")
case "processing":
statuses = append(statuses, "running")
case "succeeded":
statuses = append(statuses, "succeeded")
case "failed":
statuses = append(statuses, "failed", "cancelled")
}
}
}
return statuses
}
func klingRequestID(tasks []store.GatewayTask) string {
if len(tasks) == 0 {
return ""
}
return firstNonEmpty(tasks[0].RequestID, tasks[0].ID)
}
func klingRequestIDFromAny(items []any) string {
if len(items) == 0 {
return ""
}
item, _ := items[0].(map[string]any)
return stringFromRequestAny(item["id"])
}
func klingStringAny(value any) string {
switch typed := value.(type) {
case string:
return strings.TrimSpace(typed)
case json.Number:
return typed.String()
case float64:
return strconv.FormatFloat(typed, 'f', -1, 64)
case float32:
return strconv.FormatFloat(float64(typed), 'f', -1, 32)
case int:
return strconv.Itoa(typed)
case int64:
return strconv.FormatInt(typed, 10)
case int32:
return strconv.FormatInt(int64(typed), 10)
default:
return ""
}
}
@@ -0,0 +1,265 @@
package httpapi
import (
"context"
"io"
"log/slog"
"net/http"
"net/http/httptest"
"os"
"strconv"
"strings"
"testing"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestKlingCompatibilitySimulationFlow(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 Kling compatibility 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()
var upgradedBaseModels int
if err := db.Pool().QueryRow(ctx, `
SELECT count(*)
FROM base_model_catalog
WHERE provider_key = 'keling'
AND provider_model_name IN ('kling-video-o1', 'kling-v3-omni')
AND model_type @> '["video_generate","image_to_video","omni_video"]'::jsonb
AND capabilities ? 'video_generate'
AND capabilities ? 'image_to_video'
AND capabilities ? 'omni_video'
AND metadata->'rawModel'->'types' @> '["video_generate","image_to_video","omni_video"]'::jsonb`).Scan(&upgradedBaseModels); err != nil {
t.Fatalf("read upgraded Kling Omni base model capabilities: %v", err)
}
if upgradedBaseModels != 2 {
t.Fatalf("expected both Kling Omni base models to expose base video capabilities, got %d", upgradedBaseModels)
}
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",
CORSAllowedOrigin: "*",
}, db, slog.New(slog.NewTextHandler(io.Discard, nil))))
defer server.Close()
suffix := strconv.FormatInt(time.Now().UnixNano(), 10)
username := "kling_compat_" + suffix
password := "password123"
var registerResponse struct {
AccessToken string `json:"accessToken"`
}
doJSON(t, server.URL, http.MethodPost, "/api/v1/auth/register", "", map[string]any{
"username": username,
"email": username + "@example.com",
"password": password,
}, http.StatusCreated, &registerResponse)
var apiKeyResponse struct {
Secret string `json:"secret"`
}
doJSON(t, server.URL, http.MethodPost, "/api/v1/api-keys", registerResponse.AccessToken, map[string]any{
"name": "Kling compatibility key",
}, http.StatusCreated, &apiKeyResponse)
if _, err := db.Pool().Exec(ctx, `UPDATE gateway_users SET roles = '["admin"]'::jsonb WHERE username = $1`, username); err != nil {
t.Fatalf("promote compatibility user: %v", err)
}
var loginResponse struct {
AccessToken string `json:"accessToken"`
}
doJSON(t, server.URL, http.MethodPost, "/api/v1/auth/login", "", map[string]any{
"account": username, "password": password,
}, http.StatusOK, &loginResponse)
var platform struct {
ID string `json:"id"`
}
doJSON(t, server.URL, http.MethodPost, "/api/admin/platforms", loginResponse.AccessToken, map[string]any{
"provider": "keling",
"platformKey": "kling-compat-" + suffix,
"name": "Kling Compatibility Simulation",
"baseUrl": "https://api-beijing.klingai.com/v1",
"authType": "AccessKey-SecretKey",
"credentials": map[string]any{"accessKey": "test-ak", "secretKey": "test-sk"},
}, http.StatusCreated, &platform)
for _, model := range []string{klingO1Model, klingV3OmniModel} {
doJSON(t, server.URL, http.MethodPost, "/api/admin/platforms/"+platform.ID+"/models", loginResponse.AccessToken, map[string]any{
"canonicalModelKey": "keling:" + model,
"modelName": model,
"providerModelName": model,
"modelAlias": model,
"modelType": []string{"omni_video"},
"displayName": model,
}, http.StatusCreated, nil)
}
var upgradedPlatformModels int
if err := db.Pool().QueryRow(ctx, `
SELECT count(*)
FROM platform_models
WHERE platform_id = $1::uuid
AND model_type @> '["video_generate","image_to_video","omni_video"]'::jsonb
AND capabilities ? 'video_generate'
AND capabilities ? 'image_to_video'
AND capabilities ? 'omni_video'`, platform.ID).Scan(&upgradedPlatformModels); err != nil {
t.Fatalf("read upgraded Kling Omni platform model capabilities: %v", err)
}
if upgradedPlatformModels != 2 {
t.Fatalf("expected both Kling Omni platform models to expose base video capabilities, got %d", upgradedPlatformModels)
}
assertGenericVideoGeneration := func(name string, model string, image string, expectedModelType string) {
t.Helper()
t.Run(name, func(t *testing.T) {
request := map[string]any{
"model": model,
"prompt": "通用视频接口模拟任务",
"duration": 5,
"resolution": "720p",
"runMode": "simulation",
"simulation": true,
"simulationDurationMs": 5,
}
if image != "" {
request["image"] = image
}
var response struct {
Task struct {
ID string `json:"id"`
Status string `json:"status"`
ModelType string `json:"modelType"`
ResolvedModel string `json:"resolvedModel"`
} `json:"task"`
}
doJSONWithHeaders(t, server.URL, http.MethodPost, "/api/v1/videos/generations", apiKeyResponse.Secret, request, map[string]string{"X-Async": "true"}, http.StatusAccepted, &response)
if response.Task.ID == "" {
t.Fatal("async generic video response did not return a task id")
}
waitForTaskStatus(t, server.URL, apiKeyResponse.Secret, response.Task.ID, []string{"succeeded"}, 10*time.Second)
doJSON(t, server.URL, http.MethodGet, "/api/v1/tasks/"+response.Task.ID, apiKeyResponse.Secret, nil, http.StatusOK, &response.Task)
resolvedModel, resolved := klingV2ProviderModel(response.Task.ResolvedModel)
if response.Task.Status != "succeeded" || response.Task.ModelType != expectedModelType || !resolved || resolvedModel != model {
t.Fatalf("generic video request without modelType should use inferred capability: %+v", response.Task)
}
})
}
for _, model := range []string{klingO1Model, klingV3OmniModel} {
assertGenericVideoGeneration(model+"-text-to-video", model, "", "video_generate")
assertGenericVideoGeneration(model+"-image-to-video", model, "https://example.com/first.png", "image_to_video")
}
createV1 := func(model string, duration int, externalID string) string {
t.Helper()
var response map[string]any
doJSON(t, server.URL, http.MethodPost, "/api/v1/kling/v1/videos/omni-video", apiKeyResponse.Secret, map[string]any{
"model_name": model,
"prompt": "兼容接口模拟任务",
"duration": duration,
"mode": "std",
"external_task_id": externalID,
"runMode": "simulation",
"simulation": true,
"simulationDurationMs": 5,
}, http.StatusOK, &response)
if response["code"] != float64(0) {
t.Fatalf("unexpected V1 response: %#v", response)
}
data, _ := response["data"].(map[string]any)
taskID, _ := data["task_id"].(string)
if taskID == "" {
t.Fatalf("V1 response missing task id: %#v", response)
}
return taskID
}
o1TaskID := createV1(klingO1Model, 5, "compat-o1-"+suffix)
v3TaskID := createV1(klingV3OmniModel, 15, "compat-v3-"+suffix)
for _, taskID := range []string{o1TaskID, v3TaskID} {
waitKlingV1SimulationTask(t, server.URL, apiKeyResponse.Secret, taskID)
}
var listResponse map[string]any
doJSON(t, server.URL, http.MethodGet, "/api/v1/kling/v1/videos/omni-video?pageNum=1&pageSize=10", apiKeyResponse.Secret, nil, http.StatusOK, &listResponse)
items, _ := listResponse["data"].([]any)
if len(items) < 2 {
t.Fatalf("V1 task list did not return compatibility tasks: %#v", listResponse)
}
var duplicateResponse map[string]any
doJSON(t, server.URL, http.MethodPost, "/api/v1/kling/v1/videos/omni-video", apiKeyResponse.Secret, map[string]any{
"model_name": klingO1Model, "prompt": "duplicate", "duration": 5,
"external_task_id": "compat-o1-" + suffix,
"runMode": "simulation", "simulation": true,
}, http.StatusConflict, &duplicateResponse)
if duplicateResponse["code"] != float64(1200) || duplicateResponse["error"] != "external_task_id_reused" {
t.Fatalf("unexpected duplicate external id response: %#v", duplicateResponse)
}
var v2Response map[string]any
doJSON(t, server.URL, http.MethodPost, "/api/v1/kling/v2/omni-video/kling-v3-omni", apiKeyResponse.Secret, map[string]any{
"contents": []any{map[string]any{"type": "prompt", "text": "API 2.0 模拟任务"}},
"settings": map[string]any{"duration": 3, "resolution": "720p", "aspect_ratio": "16:9", "audio": "off"},
"options": map[string]any{"external_task_id": "compat-v2-" + suffix},
"runMode": "simulation",
"simulation": true,
"simulationDurationMs": 5,
}, http.StatusOK, &v2Response)
v2Data, _ := v2Response["data"].(map[string]any)
v2TaskID, _ := v2Data["id"].(string)
if v2TaskID == "" {
t.Fatalf("V2 response missing task id: %#v", v2Response)
}
waitKlingV2SimulationTask(t, server.URL, apiKeyResponse.Secret, v2TaskID)
}
func waitKlingV1SimulationTask(t *testing.T, baseURL string, apiKey string, taskID string) {
t.Helper()
deadline := time.Now().Add(5 * time.Second)
for time.Now().Before(deadline) {
var response map[string]any
doJSON(t, baseURL, http.MethodGet, "/api/v1/kling/v1/videos/omni-video/"+taskID, apiKey, nil, http.StatusOK, &response)
data, _ := response["data"].(map[string]any)
switch data["task_status"] {
case "succeed":
return
case "failed":
t.Fatalf("V1 simulation task failed: %#v", response)
}
time.Sleep(20 * time.Millisecond)
}
t.Fatalf("V1 simulation task %s timed out", taskID)
}
func waitKlingV2SimulationTask(t *testing.T, baseURL string, apiKey string, taskID string) {
t.Helper()
deadline := time.Now().Add(5 * time.Second)
for time.Now().Before(deadline) {
var response map[string]any
doJSON(t, baseURL, http.MethodGet, "/api/v1/kling/v2/tasks?task_ids="+taskID, apiKey, nil, http.StatusOK, &response)
items, _ := response["data"].([]any)
if len(items) == 1 {
data, _ := items[0].(map[string]any)
switch data["status"] {
case "succeeded":
return
case "failed":
t.Fatalf("V2 simulation task failed: %#v", response)
}
}
time.Sleep(20 * time.Millisecond)
}
t.Fatalf("V2 simulation task %s timed out", taskID)
}
@@ -0,0 +1,137 @@
package httpapi
import (
"encoding/json"
"net/http"
"net/http/httptest"
"strings"
"testing"
)
func TestKlingV1O1CompatibilityBody(t *testing.T) {
body, externalID, err := klingCompatTaskBody("v1", klingO1Model, map[string]any{
"prompt": "一只纸鹤飞过湖面",
"duration": json.Number("10"),
"aspect_ratio": "16:9",
"sound": "on",
"external_task_id": "client-o1-1",
"watermark_info": map[string]any{"enabled": true},
})
if err != nil {
t.Fatalf("build O1 compatibility body: %v", err)
}
if externalID != "client-o1-1" || body["model"] != klingO1Model || body["modelType"] != "omni_video" {
t.Fatalf("unexpected identity fields: %#v", body)
}
if body["mode"] != "pro" || body["resolution"] != "1080p" || body["audio"] != true {
t.Fatalf("unexpected V1 defaults: %#v", body)
}
content := mapListFromRequest(body["content"])
if len(content) != 1 || content[0]["type"] != "text" || content[0]["text"] != "一只纸鹤飞过湖面" {
t.Fatalf("unexpected canonical content: %#v", content)
}
}
func TestKlingV1V3OmniCompatibilityBody(t *testing.T) {
body, _, err := klingCompatTaskBody("v1", klingV3OmniModel, map[string]any{
"multi_shot": true,
"shot_type": "customize",
"mode": "4k",
"multi_prompt": []any{
map[string]any{"index": json.Number("1"), "prompt": "推近人物", "duration": json.Number("7")},
map[string]any{"index": json.Number("2"), "prompt": "切到城市远景", "duration": json.Number("8")},
},
"image_list": []any{
map[string]any{"image_url": "https://example.com/first.png", "type": "first_frame"},
},
})
if err != nil {
t.Fatalf("build 3.0 Omni compatibility body: %v", err)
}
if body["resolution"] != "2160p" {
t.Fatalf("4k mode was not normalized: %#v", body)
}
content := mapListFromRequest(body["content"])
if len(content) != 1 || content[0]["role"] != "first_frame" {
t.Fatalf("image input was not normalized: %#v", content)
}
}
func TestKlingV2CompatibilityBody(t *testing.T) {
body, externalID, err := klingCompatTaskBody("v2", klingV3OmniModel, map[string]any{
"contents": []any{
map[string]any{"type": "prompt", "text": "让角色向镜头挥手"},
map[string]any{"type": "first_frame", "url": "https://example.com/first.png"},
map[string]any{"type": "element", "id": json.Number("42")},
},
"settings": map[string]any{
"resolution": "1080p",
"duration": json.Number("15"),
"aspect_ratio": "9:16",
"audio": "native",
},
"options": map[string]any{
"external_task_id": "client-v2-1",
"callback_url": "https://example.com/callback",
},
})
if err != nil {
t.Fatalf("build V2 compatibility body: %v", err)
}
if externalID != "client-v2-1" || body["mode"] != "pro" || body["audio"] != true {
t.Fatalf("unexpected V2 settings: %#v", body)
}
if len(mapListFromRequest(body["image_list"])) != 1 || len(mapListFromRequest(body["element_list"])) != 1 {
t.Fatalf("unexpected V2 references: %#v", body)
}
}
func TestKlingCompatibilityValidation(t *testing.T) {
tests := []struct {
name string
model string
body map[string]any
}{
{name: "O1 duration", model: klingO1Model, body: map[string]any{"prompt": "test", "duration": 11}},
{name: "O1 text-only flexible duration", model: klingO1Model, body: map[string]any{"prompt": "test", "duration": 3}},
{name: "O1 4k", model: klingO1Model, body: map[string]any{"prompt": "test", "duration": 5, "mode": "4k"}},
{name: "O1 multi-shot", model: klingO1Model, body: map[string]any{"multi_shot": true, "multi_prompt": []any{map[string]any{"prompt": "test", "duration": 3}}}},
{name: "custom multi-shot without prompts", model: klingV3OmniModel, body: map[string]any{"multi_shot": true, "shot_type": "customize"}},
{name: "intelligence multi-shot without prompt", model: klingV3OmniModel, body: map[string]any{"multi_shot": true, "shot_type": "intelligence"}},
{name: "video with native audio", model: klingV3OmniModel, body: map[string]any{"prompt": "test", "sound": "on", "video_list": []any{map[string]any{"video_url": "https://example.com/input.mp4"}}}},
{name: "video duration too long", model: klingV3OmniModel, body: map[string]any{"prompt": "test", "duration": 15, "video_list": []any{map[string]any{"video_url": "https://example.com/input.mp4"}}}},
{name: "too many references", model: klingV3OmniModel, body: map[string]any{"prompt": "test", "image_list": []any{
map[string]any{}, map[string]any{}, map[string]any{}, map[string]any{}, map[string]any{}, map[string]any{}, map[string]any{}, map[string]any{},
}}},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
if _, _, err := klingCompatTaskBody("v1", test.model, test.body); err == nil {
t.Fatal("expected validation error")
}
})
}
if _, _, err := klingCompatTaskBody("v1", klingV3OmniModel, map[string]any{"prompt": "test", "duration": 15}); err != nil {
t.Fatalf("3.0 Omni should allow a 15-second duration: %v", err)
}
}
func TestDecodeKlingJSONRejectsTrailingData(t *testing.T) {
request := httptest.NewRequest(http.MethodPost, "/kling/v1/videos/omni-video", strings.NewReader(`{"prompt":"ok"} trailing`))
var body map[string]any
if err := decodeKlingJSON(request, &body); err == nil {
t.Fatal("expected trailing JSON error")
}
}
func TestWriteKlingCompatErrorUsesOfficialNumericEnvelope(t *testing.T) {
recorder := httptest.NewRecorder()
writeKlingCompatError(recorder, http.StatusBadRequest, "bad request", "invalid_parameter")
var body map[string]any
if err := json.Unmarshal(recorder.Body.Bytes(), &body); err != nil {
t.Fatalf("decode error response: %v", err)
}
if body["code"] != float64(1001) || body["error"] != "invalid_parameter" {
t.Fatalf("unexpected error envelope: %#v", body)
}
}
@@ -1141,6 +1141,10 @@ func canonicalCapabilityFilterValue(value string) string {
return "text_embedding"
case "rerank", "reranks":
return "text_rerank"
case "vectorize", "image_vectorizer":
return "image_vectorize"
case "video_upscale", "upscale":
return "video_enhance"
case "model":
return "model_3d"
default:
@@ -1164,6 +1168,10 @@ func capabilityFilterValueForTag(tag string) string {
return "structured_output"
case "数字人":
return "digital_human"
case "图片矢量化":
return "image_vectorize"
case "视频增强":
return "video_enhance"
case "重排序":
return "text_rerank"
case "3D 模型":
@@ -1194,11 +1202,13 @@ func capabilityLabel(value string) string {
"image_generate": "图像生成",
"image_edit": "图像编辑",
"image_analysis": "图像分析",
"image_vectorize": "图片矢量化",
"video_generate": "视频生成",
"image_to_video": "图生视频",
"text_to_video": "文生视频",
"video_edit": "视频编辑",
"video_understanding": "视频理解",
"video_enhance": "视频增强",
"audio_generate": "音频生成",
"text_to_speech": "语音合成",
"voice_clone": "音色克隆",
+34 -23
View File
@@ -7,46 +7,57 @@ import (
)
func (s *Server) platformModelResponse(ctx context.Context, model store.PlatformModel) store.PlatformModel {
return s.platformModelResponseWithRuleSets(model, s.responsePricingRuleSetConfigs(ctx, []store.PlatformModel{model}))
}
func (s *Server) platformModelResponseWithRuleSets(model store.PlatformModel, ruleSetConfigs map[string]map[string]any) store.PlatformModel {
model.Capabilities = store.EffectivePlatformModelCapabilities(model.BaseCapabilities, model.Capabilities)
model.Capabilities = enrichResponseCapabilities(model)
model = s.withEffectiveResponseBillingConfig(ctx, model)
model = withEffectiveResponseBillingConfig(model, ruleSetConfigs)
return store.FilterPlatformModelBillingConfig(model)
}
func (s *Server) platformModelResponses(ctx context.Context, models []store.PlatformModel) []store.PlatformModel {
ruleSetConfigs := s.responsePricingRuleSetConfigs(ctx, models)
items := make([]store.PlatformModel, len(models))
for i, model := range models {
items[i] = s.platformModelResponse(ctx, model)
items[i] = s.platformModelResponseWithRuleSets(model, ruleSetConfigs)
}
return items
}
func (s *Server) withEffectiveResponseBillingConfig(ctx context.Context, model store.PlatformModel) store.PlatformModel {
config := model.BillingConfig
if model.PricingRuleSetID != "" {
if ruleSetConfig, err := s.store.PricingRuleSetBillingConfig(ctx, model.PricingRuleSetID); err == nil && len(ruleSetConfig) > 0 {
config = ruleSetConfig
func (s *Server) responsePricingRuleSetConfigs(ctx context.Context, models []store.PlatformModel) map[string]map[string]any {
configs := map[string]map[string]any{}
if s.store == nil {
return configs
}
ids := map[string]bool{}
for _, model := range models {
for _, id := range []string{firstNonEmpty(model.BasePricingRuleSetID, model.PlatformPricingRuleSetID), model.PricingRuleSetID} {
if id != "" {
ids[id] = true
}
}
}
if len(model.BillingConfigOverride) > 0 {
config = mergeResponseBillingConfig(config, model.BillingConfigOverride)
for id := range ids {
if config, err := s.store.PricingRuleSetBillingConfig(ctx, id); err == nil && len(config) > 0 {
configs[id] = config
}
}
model.BillingConfig = config
return model
return configs
}
func mergeResponseBillingConfig(base map[string]any, override map[string]any) map[string]any {
if len(base) == 0 && len(override) == 0 {
return nil
}
out := make(map[string]any, len(base)+len(override))
for key, value := range base {
out[key] = value
}
for key, value := range override {
out[key] = value
}
return out
func withEffectiveResponseBillingConfig(model store.PlatformModel, ruleSetConfigs map[string]map[string]any) store.PlatformModel {
inheritedRuleSetConfig := ruleSetConfigs[firstNonEmpty(model.BasePricingRuleSetID, model.PlatformPricingRuleSetID)]
modelRuleSetConfig := ruleSetConfigs[model.PricingRuleSetID]
model.BillingConfig = store.ResolveEffectiveBillingConfig(store.EffectiveBillingConfigInput{
BaseConfig: model.BaseBillingConfig,
LegacyPlatformModelConfig: model.BillingConfig,
InheritedRuleSetConfig: inheritedRuleSetConfig,
ModelRuleSetConfig: modelRuleSetConfig,
Override: model.BillingConfigOverride,
})
return model
}
func enrichResponseCapabilities(model store.PlatformModel) map[string]any {
@@ -173,6 +173,41 @@ func TestPlatformModelResponsePreservesTextGenerateFieldsOverFallbacks(t *testin
assertStringListValue(t, textGenerate["thinkingEffortLevels"], []string{"minimal", "low", "medium"})
}
func TestPlatformModelResponseUsesBaseBillingConfigWithoutMaterializedSnapshot(t *testing.T) {
model := store.PlatformModel{
ModelName: "base-priced-model",
ModelType: store.StringList{"video_generate"},
BaseBillingConfig: map[string]any{
"video": map[string]any{"basePrice": float64(416)},
},
}
response := (&Server{}).platformModelResponse(context.Background(), model)
video, ok := response.BillingConfig["video"].(map[string]any)
if !ok || video["basePrice"] != float64(416) {
t.Fatalf("expected base billing price 416, got %#v", response.BillingConfig)
}
}
func TestEffectiveResponseBillingConfigPrefersBaseRuleOverLegacySnapshot(t *testing.T) {
model := store.PlatformModel{
BasePricingRuleSetID: "seedance-pricing",
BillingConfig: map[string]any{
"video": map[string]any{"basePrice": float64(100)},
},
}
response := withEffectiveResponseBillingConfig(model, map[string]map[string]any{
"seedance-pricing": {
"video": map[string]any{"basePrice": float64(416)},
},
})
video, ok := response.BillingConfig["video"].(map[string]any)
if !ok || video["basePrice"] != float64(416) {
t.Fatalf("expected base rule price 416, got %#v", response.BillingConfig)
}
}
func textGenerateCapabilities(t *testing.T, model store.PlatformModel) map[string]any {
t.Helper()
capabilities, ok := model.Capabilities["text_generate"].(map[string]any)
+57 -12
View File
@@ -23,8 +23,8 @@ type SkillBundleMetadataResponse struct {
Modules []string `json:"modules" example:"model-runtime"`
FileName string `json:"fileName" example:"ai-gateway-ops-management-v1.0.2.zip"`
DownloadPath string `json:"downloadPath" example:"/api/v1/public/skills/ai-gateway-ops-management/download"`
APIDocsJSONPath string `json:"apiDocsJsonPath" example:"/api-docs-json"`
APIDocsYAMLPath string `json:"apiDocsYamlPath" example:"/api-docs-yaml"`
APIDocsJSONPath string `json:"apiDocsJsonPath" example:"/api/v1/openapi.json"`
APIDocsYAMLPath string `json:"apiDocsYamlPath" example:"/api/v1/openapi.yaml"`
}
type ErrorEnvelope struct {
@@ -210,17 +210,31 @@ type PricingEstimateResponse struct {
RequestFingerprint string `json:"requestFingerprint" example:"76ef6a537de8e71bd1ca93acadc078dbdbfa9f17e45224e4f9df59f535d2886f"`
}
type DailyTokenUsageResponse struct {
Items []store.DailyTokenUsage `json:"items"`
TokenDays []store.DailyTokenUsage `json:"tokenDays"`
Summary dailyTokenUsageSummary `json:"summary"`
Range map[string]string `json:"range"`
}
type TaskRequest struct {
Model string `json:"model" example:"gpt-4o-mini"`
Messages []ChatMessage `json:"messages,omitempty"`
Input interface{} `json:"input,omitempty"`
Prompt string `json:"prompt,omitempty" example:"A watercolor robot reading a book"`
Text string `json:"text,omitempty" example:"Hello from EasyAI audio synthesis."`
TextFileID string `json:"text_file_id,omitempty" example:""`
VoiceID string `json:"voice_id,omitempty" example:"female-shaonv"`
Stream *bool `json:"stream,omitempty" example:"false"`
RunMode string `json:"runMode,omitempty" example:"simulation"`
MaxTokens *int `json:"max_tokens,omitempty" example:"512"`
Model string `json:"model" example:"gpt-4o-mini"`
Messages []ChatMessage `json:"messages,omitempty"`
Input interface{} `json:"input,omitempty"`
Query string `json:"query,omitempty" example:"Which document mentions EasyAI Gateway?"`
Documents []string `json:"documents,omitempty"`
Prompt string `json:"prompt,omitempty" example:"A watercolor robot reading a book"`
Text string `json:"text,omitempty" example:"Hello from EasyAI audio synthesis."`
TextFileID string `json:"text_file_id,omitempty" example:""`
VoiceID string `json:"voice_id,omitempty" example:"female-shaonv"`
AudioURL string `json:"audio_url,omitempty" example:"https://example.com/source-voice.mp3"`
DisplayName string `json:"display_name,omitempty" example:"EasyAI DEV acceptance"`
PromptAudioURL string `json:"prompt_audio_url,omitempty" example:"https://example.com/prompt-voice.mp3"`
PromptText string `json:"prompt_text,omitempty" example:"EasyAI voice clone prompt."`
PreviewModel string `json:"preview_model,omitempty" example:"speech-2.8-hd"`
Stream *bool `json:"stream,omitempty" example:"false"`
RunMode string `json:"runMode,omitempty" example:"simulation"`
MaxTokens *int `json:"max_tokens,omitempty" example:"512"`
// MaxCompletionTokens includes visible output and reasoning tokens.
MaxCompletionTokens *int `json:"max_completion_tokens,omitempty" example:"512"`
MaxOutputTokens *int `json:"max_output_tokens,omitempty" example:"512"`
@@ -229,6 +243,9 @@ type TaskRequest struct {
Size string `json:"size,omitempty" example:"1024x1024"`
Duration int `json:"duration,omitempty" example:"5"`
Resolution string `json:"resolution,omitempty" example:"720p"`
AspectRatio string `json:"aspect_ratio,omitempty" example:"16:9"`
Audio *bool `json:"audio,omitempty" example:"false"`
Watermark *bool `json:"watermark,omitempty" example:"false"`
MakeInstrumental bool `json:"makeInstrumental,omitempty" example:"false"`
CustomMode bool `json:"customMode,omitempty" example:"false"`
Style string `json:"style,omitempty" example:"city pop, bright synth"`
@@ -365,6 +382,34 @@ type ImageEditRequest struct {
RunMode string `json:"runMode,omitempty" example:"simulation"`
}
type ImageVectorizeSource struct {
URL string `json:"url,omitempty" example:"https://example.com/source.png"`
VectorizerTaskID string `json:"vectorizerTaskId,omitempty" example:"9f4d8f3d-5f5f-4bb7-a4be-344a9f930e25"`
}
type ImageVectorizeRequest struct {
Model string `json:"model,omitempty" example:"easy-image-vectorizer-1"`
Source ImageVectorizeSource `json:"source"`
Format string `json:"format,omitempty" example:"svg" enums:"svg,eps,pdf,dxf,png"`
MaxColors int `json:"maxColors,omitempty" example:"16" enums:"0,2,4,8,16,32"`
CleanupLevel string `json:"cleanupLevel,omitempty" example:"standard" enums:"low,standard,strong"`
}
type VideoUpscaleRequest struct {
Model string `json:"model,omitempty" example:"easy-proteus-standard-4"`
VideoURL string `json:"video_url" example:"https://example.com/source.mp4"`
Operation string `json:"operation,omitempty" example:"upscale" enums:"upscale"`
TargetResolution string `json:"target_resolution,omitempty" example:"1080p"`
OutputWidth int `json:"output_width,omitempty" example:"1920"`
OutputHeight int `json:"output_height,omitempty" example:"1080"`
PreserveAudio *bool `json:"preserve_audio,omitempty" example:"true"`
Duration float64 `json:"duration,omitempty" example:"3"`
SourceResolution string `json:"source_resolution,omitempty" example:"480p"`
SourceFrameRate float64 `json:"source_frame_rate,omitempty" example:"24"`
TargetFrameRate float64 `json:"target_frame_rate,omitempty" example:"24"`
SlowMotionRate float64 `json:"slow_motion_rate,omitempty" example:"1"`
}
type VideoGenerationRequest struct {
Model string `json:"model" example:"video-model"`
Prompt string `json:"prompt" example:"A cinematic drone shot over mountains"`
@@ -0,0 +1,329 @@
package httpapi
import (
"crypto/sha256"
"encoding/hex"
"encoding/json"
"fmt"
"io"
"net/http"
"strings"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/runner"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
const seedancePortraitAssetCategory = "seedance_portrait_asset"
// getSeedancePortraitAssetCapability godoc
// @Summary 查询 Seedance 真人资产能力
// @Description 返回当前网关是否已配置可创建、同步和引用的火山 Seedance 真人资产平台。
// @Tags portrait-assets
// @Produce json
// @Security BearerAuth
// @Success 200 {object} runner.PortraitAssetCapability
// @Router /api/v1/resource/material/seedance-portrait-assets/capability [get]
func (s *Server) getSeedancePortraitAssetCapability(w http.ResponseWriter, r *http.Request) {
capability, err := s.runner.PortraitAssetCapability(r.Context())
if err != nil {
s.logger.Error("get portrait asset capability failed", "error", err)
writeError(w, http.StatusInternalServerError, "get portrait asset capability failed")
return
}
writeJSON(w, http.StatusOK, capability)
}
// listSeedancePortraitAssets godoc
// @Summary 列出 Seedance 真人资产
// @Description 返回当前用户的真人资产;兼容 server-main material 列表响应字段。
// @Tags portrait-assets
// @Produce json
// @Security BearerAuth
// @Success 200 {object} map[string]any
// @Router /api/v1/resource/material/user/materials [get]
func (s *Server) listSeedancePortraitAssets(w http.ResponseWriter, r *http.Request) {
user, ok := auth.UserFromContext(r.Context())
if !ok || user == nil {
writeError(w, http.StatusUnauthorized, "unauthorized")
return
}
if category := strings.TrimSpace(r.URL.Query().Get("category")); category != seedancePortraitAssetCategory {
writeError(w, http.StatusNotFound, "material category not found")
return
}
items, err := s.store.ListPortraitAssets(r.Context(), user, store.PortraitAssetListFilter{
Keyword: r.URL.Query().Get("keyword"),
SourceType: firstNonEmptyQuery(r, "fileType", "sourceType"),
Page: portraitAssetQueryInt(r, "pageNumber", "page"),
PageSize: portraitAssetQueryInt(r, "pageSize"),
})
if err != nil {
s.logger.Error("list portrait assets failed", "error", err)
writeError(w, http.StatusInternalServerError, "list portrait assets failed")
return
}
responseItems := make([]any, 0, len(items.Items))
for _, item := range items.Items {
responseItems = append(responseItems, s.portraitAssetResponse(r, item))
}
writeJSON(w, http.StatusOK, map[string]any{
"data": responseItems,
"total": items.Total,
"page": items.Page,
"pageSize": items.PageSize,
})
}
// createSeedancePortraitAsset godoc
// @Summary 上传并创建 Seedance 真人资产
// @Description 文件先写入网关文件存储;仅在 private_avatar_eligible=true 时登记到火山 Assets。创建后会立即触发一次状态同步。
// @Tags portrait-assets
// @Accept multipart/form-data
// @Produce json
// @Security BearerAuth
// @Param file formData file true "真人资产源文件(图片、视频或音频)"
// @Param data formData string true "material JSONcategory 必须是 seedance_portrait_asset"
// @Success 200 {object} map[string]any
// @Router /api/v1/resource/material [post]
func (s *Server) createSeedancePortraitAsset(w http.ResponseWriter, r *http.Request) {
user, ok := auth.UserFromContext(r.Context())
if !ok || user == nil {
writeError(w, http.StatusUnauthorized, "unauthorized")
return
}
r.Body = http.MaxBytesReader(w, r.Body, maxGatewayUploadBytes)
if err := r.ParseMultipartForm(multipartTaskMemoryBytes); err != nil {
writeError(w, http.StatusBadRequest, "invalid multipart form-data body")
return
}
if r.MultipartForm != nil {
defer r.MultipartForm.RemoveAll()
}
var data map[string]any
if err := json.Unmarshal([]byte(strings.TrimSpace(r.FormValue("data"))), &data); err != nil || data == nil {
writeError(w, http.StatusBadRequest, "data must be a JSON object")
return
}
if strings.TrimSpace(portraitAssetString(data["category"])) != seedancePortraitAssetCategory {
writeError(w, http.StatusBadRequest, "category must be seedance_portrait_asset")
return
}
privateEligible, _ := data["private_avatar_eligible"].(bool)
if !privateEligible {
writeError(w, http.StatusBadRequest, "private_avatar_eligible must be true after the user confirms authorization", "portrait_asset_authorization_required")
return
}
file, header, err := r.FormFile("file")
if err != nil {
writeError(w, http.StatusBadRequest, "file is required")
return
}
defer file.Close()
payload, err := io.ReadAll(file)
if err != nil {
writeError(w, http.StatusBadRequest, "read portrait asset file failed")
return
}
contentType := strings.TrimSpace(header.Header.Get("Content-Type"))
if contentType == "" && len(payload) > 0 {
contentType = http.DetectContentType(payload)
}
sourceType := strings.ToLower(strings.TrimSpace(firstNonEmpty(portraitAssetString(data["fileType"]), portraitAssetString(data["sourceType"]))))
if !portraitAssetSourceMatchesContentType(sourceType, contentType) {
writeError(w, http.StatusBadRequest, "fileType must be image, video, or audio and match the uploaded file", "portrait_asset_unsupported_type")
return
}
upload, err := s.runner.UploadFile(r.Context(), runner.FileUploadPayload{
Bytes: payload, ContentType: contentType, FileName: header.Filename, Source: "seedance-portrait-asset", Scene: store.FileStorageSceneUpload,
})
if err != nil {
s.logger.Error("upload portrait asset failed", "error", err)
writeError(w, http.StatusBadGateway, err.Error(), clients.ErrorCode(err))
return
}
url := strings.TrimSpace(portraitAssetString(upload["url"]))
if url == "" {
writeError(w, http.StatusBadGateway, "portrait asset upload returned no URL", "portrait_asset_source_url_required")
return
}
digest := sha256.Sum256(payload)
asset, reused, err := s.runner.CreatePortraitAsset(r.Context(), user, runner.PortraitAssetCreateInput{
Name: strings.TrimSpace(portraitAssetString(data["name"])),
Description: strings.TrimSpace(portraitAssetString(data["description"])),
SourceType: sourceType,
URL: url,
Preview: firstNonEmpty(portraitAssetString(data["preview"]), url),
MimeType: contentType,
ByteSize: int64(len(payload)),
SourceSHA256: hex.EncodeToString(digest[:]),
PrivateAvatarEligible: privateEligible,
Metadata: map[string]any{
"tags": data["tags"],
"materialGroupId": data["material_group_id"],
"uploadedFileName": header.Filename,
"uploadAssetStorage": upload["assetStorage"],
},
})
if err != nil {
writePortraitAssetError(w, err)
return
}
_, _ = s.runner.SyncPortraitAssets(r.Context(), user, []string{asset.ID})
asset, _, err = s.refreshPortraitAssetForResponse(r, user, asset.ID, asset)
if err != nil {
s.logger.Error("refresh portrait asset after create failed", "error", err)
writeError(w, http.StatusInternalServerError, "refresh portrait asset failed")
return
}
response := map[string]any{"asset": s.portraitAssetResponse(r, asset)}
if reused {
response["dedupe"] = map[string]any{"reused": true, "code": "PORTRAIT_ASSET_REUSED", "reason": "same_source", "message": "已复用相同源文件的真人资产,并触发状态刷新。"}
}
writeJSON(w, http.StatusOK, response)
}
// syncSeedancePortraitAssets godoc
// @Summary 同步 Seedance 真人资产状态
// @Description 调用火山 CreateAsset/GetAsset;多次调用可把 Processing 状态刷新为 Active 或 Failed。
// @Tags portrait-assets
// @Accept json
// @Produce json
// @Security BearerAuth
// @Success 200 {object} runner.PortraitAssetSyncResponse
// @Router /api/v1/resource/material/seedance-portrait-assets/sync [post]
func (s *Server) syncSeedancePortraitAssets(w http.ResponseWriter, r *http.Request) {
user, ok := auth.UserFromContext(r.Context())
if !ok || user == nil {
writeError(w, http.StatusUnauthorized, "unauthorized")
return
}
var request struct {
IDs []string `json:"ids"`
}
if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
writeError(w, http.StatusBadRequest, "invalid json body")
return
}
if len(request.IDs) == 0 {
writeJSON(w, http.StatusOK, runner.PortraitAssetSyncResponse{SyncedIDs: []string{}, Skipped: []runner.PortraitAssetIssue{}, Failed: []runner.PortraitAssetIssue{}, Assets: []store.PortraitAsset{}})
return
}
response, err := s.runner.SyncPortraitAssets(r.Context(), user, request.IDs)
if err != nil {
s.logger.Error("sync portrait assets failed", "error", err)
writePortraitAssetError(w, err)
return
}
assets := make([]any, 0, len(response.Assets))
for _, asset := range response.Assets {
assets = append(assets, s.portraitAssetResponse(r, asset))
}
writeJSON(w, http.StatusOK, map[string]any{
"requested": response.Requested, "accepted": response.Accepted, "syncedIds": response.SyncedIDs,
"skipped": response.Skipped, "failed": response.Failed, "assets": assets,
})
}
func (s *Server) refreshPortraitAssetForResponse(r *http.Request, user *auth.User, assetID string, fallback store.PortraitAsset) (store.PortraitAsset, bool, error) {
asset, found, err := s.store.FindPortraitAssetForUser(r.Context(), user, assetID)
if err != nil || !found {
return fallback, found, err
}
return asset, true, nil
}
func (s *Server) portraitAssetResponse(r *http.Request, asset store.PortraitAsset) map[string]any {
active, total, lastError, updatedAt, err := s.store.PortraitAssetBindingSummary(r.Context(), asset.ID)
if err != nil {
active, total, lastError, updatedAt = 0, 0, asset.LastError, asset.UpdatedAt.UTC().Format(time.RFC3339Nano)
}
summaryStatus := asset.Status
if summaryStatus == "not_synced" && total == 0 {
summaryStatus = "not_synced"
}
response := map[string]any{
"id": asset.ID, "name": asset.Name, "description": asset.Description, "url": asset.URL, "preview": firstNonEmpty(asset.Preview, asset.URL),
"type": "personal", "fileType": asset.SourceType, "sourceType": asset.SourceType, "size": asset.ByteSize,
"privateAvatarEligible": asset.PrivateAvatarEligible,
"createdAt": asset.CreatedAt.UTC().Format(time.RFC3339Nano), "updatedAt": asset.UpdatedAt.UTC().Format(time.RFC3339Nano),
"seedanceAssetSummary": map[string]any{
"eligible": asset.PrivateAvatarEligible, "status": summaryStatus, "provider": "volces", "activePlatformCount": active,
"totalPlatformCount": total, "sourceType": asset.SourceType, "updatedAt": updatedAt,
},
}
if lastError != "" {
response["seedanceAssetSummary"].(map[string]any)["lastError"] = lastError
}
if asset.SourceType == "image" {
response["thumbnail"] = firstNonEmpty(asset.Preview, asset.URL)
}
return response
}
func writePortraitAssetError(w http.ResponseWriter, err error) {
status := http.StatusInternalServerError
if clientErr := clients.ErrorCode(err); clientErr != "client_error" {
switch clientErr {
case "portrait_asset_not_found":
status = http.StatusNotFound
case "portrait_asset_processing":
status = http.StatusServiceUnavailable
case "portrait_asset_authorization_required", "portrait_asset_unsupported_type", "portrait_asset_source_url_required", "portrait_asset_id_required", "portrait_asset_unsupported_model", "portrait_asset_audio_only":
status = http.StatusBadRequest
}
writeError(w, status, err.Error(), clientErr)
return
}
writeError(w, status, err.Error())
}
func portraitAssetSourceMatchesContentType(sourceType string, contentType string) bool {
contentType = strings.ToLower(strings.TrimSpace(contentType))
switch sourceType {
case "image":
return strings.HasPrefix(contentType, "image/")
case "video":
return strings.HasPrefix(contentType, "video/")
case "audio":
return strings.HasPrefix(contentType, "audio/")
default:
return false
}
}
func portraitAssetQueryInt(r *http.Request, keys ...string) int {
for _, key := range keys {
value := strings.TrimSpace(r.URL.Query().Get(key))
if value == "" {
continue
}
var parsed int
if _, err := fmt.Sscan(value, &parsed); err == nil {
return parsed
}
}
return 0
}
func firstNonEmptyQuery(r *http.Request, keys ...string) string {
for _, key := range keys {
if value := strings.TrimSpace(r.URL.Query().Get(key)); value != "" {
return value
}
}
return ""
}
func portraitAssetString(value any) string {
switch typed := value.(type) {
case string:
return typed
case fmt.Stringer:
return typed.String()
default:
return ""
}
}
@@ -161,7 +161,7 @@ func validPricingRuleSetInput(input store.PricingRuleSetInput) bool {
return false
}
switch calculator := strings.TrimSpace(rule.CalculatorType); calculator {
case "", "token_usage", "unit_weight", "duration_weight":
case "", "token_usage", "unit_weight", "duration_weight", "transition_matrix":
default:
return false
}
@@ -0,0 +1,76 @@
package httpapi
import (
"encoding/json"
"strings"
"testing"
gatewaydocs "github.com/easyai/easyai-ai-gateway/apps/api/docs"
)
func TestOpenAPIPublicRoutesUseCanonicalV1Prefix(t *testing.T) {
var document struct {
Paths map[string]any `json:"paths"`
}
if err := json.Unmarshal(gatewaydocs.SwaggerJSON, &document); err != nil {
t.Fatalf("decode embedded OpenAPI document: %v", err)
}
legacyPrefixes := []string{
"/v1/",
"/v1beta/",
"/kling/",
"/upload/",
"/api/v3/",
"/chat/",
"/images/",
"/song/",
"/music/",
"/speech/",
"/voice_clone",
"/tasks",
}
legacyExact := map[string]bool{
"/healthz": true,
"/readyz": true,
"/api-docs-json": true,
"/api-docs-yaml": true,
"/responses": true,
"/embeddings": true,
"/reranks": true,
}
for route := range document.Paths {
if legacyExact[route] {
t.Errorf("legacy public route must not be advertised in OpenAPI: %s", route)
}
for _, prefix := range legacyPrefixes {
if strings.HasPrefix(route, prefix) {
t.Errorf("legacy public route must not be advertised in OpenAPI: %s", route)
}
}
}
required := []string{
"/api/v1/healthz",
"/api/v1/readyz",
"/api/v1/openapi.json",
"/api/v1/chat/completions",
"/api/v1/responses",
"/api/v1/images/generations",
"/api/v1/images/vectorize",
"/api/v1/videos/generations",
"/api/v1/videos/upscales",
"/api/v1/pricing/estimate",
"/api/workspace/token-usage/daily",
"/api/v1/models/{model}:generateContent",
"/api/v1/videos/omni-video",
"/api/v1/kling/v1/videos/omni-video",
"/api/v1/kling/v2/omni-video/{model}",
"/api/v1/contents/generations/tasks",
}
for _, route := range required {
if _, ok := document.Paths[route]; !ok {
t.Errorf("canonical public route missing from OpenAPI: %s", route)
}
}
}
@@ -40,6 +40,18 @@ func TestReceiveSecurityEventUsesPreparedReceiverBeforeFirstActivation(t *testin
}
}
func TestReceiveSecurityEventReturnsNotFoundWithoutConfiguredReceiver(t *testing.T) {
server := &Server{identityRuntime: identityruntime.NewManager(nil, &preparedReceiverBuilder{})}
request := httptest.NewRequest(http.MethodPost, "/api/v1/security-events/ssf", nil)
response := httptest.NewRecorder()
server.receiveSecurityEvent(response, request)
if response.Code != http.StatusNotFound {
t.Fatalf("unconfigured SSF status=%d, want %d", response.Code, http.StatusNotFound)
}
}
func TestSecurityEventWriteAuditFailsClosedBeforeMutation(t *testing.T) {
server := &Server{}
request := httptest.NewRequest(http.MethodPost, "/api/admin/system/security-events/connection/verify", nil)
+40 -6
View File
@@ -127,12 +127,16 @@ func NewServerWithContext(ctx context.Context, cfg config.Config, db *store.Stor
mux := http.NewServeMux()
mux.HandleFunc("GET /healthz", server.health)
mux.HandleFunc("GET /readyz", server.ready)
mux.HandleFunc("GET /api/v1/healthz", server.health)
mux.HandleFunc("GET /api/v1/readyz", server.ready)
mux.Handle("GET /metrics", securityEventMetrics.DynamicHandler(db))
mux.HandleFunc("GET /static/simulation/{asset}", serveSimulationAsset)
mux.HandleFunc("GET /static/generated/{asset}", server.serveGeneratedStaticAsset)
mux.HandleFunc("GET /static/uploaded/{asset}", server.serveUploadedStaticAsset)
mux.HandleFunc("GET /api-docs-json", server.apiDocsJSON)
mux.HandleFunc("GET /api-docs-yaml", server.apiDocsYAML)
mux.HandleFunc("GET /api/v1/openapi.json", server.apiDocsJSON)
mux.HandleFunc("GET /api/v1/openapi.yaml", server.apiDocsYAML)
mux.HandleFunc("GET /api/v1/public/skills/ai-gateway-ops-management/metadata", server.getOpsManagementSkillMetadata)
mux.HandleFunc("GET /api/v1/public/skills/ai-gateway-ops-management/download", server.downloadOpsManagementSkill)
@@ -182,6 +186,7 @@ func NewServerWithContext(ctx context.Context, cfg config.Config, db *store.Stor
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)))
mux.Handle("POST /api/v1/api-keys/access-rules/batch", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.batchAPIKeyAccessRules)))
mux.Handle("GET /api/v1/api-keys/assignable-models", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.listAPIKeyAssignableModels)))
mux.Handle("PATCH /api/v1/api-keys/{apiKeyID}/scopes", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.updateAPIKeyScopes)))
mux.Handle("PATCH /api/v1/api-keys/{apiKeyID}/disable", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.disableAPIKey)))
mux.Handle("DELETE /api/v1/api-keys/{apiKeyID}", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.deleteAPIKey)))
@@ -190,6 +195,7 @@ func NewServerWithContext(ctx context.Context, cfg config.Config, db *store.Stor
mux.Handle("GET /api/workspace/user-groups", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.listCurrentUserGroups)))
mux.Handle("GET /api/workspace/wallet", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.getWallet)))
mux.Handle("GET /api/workspace/wallet/transactions", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.listWalletTransactions)))
mux.Handle("GET /api/workspace/token-usage/daily", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.dailyTokenUsage)))
mux.Handle("GET /api/workspace/tasks", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.listTasks)))
mux.Handle("GET /api/workspace/tasks/{taskID}", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.getTask)))
mux.Handle("POST /api/workspace/tasks/{taskID}/cancel", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.cancelTask)))
@@ -251,12 +257,17 @@ func NewServerWithContext(ctx context.Context, cfg config.Config, db *store.Stor
mux.Handle("GET /api/admin/runtime/rate-limit-windows", server.requireAdmin(auth.PermissionPower, http.HandlerFunc(server.listRateLimitWindows)))
mux.Handle("GET /api/admin/runtime/model-rate-limits", server.requireAdmin(auth.PermissionPower, http.HandlerFunc(server.listModelRateLimitStatuses)))
mux.Handle("POST /api/v1/chat/completions", server.requireUser(auth.PermissionBasic, server.createAPIV1ChatCompletions()))
mux.Handle("POST /api/v1/responses", server.requireUser(auth.PermissionBasic, server.createTask("responses", false)))
mux.Handle("POST /api/v1/embeddings", server.requireUser(auth.PermissionBasic, server.createTask("embeddings", false)))
mux.Handle("POST /api/v1/reranks", server.requireUser(auth.PermissionBasic, server.createTask("reranks", false)))
mux.Handle("POST /api/v1/images/generations", server.requireUser(auth.PermissionBasic, server.createTask("images.generations", false)))
mux.Handle("POST /api/v1/images/edits", server.requireUser(auth.PermissionBasic, server.createTask("images.edits", false)))
mux.Handle("POST /api/v1/videos/generations", server.requireUser(auth.PermissionBasic, server.createTask("videos.generations", false)))
mux.Handle("POST /api/v1/responses", server.requireUser(auth.PermissionBasic, server.createTask("responses", true)))
mux.Handle("POST /api/v1/embeddings", server.requireUser(auth.PermissionBasic, server.createTask("embeddings", true)))
mux.Handle("POST /api/v1/reranks", server.requireUser(auth.PermissionBasic, server.createTask("reranks", true)))
mux.Handle("POST /api/v1/images/generations", server.requireUser(auth.PermissionBasic, server.createTask("images.generations", true)))
mux.Handle("POST /api/v1/images/edits", server.requireUser(auth.PermissionBasic, server.createTask("images.edits", true)))
mux.Handle("POST /api/v1/images/vectorize", server.requireUser(auth.PermissionBasic, server.createImageVectorizeTask()))
mux.Handle("POST /api/v1/videos/generations", server.requireUser(auth.PermissionBasic, server.createTask("videos.generations", true)))
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("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)))
@@ -264,9 +275,24 @@ func NewServerWithContext(ctx context.Context, cfg config.Config, db *store.Stor
mux.Handle("GET /api/v1/voice_clone/voices", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.listClonedVoices)))
mux.Handle("DELETE /api/v1/voice_clone/voices/{voiceID}", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.deleteClonedVoice)))
mux.Handle("POST /api/v1/files/upload", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.uploadFile)))
mux.Handle("GET /api/v1/resource/material/seedance-portrait-assets/capability", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.getSeedancePortraitAssetCapability)))
mux.Handle("GET /api/v1/resource/material/user/materials", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.listSeedancePortraitAssets)))
mux.Handle("POST /api/v1/resource/material", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.createSeedancePortraitAsset)))
mux.Handle("POST /api/v1/resource/material/seedance-portrait-assets/sync", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.syncSeedancePortraitAssets)))
mux.Handle("POST /api/v3/contents/generations/tasks", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.createVolcesContentsGenerationTask)))
mux.Handle("GET /api/v3/contents/generations/tasks", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.listVolcesContentsGenerationTasks)))
mux.Handle("GET /api/v3/contents/generations/tasks/{taskID}", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.getVolcesContentsGenerationTask)))
mux.Handle("DELETE /api/v3/contents/generations/tasks/{taskID}", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.deleteVolcesContentsGenerationTask)))
mux.Handle("POST /api/v1/contents/generations/tasks", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.createVolcesContentsGenerationTask)))
mux.Handle("GET /api/v1/contents/generations/tasks", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.listVolcesContentsGenerationTasks)))
mux.Handle("GET /api/v1/contents/generations/tasks/{taskID}", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.getVolcesContentsGenerationTask)))
mux.Handle("DELETE /api/v1/contents/generations/tasks/{taskID}", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.deleteVolcesContentsGenerationTask)))
server.registerGeminiGenerateContentRoutes(mux)
server.registerKlingCompatibilityRoutes(mux)
mux.Handle("POST /upload/{version}/files", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.geminiFilesUpload)))
mux.Handle("POST /upload/{version}/files/{uploadID}", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.geminiFilesUploadFinalize)))
mux.Handle("POST /api/v1/gemini/upload/{version}/files", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.geminiFilesUpload)))
mux.Handle("POST /api/v1/gemini/upload/{version}/files/{uploadID}", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.geminiFilesUploadFinalize)))
mux.Handle("GET /api/v1/tasks", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.listTasks)))
mux.Handle("GET /api/v1/tasks/{taskID}", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.getTask)))
mux.Handle("POST /api/v1/tasks/{taskID}/cancel", server.requireUser(auth.PermissionBasic, http.HandlerFunc(server.cancelTask)))
@@ -289,6 +315,14 @@ func NewServerWithContext(ctx context.Context, cfg config.Config, db *store.Stor
mux.Handle("POST /v1/images/generations", server.requireUser(auth.PermissionBasic, server.createTask("images.generations", true)))
mux.Handle("POST /images/edits", server.requireUser(auth.PermissionBasic, server.createTask("images.edits", true)))
mux.Handle("POST /v1/images/edits", server.requireUser(auth.PermissionBasic, server.createTask("images.edits", true)))
mux.Handle("POST /images/vectorize", server.requireUser(auth.PermissionBasic, server.createImageVectorizeTask()))
mux.Handle("POST /v1/images/vectorize", server.requireUser(auth.PermissionBasic, server.createImageVectorizeTask()))
mux.Handle("POST /video/upscale", server.requireUser(auth.PermissionBasic, server.createVideoUpscaleTask()))
mux.Handle("POST /v1/video/upscale", server.requireUser(auth.PermissionBasic, server.createVideoUpscaleTask()))
mux.Handle("POST /v1/videos/omni-video", server.requireKelingAPIKey(http.HandlerFunc(server.createKelingOmniVideo)))
mux.Handle("GET /v1/videos/omni-video/{taskID}", server.requireKelingAPIKey(http.HandlerFunc(server.getKelingOmniVideo)))
mux.Handle("POST /api/v1/videos/omni-video", server.requireKelingAPIKey(http.HandlerFunc(server.createKelingOmniVideo)))
mux.Handle("GET /api/v1/videos/omni-video/{taskID}", server.requireKelingAPIKey(http.HandlerFunc(server.getKelingOmniVideo)))
mux.Handle("POST /song/generations", server.requireUser(auth.PermissionBasic, server.createTask("song.generations", true)))
mux.Handle("POST /v1/song/generations", server.requireUser(auth.PermissionBasic, server.createTask("song.generations", true)))
mux.Handle("POST /music/generations", server.requireUser(auth.PermissionBasic, server.createTask("music.generations", true)))
@@ -16,8 +16,6 @@ import (
// @Failure 401 {object} ErrorEnvelope
// @Failure 500 {object} ErrorEnvelope
// @Router /api/v1/voice_clone/voices [get]
// @Router /v1/voice_clone/voices [get]
// @Router /voice_clone/voices [get]
func (s *Server) listClonedVoices(w http.ResponseWriter, r *http.Request) {
user, ok := auth.UserFromContext(r.Context())
if !ok {
@@ -51,8 +49,6 @@ func (s *Server) listClonedVoices(w http.ResponseWriter, r *http.Request) {
// @Failure 404 {object} ErrorEnvelope
// @Failure 502 {object} ErrorEnvelope
// @Router /api/v1/voice_clone/voices/{voiceID} [delete]
// @Router /v1/voice_clone/voices/{voiceID} [delete]
// @Router /voice_clone/voices/{voiceID} [delete]
func (s *Server) deleteClonedVoice(w http.ResponseWriter, r *http.Request) {
user, ok := auth.UserFromContext(r.Context())
if !ok {
@@ -0,0 +1,333 @@
package httpapi
import (
"encoding/json"
"errors"
"net/http"
"strconv"
"strings"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/runner"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
const volcesContentsCompatibilityMarker = "volces_contents_generations_v3"
// createVolcesContentsGenerationTask godoc
// @Summary 创建火山内容生成任务
// @Description 统一公开入口兼容火山方舟内容生成任务。网关 task id 是查询与取消用的公开 id;上游 id 另以 upstream_task_id 保留。
// @Tags volces-compatible
// @Accept json
// @Produce json
// @Security BearerAuth
// @Success 200 {object} map[string]any
// @Router /api/v1/contents/generations/tasks [post]
func (s *Server) createVolcesContentsGenerationTask(w http.ResponseWriter, r *http.Request) {
user, ok := auth.UserFromContext(r.Context())
if !ok || user == nil {
writeError(w, http.StatusUnauthorized, "unauthorized")
return
}
body, err := s.decodeTaskRequestBody(r.Context(), w, r, "videos.generations")
if err != nil {
writeError(w, http.StatusBadRequest, err.Error(), clients.ErrorCode(err))
return
}
task, err := s.createVolcesCompatibleTask(r, user, body)
if err != nil {
writeVolcesCompatibleTaskError(w, err)
return
}
writeJSON(w, http.StatusOK, volcesCompatibleTask(task))
}
// getVolcesContentsGenerationTask godoc
// @Summary 查询火山内容生成任务
// @Tags volces-compatible
// @Produce json
// @Security BearerAuth
// @Param taskID path string true "任务 ID"
// @Success 200 {object} map[string]any
// @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
}
writeJSON(w, http.StatusOK, volcesCompatibleTask(task))
}
// listVolcesContentsGenerationTasks godoc
// @Summary 列出火山内容生成任务
// @Tags volces-compatible
// @Produce json
// @Security BearerAuth
// @Success 200 {object} map[string]any
// @Router /api/v1/contents/generations/tasks [get]
func (s *Server) listVolcesContentsGenerationTasks(w http.ResponseWriter, r *http.Request) {
user, ok := auth.UserFromContext(r.Context())
if !ok || user == nil {
writeError(w, http.StatusUnauthorized, "unauthorized")
return
}
page := portraitAssetQueryInt(r, "page_num", "pageNumber", "page")
pageSize := portraitAssetQueryInt(r, "page_size", "pageSize")
tasks, err := s.store.ListVolcesCompatibleTasks(r.Context(), user, store.VolcesCompatibleTaskListFilter{
CompatibilityMarker: volcesContentsCompatibilityMarker,
Status: r.URL.Query().Get("filter.status"),
Model: r.URL.Query().Get("filter.model"),
TaskIDs: r.URL.Query()["filter.task_ids"],
Page: page,
PageSize: pageSize,
})
if err != nil {
s.logger.Error("list Volces-compatible tasks failed", "error", err)
writeError(w, http.StatusInternalServerError, "list tasks failed")
return
}
items := make([]any, 0)
for _, task := range tasks.Items {
items = append(items, volcesCompatibleTask(task))
}
writeJSON(w, http.StatusOK, map[string]any{
"items": items, "total": tasks.Total,
"page_num": tasks.Page, "page_size": tasks.PageSize,
// data/page are retained as additive gateway fields for existing callers.
"data": items, "page": tasks.Page,
})
}
// deleteVolcesContentsGenerationTask godoc
// @Summary 取消火山内容生成任务
// @Description 取消网关任务;对于已提交且保存了上游任务标识的 Volces 视频任务,同时调用火山 DELETE 接口并持久化取消状态。
// @Tags volces-compatible
// @Produce json
// @Security BearerAuth
// @Param taskID path string true "任务 ID"
// @Success 200 {object} map[string]any
// @Router /api/v1/contents/generations/tasks/{taskID} [delete]
func (s *Server) deleteVolcesContentsGenerationTask(w http.ResponseWriter, r *http.Request) {
task, ok := s.volcesCompatibleTaskForUser(w, r)
if !ok {
return
}
user, _ := auth.UserFromContext(r.Context())
result, err := s.runner.CancelVolcesVideoTask(r.Context(), task, user)
if err != nil {
if errors.Is(err, runner.ErrTaskAccessDenied) {
writeError(w, http.StatusNotFound, "task not found")
return
}
s.logger.Error("cancel Volces-compatible task failed", "error", err)
writeError(w, http.StatusInternalServerError, "cancel task failed")
return
}
updated, err := s.store.GetTask(r.Context(), task.ID)
if err != nil {
writeError(w, http.StatusInternalServerError, "get cancelled task failed")
return
}
response := volcesCompatibleTask(updated)
response["cancelled"] = result.Cancelled
response["cancellable"] = result.Cancellable
response["message"] = result.Message
writeJSON(w, http.StatusOK, response)
}
// createLegacyVolcesVideoGeneration godoc
// @Summary 创建 server-main 兼容视频任务
// @Description 兼容 server-main 的 /api/v1/video/generations,返回 submitted 和 task_id;额外保留火山任务字段。
// @Tags volces-compatible
// @Accept json
// @Produce json
// @Security BearerAuth
// @Success 200 {object} map[string]any
// @Router /api/v1/video/generations [post]
func (s *Server) createLegacyVolcesVideoGeneration(w http.ResponseWriter, r *http.Request) {
user, ok := auth.UserFromContext(r.Context())
if !ok || user == nil {
writeError(w, http.StatusUnauthorized, "unauthorized")
return
}
body, err := s.decodeTaskRequestBody(r.Context(), w, r, "videos.generations")
if err != nil {
writeError(w, http.StatusBadRequest, err.Error(), clients.ErrorCode(err))
return
}
task, err := s.createVolcesCompatibleTask(r, user, body)
if err != nil {
writeVolcesCompatibleTaskError(w, err)
return
}
response := volcesCompatibleTask(task)
response["status"] = "submitted"
response["task_id"] = task.ID
writeJSON(w, http.StatusOK, response)
}
// getLegacyVolcesVideoResult godoc
// @Summary 查询 server-main 兼容视频结果
// @Tags volces-compatible
// @Produce json
// @Security BearerAuth
// @Param taskID path string true "任务 ID"
// @Success 200 {object} map[string]any
// @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 {
return
}
compat := volcesCompatibleTask(task)
legacyStatus := "process"
switch compat["status"] {
case "succeeded":
legacyStatus = "success"
case "failed", "cancelled":
legacyStatus = "failed"
}
writeJSON(w, http.StatusOK, map[string]any{
"status": legacyStatus, "task_id": task.ID, "data": compat["content"], "result": compat,
})
}
func (s *Server) createVolcesCompatibleTask(r *http.Request, user *auth.User, body map[string]any) (store.GatewayTask, error) {
model := strings.TrimSpace(volcesCompatString(body["model"]))
if model == "" {
return store.GatewayTask{}, &clients.ClientError{Code: "invalid_parameter", Message: "model is required", StatusCode: http.StatusBadRequest, Retryable: false}
}
if !apiKeyScopeAllowed(user, "videos.generations") {
return store.GatewayTask{}, &clients.ClientError{Code: "forbidden", Message: "api key scope does not allow video generation", StatusCode: http.StatusForbidden, Retryable: false}
}
body["_gateway_compatibility"] = volcesContentsCompatibilityMarker
task, err := s.prepareAndCreateGatewayTask(r.Context(), r, user, "videos.generations", model, body, true)
if err != nil {
return store.GatewayTask{}, err
}
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}
}
return task, nil
}
func (s *Server) volcesCompatibleTaskForUser(w http.ResponseWriter, r *http.Request) (store.GatewayTask, bool) {
user, ok := auth.UserFromContext(r.Context())
if !ok || user == nil {
writeError(w, http.StatusUnauthorized, "unauthorized")
return store.GatewayTask{}, false
}
task, err := s.store.GetTask(r.Context(), strings.TrimSpace(r.PathValue("taskID")))
if err != nil {
if store.IsNotFound(err) {
writeError(w, http.StatusNotFound, "task not found")
return store.GatewayTask{}, false
}
s.logger.Error("get Volces-compatible task failed", "error", err)
writeError(w, http.StatusInternalServerError, "get task failed")
return store.GatewayTask{}, false
}
if !isVolcesCompatibleTask(task) || !kelingCompatTaskOwnedBy(task, user) {
writeError(w, http.StatusNotFound, "task not found")
return store.GatewayTask{}, false
}
return task, true
}
func isVolcesCompatibleTask(task store.GatewayTask) bool {
return task.Kind == "videos.generations" && strings.TrimSpace(volcesCompatString(task.Request["_gateway_compatibility"])) == volcesContentsCompatibilityMarker
}
func volcesCompatibleTask(task store.GatewayTask) map[string]any {
response := cloneVolcesCompatibleMap(task.Result)
if len(response) == 0 {
response = cloneVolcesCompatibleMap(task.RemoteTaskPayload)
}
if response == nil {
response = map[string]any{}
}
response["id"] = task.ID
response["model"] = firstNonEmpty(volcesCompatString(response["model"]), task.Model)
response["status"] = volcesCompatibleTaskStatus(task.Status)
response["created_at"] = task.CreatedAt.Unix()
response["updated_at"] = task.UpdatedAt.Unix()
if task.RemoteTaskID != "" {
response["upstream_task_id"] = task.RemoteTaskID
}
for _, key := range []string{"content", "seed", "resolution", "ratio", "duration", "frames", "framespersecond"} {
if response[key] == nil && task.Request[key] != nil {
response[key] = task.Request[key]
}
}
if len(task.Usage) > 0 && response["usage"] == nil {
response["usage"] = task.Usage
}
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)}
}
response["gateway_task_id"] = task.ID
response["gateway_status"] = task.Status
response["billings"] = task.Billings
response["billing_summary"] = task.BillingSummary
response["final_charge_amount"] = task.FinalChargeAmount
return response
}
func volcesCompatibleTaskStatus(status string) string {
switch strings.ToLower(strings.TrimSpace(status)) {
case "succeeded", "success", "completed":
return "succeeded"
case "failed":
return "failed"
case "cancelled", "canceled":
return "cancelled"
case "running", "processing":
return "running"
default:
return "queued"
}
}
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
if errors.As(err, &staged) {
err = staged.Err
}
var clientErr *clients.ClientError
if errors.As(err, &clientErr) && clientErr.StatusCode > 0 {
status = clientErr.StatusCode
} else if errors.As(err, &clientErr) {
status = http.StatusBadRequest
}
writeError(w, status, err.Error(), clients.ErrorCode(err))
}
func volcesCompatString(value any) string {
switch typed := value.(type) {
case string:
return strings.TrimSpace(typed)
case json.Number:
return typed.String()
case float64:
return strconv.FormatFloat(typed, 'f', -1, 64)
default:
return ""
}
}
@@ -0,0 +1,29 @@
package httpapi
import (
"testing"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestVolcesCompatibleTaskPreservesOfficialFieldsAndGatewayBilling(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{
"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},
},
Billings: []any{map[string]any{"amount": 3}}, BillingSummary: map[string]any{"currency": "resource"}, FinalChargeAmount: 3,
}
got := volcesCompatibleTask(task)
if got["id"] != task.ID || got["upstream_task_id"] != task.RemoteTaskID || got["status"] != "succeeded" {
t.Fatalf("unexpected compatibility identity/status: %+v", got)
}
content, _ := got["content"].(map[string]any)
if content["video_url"] != "https://example.com/out.mp4" || got["usage"] == nil || got["billings"] == nil {
t.Fatalf("official or billing fields were lost: %+v", got)
}
}
+4 -1
View File
@@ -105,7 +105,10 @@ func (manager *Manager) SecurityEventReceiver() http.Handler {
if runtime := manager.Current(); runtime != nil && runtime.SecurityEvents != nil {
return runtime.SecurityEvents
}
return manager.SecurityEventManager()
if securityEventManager := manager.SecurityEventManager(); securityEventManager != nil {
return securityEventManager
}
return nil
}
// SecurityEventManager resolves the manager used by administrative recovery
@@ -630,6 +630,14 @@ func TestSecurityEventManagerExposesPreparedRecoveryManagerWithoutActiveRuntime(
}
}
func TestSecurityEventReceiverReturnsNilWithoutConfiguredManager(t *testing.T) {
manager := NewManager(&runtimeRepositoryFake{}, &runtimeBuilderFake{})
if manager.SecurityEventReceiver() != nil {
t.Fatal("unconfigured security event receiver should be nil")
}
}
func TestSecurityEventManagerPrefersActiveRuntime(t *testing.T) {
active := &securityevents.ConnectionManager{}
prepared := &securityevents.ConnectionManager{}
@@ -0,0 +1,32 @@
package runner
import (
"strings"
"testing"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestCandidateWithEnvironmentCredentialsResolvesNamesAtRuntime(t *testing.T) {
t.Setenv("TEST_GATEWAY_PROVIDER_SECRET", "secret-value")
candidate, err := candidateWithEnvironmentCredentials(store.RuntimeModelCandidate{
Credentials: map[string]any{"safe": "existing"},
PlatformConfig: map[string]any{"credentialEnv": map[string]any{
"apiKey": "TEST_GATEWAY_PROVIDER_SECRET",
}},
})
if err != nil {
t.Fatal(err)
}
if candidate.Credentials["apiKey"] != "secret-value" || candidate.Credentials["safe"] != "existing" {
t.Fatalf("credentials not resolved: %+v", candidate.Credentials)
}
}
func TestCandidateWithEnvironmentCredentialsDoesNotLeakMissingValue(t *testing.T) {
candidate := store.RuntimeModelCandidate{PlatformConfig: map[string]any{"credentialEnv": map[string]any{"apiKey": "MISSING_GATEWAY_PROVIDER_SECRET"}}}
_, err := candidateWithEnvironmentCredentials(candidate)
if err == nil || strings.Contains(err.Error(), "MISSING_GATEWAY_PROVIDER_SECRET") {
t.Fatalf("missing credential error should be generic: %v", err)
}
}
@@ -0,0 +1,27 @@
package runner
import (
"testing"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestKelingO1GeneratedAudioIsRejectedInsteadOfSilentlyRemoved(t *testing.T) {
result := preprocessRequestWithLog("videos.generations", map[string]any{
"model": "kling-o1",
"audio": true,
}, store.RuntimeModelCandidate{
Provider: "keling",
ProviderModelName: "kling-video-o1",
ModelType: "video_generate",
Capabilities: map[string]any{
"video_generate": map[string]any{"output_audio": false},
},
})
if result.Err == nil {
t.Fatal("Keling O1 audio=true must be rejected")
}
if len(result.Log.Changes) == 0 || result.Log.Changes[len(result.Log.Changes)-1].Action != "reject" {
t.Fatalf("expected an auditable reject change, got %+v", result.Log.Changes)
}
}
@@ -4,6 +4,8 @@ import (
"fmt"
"math"
"strings"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
type resolutionNormalizeProcessor struct{}
@@ -691,6 +693,16 @@ func (audioProcessor) ShouldProcess(params map[string]any, modelType string, con
}
func (audioProcessor) Process(params map[string]any, modelType string, context *paramProcessContext) bool {
if context != nil && kelingO1GeneratedAudioRequested(params, context.candidate) {
return context.reject(
"AudioProcessor",
"audio",
params["audio"],
"kling-video-o1 does not support generated audio",
capabilityPath(modelType, "output_audio"),
capabilityValue(context.modelCapability, modelType, "output_audio"),
)
}
capability := capabilityForType(context.modelCapability, modelType)
if capability == nil || !boolFromAny(capability["output_audio"]) {
for _, key := range []string{"audio", "output_audio"} {
@@ -712,6 +724,17 @@ func (audioProcessor) Process(params map[string]any, modelType string, context *
return true
}
func kelingO1GeneratedAudioRequested(params map[string]any, candidate store.RuntimeModelCandidate) bool {
if !strings.EqualFold(strings.TrimSpace(candidate.Provider), "keling") {
return false
}
model := strings.ToLower(strings.TrimSpace(candidate.ProviderModelName))
if model != "kling-o1" && model != "kling-video-o1" {
return false
}
return boolFromAny(params["audio"]) || boolFromAny(params["output_audio"])
}
type imageCountProcessor struct{}
func (imageCountProcessor) Name() string { return "ImageCountProcessor" }
@@ -39,6 +39,20 @@ func TestVideoModelTypeInferenceReadsContentArray(t *testing.T) {
}
}
func TestAdvancedMediaModelTypeInference(t *testing.T) {
if got := modelTypeFromKind("images.vectorize", nil); got != "image_vectorize" {
t.Fatalf("images.vectorize model type = %s", got)
}
if got := modelTypeFromKind("videos.upscales", nil); got != "video_enhance" {
t.Fatalf("videos.upscales model type = %s", got)
}
for _, alias := range []string{"vectorize", "image-vectorizer", "image_vectorize"} {
if got := canonicalModelType(alias); got != "image_vectorize" {
t.Fatalf("canonical image vectorize alias %q = %q", alias, got)
}
}
}
func TestVideoContentTextContributesToTokenEstimate(t *testing.T) {
tokens := estimateRequestTokens(map[string]any{
"model": "demo-video",
+528
View File
@@ -0,0 +1,528 @@
package runner
import (
"context"
"crypto/sha256"
"encoding/hex"
"fmt"
"net/http"
"net/url"
"regexp"
"strings"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
var portraitAssetPlaceholderPattern = regexp.MustCompile(`(?i)<<<[[:space:]]*portrait[_-]?asset_([0-9]+)[[:space:]]*>>>|@portrait_asset([0-9]+)|@人像资产([0-9]+)`)
type PortraitAssetCapability struct {
Enabled bool `json:"enabled"`
CanUse bool `json:"canUse"`
CanCreate bool `json:"canCreate"`
CanUseAsPortraitAsset bool `json:"canUseAsPortraitAsset"`
CanUseAsPlainMaterial bool `json:"canUseAsPlainMaterial"`
AvailablePlatformIDs []string `json:"availablePlatformIds"`
CreationPlatformIDs []string `json:"creationPlatformIds"`
CanReferenceTencentAsset bool `json:"canReferenceTencentAssetUri"`
Reason string `json:"reason,omitempty"`
}
type PortraitAssetCreateInput struct {
Name string
Description string
SourceType string
URL string
Preview string
MimeType string
ByteSize int64
SourceSHA256 string
PrivateAvatarEligible bool
Metadata map[string]any
}
type PortraitAssetSyncResponse struct {
Requested int `json:"requested"`
Accepted int `json:"accepted"`
SyncedIDs []string `json:"syncedIds"`
Skipped []PortraitAssetIssue `json:"skipped"`
Failed []PortraitAssetIssue `json:"failed"`
Assets []store.PortraitAsset `json:"assets"`
}
type PortraitAssetIssue struct {
ID string `json:"id"`
Reason string `json:"reason"`
}
type portraitAssetPlatformSettings struct {
ProjectName string
AssetGroupID string
Credentials clients.VolcesAssetCredentials
}
func (s *Service) PortraitAssetCapability(ctx context.Context) (PortraitAssetCapability, error) {
platforms, err := s.store.ListPortraitAssetPlatforms(ctx)
if err != nil {
return PortraitAssetCapability{}, err
}
ids := make([]string, 0, len(platforms))
for _, platform := range platforms {
if _, ok := portraitAssetSettings(platform); ok {
ids = append(ids, platform.PlatformID)
}
}
capability := PortraitAssetCapability{
Enabled: len(ids) > 0,
CanUse: len(ids) > 0,
CanCreate: len(ids) > 0,
CanUseAsPortraitAsset: len(ids) > 0,
CanUseAsPlainMaterial: true,
AvailablePlatformIDs: ids,
CreationPlatformIDs: ids,
}
if len(ids) == 0 {
capability.Reason = "未配置可用的火山 Seedance 人像资产平台;请在 Volces 平台 config.seedancePrivateAsset 中配置 enabled、accessKey、secretKey、projectName、assetGroupId。"
}
return capability, nil
}
func (s *Service) CreatePortraitAsset(ctx context.Context, user *auth.User, input PortraitAssetCreateInput) (store.PortraitAsset, bool, error) {
if s.store == nil {
return store.PortraitAsset{}, false, fmt.Errorf("portrait asset store is unavailable")
}
if !validPortraitAssetSourceType(input.SourceType) {
return store.PortraitAsset{}, false, &clients.ClientError{Code: "portrait_asset_unsupported_type", Message: "source type must be image, video, or audio", StatusCode: http.StatusBadRequest, Retryable: false}
}
if strings.TrimSpace(input.URL) == "" {
return store.PortraitAsset{}, false, &clients.ClientError{Code: "portrait_asset_source_url_required", Message: "portrait asset source URL is required", StatusCode: http.StatusBadRequest, Retryable: false}
}
if !input.PrivateAvatarEligible {
return store.PortraitAsset{}, false, &clients.ClientError{Code: "portrait_asset_authorization_required", Message: "private_avatar_eligible must be true after the user confirms authorization", StatusCode: http.StatusBadRequest, Retryable: false}
}
if existing, found, err := s.store.FindPortraitAssetBySourceHash(ctx, user, input.SourceSHA256); err != nil {
return store.PortraitAsset{}, false, err
} else if found {
return existing, true, nil
}
gatewayUserID, userID := portraitAssetUserKeys(user)
if user == nil || userID == "" {
return store.PortraitAsset{}, false, store.ErrLocalUserRequired
}
asset, err := s.store.CreatePortraitAsset(ctx, store.PortraitAssetInput{
GatewayUserID: gatewayUserID,
UserID: userID,
GatewayTenantID: strings.TrimSpace(user.GatewayTenantID),
TenantID: strings.TrimSpace(user.TenantID),
TenantKey: strings.TrimSpace(user.TenantKey),
Name: strings.TrimSpace(input.Name),
Description: strings.TrimSpace(input.Description),
SourceType: strings.ToLower(strings.TrimSpace(input.SourceType)),
URL: strings.TrimSpace(input.URL),
Preview: firstNonEmptyString(strings.TrimSpace(input.Preview), strings.TrimSpace(input.URL)),
MimeType: strings.TrimSpace(input.MimeType),
ByteSize: input.ByteSize,
SourceSHA256: strings.TrimSpace(input.SourceSHA256),
PrivateAvatarEligible: input.PrivateAvatarEligible,
Metadata: input.Metadata,
})
return asset, false, err
}
func (s *Service) SyncPortraitAssets(ctx context.Context, user *auth.User, ids []string) (PortraitAssetSyncResponse, error) {
response := PortraitAssetSyncResponse{
Requested: len(ids), SyncedIDs: make([]string, 0), Skipped: make([]PortraitAssetIssue, 0), Failed: make([]PortraitAssetIssue, 0), Assets: make([]store.PortraitAsset, 0),
}
platforms, err := s.store.ListPortraitAssetPlatforms(ctx)
if err != nil {
return response, err
}
configured := make([]store.PortraitAssetPlatform, 0, len(platforms))
for _, platform := range platforms {
if _, ok := portraitAssetSettings(platform); ok {
configured = append(configured, platform)
}
}
seen := map[string]bool{}
for _, value := range ids {
assetID := strings.TrimSpace(value)
if assetID == "" || seen[assetID] {
continue
}
seen[assetID] = true
asset, found, err := s.store.FindPortraitAssetForUser(ctx, user, assetID)
if err != nil {
return response, err
}
if !found {
response.Skipped = append(response.Skipped, PortraitAssetIssue{ID: assetID, Reason: "portrait asset not found"})
continue
}
if !asset.PrivateAvatarEligible {
response.Skipped = append(response.Skipped, PortraitAssetIssue{ID: asset.ID, Reason: "portrait asset authorization is required"})
continue
}
if len(configured) == 0 {
_ = s.store.UpdatePortraitAssetStatus(ctx, asset.ID, "not_configured", "no configured Volces portrait asset platform")
asset.Status = "not_configured"
asset.LastError = "no configured Volces portrait asset platform"
response.Skipped = append(response.Skipped, PortraitAssetIssue{ID: asset.ID, Reason: asset.LastError})
response.Assets = append(response.Assets, asset)
continue
}
response.Accepted++
assetFailed := false
for _, platform := range configured {
if err := s.syncPortraitAssetToPlatform(ctx, asset, platform); err != nil {
assetFailed = true
response.Failed = append(response.Failed, PortraitAssetIssue{ID: asset.ID, Reason: platform.PlatformID + ": " + err.Error()})
}
}
updated, _, err := s.refreshPortraitAssetStatus(ctx, user, asset.ID)
if err != nil {
return response, err
}
response.Assets = append(response.Assets, updated)
if !assetFailed {
response.SyncedIDs = append(response.SyncedIDs, updated.ID)
}
}
return response, nil
}
func (s *Service) syncPortraitAssetToPlatform(ctx context.Context, asset store.PortraitAsset, platform store.PortraitAssetPlatform) error {
settings, ok := portraitAssetSettings(platform)
if !ok {
return &clients.ClientError{Code: "portrait_asset_not_configured", Message: "platform portrait asset configuration is incomplete", Retryable: false}
}
binding, found, err := s.store.GetPortraitAssetBinding(ctx, asset.ID, platform.PlatformID)
if err != nil {
return err
}
if !found {
binding = store.PortraitAssetBinding{AssetID: asset.ID, PlatformID: platform.PlatformID, ProjectName: settings.ProjectName, AssetGroupID: settings.AssetGroupID, Status: "pending"}
}
if !portraitAssetHasPublicURL(asset.URL) {
return s.recordPortraitAssetBindingFailure(ctx, binding, settings, &clients.ClientError{
Code: "portrait_asset_public_url_required",
Message: "portrait asset URL must be an absolute http(s) URL reachable by Volces",
StatusCode: http.StatusBadRequest,
Retryable: false,
})
}
client := clients.VolcesAssetClient{HTTPClient: s.portraitAssetHTTPClient()}
remoteID := strings.TrimSpace(binding.RemoteAssetID)
if remoteID == "" {
created, _, createErr := client.CreateAsset(ctx, settings.Credentials, map[string]any{
"GroupId": settings.AssetGroupID, "URL": asset.URL, "Name": asset.Name,
"AssetType": volcesPortraitAssetType(asset.SourceType), "ProjectName": settings.ProjectName,
})
if createErr != nil {
return s.recordPortraitAssetBindingFailure(ctx, binding, settings, createErr)
}
remoteID = strings.TrimSpace(created.ID)
if remoteID == "" {
return s.recordPortraitAssetBindingFailure(ctx, binding, settings, &clients.ClientError{Code: "invalid_response", Message: "volces CreateAsset returned no asset id", Retryable: false})
}
binding.RemoteAssetID = remoteID
}
remote, _, getErr := client.GetAsset(ctx, settings.Credentials, map[string]any{"Id": remoteID, "ProjectName": settings.ProjectName})
if getErr != nil {
return s.recordPortraitAssetBindingFailure(ctx, binding, settings, getErr)
}
binding.ProjectName = settings.ProjectName
binding.AssetGroupID = settings.AssetGroupID
binding.RemoteAssetID = firstNonEmptyString(remote.ID, remoteID)
binding.RemoteAssetURI = "asset://" + binding.RemoteAssetID
binding.Status = portraitAssetBindingStatus(remote.Status)
binding.LastErrorCode = strings.TrimSpace(stringFromMap(remote.Error, "Code"))
binding.LastErrorMessage = strings.TrimSpace(stringFromMap(remote.Error, "Message"))
if binding.Status == "failed" && binding.LastErrorMessage == "" {
binding.LastErrorMessage = "volces portrait asset processing failed"
}
_, err = s.store.UpsertPortraitAssetBinding(ctx, binding)
return err
}
func (s *Service) recordPortraitAssetBindingFailure(ctx context.Context, binding store.PortraitAssetBinding, settings portraitAssetPlatformSettings, cause error) error {
binding.ProjectName = settings.ProjectName
binding.AssetGroupID = settings.AssetGroupID
binding.Status = "failed"
binding.LastErrorCode = clients.ErrorCode(cause)
binding.LastErrorMessage = cause.Error()
_, err := s.store.UpsertPortraitAssetBinding(ctx, binding)
if err != nil {
return err
}
return cause
}
func (s *Service) refreshPortraitAssetStatus(ctx context.Context, user *auth.User, assetID string) (store.PortraitAsset, bool, error) {
asset, found, err := s.store.FindPortraitAssetForUser(ctx, user, assetID)
if err != nil || !found {
return asset, found, err
}
active, total, lastError, _, err := s.store.PortraitAssetBindingSummary(ctx, asset.ID)
if err != nil {
return asset, true, err
}
status := "not_synced"
if total == 0 {
status = "not_synced"
} else if active > 0 {
status = "active"
if active < total {
status = "partial"
}
} else if lastError != "" {
status = "failed"
} else {
status = "pending"
}
if err := s.store.UpdatePortraitAssetStatus(ctx, asset.ID, status, lastError); err != nil {
return asset, true, err
}
asset.Status = status
asset.LastError = lastError
return asset, true, nil
}
func (s *Service) compilePortraitAssetReferences(ctx context.Context, user *auth.User, kind string, body map[string]any, candidate store.RuntimeModelCandidate) (map[string]any, error) {
entries := portraitAssetList(body["portrait_asset_list"])
if len(entries) == 0 {
return body, nil
}
if kind != "videos.generations" || !isVolcesPortraitAssetCandidate(candidate) {
return nil, &clients.ClientError{Code: "portrait_asset_unsupported_model", Message: "portrait assets require a configured Volces Seedance omni video model", StatusCode: http.StatusBadRequest, Retryable: false}
}
if !candidateSupportsPortraitAssets(candidate) {
return nil, &clients.ClientError{Code: "portrait_asset_unsupported_model", Message: "selected model does not enable supports_portrait_asset_reference", StatusCode: http.StatusBadRequest, Retryable: false}
}
out := cloneMap(body)
content := contentItems(out["content"])
labels := make([]string, len(entries))
nonAudioAssets := 0
for index, entry := range entries {
assetID := firstNonEmptyString(stringFromMap(entry, "id"), stringFromMap(entry, "easyai_portrait_asset_id"))
if assetID == "" {
return nil, &clients.ClientError{Code: "portrait_asset_id_required", Message: fmt.Sprintf("portrait_asset_list[%d].id is required", index), StatusCode: http.StatusBadRequest, Retryable: false}
}
asset, found, err := s.store.FindPortraitAssetForUser(ctx, user, assetID)
if err != nil {
return nil, err
}
if !found {
return nil, &clients.ClientError{Code: "portrait_asset_not_found", Message: "portrait asset not found", StatusCode: http.StatusNotFound, Retryable: false}
}
if !asset.PrivateAvatarEligible {
return nil, &clients.ClientError{Code: "portrait_asset_authorization_required", Message: "portrait asset authorization is required", StatusCode: http.StatusBadRequest, Retryable: false}
}
binding, bound, err := s.store.GetPortraitAssetBinding(ctx, asset.ID, candidate.PlatformID)
if err != nil {
return nil, err
}
if !bound || binding.Status != "active" || strings.TrimSpace(binding.RemoteAssetURI) == "" {
return nil, &clients.ClientError{Code: "portrait_asset_processing", Message: "portrait asset is not active for the selected Volces platform; sync it and retry", StatusCode: http.StatusServiceUnavailable, Retryable: true}
}
labels[index] = firstNonEmptyString(strings.TrimSpace(stringFromMap(entry, "name")), asset.Name, "portrait asset "+fmt.Sprint(index+1))
if asset.SourceType != "audio" {
nonAudioAssets++
}
content = append(content, portraitAssetContent(asset.SourceType, binding.RemoteAssetURI))
}
if nonAudioAssets == 0 {
return nil, &clients.ClientError{Code: "portrait_asset_audio_only", Message: "portrait_asset_list cannot contain audio-only assets", StatusCode: http.StatusBadRequest, Retryable: false}
}
for index := range content {
if strings.ToLower(strings.TrimSpace(stringFromAny(content[index]["type"]))) != "text" {
continue
}
content[index]["text"] = replacePortraitAssetPlaceholders(stringFromAny(content[index]["text"]), labels)
}
out["content"] = mapsToAnySlice(content)
delete(out, "portrait_asset_list")
return out, nil
}
func (s *Service) portraitAssetHTTPClient() *http.Client {
if s.httpClients != nil && s.httpClients.none != nil {
return s.httpClients.none
}
return http.DefaultClient
}
func portraitAssetSettings(platform store.PortraitAssetPlatform) (portraitAssetPlatformSettings, bool) {
config := portraitAssetNestedConfig(platform.Config)
accessKey := firstNonEmptyString(portraitAssetValue(config, "accessKey", "access_key"), portraitAssetValue(platform.Credentials, "accessKey", "access_key"))
secretKey := firstNonEmptyString(portraitAssetValue(config, "secretKey", "secret_key"), portraitAssetValue(platform.Credentials, "secretKey", "secret_key"))
projectName := firstNonEmptyString(portraitAssetValue(config, "projectName", "project_name"), "default")
assetGroupID := portraitAssetValue(config, "assetGroupId", "asset_group_id")
endpoint := firstNonEmptyString(portraitAssetValue(config, "assetEndpoint", "asset_endpoint", "volcesAssetEndpoint", "volces_asset_endpoint"), clientsVolcesAssetDefaultEndpoint())
if accessKey == "" || secretKey == "" || projectName == "" || assetGroupID == "" {
return portraitAssetPlatformSettings{}, false
}
if enabled, present := portraitAssetBool(config, "enabled"); present && !enabled {
return portraitAssetPlatformSettings{}, false
}
return portraitAssetPlatformSettings{ProjectName: projectName, AssetGroupID: assetGroupID, Credentials: clients.VolcesAssetCredentials{AccessKey: accessKey, SecretKey: secretKey, Endpoint: endpoint}}, true
}
func portraitAssetNestedConfig(config map[string]any) map[string]any {
for _, key := range []string{"seedancePrivateAsset", "seedance_private_asset", "portraitAsset", "portrait_asset"} {
if nested, ok := config[key].(map[string]any); ok {
return nested
}
}
return config
}
func portraitAssetValue(values map[string]any, keys ...string) string {
for _, key := range keys {
if value := strings.TrimSpace(stringFromAny(values[key])); value != "" {
return value
}
}
return ""
}
func portraitAssetBool(values map[string]any, key string) (bool, bool) {
value, ok := values[key]
if !ok {
return false, false
}
switch typed := value.(type) {
case bool:
return typed, true
case string:
return strings.EqualFold(strings.TrimSpace(typed), "true"), true
default:
return false, false
}
}
func validPortraitAssetSourceType(value string) bool {
switch strings.ToLower(strings.TrimSpace(value)) {
case "image", "video", "audio":
return true
default:
return false
}
}
func portraitAssetHasPublicURL(value string) bool {
parsed, err := url.Parse(strings.TrimSpace(value))
if err != nil || parsed.Host == "" {
return false
}
return strings.EqualFold(parsed.Scheme, "http") || strings.EqualFold(parsed.Scheme, "https")
}
func volcesPortraitAssetType(sourceType string) string {
switch strings.ToLower(strings.TrimSpace(sourceType)) {
case "video":
return "Video"
case "audio":
return "Audio"
default:
return "Image"
}
}
func portraitAssetBindingStatus(value string) string {
switch strings.ToLower(strings.TrimSpace(value)) {
case "active", "succeeded", "success":
return "active"
case "failed", "error":
return "failed"
default:
return "processing"
}
}
func portraitAssetList(value any) []map[string]any {
switch typed := value.(type) {
case []any:
out := make([]map[string]any, 0, len(typed))
for _, item := range typed {
if object, ok := item.(map[string]any); ok {
out = append(out, object)
}
}
return out
case []map[string]any:
return typed
default:
return nil
}
}
func portraitAssetContent(sourceType string, assetURI string) map[string]any {
switch strings.ToLower(strings.TrimSpace(sourceType)) {
case "video":
return map[string]any{"type": "video_url", "role": "reference_video", "video_url": map[string]any{"url": assetURI}}
case "audio":
return map[string]any{"type": "audio_url", "role": "reference_audio", "audio_url": map[string]any{"url": assetURI}}
default:
return map[string]any{"type": "image_url", "role": "reference_image", "image_url": map[string]any{"url": assetURI}}
}
}
func replacePortraitAssetPlaceholders(value string, labels []string) string {
return portraitAssetPlaceholderPattern.ReplaceAllStringFunc(value, func(match string) string {
parts := portraitAssetPlaceholderPattern.FindStringSubmatch(match)
for index := 1; index < len(parts); index++ {
if parts[index] == "" {
continue
}
position := int(parts[index][0] - '0')
if len(parts[index]) > 1 {
position = 0
for _, r := range parts[index] {
position = position*10 + int(r-'0')
}
}
if position > 0 && position <= len(labels) && strings.TrimSpace(labels[position-1]) != "" {
return labels[position-1]
}
}
return match
})
}
func isVolcesPortraitAssetCandidate(candidate store.RuntimeModelCandidate) bool {
provider := strings.ToLower(strings.TrimSpace(candidate.Provider))
return provider == "volces" || provider == "volces-openai"
}
func candidateSupportsPortraitAssets(candidate store.RuntimeModelCandidate) bool {
capabilities := effectiveModelCapability(candidate)
for _, key := range []string{candidate.ModelType, "omni_video", "omni", "video_generate"} {
if capability, ok := capabilities[key].(map[string]any); ok {
if enabled, present := portraitAssetBool(capability, "supports_portrait_asset_reference"); present {
return enabled
}
}
}
return false
}
func portraitAssetUserKeys(user *auth.User) (string, string) {
if user == nil {
return "", ""
}
gatewayUserID := strings.TrimSpace(user.GatewayUserID)
if gatewayUserID == "" && user.Source == "gateway" {
gatewayUserID = strings.TrimSpace(user.ID)
}
return gatewayUserID, strings.TrimSpace(user.ID)
}
func portraitAssetSHA256(payload []byte) string {
digest := sha256.Sum256(payload)
return hex.EncodeToString(digest[:])
}
func clientsVolcesAssetDefaultEndpoint() string { return "https://ark.cn-beijing.volcengineapi.com" }
@@ -0,0 +1,50 @@
package runner
import (
"testing"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestReplacePortraitAssetPlaceholders(t *testing.T) {
got := replacePortraitAssetPlaceholders("让 <<<portrait_asset_1>>> 和 @portrait_asset2、@人像资产3 出镜", []string{"Alice", "Bob", "Carol"})
want := "让 Alice 和 Bob、Carol 出镜"
if got != want {
t.Fatalf("placeholder replacement = %q, want %q", got, want)
}
}
func TestPortraitAssetContentUsesAssetURI(t *testing.T) {
item := portraitAssetContent("video", "asset://volces-video-1")
video, _ := item["video_url"].(map[string]any)
if item["type"] != "video_url" || item["role"] != "reference_video" || video["url"] != "asset://volces-video-1" {
t.Fatalf("unexpected portrait asset content: %+v", item)
}
}
func TestPortraitAssetSettingsRequireConfiguredVolcesAssetGroup(t *testing.T) {
settings, ok := portraitAssetSettings(store.PortraitAssetPlatform{Config: map[string]any{
"seedancePrivateAsset": map[string]any{
"enabled": true, "accessKey": "ak", "secretKey": "sk", "projectName": "project", "assetGroupId": "group",
},
}})
if !ok || settings.ProjectName != "project" || settings.AssetGroupID != "group" || settings.Credentials.AccessKey != "ak" {
t.Fatalf("unexpected configured portrait asset settings: %+v ok=%v", settings, ok)
}
if _, ok := portraitAssetSettings(store.PortraitAssetPlatform{Config: map[string]any{"seedancePrivateAsset": map[string]any{"enabled": true, "accessKey": "ak"}}}); ok {
t.Fatal("incomplete platform config must not enable portrait assets")
}
}
func TestPortraitAssetHasPublicURL(t *testing.T) {
for _, value := range []string{"https://assets.example.com/portrait.png", "http://assets.example.com/portrait.mp4"} {
if !portraitAssetHasPublicURL(value) {
t.Fatalf("expected public URL: %q", value)
}
}
for _, value := range []string{"/uploads/portrait.png", "file:///tmp/portrait.png", "asset://portrait-id"} {
if portraitAssetHasPublicURL(value) {
t.Fatalf("expected non-public URL: %q", value)
}
}
}
+197 -19
View File
@@ -3,6 +3,7 @@ package runner
import (
"context"
"math"
"strconv"
"strings"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
@@ -10,6 +11,8 @@ import (
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
const videoBillingUnitSeconds = 5
type EstimateResult struct {
Items []any `json:"items"`
Resolver string `json:"resolver"`
@@ -129,13 +132,27 @@ func (s *Service) billings(ctx context.Context, user *auth.User, kind string, bo
resource = "image_edit"
baseKey = "editBase"
}
if kind == "images.vectorize" {
resource = "image_vectorize"
unit = "conversion"
baseKey = "vectorizeBase"
}
if kind == "videos.upscales" {
resource = "video_enhance"
unit = "5s_video"
baseKey = "videoEnhanceBase"
inputs := resolveVideoEnhancePricingInputs(body, response)
amount, details := videoEnhanceAmount(config, inputs, float64(count), resourcePrice(config, resource, baseKey, "basePrice"))
amount = math.Ceil(amount*1e8) / 1e8 * discount
return []any{billingLineWithDetails(candidate, resource, unit, float64(count)*inputs.DurationUnits, roundPrice(amount), discount, simulated, details)}
}
if kind == "videos.generations" {
resource = "video"
unit = "5s_video"
baseKey = "videoBase"
duration, durationSource := billingDurationSeconds(body, response)
audioEnabled, audioSource := billingAudioEnabled(body, response)
durationUnits := math.Max(1, math.Ceil(duration/5))
durationUnits := videoDurationUnits(duration)
amount := float64(count) *
durationUnits *
resourcePrice(config, resource, baseKey, "basePrice") *
@@ -144,7 +161,7 @@ func (s *Service) billings(ctx context.Context, user *auth.User, kind string, bo
resourceWeight(config, resource, "referenceVideoWeights", boolWeightKey(requestHasReferenceVideo(body))) *
resourceWeight(config, resource, "voiceSpecifiedWeights", boolWeightKey(requestHasVoiceID(body, audioEnabled))) *
discount
return []any{billingLineWithDetails(candidate, resource, unit, count*int(durationUnits), roundPrice(amount), discount, simulated, map[string]any{
return []any{billingLineWithDetails(candidate, resource, unit, videoDurationQuantity(duration, count), roundPrice(amount), discount, simulated, map[string]any{
"count": count,
"audio": audioEnabled,
"audioSource": audioSource,
@@ -191,25 +208,180 @@ func (s *Service) billings(ctx context.Context, user *auth.User, kind string, bo
return []any{billingLine(candidate, resource, unit, count, roundPrice(amount), discount, simulated)}
}
type videoEnhancePricingInputs struct {
DurationSeconds, DurationUnits, SourceFrameRate, TargetFrameRate, SlowMotionRate float64
SourceResolution, TargetResolution, Model string
}
var videoEnhanceResolutionPixels = map[string]float64{
"480p": 854 * 480, "720p": 1280 * 720, "1080p": 1920 * 1080,
"1440p": 2560 * 1440, "2160p": 3840 * 2160,
}
var videoEnhanceResolutionBaseline = map[string]float64{"480p": 1, "720p": 1.4, "1080p": 1.8, "1440p": 2.1, "2160p": 2.6}
func resolveVideoEnhancePricingInputs(body map[string]any, response clients.Response) videoEnhancePricingInputs {
merged := cloneMap(body)
if data, ok := response.Result["data"].([]any); ok && len(data) > 0 {
if item, ok := data[0].(map[string]any); ok {
for _, key := range []string{"duration", "source_resolution", "target_resolution", "source_frame_rate", "target_frame_rate", "slow_motion_rate"} {
if merged[key] == nil && item[key] != nil {
merged[key] = item[key]
}
}
}
}
duration := floatFromAny(merged["duration"])
if duration <= 0 {
duration = 5
}
targetValue := firstNonEmptyString(stringFromMap(merged, "target_resolution"), stringFromMap(merged, "output_resolution"), stringFromMap(merged, "resolution"))
if dimensions := videoEnhanceDimensionsValue(merged, "output_width", "output_height"); dimensions != "" {
targetValue = dimensions
}
target := normalizeVideoEnhanceResolution(targetValue, "1080p")
sourceValue := firstNonEmptyString(stringFromMap(merged, "source_resolution"), stringFromMap(merged, "resolution"))
if dimensions := videoEnhanceDimensionsValue(merged, "source_width", "source_height"); dimensions != "" {
sourceValue = dimensions
}
source := normalizeVideoEnhanceResolution(sourceValue, target)
sourceFPS := floatFromAny(firstPresentValue(merged, "source_frame_rate", "frame_rate", "frameRate"))
if sourceFPS <= 0 {
sourceFPS = 24
}
targetFPS := floatFromAny(firstPresentValue(merged, "target_frame_rate", "output_frame_rate"))
if targetFPS <= 0 {
targetFPS = sourceFPS
}
slowMotion := floatFromAny(merged["slow_motion_rate"])
if slowMotion <= 0 {
slowMotion = 1
}
return videoEnhancePricingInputs{
DurationSeconds: duration, DurationUnits: math.Max(1, math.Round(duration)/5),
SourceResolution: source, TargetResolution: target,
SourceFrameRate: sourceFPS, TargetFrameRate: targetFPS, SlowMotionRate: slowMotion,
Model: firstNonEmptyString(stringFromMap(merged, "model"), stringFromMap(merged, "enhancement_model")),
}
}
func videoEnhanceDimensionsValue(values map[string]any, widthKey, heightKey string) string {
width := int(math.Round(floatFromAny(values[widthKey])))
height := int(math.Round(floatFromAny(values[heightKey])))
if width <= 0 || height <= 0 {
return ""
}
return strconv.Itoa(width) + "x" + strconv.Itoa(height)
}
func normalizeVideoEnhanceResolution(value, fallback string) string {
normalized := strings.TrimSuffix(strings.ToLower(strings.TrimSpace(value)), "_upscale")
if normalized == "2k" {
normalized = "1440p"
}
if normalized == "4k" {
normalized = "2160p"
}
if width, height, ok := parseVideoEnhanceDimensions(normalized); ok {
pixels := float64(width * height)
for _, bucket := range []string{"480p", "720p", "1080p", "1440p", "2160p"} {
if pixels <= videoEnhanceResolutionPixels[bucket] {
return bucket
}
}
return "2160p"
}
if _, ok := videoEnhanceResolutionPixels[normalized]; ok {
return normalized
}
if _, ok := videoEnhanceResolutionPixels[fallback]; ok {
return fallback
}
return "1080p"
}
func parseVideoEnhanceDimensions(value string) (int, int, bool) {
parts := strings.Split(strings.ReplaceAll(value, " ", ""), "x")
if len(parts) != 2 {
return 0, 0, false
}
width, widthErr := strconv.Atoi(parts[0])
height, heightErr := strconv.Atoi(parts[1])
return width, height, widthErr == nil && heightErr == nil && width > 0 && height > 0
}
func videoEnhanceAmount(config map[string]any, input videoEnhancePricingInputs, count, basePrice float64) (float64, map[string]any) {
transitionKey := input.SourceResolution + "->" + input.TargetResolution
minimumResolutionWeight := videoEnhanceResolutionBaseline[input.TargetResolution]
resolutionWeight, configuredResolution := pricingDynamicWeight(config, "video_enhance", "resolutionTransitions", transitionKey)
if !configuredResolution {
resolutionWeight = math.Max(minimumResolutionWeight, videoEnhanceResolutionPixels[input.TargetResolution]/videoEnhanceResolutionPixels[input.SourceResolution])
} else {
resolutionWeight = math.Max(minimumResolutionWeight, resolutionWeight)
}
sourceFPS := math.Max(1, math.Round(input.SourceFrameRate))
targetFPS := math.Max(1, math.Round(input.TargetFrameRate))
frameKey := strconv.Itoa(int(sourceFPS)) + "->" + strconv.Itoa(int(targetFPS))
minimumFrameWeight := math.Max(1, targetFPS/24)
frameWeight, configuredFrame := pricingDynamicWeight(config, "video_enhance", "frameRateTransitions", frameKey)
if !configuredFrame {
frameWeight = minimumFrameWeight
} else {
frameWeight = math.Max(minimumFrameWeight, frameWeight)
}
modelWeight, configuredModel := pricingDynamicWeight(config, "video_enhance", "modelWeights", input.Model)
if !configuredModel {
modelWeight = 1
}
markup, foundMarkup := pricingDynamicWeight(config, "video_enhance", "", "markup")
if !foundMarkup {
markup = 1
}
amount := count * input.DurationUnits * basePrice * resolutionWeight * frameWeight * input.SlowMotionRate * modelWeight * markup
return amount, map[string]any{
"count": count, "durationSeconds": input.DurationSeconds, "durationUnit": "5s", "durationUnitCount": input.DurationUnits,
"sourceResolution": input.SourceResolution, "targetResolution": input.TargetResolution,
"resolutionTransitionKey": transitionKey, "resolutionTransitionWeight": resolutionWeight, "resolutionTransitionFallback": !configuredResolution,
"sourceFrameRate": input.SourceFrameRate, "targetFrameRate": input.TargetFrameRate,
"frameRateTransitionKey": frameKey, "frameRateTransitionWeight": frameWeight, "frameRateTransitionFallback": !configuredFrame,
"slowMotionRate": input.SlowMotionRate, "model": input.Model, "modelWeight": modelWeight, "markup": markup,
}
}
func pricingDynamicWeight(config map[string]any, resource, group, key string) (float64, bool) {
resourceConfig, _ := config[resource].(map[string]any)
dynamic, _ := resourceConfig["dynamicWeight"].(map[string]any)
if group == "" {
value, ok := dynamic[key]
weight := floatFromAny(value)
return weight, ok && weight > 0
}
values, _ := dynamic[group].(map[string]any)
value, ok := values[key]
weight := floatFromAny(value)
return weight, ok && weight > 0
}
func (s *Service) effectiveBillingConfig(ctx context.Context, candidate store.RuntimeModelCandidate) map[string]any {
base := candidate.BaseBillingConfig
if ruleSetID := firstNonEmptyString(candidate.BasePricingRuleSetID, candidate.PlatformPricingRuleSetID); ruleSetID != "" {
var inheritedRuleSetConfig map[string]any
if ruleSetID := firstNonEmptyString(candidate.BasePricingRuleSetID, candidate.PlatformPricingRuleSetID); ruleSetID != "" && s.store != nil {
if ruleSetConfig, err := s.store.PricingRuleSetBillingConfig(ctx, ruleSetID); err == nil && len(ruleSetConfig) > 0 {
base = ruleSetConfig
inheritedRuleSetConfig = ruleSetConfig
}
}
if len(candidate.BillingConfig) > 0 {
base = candidate.BillingConfig
}
if candidate.ModelPricingRuleSetID != "" {
var modelRuleSetConfig map[string]any
if candidate.ModelPricingRuleSetID != "" && s.store != nil {
if ruleSetConfig, err := s.store.PricingRuleSetBillingConfig(ctx, candidate.ModelPricingRuleSetID); err == nil && len(ruleSetConfig) > 0 {
base = ruleSetConfig
modelRuleSetConfig = ruleSetConfig
}
}
if len(candidate.BillingConfigOverride) > 0 {
base = mergeMap(base, candidate.BillingConfigOverride)
}
return base
return store.ResolveEffectiveBillingConfig(store.EffectiveBillingConfigInput{
BaseConfig: candidate.BaseBillingConfig,
LegacyPlatformModelConfig: candidate.BillingConfig,
InheritedRuleSetConfig: inheritedRuleSetConfig,
ModelRuleSetConfig: modelRuleSetConfig,
Override: candidate.BillingConfigOverride,
})
}
func effectiveDiscount(ctx context.Context, db *store.Store, user *auth.User, candidate store.RuntimeModelCandidate) float64 {
@@ -416,6 +588,16 @@ func weightValueAliases(key string, name string) []string {
}
}
func videoDurationUnits(durationSeconds float64) float64 {
return videoDurationQuantity(durationSeconds, 1)
}
func videoDurationQuantity(durationSeconds float64, count int) float64 {
const durationPrecision = 1_000_000_000
scaledDuration := math.Round(durationSeconds * durationPrecision)
return scaledDuration * float64(count) / (durationPrecision * videoBillingUnitSeconds)
}
func requestOutputCount(body map[string]any) int {
for _, key := range []string{"n", "count", "batch_size", "batchSize"} {
if value := int(math.Ceil(floatFromAny(body[key]))); value > 0 {
@@ -476,11 +658,7 @@ func generatedVideoDurationSeconds(result map[string]any) (float64, bool) {
if duration <= 0 {
continue
}
rounded := math.Round(duration)
if rounded <= 0 {
rounded = 1
}
return rounded, true
return duration, true
}
return 0, false
}
+65 -7
View File
@@ -38,7 +38,7 @@ func TestImageBillingEstimateUsesCountResolutionAndQuality(t *testing.T) {
}
}
func TestVideoBillingEstimateUsesFiveSecondUnitsAndDynamicWeights(t *testing.T) {
func TestVideoBillingEstimateProratesFiveSecondUnitsAndDynamicWeights(t *testing.T) {
service := &Service{}
candidate := store.RuntimeModelCandidate{
ModelName: "video-model",
@@ -67,13 +67,13 @@ func TestVideoBillingEstimateUsesFiveSecondUnitsAndDynamicWeights(t *testing.T)
}, candidate, clients.Response{}, true)
line := firstBillingLine(t, items)
if got, want := floatFromAny(line["amount"]), 1620.0; got != want {
if got, want := floatFromAny(line["amount"]), 1296.0; got != want {
t.Fatalf("video estimated amount = %v, want %v", got, want)
}
if got, want := floatFromAny(line["durationUnitCount"]), 3.0; got != want {
if got, want := floatFromAny(line["durationUnitCount"]), 2.4; got != want {
t.Fatalf("video duration units = %v, want %v", got, want)
}
if got, want := line["quantity"], 3; got != want {
if got, want := floatFromAny(line["quantity"]), 2.4; got != want {
t.Fatalf("video quantity = %v, want %v", got, want)
}
if got, want := line["durationSource"], "preprocessed_request"; got != want {
@@ -93,6 +93,64 @@ func TestVideoBillingEstimateUsesFiveSecondUnitsAndDynamicWeights(t *testing.T)
}
}
func TestAdvancedMediaBillingUsesConversionAndTransitionMatrix(t *testing.T) {
service := &Service{}
vectorCandidate := store.RuntimeModelCandidate{ModelName: "easy-image-vectorizer-1", BaseBillingConfig: map[string]any{
"image_vectorize": map[string]any{"basePrice": 200},
}}
vectorLine := firstBillingLine(t, service.billings(context.Background(), nil, "images.vectorize", map[string]any{"count": 1}, vectorCandidate, clients.Response{}, true))
if vectorLine["resourceType"] != "image_vectorize" || vectorLine["unit"] != "conversion" || floatFromAny(vectorLine["amount"]) != 200 {
t.Fatalf("unexpected vectorizer billing: %+v", vectorLine)
}
topazCandidate := store.RuntimeModelCandidate{ModelName: "easy-proteus-standard-4", BaseBillingConfig: map[string]any{
"video_enhance": map[string]any{
"basePrice": 100,
"dynamicWeight": map[string]any{
"resolutionTransitions": map[string]any{"720p->1080p": 1.8},
"frameRateTransitions": map[string]any{"24->60": 2.5},
"markup": 1,
},
},
}}
topazLine := firstBillingLine(t, service.billings(context.Background(), nil, "videos.upscales", map[string]any{
"count": 1, "duration": 5, "source_resolution": "720p", "target_resolution": "1080p",
"source_frame_rate": 24, "target_frame_rate": 60, "slow_motion_rate": 2,
"model": "easy-proteus-standard-4",
}, topazCandidate, clients.Response{}, true))
if got, want := floatFromAny(topazLine["amount"]), 900.0; got != want {
t.Fatalf("video enhance amount=%v, want=%v, line=%+v", got, want, topazLine)
}
if topazLine["resolutionTransitionKey"] != "720p->1080p" || topazLine["frameRateTransitionKey"] != "24->60" {
t.Fatalf("missing transition evidence: %+v", topazLine)
}
}
func TestVideoEnhancePricingFallbackUsesPixelsAndBaseline(t *testing.T) {
amount, details := videoEnhanceAmount(map[string]any{"video_enhance": map[string]any{"dynamicWeight": map[string]any{}}}, videoEnhancePricingInputs{
DurationSeconds: 5, DurationUnits: 1, SourceResolution: "720p", TargetResolution: "2160p",
SourceFrameRate: 30, TargetFrameRate: 60, SlowMotionRate: 1,
}, 1, 100)
if amount != 2250 {
t.Fatalf("fallback amount=%v, want=2250 details=%+v", amount, details)
}
}
func TestVideoEnhancePricingPrefersProbedDimensions(t *testing.T) {
inputs := resolveVideoEnhancePricingInputs(map[string]any{
"duration": 3,
"source_resolution": "180p",
"source_width": 320,
"source_height": 180,
"target_resolution": "720p",
"output_width": 1280,
"output_height": 720,
}, clients.Response{})
if inputs.SourceResolution != "480p" || inputs.TargetResolution != "720p" {
t.Fatalf("dimension buckets = %s->%s, want 480p->720p", inputs.SourceResolution, inputs.TargetResolution)
}
}
func TestMusicBillingUsesSongResourceAndOutputCount(t *testing.T) {
service := &Service{}
candidate := store.RuntimeModelCandidate{
@@ -172,13 +230,13 @@ func TestVideoBillingPrefersGeneratedDuration(t *testing.T) {
}, false)
line := firstBillingLine(t, items)
if got, want := floatFromAny(line["durationSeconds"]), 7.0; got != want {
if got, want := floatFromAny(line["durationSeconds"]), 6.6; got != want {
t.Fatalf("video generated duration = %v, want %v", got, want)
}
if got, want := floatFromAny(line["durationUnitCount"]), 2.0; got != want {
if got, want := floatFromAny(line["durationUnitCount"]), 1.32; got != want {
t.Fatalf("video generated duration units = %v, want %v", got, want)
}
if got, want := floatFromAny(line["amount"]), 200.0; got != want {
if got, want := floatFromAny(line["amount"]), 132.0; got != want {
t.Fatalf("video generated duration amount = %v, want %v", got, want)
}
if got, want := line["durationSource"], "generated_video"; got != want {
+51 -4
View File
@@ -184,6 +184,22 @@ func multiplyFixedAmountRatio(amount fixedAmount, numerator int, denominator int
return fixedAmountFromBigInt(roundBigIntRatio(product, big.NewInt(int64(denominator))))
}
func multiplyFixedProductRatio(base fixedAmount, integerFactors []int, fixedFactors []fixedAmount, denominator int) (fixedAmount, error) {
if denominator == 0 {
return 0, fmt.Errorf("division by zero")
}
product := big.NewInt(int64(base))
for _, factor := range integerFactors {
product.Mul(product, big.NewInt(int64(factor)))
}
divisor := big.NewInt(int64(denominator))
for _, factor := range fixedFactors {
product.Mul(product, big.NewInt(int64(factor)))
divisor.Mul(divisor, big.NewInt(fixedScale))
}
return fixedAmountFromBigInt(roundBigIntRatio(product, divisor))
}
func fixedAmountFromBigInt(value *big.Int) (fixedAmount, error) {
if value == nil || !value.IsInt64() {
return 0, errFixedAmountOverflow
@@ -643,13 +659,39 @@ func (s *Service) billingsWithResolvedPricingV2(
resource = "image_edit"
baseKey = "editBase"
}
if kind == "images.vectorize" {
resource = "image_vectorize"
unit = "conversion"
baseKey = "vectorizeBase"
}
if kind == "videos.upscales" {
resource = "video_enhance"
unit = "5s_video"
baseKey = "videoEnhanceBase"
inputs := resolveVideoEnhancePricingInputs(body, response)
price, priceErr := pricing.requiredPrice(resource, baseKey, "basePrice")
if priceErr != nil {
return nil, 0, resolvedPricing{}, priceErr
}
rawAmount, details := videoEnhanceAmount(pricing.Config, inputs, float64(count), price.Float64())
rawAmount = math.Ceil(rawAmount*1e8) / 1e8
amount, amountErr := fixedAmountFromAny(rawAmount * discount.Float64())
if amountErr != nil {
return nil, 0, resolvedPricing{}, pricing.calculationError(resource, amountErr)
}
return []any{buildLine(resource, unit, float64(count)*inputs.DurationUnits, amount, details)}, amount, pricing, nil
}
if kind == "videos.generations" {
resource = "video"
unit = "5s_video"
baseKey = "videoBase"
duration, durationSource := billingDurationSeconds(body, response)
audioEnabled, audioSource := billingAudioEnabled(body, response)
durationUnits := int(math.Max(1, math.Ceil(duration/5)))
durationUnits := videoDurationUnits(duration)
durationFixed, durationErr := fixedAmountFromAny(duration)
if durationErr != nil {
return nil, 0, resolvedPricing{}, pricing.calculationError(resource, durationErr)
}
price, priceErr := pricing.requiredPrice(resource, baseKey, "basePrice")
if priceErr != nil {
return nil, 0, resolvedPricing{}, priceErr
@@ -670,11 +712,16 @@ func (s *Service) billingsWithResolvedPricingV2(
if weightErr != nil {
return nil, 0, resolvedPricing{}, weightErr
}
amount, calculationErr := pricing.calculate(resource, price, []int{count, durationUnits}, resolutionWeight, audioWeight, referenceVideoWeight, voiceWeight, discount)
amount, calculationErr := multiplyFixedProductRatio(
price,
[]int{count},
[]fixedAmount{durationFixed, resolutionWeight, audioWeight, referenceVideoWeight, voiceWeight, discount},
videoBillingUnitSeconds,
)
if calculationErr != nil {
return nil, 0, resolvedPricing{}, calculationErr
return nil, 0, resolvedPricing{}, pricing.calculationError(resource, calculationErr)
}
item := buildLine(resource, unit, count*durationUnits, amount, map[string]any{
item := buildLine(resource, unit, videoDurationQuantity(duration, count), amount, map[string]any{
"count": count, "audio": audioEnabled, "audioSource": audioSource,
"durationSeconds": duration, "durationSource": durationSource,
"durationUnit": "5s", "durationUnitCount": durationUnits,
+107
View File
@@ -1,10 +1,12 @@
package runner
import (
"context"
"errors"
"math"
"testing"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
@@ -39,6 +41,9 @@ func TestFixedAmountOperationsRejectOverflow(t *testing.T) {
if _, err := pricing.calculate("image", maximum, []int{2}); !isPricingUnavailable(err) {
t.Fatalf("pricing overflow should be unavailable: %v", err)
}
if _, err := multiplyFixedProductRatio(maximum, []int{2}, nil, 1); !errors.Is(err, errFixedAmountOverflow) {
t.Fatalf("product ratio overflow error=%v", err)
}
}
func TestEstimatedOutputTokensUsesAliasesAndCapabilityFallback(t *testing.T) {
@@ -138,6 +143,108 @@ func TestPricingWeightsUseFixedPrecisionAndRejectInvalidValues(t *testing.T) {
}
}
func TestVideoBillingV2ProratesFiveSecondPriceByActualDuration(t *testing.T) {
service := &Service{}
candidate := store.RuntimeModelCandidate{ModelName: "video-model"}
pricing := resolvedPricing{
Config: map[string]any{
"video": map[string]any{
"basePrice": 100,
"dynamicWeight": map[string]any{
"audioWeights": map[string]any{"true": 2},
},
},
},
Currency: "resource",
}
tests := []struct {
name string
duration float64
wantUnits float64
wantAmount float64
}{
{name: "three seconds uses zero point six units", duration: 3, wantUnits: 0.6, wantAmount: 120},
{name: "five seconds uses one unit", duration: 5, wantUnits: 1, wantAmount: 200},
{name: "six seconds uses one point two units", duration: 6, wantUnits: 1.2, wantAmount: 240},
{name: "fractional seconds retain fixed amount precision", duration: 6.000000001, wantUnits: 1.2000000002, wantAmount: 240.00000004},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
items, total, _, err := service.billingsWithResolvedPricingV2(
context.Background(), nil, "videos.generations",
map[string]any{"duration": test.duration, "audio": true},
candidate, clients.Response{}, true, pricing,
)
if err != nil {
t.Fatalf("bill video: %v", err)
}
line := firstBillingLine(t, items)
if got := total.Float64(); math.Abs(got-test.wantAmount) > 1e-9 {
t.Fatalf("total amount=%v, want %v", got, test.wantAmount)
}
if got := floatFromAny(line["amount"]); math.Abs(got-test.wantAmount) > 1e-9 {
t.Fatalf("line amount=%v, want %v", got, test.wantAmount)
}
if got := floatFromAny(line["quantity"]); math.Abs(got-test.wantUnits) > 1e-12 {
t.Fatalf("quantity=%v, want %v", got, test.wantUnits)
}
if got := floatFromAny(line["durationUnitCount"]); math.Abs(got-test.wantUnits) > 1e-12 {
t.Fatalf("duration units=%v, want %v", got, test.wantUnits)
}
})
}
}
func TestVideoBillingV2RoundsOnlyAfterApplyingDurationCountAndWeights(t *testing.T) {
service := &Service{}
candidate := store.RuntimeModelCandidate{ModelName: "video-model"}
tests := []struct {
name string
body map[string]any
pricing resolvedPricing
wantAmount string
}{
{
name: "count preserves a sub-nano duration share",
body: map[string]any{"duration": 1, "count": 5},
pricing: resolvedPricing{Config: map[string]any{
"video": map[string]any{"basePrice": "0.000000001"},
}},
wantAmount: "0.000000001",
},
{
name: "weight does not amplify a rounded duration share",
body: map[string]any{"duration": 3, "audio": true},
pricing: resolvedPricing{Config: map[string]any{
"video": map[string]any{
"basePrice": "0.000000001",
"dynamicWeight": map[string]any{
"audioWeights": map[string]any{"true": 2},
},
},
}},
wantAmount: "0.000000001",
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
_, total, _, err := service.billingsWithResolvedPricingV2(
context.Background(), nil, "videos.generations", test.body,
candidate, clients.Response{}, true, test.pricing,
)
if err != nil {
t.Fatalf("bill video: %v", err)
}
if got := total.String(); got != test.wantAmount {
t.Fatalf("total amount=%s, want %s", got, test.wantAmount)
}
})
}
}
func mustFixedAmount(t *testing.T, value string) fixedAmount {
t.Helper()
amount, err := parseFixedAmount(value)
+150 -2
View File
@@ -5,6 +5,8 @@ import (
"errors"
"fmt"
"log/slog"
"os"
"regexp"
"strconv"
"strings"
"time"
@@ -90,6 +92,8 @@ func New(cfg config.Config, db *store.Store, logger *slog.Logger, observers ...b
"volces": clients.VolcesClient{HTTPClient: httpClients.none},
"keling": clients.KelingClient{HTTPClient: httpClients.none},
"kling": clients.KelingClient{HTTPClient: httpClients.none},
"vectorizer": clients.VectorizerClient{HTTPClient: httpClients.none},
"topaz": clients.TopazClient{HTTPClient: httpClients.none},
"universal": clients.UniversalClient{HTTPClient: httpClients.none, ScriptExecutor: scriptExecutor},
"simulation": clients.SimulationClient{},
},
@@ -526,6 +530,20 @@ 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)
if err != nil && isVolcesRemoteTaskCancellation(candidate, err) {
cancelled, changed, cancelErr := s.store.CancelSubmittedTask(ctx, task.ID, task.ExecutionToken, "任务已由火山引擎取消")
if cancelErr != nil {
return Result{}, cancelErr
}
if changed {
// CancelSubmittedTask atomically transfers any reservation to the release Outbox.
walletReservationFinalized = true
if emitErr := s.emit(ctx, task.ID, "task.cancelled", "cancelled", "cancelled", 1, "任务已由火山引擎取消", map[string]any{"taskId": task.ID, "reason": "upstream_cancelled"}, isSimulation(task, candidate)); emitErr != nil {
return Result{}, emitErr
}
return Result{Task: cancelled, Output: cancelled.Result}, nil
}
}
if err == nil {
attemptNo = nextAttemptNo
var billings []any
@@ -592,6 +610,13 @@ candidatesLoop:
ResponseDurationMS: record.ResponseDurationMS,
})
if finishErr != nil {
if errors.Is(finishErr, store.ErrTaskExecutionLeaseLost) {
latest, latestErr := s.store.GetTask(ctx, task.ID)
if latestErr == nil && latest.Status == "cancelled" {
walletReservationFinalized = true
return Result{Task: latest, Output: latest.Result}, nil
}
}
return Result{}, finishErr
}
walletReservationFinalized = true
@@ -836,6 +861,11 @@ func billingItemsFixedTotal(items []any) fixedAmount {
}
func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user *auth.User, body map[string]any, preprocessing parameterPreprocessingLog, candidate store.RuntimeModelCandidate, pricing resolvedPricing, attemptNo int, onDelta clients.StreamDelta, responseExecution responseExecutionContext, singleSourceProtected bool, cacheAffinityPolicy map[string]any, cacheAffinityRecordKeys []string) (clients.Response, error) {
var err error
candidate, err = candidateWithEnvironmentCredentials(candidate)
if err != nil {
return clients.Response{}, err
}
simulated := isSimulation(task, candidate)
baseAttemptMetrics := mergeMetrics(attemptMetrics(candidate, attemptNo, simulated), parameterPreprocessingMetrics(preprocessing))
reservations := s.rateLimitReservations(ctx, user, candidate, body)
@@ -916,7 +946,28 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
return clients.Response{}, fmt.Errorf("prepare http client: %w", err)
}
client := s.clientFor(candidate, simulated)
providerBody, err := s.hydrateProviderRequestAssets(ctx, body, candidate)
providerBody, err := s.compilePortraitAssetReferences(ctx, user, task.Kind, body, 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
}
providerBody, err = s.preparePrivateProviderRequest(ctx, user, task.Kind, providerBody)
if err != nil {
_ = s.store.FinishTaskAttempt(ctx, store.FinishTaskAttemptInput{
AttemptID: attemptID, Status: "failed", Retryable: false,
Metrics: mergeMetrics(baseAttemptMetrics, map[string]any{"error": err.Error(), "retryable": false}),
ErrorCode: clients.ErrorCode(err), ErrorMessage: err.Error(),
})
return clients.Response{}, err
}
providerBody, err = s.hydrateProviderRequestAssets(ctx, providerBody, candidate)
if err != nil {
_ = s.store.FinishTaskAttempt(ctx, store.FinishTaskAttemptInput{
AttemptID: attemptID,
@@ -953,6 +1004,12 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
}
return s.store.SetTaskRemoteTask(context.WithoutCancel(ctx), task.ID, task.ExecutionToken, attemptID, remoteTaskID, payload)
},
OnRemoteTaskPolled: func(remoteTaskID string, payload map[string]any) error {
if strings.TrimSpace(remoteTaskID) == "" {
return nil
}
return s.store.SetTaskRemoteTask(context.WithoutCancel(ctx), task.ID, task.ExecutionToken, attemptID, remoteTaskID, payload)
},
Stream: boolFromMap(providerBody, "stream"),
StreamDelta: onDelta,
UpstreamProtocol: candidate.ResponseProtocol,
@@ -1187,12 +1244,19 @@ func (s *Service) failTask(ctx context.Context, taskID string, executionToken st
if err != nil {
return store.GatewayTask{}, err
}
if failed.Status == "cancelled" {
return failed, nil
}
if eventErr := s.emit(ctx, taskID, "task.failed", "failed", "failed", 1, message, map[string]any{"code": code, "requestId": requestID, "metrics": metrics}, simulated); eventErr != nil {
return store.GatewayTask{}, eventErr
}
return failed, nil
}
func isVolcesRemoteTaskCancellation(candidate store.RuntimeModelCandidate, err error) bool {
return isVolcesCancellationCandidate(candidate) && strings.EqualFold(clients.ErrorCode(err), "volces_task_cancelled")
}
type failedAttemptRecord struct {
Task store.GatewayTask
Body map[string]any
@@ -1375,6 +1439,8 @@ func modelTypeFromKind(kind string, body map[string]any) string {
return "image_edit"
}
return "image_generate"
case "images.vectorize":
return "image_vectorize"
case "videos.generations":
if videoRequestHasVideoOrAudioReference(body) {
return "omni_video"
@@ -1383,6 +1449,8 @@ func modelTypeFromKind(kind string, body map[string]any) string {
return "image_to_video"
}
return "video_generate"
case "videos.upscales":
return "video_enhance"
case "song.generations", "music.generations":
return "audio_generate"
case "speech.generations":
@@ -1417,6 +1485,10 @@ func canonicalModelType(value string) string {
return "text_to_speech"
case "voice", "voice_clone", "voiceclone", "voice.cloning":
return "voice_clone"
case "vectorize", "image_vectorizer", "image_vectorize":
return "image_vectorize"
case "video_upscale", "video_enhance", "upscale":
return "video_enhance"
default:
return normalized
}
@@ -1424,7 +1496,7 @@ func canonicalModelType(value string) string {
func isKnownModelType(value string) bool {
switch value {
case "text_generate", "text_embedding", "text_rerank", "image_generate", "image_edit", "video_generate", "image_to_video", "text_to_video", "video_edit", "video_reference", "video_first_last_frame", "omni_video", "omni", "audio_generate", "text_to_speech", "voice_clone":
case "text_generate", "text_embedding", "text_rerank", "image_generate", "image_edit", "image_vectorize", "video_generate", "video_enhance", "image_to_video", "text_to_video", "video_edit", "video_reference", "video_first_last_frame", "omni_video", "omni", "audio_generate", "text_to_speech", "voice_clone":
return true
default:
return false
@@ -1657,6 +1729,14 @@ func validateRequest(kind string, body map[string]any) error {
if strings.TrimSpace(stringFromMap(body, "prompt")) == "" {
return errors.New("prompt is required")
}
case "images.vectorize":
if vectorizerSourceURL(body) == "" && vectorizerSourceTaskID(body) == "" {
return &clients.ClientError{Code: "invalid_parameter", Message: "source.url or source.vectorizerTaskId is required", Param: "source", StatusCode: 400}
}
case "videos.upscales":
if firstNonEmptyString(stringFromMap(body, "video_url"), stringFromMap(body, "videoUrl")) == "" {
return &clients.ClientError{Code: "invalid_parameter", Message: "video_url is required", Param: "video_url", StatusCode: 400}
}
case "song.generations", "music.generations":
if strings.TrimSpace(stringFromMap(body, "prompt")) == "" {
return errors.New("prompt is required")
@@ -1676,6 +1756,74 @@ func validateRequest(kind string, body map[string]any) error {
return nil
}
var credentialEnvironmentName = regexp.MustCompile(`^[A-Z_][A-Z0-9_]*$`)
func candidateWithEnvironmentCredentials(candidate store.RuntimeModelCandidate) (store.RuntimeModelCandidate, error) {
references, _ := candidate.PlatformConfig["credentialEnv"].(map[string]any)
if len(references) == 0 {
references, _ = candidate.PlatformConfig["credential_env"].(map[string]any)
}
if len(references) == 0 {
return candidate, nil
}
credentials := cloneMap(candidate.Credentials)
for credentialName, rawEnvironmentName := range references {
environmentName := strings.TrimSpace(fmt.Sprint(rawEnvironmentName))
if !credentialEnvironmentName.MatchString(environmentName) {
return candidate, &clients.ClientError{Code: "missing_credentials", Message: "platform credential environment reference is invalid", Retryable: false}
}
value, ok := os.LookupEnv(environmentName)
if !ok || strings.TrimSpace(value) == "" {
return candidate, &clients.ClientError{Code: "missing_credentials", Message: "platform credential environment variable is not configured", Retryable: false}
}
credentials[credentialName] = value
}
candidate.Credentials = credentials
return candidate, nil
}
func (s *Service) preparePrivateProviderRequest(ctx context.Context, user *auth.User, kind string, body map[string]any) (map[string]any, error) {
if kind != "images.vectorize" {
return body, nil
}
taskID := vectorizerSourceTaskID(body)
if taskID == "" {
return body, nil
}
sourceTask, err := s.store.GetTask(ctx, taskID)
if err != nil || !taskAccessibleToUser(sourceTask, user) || sourceTask.Kind != "images.vectorize" {
return nil, &clients.ClientError{Code: "not_found", Message: "source vectorizer task not found", StatusCode: 404, Retryable: false}
}
imageToken := strings.TrimSpace(fmt.Sprint(sourceTask.RemoteTaskPayload["imageToken"]))
if imageToken == "" {
return nil, &clients.ClientError{Code: "invalid_parameter", Message: "source vectorizer task cannot be reused", Param: "source.vectorizerTaskId", StatusCode: 400, Retryable: false}
}
privateBody := cloneMap(body)
privateBody["_vectorizer_image_token"] = imageToken
if receipt := strings.TrimSpace(fmt.Sprint(sourceTask.RemoteTaskPayload["receipt"])); receipt != "" {
privateBody["_vectorizer_receipt"] = receipt
}
return privateBody, nil
}
func vectorizerSourceURL(body map[string]any) string {
if source, ok := body["source"].(map[string]any); ok {
return strings.TrimSpace(firstNonEmptyString(stringFromMap(source, "url"), stringFromMap(source, "image_url"), stringFromMap(source, "imageUrl")))
}
return strings.TrimSpace(firstNonEmptyString(stringFromMap(body, "image_url"), stringFromMap(body, "imageUrl")))
}
func vectorizerSourceTaskID(body map[string]any) string {
for _, key := range []string{"source", "sourceContext"} {
if source, ok := body[key].(map[string]any); ok {
if taskID := strings.TrimSpace(firstNonEmptyString(stringFromMap(source, "vectorizerTaskId"), stringFromMap(source, "vectorizer_task_id"), stringFromMap(source, "taskId"), stringFromMap(source, "task_id"))); taskID != "" {
return taskID
}
}
}
return strings.TrimSpace(firstNonEmptyString(stringFromMap(body, "vectorizerTaskId"), stringFromMap(body, "vectorizer_task_id")))
}
func hasRerankDocuments(value any) bool {
switch typed := value.(type) {
case []any:
+63
View File
@@ -6,6 +6,7 @@ import (
"strings"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
"github.com/riverqueue/river/rivertype"
)
@@ -104,6 +105,61 @@ func (s *Service) CancelTask(ctx context.Context, taskID string, user *auth.User
}, nil
}
// CancelVolcesVideoTask extends local queue cancellation with the official
// Volces DELETE call once a video task has a persisted remote task id.
func (s *Service) CancelVolcesVideoTask(ctx context.Context, task store.GatewayTask, user *auth.User) (TaskCancelResult, error) {
local, err := s.CancelTask(ctx, task.ID, user)
if err != nil || local.Cancelled || strings.TrimSpace(task.RemoteTaskID) == "" {
return local, err
}
if taskCancelTerminalStatus(task.Status) {
return local, nil
}
var latest store.TaskAttempt
for _, attempt := range task.Attempts {
if attempt.PlatformModelID != "" && (latest.AttemptNo == 0 || attempt.AttemptNo >= latest.AttemptNo) {
latest = attempt
}
}
candidate, found, err := s.store.GetRuntimeModelCandidateForRemoteTask(ctx, latest.PlatformModelID, latest.PlatformID)
if err != nil {
return TaskCancelResult{}, err
}
if !found || !isVolcesCancellationCandidate(candidate) {
return local, nil
}
httpClient, err := s.httpClientForCandidate(candidate, false)
if err != nil {
return TaskCancelResult{}, err
}
_, _, err = (clients.VolcesClient{HTTPClient: httpClient}).DeleteVideoTask(ctx, clients.Request{
Kind: "videos.generations", Candidate: candidate, HTTPClient: httpClient, RemoteTaskID: task.RemoteTaskID,
})
if err != nil {
return TaskCancelResult{}, err
}
cancelledTask, cancelled, err := s.store.CancelSubmittedTask(ctx, task.ID, task.ExecutionToken, "任务已由火山引擎取消")
if err != nil {
return TaskCancelResult{}, err
}
if !cancelled {
latestTask, latestErr := s.store.GetTask(ctx, task.ID)
if latestErr == nil {
return taskCancelUnavailable(latestTask, "任务状态已变化,未覆盖本地最终状态"), nil
}
return local, nil
}
if err := s.emit(ctx, cancelledTask.ID, "task.cancelled", "cancelled", "cancelled", 1, "任务已由火山引擎取消", map[string]any{"taskId": cancelledTask.ID, "reason": "upstream_cancel"}, cancelledTask.RunMode == "simulation"); err != nil {
return TaskCancelResult{}, err
}
return TaskCancelResult{TaskID: cancelledTask.ID, Cancelled: true, Cancellable: true, Submitted: true, Message: "任务已由火山引擎取消"}, nil
}
func isVolcesCancellationCandidate(candidate store.RuntimeModelCandidate) bool {
provider := strings.ToLower(strings.TrimSpace(candidate.Provider))
return provider == "volces" || provider == "volces-openai"
}
func taskCancelUnavailable(task store.GatewayTask, message string) TaskCancelResult {
return TaskCancelResult{
TaskID: task.ID,
@@ -155,6 +211,13 @@ func taskAccessibleToUser(task store.GatewayTask, user *auth.User) bool {
return userID != "" && strings.TrimSpace(task.UserID) == userID
}
// TaskAccessibleToUser applies the same ownership boundary to every public
// task surface. API-key requests are deliberately scoped to the exact key,
// while JWT requests may see the current gateway user's tasks.
func TaskAccessibleToUser(task store.GatewayTask, user *auth.User) bool {
return taskAccessibleToUser(task, user)
}
func gatewayUserIDForAuth(user *auth.User) string {
if user == nil {
return ""
+80 -3
View File
@@ -8,10 +8,12 @@ import (
"encoding/base64"
"encoding/hex"
"encoding/json"
"errors"
"fmt"
"io"
"mime"
"mime/multipart"
"net"
"net/http"
"net/textproto"
"net/url"
@@ -85,6 +87,11 @@ func (s *Service) uploadGeneratedAssets(ctx context.Context, taskID string, task
if err != nil {
return nil, &clients.ClientError{Code: "upload_config_failed", Message: err.Error(), Retryable: true}
}
// Topaz download URLs are short-lived. Persist them before the task can be
// marked succeeded even when the global policy normally keeps URL media.
if taskKind == "videos.upscales" {
policy.UploadURLMedia = true
}
decisions := make([]generatedAssetDecision, len(data))
needsUpload := false
changed := false
@@ -723,7 +730,8 @@ func (s *Service) readGeneratedURLAsset(ctx context.Context, asset *generatedURL
if err != nil {
return nil, "", err
}
resp, err := http.DefaultClient.Do(req)
allowLocal := strings.HasPrefix(strings.TrimSpace(asset.URL), "/")
resp, err := generatedAssetHTTPClient(allowLocal).Do(req)
if err != nil {
return nil, "", &clients.ClientError{Code: "upload_source_fetch_failed", Message: err.Error(), Retryable: true}
}
@@ -752,6 +760,47 @@ func (s *Service) readGeneratedURLAsset(ctx context.Context, asset *generatedURL
return payload, strings.TrimSpace(strings.Split(contentType, ";")[0]), nil
}
func generatedAssetHTTPClient(allowLocal bool) *http.Client {
transport := http.DefaultTransport.(*http.Transport).Clone()
transport.Proxy = nil
transport.DialContext = func(ctx context.Context, network, address string) (net.Conn, error) {
host, port, err := net.SplitHostPort(address)
if err != nil {
return nil, err
}
addresses, err := net.DefaultResolver.LookupIPAddr(ctx, host)
if err != nil || len(addresses) == 0 {
return nil, errors.New("generated media DNS resolution failed")
}
for _, address := range addresses {
if generatedAssetBlockedAddress(address.IP) && !(allowLocal && address.IP.IsLoopback()) {
return nil, errors.New("generated media resolved to a blocked network")
}
}
dialer := &net.Dialer{Timeout: 10 * time.Second}
var attempts []error
for _, resolved := range addresses {
connection, dialErr := dialer.DialContext(ctx, network, net.JoinHostPort(resolved.IP.String(), port))
if dialErr == nil {
return connection, nil
}
attempts = append(attempts, dialErr)
}
return nil, errors.Join(attempts...)
}
return &http.Client{
Transport: transport,
Timeout: 10 * time.Minute,
CheckRedirect: func(*http.Request, []*http.Request) error {
return http.ErrUseLastResponse
},
}
}
func generatedAssetBlockedAddress(ip net.IP) bool {
return ip == nil || ip.IsLoopback() || ip.IsPrivate() || ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() || ip.IsUnspecified() || ip.IsMulticast()
}
func (s *Service) generatedAssetFetchURL(raw string) (string, error) {
value := strings.TrimSpace(raw)
if value == "" {
@@ -1304,6 +1353,9 @@ func mediaContentTypeFromItem(item map[string]any) string {
func mediaKindForAsset(taskKind string, item map[string]any, sourceKey string, contentType string) string {
contentType = strings.ToLower(strings.TrimSpace(contentType))
if generatedContentTypeIsDocument(contentType) {
return "file"
}
if strings.HasPrefix(contentType, "image/") {
return "image"
}
@@ -1314,6 +1366,9 @@ func mediaKindForAsset(taskKind string, item map[string]any, sourceKey string, c
return "audio"
}
itemType := strings.ToLower(strings.TrimSpace(stringFromAny(item["type"])))
if itemType == "file" || strings.Contains(itemType, "document") {
return "file"
}
if strings.Contains(itemType, "video") {
return "video"
}
@@ -1346,6 +1401,8 @@ func defaultContentTypeForGeneratedAsset(kind string) string {
return "video/mp4"
case "audio":
return "audio/mpeg"
case "file":
return "application/octet-stream"
default:
return "image/png"
}
@@ -1354,10 +1411,10 @@ func defaultContentTypeForGeneratedAsset(kind string) string {
func resolvedGeneratedAssetContentType(declared string, kind string, payload []byte) string {
declared = normalizeGeneratedContentType(declared)
detected := detectGeneratedAssetContentType(payload)
if generatedContentTypeIsMedia(detected) {
if generatedContentTypeIsMedia(detected) || generatedContentTypeIsDocument(detected) {
return detected
}
if generatedContentTypeIsMedia(declared) {
if generatedContentTypeIsMedia(declared) || generatedContentTypeIsDocument(declared) {
return declared
}
return defaultContentTypeForGeneratedAsset(kind)
@@ -1380,6 +1437,15 @@ func generatedContentTypeIsMedia(contentType string) bool {
strings.HasPrefix(contentType, "audio/")
}
func generatedContentTypeIsDocument(contentType string) bool {
switch normalizeGeneratedContentType(contentType) {
case "application/pdf", "application/postscript", "application/eps", "application/dxf", "image/vnd.dxf":
return true
default:
return false
}
}
func generatedAssetKindFromContentType(fallback string, contentType string) string {
contentType = normalizeGeneratedContentType(contentType)
if strings.HasPrefix(contentType, "image/") {
@@ -1391,6 +1457,9 @@ func generatedAssetKindFromContentType(fallback string, contentType string) stri
if strings.HasPrefix(contentType, "audio/") {
return "audio"
}
if generatedContentTypeIsDocument(contentType) {
return "file"
}
fallback = strings.ToLower(strings.TrimSpace(fallback))
if fallback != "" {
return fallback
@@ -1434,6 +1503,14 @@ func randomHexSuffix(byteCount int) string {
func fileExtensionForContentType(contentType string, kind string) string {
normalized := strings.ToLower(strings.TrimSpace(strings.Split(contentType, ";")[0]))
switch normalized {
case "application/pdf":
return ".pdf"
case "application/postscript", "application/eps":
return ".eps"
case "application/dxf", "image/vnd.dxf":
return ".dxf"
}
switch normalized {
case "image/jpeg", "image/jpg":
return ".jpg"
case "image/webp":
+19
View File
@@ -52,6 +52,25 @@ func TestGeneratedAssetDecisionUploadsInlineImageBase64(t *testing.T) {
}
}
func TestGeneratedAssetDecisionUploadsVectorDocumentWithoutChangingType(t *testing.T) {
item := map[string]any{
"type": "file",
"b64_json": base64.StdEncoding.EncodeToString([]byte("%PDF-1.7\n")),
"mime_type": "application/pdf",
}
decision, err := generatedAssetDecisionForItem("images.vectorize", item, defaultGeneratedAssetUploadPolicy())
if err != nil {
t.Fatal(err)
}
if decision.Inline == nil || decision.Inline.Kind != "file" || decision.Inline.ContentType != "application/pdf" {
t.Fatalf("unexpected vector document decision: %+v", decision)
}
contentType := resolvedGeneratedAssetContentType(decision.Inline.ContentType, decision.Inline.Kind, decision.Inline.Bytes)
if contentType != "application/pdf" || fileExtensionForContentType(contentType, "file") != ".pdf" {
t.Fatalf("vector document type changed: contentType=%s", contentType)
}
}
func TestGeneratedAssetDecisionUploadsInlineVideoBuffer(t *testing.T) {
item := map[string]any{
"type": "video",
+4
View File
@@ -215,6 +215,10 @@ func (s *Service) DeleteClonedVoice(ctx context.Context, user *auth.User, rawID
if !ok {
return DeletedClonedVoiceResult{}, &clients.ClientError{Code: "cloned_voice_platform_unavailable", Message: "cloned voice platform binding is unavailable", StatusCode: 400, Retryable: false}
}
candidate, err = candidateWithEnvironmentCredentials(candidate)
if err != nil {
return DeletedClonedVoiceResult{}, err
}
requestHTTPClient, err := s.httpClientForCandidate(candidate, false)
if err != nil {
return DeletedClonedVoiceResult{}, err
@@ -18,7 +18,7 @@ Use this skill to operate AI Gateway administration APIs through documented, evi
- Reuse existing pricing rules, runtime policy sets, providers, protocol clients, base models, and platforms whenever their effective behavior satisfies the target. Do not create a near-duplicate resource merely because the upstream account, base URL, or provider-side model name differs.
- Prefer a supported standard client before using `universal` scripts. Use custom scripts only when the upstream contract cannot be represented by the existing OpenAI, Gemini, or provider-specific clients.
- Do not invent platform config fields or assume an arbitrary config key is enforced. For `universal`, use only the recognized keys documented in `references/model-universal-platforms.md`; treat any extra key as script-owned data available through `context.env`.
- Use the module references as the primary API source. Only when the required API is absent, inspect `<gateway-api-base-url>/api-docs-json`; continue only when path, method, schema, authentication, permission, and side effects are unambiguous.
- Use the module references as the primary API source. Only when the required API is absent, inspect `<gateway-origin>/api/v1/openapi.json`; continue only when path, method, schema, authentication, permission, and side effects are unambiguous.
## Module Routing
@@ -2,7 +2,7 @@
## Required Inputs
- Gateway API base URL. When using the bundled Web deployment this commonly includes `/gateway-api`; direct API access commonly uses port `8088`.
- Gateway origin and public API base URL. Public API access always ends with `/api/v1`; direct local access commonly uses `http://127.0.0.1:8088/api/v1`.
- Administrator JWT with the `manager` or `admin` role.
- Target provider documentation and authorization material.
- Clear requested outcome and whether real upstream calls are allowed.
@@ -10,7 +10,8 @@
Do not place credentials in files or reusable commands. Use shell environment variables:
```bash
export GATEWAY_BASE_URL='https://gateway.example.com/gateway-api'
export GATEWAY_ORIGIN='https://gateway.example.com'
export GATEWAY_PUBLIC_API_BASE="$GATEWAY_ORIGIN/api/v1"
export GATEWAY_ADMIN_TOKEN='<administrator-jwt>'
```
@@ -24,7 +25,7 @@ For standalone or hybrid deployments, local login can return a JWT:
curl --fail-with-body \
-H 'Content-Type: application/json' \
-d '{"account":"<admin-account>","password":"<admin-password>"}' \
"$GATEWAY_BASE_URL/api/v1/auth/login"
"$GATEWAY_PUBLIC_API_BASE/auth/login"
```
Do not use local login when the deployment requires OIDC or server-main identity. Obtain the deployment's administrator access token instead.
@@ -34,7 +35,7 @@ Verify identity and role before writes:
```bash
curl --fail-with-body \
-H "Authorization: Bearer $GATEWAY_ADMIN_TOKEN" \
"$GATEWAY_BASE_URL/api/v1/me"
"$GATEWAY_PUBLIC_API_BASE/me"
```
## Request Pattern
@@ -47,7 +48,7 @@ curl --fail-with-body \
-H 'Content-Type: application/json' \
-X POST \
-d '<json-body>' \
"$GATEWAY_BASE_URL/api/admin/<resource>"
"$GATEWAY_ORIGIN/api/admin/<resource>"
```
Always read current state before PATCH, DELETE, reset, disable, or full replacement. PATCH handlers for providers, base models, pricing rule sets, runtime policy sets, runner policy, and platforms write complete resource shapes rather than merging every omitted field.
@@ -67,7 +68,7 @@ Obtain explicit confirmation after showing the current snapshot and impact befor
The live machine-readable documents are:
- `<gateway-api-base-url>/api-docs-json`
- `<gateway-api-base-url>/api-docs-yaml`
- `<gateway-origin>/api/v1/openapi.json`
- `<gateway-origin>/api/v1/openapi.yaml`
Use them only when this Skill does not document the required operation. Before acting, confirm the exact path, method, body, authentication, permission, response, and side effect. Do not infer a write operation from a similarly named endpoint.
@@ -33,7 +33,7 @@ Use an authorized user JWT:
```bash
curl --fail-with-body \
-H "Authorization: Bearer <user-jwt>" \
"$GATEWAY_BASE_URL/api/v1/model-catalog"
"$GATEWAY_PUBLIC_API_BASE/model-catalog"
```
Confirm model alias, model types, provider source, effective capabilities, pricing summary, rate limits, permissions, and enabled state.
@@ -57,7 +57,7 @@ curl --fail-with-body \
"simulation": true,
"stream": false
}' \
"$GATEWAY_BASE_URL/v1/chat/completions"
"$GATEWAY_PUBLIC_API_BASE/chat/completions"
```
Simulation verifies Gateway routing, permissions, parameter normalization, pricing, and task behavior, but it does not execute the real universal submit or poll scripts. Validate universal scripts separately against a local mock or approved provider test environment before enabling production traffic.
@@ -82,4 +82,4 @@ With explicit approval, run one real minimal request and verify upstream request
## Final Report
Report resource IDs, before/after behavior, requests used for verification, simulation or real mode, billing evidence, remaining risks, rollback readiness, and whether `/api-docs-json` was used. Never include credentials or raw secret-bearing payloads.
Report resource IDs, before/after behavior, requests used for verification, simulation or real mode, billing evidence, remaining risks, rollback readiness, and whether `/api/v1/openapi.json` was used. Never include credentials or raw secret-bearing payloads.
@@ -166,16 +166,16 @@ The response has this shape:
"platformModelId": "<platform-model-id>",
"resourceType": "video",
"unit": "5s_video",
"quantity": 3,
"quantity": 2.4,
"amount": 12.5,
"currency": "resource",
"discountFactor": 0.8,
"simulated": true,
"durationSeconds": 12,
"durationUnitCount": 3
"durationUnitCount": 2.4
}
],
"resolver": "effective-pricing-v1",
"resolver": "effective-pricing-v2",
"totalAmount": 12.5,
"currency": "resource"
}
@@ -197,7 +197,7 @@ Useful calculation checks:
text input = input tokens / 1000 × input price × discount
text output = output tokens / 1000 × output price × discount
image = count × base price × quality/size/resolution weights × discount
video = count × ceil(duration seconds / 5) × base price × applicable weights × discount
video = count × (duration seconds / 5) × base price × applicable weights × discount
speech = Unicode character count × audio price × discount
```
+40 -3
View File
@@ -283,6 +283,21 @@ func (s *Store) filterCandidatesByAccessRules(ctx context.Context, user *auth.Us
}
func (s *Store) ListAccessiblePlatformModels(ctx context.Context, user *auth.User) ([]PlatformModel, error) {
return s.listPlatformModelsForAccessRules(ctx, user, nil)
}
// ListAPIKeyAssignablePlatformModels returns the enabled models that the
// current user may delegate to their API keys. API-key rules are deliberately
// excluded here: they restrict individual credentials and must not shrink the
// resource pool that the owning user can manage.
func (s *Store) ListAPIKeyAssignablePlatformModels(ctx context.Context, user *auth.User) ([]PlatformModel, error) {
if localGatewayUserID(user) == "" {
return nil, ErrLocalUserRequired
}
return s.listPlatformModelsForAccessRules(ctx, user, map[string]bool{"api_key": true})
}
func (s *Store) listPlatformModelsForAccessRules(ctx context.Context, user *auth.User, excludedSubjectTypes map[string]bool) ([]PlatformModel, error) {
accessUser, err := s.resolveCurrentAccessUser(ctx, user)
if err != nil {
return nil, err
@@ -307,7 +322,7 @@ func (s *Store) ListAccessiblePlatformModels(ctx context.Context, user *auth.Use
enabled = append(enabled, model)
}
}
return s.filterPlatformModelsByAccessRules(ctx, accessUser, enabled)
return s.filterPlatformModelsByAccessRules(ctx, accessUser, enabled, excludedSubjectTypes)
}
func (s *Store) ensureAPIKeyAccessRuleResourcesAllowed(ctx context.Context, user *auth.User, resources []AccessRuleResourceInput) error {
@@ -328,7 +343,7 @@ func (s *Store) ensureAPIKeyAccessRuleResourcesAllowed(ctx context.Context, user
}
func (s *Store) accessibleAccessRuleResources(ctx context.Context, user *auth.User) (map[string]bool, error) {
models, err := s.ListAccessiblePlatformModels(ctx, user)
models, err := s.ListAPIKeyAssignablePlatformModels(ctx, user)
if err != nil {
return nil, err
}
@@ -383,7 +398,12 @@ WHERE id = $1::uuid
return &next, nil
}
func (s *Store) filterPlatformModelsByAccessRules(ctx context.Context, user *auth.User, models []PlatformModel) ([]PlatformModel, error) {
func (s *Store) filterPlatformModelsByAccessRules(
ctx context.Context,
user *auth.User,
models []PlatformModel,
excludedSubjectTypes map[string]bool,
) ([]PlatformModel, error) {
if len(models) == 0 {
return models, nil
}
@@ -398,6 +418,12 @@ func (s *Store) filterPlatformModelsByAccessRules(ctx context.Context, user *aut
if len(rules) == 0 {
return models, nil
}
if len(excludedSubjectTypes) > 0 {
rules = filterAccessRulesBySubjectType(rules, excludedSubjectTypes)
if len(rules) == 0 {
return models, nil
}
}
subjects := accessRuleSubjects(user)
level := 0
if user != nil {
@@ -412,6 +438,17 @@ func (s *Store) filterPlatformModelsByAccessRules(ctx context.Context, user *aut
return filtered, nil
}
func filterAccessRulesBySubjectType(rules []AccessRule, excludedSubjectTypes map[string]bool) []AccessRule {
filtered := make([]AccessRule, 0, len(rules))
for _, rule := range rules {
if excludedSubjectTypes[rule.SubjectType] {
continue
}
filtered = append(filtered, rule)
}
return filtered
}
func (s *Store) listActiveAccessRulesForResources(ctx context.Context, resources []accessRuleResource) ([]AccessRule, error) {
values := make([]string, 0, len(resources))
for _, resource := range resources {
@@ -0,0 +1,23 @@
package store
import "testing"
func TestFilterAccessRulesBySubjectTypeExcludesAPIKeyRulesOnly(t *testing.T) {
rules := []AccessRule{
{ID: "api-key-allow", SubjectType: "api_key", Effect: "allow"},
{ID: "api-key-deny", SubjectType: "api_key", Effect: "deny"},
{ID: "user-group-allow", SubjectType: "user_group", Effect: "allow"},
{ID: "user-deny", SubjectType: "user", Effect: "deny"},
{ID: "tenant-allow", SubjectType: "tenant", Effect: "allow"},
}
filtered := filterAccessRulesBySubjectType(rules, map[string]bool{"api_key": true})
if len(filtered) != 3 {
t.Fatalf("filtered rule count = %d, want 3: %+v", len(filtered), filtered)
}
for _, rule := range filtered {
if rule.SubjectType == "api_key" {
t.Fatalf("api-key rule should not affect the owning user's assignable resources: %+v", rule)
}
}
}
@@ -0,0 +1,99 @@
package store
import (
"context"
"sync"
"testing"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
"github.com/jackc/pgx/v5/pgxpool"
)
func TestVerifyLocalAPIKeyWorksWithSingleConnectionPool(t *testing.T) {
db, verificationStore, created, user := newLocalAPIKeyVerificationFixture(t, 1)
ctx := context.Background()
verifyCtx, cancel := context.WithTimeout(ctx, 2*time.Second)
defer cancel()
verified, err := verificationStore.VerifyLocalAPIKey(verifyCtx, created.Secret)
if err != nil {
t.Fatalf("verify local API key with one connection: %v", err)
}
if verified.APIKeyID != created.APIKey.ID || verified.GatewayUserID != user.ID {
t.Fatalf("verified identity = %+v, want API key %q and user %q", verified, created.APIKey.ID, user.ID)
}
var lastUsedAt *time.Time
if err := db.pool.QueryRow(ctx, `SELECT last_used_at FROM gateway_api_keys WHERE id=$1::uuid`, created.APIKey.ID).Scan(&lastUsedAt); err != nil {
t.Fatalf("read API key last_used_at: %v", err)
}
if lastUsedAt == nil {
t.Fatal("successful API key verification did not update last_used_at")
}
}
func TestVerifyLocalAPIKeyHandlesEightConcurrentRequestsWithFourConnections(t *testing.T) {
_, verificationStore, created, _ := newLocalAPIKeyVerificationFixture(t, 4)
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
const requestCount = 8
start := make(chan struct{})
errorsByRequest := make(chan error, requestCount)
var requests sync.WaitGroup
requests.Add(requestCount)
for range requestCount {
go func() {
defer requests.Done()
<-start
_, err := verificationStore.VerifyLocalAPIKey(ctx, created.Secret)
errorsByRequest <- err
}()
}
close(start)
requests.Wait()
close(errorsByRequest)
for err := range errorsByRequest {
if err != nil {
t.Fatalf("concurrent API key verification failed: %v", err)
}
}
if acquired := verificationStore.pool.Stat().AcquiredConns(); acquired != 0 {
t.Fatalf("API key verification left %d connections acquired", acquired)
}
}
func newLocalAPIKeyVerificationFixture(t *testing.T, maxConnections int32) (*Store, *Store, CreatedAPIKey, GatewayUser) {
t.Helper()
db := newIdentityPairingPostgresTestStore(t)
ctx := context.Background()
user, err := db.RegisterLocalUser(ctx, LocalRegisterInput{
Username: "api-key-verification-user",
Password: "api-key-verification-password",
})
if err != nil {
t.Fatalf("register API key verification user: %v", err)
}
created, err := db.CreateAPIKey(ctx, CreateAPIKeyInput{Name: "API key verification fixture"}, &auth.User{
ID: user.ID,
GatewayUserID: user.ID,
GatewayTenantID: user.GatewayTenantID,
TenantID: user.TenantID,
TenantKey: user.TenantKey,
})
if err != nil {
t.Fatalf("create API key verification fixture: %v", err)
}
config := db.pool.Config()
config.MaxConns = maxConnections
config.MinConns = 0
pool, err := pgxpool.NewWithConfig(ctx, config)
if err != nil {
t.Fatalf("create verification pool with %d connections: %v", maxConnections, err)
}
t.Cleanup(pool.Close)
return db, &Store{pool: pool}, created, user
}
@@ -0,0 +1,169 @@
package store
import (
"context"
"errors"
"testing"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgconn"
"golang.org/x/crypto/bcrypt"
)
func TestVerifyLocalAPIKeyClosesCandidateRowsBeforeUpdatingUsage(t *testing.T) {
secret := "sk-gw-matching-secret"
wrongHash, err := bcrypt.GenerateFromPassword([]byte("sk-gw-different-secret"), bcrypt.MinCost)
if err != nil {
t.Fatalf("hash non-matching API key: %v", err)
}
matchingHash, err := bcrypt.GenerateFromPassword([]byte(secret), bcrypt.MinCost)
if err != nil {
t.Fatalf("hash matching API key: %v", err)
}
rows := &fakeLocalAPIKeyRows{candidates: []localAPIKeyCandidate{
{apiKeyID: "wrong-key", hash: string(wrongHash), keyPrefix: apiKeyPrefix(secret)},
{
apiKeyID: "matching-key",
hash: string(matchingHash),
keyPrefix: apiKeyPrefix(secret),
keyName: "Matching key",
scopesBytes: []byte(`["chat"]`),
userGroupID: "group-id",
gatewayUserID: "user-id",
username: "api-key-user",
rolesBytes: []byte(`["user"]`),
gatewayTenantID: "gateway-tenant-id",
tenantID: "tenant-id",
tenantKey: "tenant-key",
},
}}
database := &fakeLocalAPIKeyDatabase{rows: rows}
user, err := verifyLocalAPIKey(context.Background(), database, secret)
if err != nil {
t.Fatalf("verify local API key: %v", err)
}
if !rows.closed {
t.Fatal("candidate rows remained open after API key verification")
}
if database.updatedAPIKeyID != "matching-key" {
t.Fatalf("updated API key = %q, want matching-key", database.updatedAPIKeyID)
}
if user.APIKeyID != "matching-key" || user.GatewayUserID != "user-id" {
t.Fatalf("verified user = %+v", user)
}
}
func TestVerifyLocalAPIKeyReturnsUnauthorizedAfterClosingCandidateRows(t *testing.T) {
secret := "sk-gw-unknown-secret"
wrongHash, err := bcrypt.GenerateFromPassword([]byte("sk-gw-different-secret"), bcrypt.MinCost)
if err != nil {
t.Fatalf("hash non-matching API key: %v", err)
}
rows := &fakeLocalAPIKeyRows{candidates: []localAPIKeyCandidate{{
apiKeyID: "wrong-key",
hash: string(wrongHash),
}}}
database := &fakeLocalAPIKeyDatabase{rows: rows}
_, err = verifyLocalAPIKey(context.Background(), database, secret)
if !errors.Is(err, auth.ErrUnauthorized) {
t.Fatalf("verify error = %v, want unauthorized", err)
}
if !rows.closed {
t.Fatal("candidate rows remained open after unsuccessful API key verification")
}
if database.updatedAPIKeyID != "" {
t.Fatalf("unexpected API key usage update for %q", database.updatedAPIKeyID)
}
}
type fakeLocalAPIKeyDatabase struct {
rows *fakeLocalAPIKeyRows
updatedAPIKeyID string
}
func (database *fakeLocalAPIKeyDatabase) Query(context.Context, string, ...any) (pgx.Rows, error) {
return database.rows, nil
}
func (database *fakeLocalAPIKeyDatabase) Exec(_ context.Context, _ string, arguments ...any) (pgconn.CommandTag, error) {
if !database.rows.closed {
return pgconn.CommandTag{}, errors.New("API key usage update started before candidate rows closed")
}
database.updatedAPIKeyID, _ = arguments[0].(string)
return pgconn.NewCommandTag("UPDATE 1"), nil
}
type fakeLocalAPIKeyRows struct {
candidates []localAPIKeyCandidate
current int
closed bool
}
func (rows *fakeLocalAPIKeyRows) Close() {
rows.closed = true
}
func (rows *fakeLocalAPIKeyRows) Err() error {
return nil
}
func (rows *fakeLocalAPIKeyRows) CommandTag() pgconn.CommandTag {
return pgconn.CommandTag{}
}
func (rows *fakeLocalAPIKeyRows) FieldDescriptions() []pgconn.FieldDescription {
return nil
}
func (rows *fakeLocalAPIKeyRows) Next() bool {
if rows.current >= len(rows.candidates) {
rows.Close()
return false
}
rows.current++
return true
}
func (rows *fakeLocalAPIKeyRows) Scan(destinations ...any) error {
candidate := rows.candidates[rows.current-1]
values := []any{
candidate.apiKeyID,
candidate.hash,
candidate.keyPrefix,
candidate.keyName,
candidate.scopesBytes,
candidate.userGroupID,
candidate.gatewayUserID,
candidate.username,
candidate.rolesBytes,
candidate.gatewayTenantID,
candidate.tenantID,
candidate.tenantKey,
}
for index, value := range values {
switch destination := destinations[index].(type) {
case *string:
*destination = value.(string)
case *[]byte:
*destination = value.([]byte)
default:
return errors.New("unsupported fake row destination")
}
}
return nil
}
func (rows *fakeLocalAPIKeyRows) Values() ([]any, error) {
return nil, errors.New("not implemented")
}
func (rows *fakeLocalAPIKeyRows) RawValues() [][]byte {
return nil
}
func (rows *fakeLocalAPIKeyRows) Conn() *pgx.Conn {
return nil
}
+6
View File
@@ -486,8 +486,14 @@ func modelTypeAliases(value string) []string {
return []string{"image_generate"}
case "images.edits":
return []string{"image_edit"}
case "images.vectorize", "vectorize":
return []string{"image_vectorize"}
case "video", "videos.generations":
return []string{"video_generate"}
case "videos.upscales", "video_upscale":
return []string{"video_enhance"}
case "omni_video":
return []string{"video_generate", "image_to_video", "omni_video"}
case "song", "music", "song.generations", "music.generations", "music_generate":
return []string{"audio_generate"}
case "speech", "speech.generations", "tts":
+32
View File
@@ -0,0 +1,32 @@
package store
// EffectiveBillingConfigInput describes the billing layers used by runtime and
// catalog responses. LegacyPlatformModelConfig is retained only as a fallback
// for models that do not have an effective pricing rule set.
type EffectiveBillingConfigInput struct {
BaseConfig map[string]any
LegacyPlatformModelConfig map[string]any
InheritedRuleSetConfig map[string]any
ModelRuleSetConfig map[string]any
Override map[string]any
}
// ResolveEffectiveBillingConfig keeps inherited pricing rules authoritative over
// the legacy materialized snapshot. Explicit model rules and overrides retain
// their higher-priority exception semantics.
func ResolveEffectiveBillingConfig(input EffectiveBillingConfigInput) map[string]any {
config := mergeObjects(input.BaseConfig, nil)
if len(input.InheritedRuleSetConfig) > 0 {
// Rule sets are allowed to cover only a subset of resource types. Keep
// base-model prices for resources that the inherited rule set does not
// define, while letting the rule set remain authoritative for matching
// top-level keys.
config = mergeObjects(config, input.InheritedRuleSetConfig)
} else if len(input.LegacyPlatformModelConfig) > 0 {
config = mergeObjects(config, input.LegacyPlatformModelConfig)
}
if len(input.ModelRuleSetConfig) > 0 {
config = mergeObjects(config, input.ModelRuleSetConfig)
}
return mergeObjects(config, input.Override)
}
@@ -0,0 +1,94 @@
package store
import "testing"
func TestResolveEffectiveBillingConfigKeepsPricingRulesAuthoritative(t *testing.T) {
tests := []struct {
name string
input EffectiveBillingConfigInput
want float64
}{
{
name: "inherited rule replaces stale platform snapshot",
input: EffectiveBillingConfigInput{
BaseConfig: videoBillingConfig(100),
LegacyPlatformModelConfig: videoBillingConfig(100),
InheritedRuleSetConfig: videoBillingConfig(416),
},
want: 416,
},
{
name: "legacy snapshot remains a fallback without a rule",
input: EffectiveBillingConfigInput{
BaseConfig: videoBillingConfig(100),
LegacyPlatformModelConfig: videoBillingConfig(125),
},
want: 125,
},
{
name: "model rule remains an explicit pricing exception",
input: EffectiveBillingConfigInput{
BaseConfig: videoBillingConfig(100),
LegacyPlatformModelConfig: videoBillingConfig(125),
InheritedRuleSetConfig: videoBillingConfig(416),
ModelRuleSetConfig: videoBillingConfig(500),
},
want: 500,
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
config := ResolveEffectiveBillingConfig(test.input)
video, ok := config["video"].(map[string]any)
if !ok {
t.Fatalf("expected video billing config, got %#v", config)
}
if got := video["basePrice"]; got != test.want {
t.Fatalf("video base price = %#v, want %v", got, test.want)
}
})
}
}
func TestResolveEffectiveBillingConfigAppliesOverrideLast(t *testing.T) {
config := ResolveEffectiveBillingConfig(EffectiveBillingConfigInput{
InheritedRuleSetConfig: videoBillingConfig(416),
Override: videoBillingConfig(600),
})
video, ok := config["video"].(map[string]any)
if !ok || video["basePrice"] != float64(600) {
t.Fatalf("expected override price 600, got %#v", config)
}
}
func TestResolveEffectiveBillingConfigPreservesBaseResourcesMissingFromRuleSet(t *testing.T) {
config := ResolveEffectiveBillingConfig(EffectiveBillingConfigInput{
BaseConfig: map[string]any{
"music": map[string]any{"basePrice": float64(20)},
"audio": map[string]any{"basePrice": float64(1)},
"video": map[string]any{"basePrice": float64(100)},
},
InheritedRuleSetConfig: map[string]any{
"video": map[string]any{"basePrice": float64(416)},
},
})
assertBillingBasePrice(t, config, "music", 20)
assertBillingBasePrice(t, config, "audio", 1)
assertBillingBasePrice(t, config, "video", 416)
}
func assertBillingBasePrice(t *testing.T, config map[string]any, resource string, want float64) {
t.Helper()
resourceConfig, ok := config[resource].(map[string]any)
if !ok || resourceConfig["basePrice"] != want {
t.Fatalf("%s base price = %#v, want %v", resource, config[resource], want)
}
}
func videoBillingConfig(basePrice float64) map[string]any {
return map[string]any{
"video": map[string]any{"basePrice": basePrice},
}
}
@@ -9,6 +9,19 @@ func TestNormalizeModelMatchKeyRemovesWhitespace(t *testing.T) {
}
}
func TestNormalizeModelTypeListExpandsOmniVideoBaseCapabilities(t *testing.T) {
got := normalizeModelTypeList([]string{"omni_video"})
want := StringList{"video_generate", "image_to_video", "omni_video"}
if len(got) != len(want) {
t.Fatalf("omni_video should include text-to-video and image-to-video capabilities: got=%v want=%v", got, want)
}
for index := range want {
if got[index] != want[index] {
t.Fatalf("omni_video capability mismatch at %d: got=%v want=%v", index, got, want)
}
}
}
func TestTaskBillingModelIdentityKeepsRequestedModelPrimary(t *testing.T) {
identity := taskBillingModelIdentity(GatewayTask{
Model: "doubao-5.0 图像编辑",
+68
View File
@@ -0,0 +1,68 @@
package store
import (
"context"
"strings"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
)
type DailyTokenUsage struct {
Date string `json:"date"`
TotalTokens int64 `json:"totalTokens"`
InputTokens int64 `json:"inputTokens"`
OutputTokens int64 `json:"outputTokens"`
CachedInputTokens int64 `json:"cachedInputTokens"`
ResourcePoints float64 `json:"resourcePoints"`
TaskCount int64 `json:"taskCount"`
}
func (s *Store) ListDailyTokenUsage(ctx context.Context, user *auth.User, from, to time.Time, timezone string) ([]DailyTokenUsage, error) {
gatewayUserID := localGatewayUserID(user)
apiKeyID := ""
userID := ""
if user != nil {
apiKeyID = strings.TrimSpace(user.APIKeyID)
userID = strings.TrimSpace(user.ID)
}
if gatewayUserID == "" && userID == "" {
return nil, ErrLocalUserRequired
}
rows, err := s.pool.Query(ctx, `
SELECT to_char(created_at AT TIME ZONE $6, 'YYYY-MM-DD') AS usage_day,
COALESCE(sum(
CASE WHEN jsonb_typeof(usage->'totalTokens') = 'number' THEN (usage->>'totalTokens')::numeric
WHEN jsonb_typeof(usage->'total_tokens') = 'number' THEN (usage->>'total_tokens')::numeric
ELSE COALESCE(CASE WHEN jsonb_typeof(usage->'inputTokens') = 'number' THEN (usage->>'inputTokens')::numeric ELSE 0 END, 0)
+ COALESCE(CASE WHEN jsonb_typeof(usage->'outputTokens') = 'number' THEN (usage->>'outputTokens')::numeric ELSE 0 END, 0)
END
), 0)::bigint AS total_tokens,
COALESCE(sum(CASE WHEN jsonb_typeof(usage->'inputTokens') = 'number' THEN (usage->>'inputTokens')::numeric WHEN jsonb_typeof(usage->'input_tokens') = 'number' THEN (usage->>'input_tokens')::numeric ELSE 0 END), 0)::bigint,
COALESCE(sum(CASE WHEN jsonb_typeof(usage->'outputTokens') = 'number' THEN (usage->>'outputTokens')::numeric WHEN jsonb_typeof(usage->'output_tokens') = 'number' THEN (usage->>'output_tokens')::numeric ELSE 0 END), 0)::bigint,
COALESCE(sum(CASE WHEN jsonb_typeof(usage->'cachedInputTokens') = 'number' THEN (usage->>'cachedInputTokens')::numeric WHEN jsonb_typeof(usage->'cached_input_tokens') = 'number' THEN (usage->>'cached_input_tokens')::numeric ELSE 0 END), 0)::bigint,
COALESCE(sum(final_charge_amount), 0)::float8,
count(*)::bigint
FROM gateway_tasks
WHERE status = 'succeeded'
AND ((NULLIF($1, '')::uuid IS NOT NULL AND gateway_user_id = NULLIF($1, '')::uuid)
OR (NULLIF($1, '')::uuid IS NULL AND NULLIF($2, '') IS NOT NULL AND user_id = $2))
AND (NULLIF($3, '') IS NULL OR api_key_id = $3)
AND created_at >= ($4::date::timestamp AT TIME ZONE $6)
AND created_at < (($5::date + 1)::timestamp AT TIME ZONE $6)
GROUP BY usage_day
ORDER BY usage_day`, gatewayUserID, userID, apiKeyID, from.Format("2006-01-02"), to.Format("2006-01-02"), timezone)
if err != nil {
return nil, err
}
defer rows.Close()
items := make([]DailyTokenUsage, 0)
for rows.Next() {
var item DailyTokenUsage
if err := rows.Scan(&item.Date, &item.TotalTokens, &item.InputTokens, &item.OutputTokens, &item.CachedInputTokens, &item.ResourcePoints, &item.TaskCount); err != nil {
return nil, err
}
items = append(items, item)
}
return items, rows.Err()
}
@@ -0,0 +1,79 @@
package store
import (
"context"
"fmt"
"os"
"strings"
"testing"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
)
func TestListDailyTokenUsageScopesAPIKeyAndAppliesIANATimezone(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 daily usage PostgreSQL integration tests")
}
ctx := context.Background()
db, err := Connect(ctx, databaseURL)
if err != nil {
t.Fatal(err)
}
defer db.Close()
var databaseName string
if err := db.pool.QueryRow(ctx, `SELECT current_database()`).Scan(&databaseName); err != nil {
t.Fatal(err)
}
if !strings.Contains(strings.ToLower(databaseName), "test") {
t.Fatalf("refusing to use non-test database %q", databaseName)
}
suffix := fmt.Sprint(time.Now().UnixNano())
var userID, otherUserID string
if err := db.pool.QueryRow(ctx, `INSERT INTO gateway_users (user_key, username) VALUES ($1, $2) RETURNING id::text`, "daily-user-"+suffix, "daily-user-"+suffix).Scan(&userID); err != nil {
t.Fatal(err)
}
if err := db.pool.QueryRow(ctx, `INSERT INTO gateway_users (user_key, username) VALUES ($1, $2) RETURNING id::text`, "daily-other-"+suffix, "daily-other-"+suffix).Scan(&otherUserID); err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_tasks WHERE gateway_user_id IN ($1::uuid, $2::uuid)`, userID, otherUserID)
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_users WHERE id IN ($1::uuid, $2::uuid)`, userID, otherUserID)
})
insertTask := func(ownerID, keyID, createdAt string, totalTokens int, points float64) {
t.Helper()
_, err := db.pool.Exec(ctx, `
INSERT INTO gateway_tasks (kind, user_id, gateway_user_id, api_key_id, model, status, usage, final_charge_amount, created_at, finished_at)
VALUES ('chat.completions', $1::text, ($1::text)::uuid, $2::text, 'daily-test-model', 'succeeded', jsonb_build_object('totalTokens', $3::bigint, 'inputTokens', $3::bigint - 1, 'outputTokens', 1), $4::numeric, $5::timestamptz, $5::timestamptz)`, ownerID, keyID, totalTokens, points, createdAt)
if err != nil {
t.Fatal(err)
}
}
insertTask(userID, "daily-key-a-"+suffix, "2026-07-20T15:30:00Z", 10, 1)
insertTask(userID, "daily-key-a-"+suffix, "2026-07-21T16:30:00Z", 20, 2)
insertTask(userID, "daily-key-b-"+suffix, "2026-07-20T16:30:00Z", 100, 10)
insertTask(otherUserID, "daily-key-a-"+suffix, "2026-07-20T16:30:00Z", 999, 99)
location, err := time.LoadLocation("Asia/Shanghai")
if err != nil {
t.Fatal(err)
}
from := time.Date(2026, 7, 20, 0, 0, 0, 0, location)
to := time.Date(2026, 7, 22, 0, 0, 0, 0, location)
keyItems, err := db.ListDailyTokenUsage(ctx, &auth.User{GatewayUserID: userID, APIKeyID: "daily-key-a-" + suffix}, from, to, "Asia/Shanghai")
if err != nil {
t.Fatal(err)
}
if len(keyItems) != 2 || keyItems[0].Date != "2026-07-20" || keyItems[0].TotalTokens != 10 || keyItems[1].Date != "2026-07-22" || keyItems[1].TotalTokens != 20 {
t.Fatalf("API key usage leaked another key/user or ignored timezone: %+v", keyItems)
}
userItems, err := db.ListDailyTokenUsage(ctx, &auth.User{GatewayUserID: userID}, from, to, "Asia/Shanghai")
if err != nil {
t.Fatal(err)
}
if len(userItems) != 3 || userItems[1].Date != "2026-07-21" || userItems[1].TotalTokens != 100 {
t.Fatalf("JWT user usage did not include all owned API keys: %+v", userItems)
}
}
@@ -685,11 +685,13 @@ func newIdentityPairingPostgresTestStore(t *testing.T) *Store {
migrationDirectory := filepath.Join(filepath.Dir(filename), "..", "..", "migrations")
for _, migrationName := range []string{
"0001_init.sql",
"0017_task_record_enrichment.sql",
"0061_oidc_server_sessions.sql",
"0065_identity_configuration_revisions.sql",
"0066_identity_onboarding_exchanges.sql",
"0067_identity_secret_cleanup_queue.sql",
"0068_identity_pairing_start_reservation.sql",
"0069_billing_correctness_v2.sql",
} {
migration, err := os.ReadFile(filepath.Join(migrationDirectory, migrationName))
if err != nil {
@@ -54,9 +54,13 @@ func billingResourcesForModelTypes(modelTypes []string) map[string]bool {
resources["image"] = true
case "images.edits", "image_edit":
resources["image_edit"] = true
case "image_vectorize", "images.vectorize":
resources["image_vectorize"] = true
case "video", "videos.generations", "video_generate", "image_to_video", "text_to_video",
"video_edit", "omni_video", "video_reference", "video_first_last_frame":
resources["video"] = true
case "video_enhance", "videos.upscales", "video_upscale":
resources["video_enhance"] = true
case "audio", "text_to_speech", "speech", "voice_clone":
resources["audio"] = true
case "music", "music_generate", "audio_generate":
@@ -100,8 +104,12 @@ func billingConfigKeyAllowed(key string, resources map[string]bool) bool {
return resources["image"]
case "image_edit", "imageedit", "editbase":
return resources["image_edit"]
case "image_vectorize", "imagevectorize", "vectorizebase":
return resources["image_vectorize"]
case "video", "videobase":
return resources["video"]
case "video_enhance", "videoenhance", "videoenhancebase":
return resources["video_enhance"]
case "audio", "audiobase":
return resources["audio"]
case "music", "musicbase":
+8 -4
View File
@@ -23,6 +23,7 @@ type modelCatalogSnapshot struct {
DisplayName string
Capabilities map[string]any
BaseBillingConfig map[string]any
PricingRuleSetID string
DefaultRateLimitPolicy map[string]any
RuntimePolicySetID string
RuntimePolicyOverride map[string]any
@@ -121,10 +122,10 @@ func (s *Store) createPlatformModel(ctx context.Context, q platformModelQuerier,
if err := validateEnabledVolcesTextModelCapabilities(ctx, q, input, capabilities); err != nil {
return PlatformModel{}, err
}
// billing_config is a legacy, explicitly supplied compatibility field. Do
// not materialize base-model pricing into it: copied prices become stale as
// soon as the base pricing rule changes and can mask the authoritative rule.
billingConfig := input.BillingConfig
if len(billingConfig) == 0 {
billingConfig = mergeObjects(base.BaseBillingConfig, input.BillingConfigOverride)
}
explicitRuntimePolicySetID := strings.TrimSpace(input.RuntimePolicySetID)
rateLimitPolicy := input.RateLimitPolicy
if len(rateLimitPolicy) == 0 && explicitRuntimePolicySetID == "" {
@@ -260,6 +261,8 @@ RETURNING id::text, platform_id::text, COALESCE(base_model_id::text, ''), model_
model.ModelType = decodeStringArray(modelTypeBytes)
model.BillingConfigOverride = decodeObject(billingOverrideBytes)
model.BillingConfig = decodeObject(billingBytes)
model.BaseBillingConfig = base.BaseBillingConfig
model.BasePricingRuleSetID = base.PricingRuleSetID
model.PermissionConfig = decodeObject(permissionBytes)
model.RetryPolicy = decodeObject(retryPolicyBytes)
model.RateLimitPolicy = decodeObject(rateLimitPolicyBytes)
@@ -368,7 +371,7 @@ func (s *Store) lookupBaseModel(ctx context.Context, q platformModelQuerier, id
var modelTypeBytes []byte
err := q.QueryRow(ctx, `
SELECT id::text, provider_key, canonical_model_key, provider_model_name, model_type, display_name,
capabilities, base_billing_config, default_rate_limit_policy,
capabilities, base_billing_config, COALESCE(pricing_rule_set_id::text, ''), default_rate_limit_policy,
COALESCE(runtime_policy_set_id::text, ''), runtime_policy_override
FROM base_model_catalog
WHERE ($1 <> '' AND id = NULLIF($1, '')::uuid)
@@ -384,6 +387,7 @@ LIMIT 1`, strings.TrimSpace(id), strings.TrimSpace(canonicalKey), strings.TrimSp
&item.DisplayName,
&capabilities,
&billingConfig,
&item.PricingRuleSetID,
&rateLimitPolicy,
&item.RuntimePolicySetID,
&runtimePolicyOverride,
@@ -0,0 +1,37 @@
package store
import (
"context"
"os"
"strings"
"testing"
)
func TestListModelsLoadsEffectiveBillingSources(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 platform-model billing source integration test")
}
ctx := context.Background()
db, err := Connect(ctx, databaseURL)
if err != nil {
t.Fatalf("connect store: %v", err)
}
defer db.Close()
models, err := db.ListModels(ctx)
if err != nil {
t.Fatalf("list models with effective billing sources: %v", err)
}
for _, model := range models {
if model.BaseModelID == "" {
continue
}
if model.BaseBillingConfig == nil {
t.Fatalf("platform model %s did not load base billing config", model.ID)
}
return
}
t.Skip("database has no base-model-backed platform model")
}
+326
View File
@@ -0,0 +1,326 @@
package store
import (
"context"
"encoding/json"
"strings"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
)
type PortraitAsset struct {
ID string `json:"id"`
GatewayUserID string `json:"gatewayUserId,omitempty"`
UserID string `json:"userId"`
GatewayTenantID string `json:"gatewayTenantId,omitempty"`
TenantID string `json:"tenantId,omitempty"`
TenantKey string `json:"tenantKey,omitempty"`
Name string `json:"name"`
Description string `json:"description,omitempty"`
SourceType string `json:"sourceType"`
URL string `json:"url"`
Preview string `json:"preview,omitempty"`
MimeType string `json:"mimeType,omitempty"`
ByteSize int64 `json:"size,omitempty"`
SourceSHA256 string `json:"sourceSha256,omitempty"`
PrivateAvatarEligible bool `json:"privateAvatarEligible"`
Status string `json:"status"`
LastError string `json:"lastError,omitempty"`
Metadata map[string]any `json:"metadata,omitempty"`
CreatedAt time.Time `json:"createdAt"`
UpdatedAt time.Time `json:"updatedAt"`
}
type PortraitAssetBinding struct {
ID string `json:"id"`
AssetID string `json:"assetId"`
PlatformID string `json:"platformId"`
ProjectName string `json:"projectName,omitempty"`
AssetGroupID string `json:"assetGroupId,omitempty"`
RemoteAssetID string `json:"remoteAssetId,omitempty"`
RemoteAssetURI string `json:"remoteAssetUri,omitempty"`
Status string `json:"status"`
LastErrorCode string `json:"lastErrorCode,omitempty"`
LastErrorMessage string `json:"lastErrorMessage,omitempty"`
LastSyncedAt string `json:"lastSyncedAt,omitempty"`
CreatedAt time.Time `json:"createdAt"`
UpdatedAt time.Time `json:"updatedAt"`
}
type PortraitAssetInput struct {
GatewayUserID string
UserID string
GatewayTenantID string
TenantID string
TenantKey string
Name string
Description string
SourceType string
URL string
Preview string
MimeType string
ByteSize int64
SourceSHA256 string
PrivateAvatarEligible bool
Metadata map[string]any
}
type PortraitAssetListFilter struct {
Keyword string
SourceType string
Page int
PageSize int
}
type PortraitAssetListResult struct {
Items []PortraitAsset
Total int
Page int
PageSize int
}
type PortraitAssetPlatform struct {
PlatformID string
PlatformKey string
Provider string
Credentials map[string]any
Config map[string]any
}
const portraitAssetColumns = `
a.id::text, COALESCE(a.gateway_user_id::text, ''), a.user_id,
COALESCE(a.gateway_tenant_id::text, ''), COALESCE(a.tenant_id, ''), COALESCE(a.tenant_key, ''),
a.name, a.description, a.source_type, a.url, a.preview, a.mime_type, a.byte_size,
a.source_sha256, a.private_avatar_eligible, a.status, a.last_error, a.metadata, a.created_at, a.updated_at`
func (s *Store) CreatePortraitAsset(ctx context.Context, input PortraitAssetInput) (PortraitAsset, error) {
metadata, _ := json.Marshal(emptyObjectIfNil(input.Metadata))
return scanPortraitAsset(s.pool.QueryRow(ctx, `
INSERT INTO gateway_portrait_assets (
gateway_user_id, user_id, gateway_tenant_id, tenant_id, tenant_key,
name, description, source_type, url, preview, mime_type, byte_size, source_sha256,
private_avatar_eligible, status, metadata
)
VALUES (
NULLIF($1, '')::uuid, $2, NULLIF($3, '')::uuid, NULLIF($4, ''), NULLIF($5, ''),
$6, $7, $8, $9, $10, $11, $12, $13, $14, 'not_synced', $15::jsonb
)
RETURNING `+portraitAssetColumns,
input.GatewayUserID, input.UserID, input.GatewayTenantID, input.TenantID, input.TenantKey,
strings.TrimSpace(input.Name), strings.TrimSpace(input.Description), strings.TrimSpace(input.SourceType),
strings.TrimSpace(input.URL), strings.TrimSpace(input.Preview), strings.TrimSpace(input.MimeType), input.ByteSize,
strings.TrimSpace(input.SourceSHA256), input.PrivateAvatarEligible, string(metadata),
))
}
func (s *Store) FindPortraitAssetBySourceHash(ctx context.Context, user *auth.User, sourceSHA256 string) (PortraitAsset, bool, error) {
sourceSHA256 = strings.TrimSpace(sourceSHA256)
if sourceSHA256 == "" {
return PortraitAsset{}, false, nil
}
gatewayUserID, userID := portraitAssetUserKeys(user)
asset, err := scanPortraitAsset(s.pool.QueryRow(ctx, `
SELECT `+portraitAssetColumns+`
FROM gateway_portrait_assets a
WHERE ((NULLIF($1, '')::uuid IS NOT NULL AND a.gateway_user_id = NULLIF($1, '')::uuid)
OR (NULLIF($2, '') IS NOT NULL AND a.user_id = $2))
AND a.source_sha256 = $3
ORDER BY a.created_at DESC
LIMIT 1`, gatewayUserID, userID, sourceSHA256))
if IsNotFound(err) {
return PortraitAsset{}, false, nil
}
return asset, err == nil, err
}
func (s *Store) FindPortraitAssetForUser(ctx context.Context, user *auth.User, assetID string) (PortraitAsset, bool, error) {
assetID = strings.TrimSpace(assetID)
if assetID == "" {
return PortraitAsset{}, false, nil
}
gatewayUserID, userID := portraitAssetUserKeys(user)
asset, err := scanPortraitAsset(s.pool.QueryRow(ctx, `
SELECT `+portraitAssetColumns+`
FROM gateway_portrait_assets a
WHERE a.id = NULLIF($3, '')::uuid
AND ((NULLIF($1, '')::uuid IS NOT NULL AND a.gateway_user_id = NULLIF($1, '')::uuid)
OR (NULLIF($2, '') IS NOT NULL AND a.user_id = $2))`, gatewayUserID, userID, assetID))
if IsNotFound(err) {
return PortraitAsset{}, false, nil
}
return asset, err == nil, err
}
func (s *Store) ListPortraitAssets(ctx context.Context, user *auth.User, filter PortraitAssetListFilter) (PortraitAssetListResult, error) {
page := filter.Page
if page < 1 {
page = 1
}
pageSize := filter.PageSize
if pageSize < 1 {
pageSize = 20
}
if pageSize > 100 {
pageSize = 100
}
gatewayUserID, userID := portraitAssetUserKeys(user)
keyword := strings.TrimSpace(filter.Keyword)
if keyword != "" {
keyword = "%" + keyword + "%"
}
where := `
WHERE ((NULLIF($1, '')::uuid IS NOT NULL AND a.gateway_user_id = NULLIF($1, '')::uuid)
OR (NULLIF($2, '') IS NOT NULL AND a.user_id = $2))
AND (NULLIF($3, '') IS NULL OR a.source_type = $3)
AND (NULLIF($4, '') IS NULL OR a.name ILIKE $4 OR a.description ILIKE $4)`
args := []any{gatewayUserID, userID, strings.TrimSpace(filter.SourceType), keyword}
var total int
if err := s.pool.QueryRow(ctx, `SELECT count(*) FROM gateway_portrait_assets a `+where, args...).Scan(&total); err != nil {
return PortraitAssetListResult{}, err
}
args = append(args, pageSize, (page-1)*pageSize)
rows, err := s.pool.Query(ctx, `SELECT `+portraitAssetColumns+`
FROM gateway_portrait_assets a `+where+`
ORDER BY a.created_at DESC
LIMIT $5 OFFSET $6`, args...)
if err != nil {
return PortraitAssetListResult{}, err
}
defer rows.Close()
items := make([]PortraitAsset, 0)
for rows.Next() {
asset, err := scanPortraitAsset(rows)
if err != nil {
return PortraitAssetListResult{}, err
}
items = append(items, asset)
}
if err := rows.Err(); err != nil {
return PortraitAssetListResult{}, err
}
return PortraitAssetListResult{Items: items, Total: total, Page: page, PageSize: pageSize}, nil
}
func (s *Store) GetPortraitAssetBinding(ctx context.Context, assetID string, platformID string) (PortraitAssetBinding, bool, error) {
binding, err := scanPortraitAssetBinding(s.pool.QueryRow(ctx, `
SELECT id::text, asset_id::text, platform_id::text, project_name, asset_group_id,
remote_asset_id, remote_asset_uri, status, last_error_code, last_error_message,
COALESCE(last_synced_at::text, ''), created_at, updated_at
FROM gateway_portrait_asset_bindings
WHERE asset_id = $1::uuid AND platform_id = $2::uuid`, assetID, platformID))
if IsNotFound(err) {
return PortraitAssetBinding{}, false, nil
}
return binding, err == nil, err
}
func (s *Store) UpsertPortraitAssetBinding(ctx context.Context, binding PortraitAssetBinding) (PortraitAssetBinding, error) {
return scanPortraitAssetBinding(s.pool.QueryRow(ctx, `
INSERT INTO gateway_portrait_asset_bindings (
asset_id, platform_id, project_name, asset_group_id, remote_asset_id, remote_asset_uri,
status, last_error_code, last_error_message, last_synced_at
)
VALUES ($1::uuid, $2::uuid, $3, $4, $5, $6, $7, $8, $9, now())
ON CONFLICT (asset_id, platform_id) DO UPDATE SET
project_name = EXCLUDED.project_name,
asset_group_id = EXCLUDED.asset_group_id,
remote_asset_id = CASE WHEN EXCLUDED.remote_asset_id <> '' THEN EXCLUDED.remote_asset_id ELSE gateway_portrait_asset_bindings.remote_asset_id END,
remote_asset_uri = CASE WHEN EXCLUDED.remote_asset_uri <> '' THEN EXCLUDED.remote_asset_uri ELSE gateway_portrait_asset_bindings.remote_asset_uri END,
status = EXCLUDED.status,
last_error_code = EXCLUDED.last_error_code,
last_error_message = EXCLUDED.last_error_message,
last_synced_at = now(),
updated_at = now()
RETURNING id::text, asset_id::text, platform_id::text, project_name, asset_group_id,
remote_asset_id, remote_asset_uri, status, last_error_code, last_error_message,
COALESCE(last_synced_at::text, ''), created_at, updated_at`,
binding.AssetID, binding.PlatformID, strings.TrimSpace(binding.ProjectName), strings.TrimSpace(binding.AssetGroupID),
strings.TrimSpace(binding.RemoteAssetID), strings.TrimSpace(binding.RemoteAssetURI), strings.TrimSpace(binding.Status),
strings.TrimSpace(binding.LastErrorCode), strings.TrimSpace(binding.LastErrorMessage),
))
}
func (s *Store) UpdatePortraitAssetStatus(ctx context.Context, assetID string, status string, lastError string) error {
_, err := s.pool.Exec(ctx, `
UPDATE gateway_portrait_assets
SET status = $2, last_error = $3, updated_at = now()
WHERE id = $1::uuid`, assetID, strings.TrimSpace(status), strings.TrimSpace(lastError))
return err
}
func (s *Store) PortraitAssetBindingSummary(ctx context.Context, assetID string) (active int, total int, latestError string, updatedAt string, err error) {
err = s.pool.QueryRow(ctx, `
SELECT COUNT(*) FILTER (WHERE status = 'active'), COUNT(*),
COALESCE((ARRAY_AGG(NULLIF(last_error_message, '') ORDER BY updated_at DESC) FILTER (WHERE NULLIF(last_error_message, '') IS NOT NULL))[1], ''),
COALESCE(MAX(updated_at)::text, '')
FROM gateway_portrait_asset_bindings
WHERE asset_id = $1::uuid`, assetID).Scan(&active, &total, &latestError, &updatedAt)
return
}
func (s *Store) ListPortraitAssetPlatforms(ctx context.Context) ([]PortraitAssetPlatform, error) {
rows, err := s.pool.Query(ctx, `
SELECT p.id::text, p.platform_key, p.provider, p.credentials, p.config
FROM integration_platforms p
WHERE p.deleted_at IS NULL
AND p.status = 'enabled'
AND LOWER(p.provider) IN ('volces', 'volces-openai')
ORDER BY COALESCE(p.dynamic_priority, p.priority), p.created_at`)
if err != nil {
return nil, err
}
defer rows.Close()
items := make([]PortraitAssetPlatform, 0)
for rows.Next() {
var item PortraitAssetPlatform
var credentials, config []byte
if err := rows.Scan(&item.PlatformID, &item.PlatformKey, &item.Provider, &credentials, &config); err != nil {
return nil, err
}
item.Credentials = decodeObject(credentials)
item.Config = decodeObject(config)
items = append(items, item)
}
return items, rows.Err()
}
func portraitAssetUserKeys(user *auth.User) (string, string) {
if user == nil {
return "", ""
}
gatewayUserID := strings.TrimSpace(user.GatewayUserID)
if gatewayUserID == "" && user.Source == "gateway" {
gatewayUserID = strings.TrimSpace(user.ID)
}
return gatewayUserID, strings.TrimSpace(user.ID)
}
type portraitAssetScanner interface{ Scan(dest ...any) error }
func scanPortraitAsset(scanner portraitAssetScanner) (PortraitAsset, error) {
var asset PortraitAsset
var metadata []byte
err := scanner.Scan(
&asset.ID, &asset.GatewayUserID, &asset.UserID, &asset.GatewayTenantID, &asset.TenantID, &asset.TenantKey,
&asset.Name, &asset.Description, &asset.SourceType, &asset.URL, &asset.Preview, &asset.MimeType, &asset.ByteSize,
&asset.SourceSHA256, &asset.PrivateAvatarEligible, &asset.Status, &asset.LastError, &metadata, &asset.CreatedAt, &asset.UpdatedAt,
)
if err != nil {
return PortraitAsset{}, err
}
asset.Metadata = decodeObject(metadata)
return asset, nil
}
func scanPortraitAssetBinding(scanner portraitAssetScanner) (PortraitAssetBinding, error) {
var binding PortraitAssetBinding
if err := scanner.Scan(
&binding.ID, &binding.AssetID, &binding.PlatformID, &binding.ProjectName, &binding.AssetGroupID,
&binding.RemoteAssetID, &binding.RemoteAssetURI, &binding.Status, &binding.LastErrorCode, &binding.LastErrorMessage,
&binding.LastSyncedAt, &binding.CreatedAt, &binding.UpdatedAt,
); err != nil {
return PortraitAssetBinding{}, err
}
return binding, nil
}
+156 -71
View File
@@ -7,6 +7,7 @@ import (
"encoding/base64"
"encoding/json"
"errors"
"net"
"strings"
"time"
"unicode"
@@ -22,8 +23,13 @@ type Store struct {
pool *pgxpool.Pool
}
const (
postgresApplicationName = "easyai-ai-gateway"
postgresConnectTimeout = 5 * time.Second
)
func defaultAPIKeyScopes() []string {
return []string{"chat", "embedding", "rerank", "image", "video", "music", "audio", "voice_clone"}
return []string{"chat", "embedding", "rerank", "image", "image_vectorize", "video", "video_enhance", "music", "audio", "voice_clone"}
}
func normalizeAPIKeyScopes(scopes []string) []string {
@@ -57,6 +63,7 @@ var (
ErrBalanceBelowFrozen = errors.New("wallet balance cannot be below frozen balance")
ErrInvalidWalletAmount = errors.New("wallet amount must be a decimal with at most nine fractional digits")
ErrIdempotencyKeyReused = errors.New("idempotency key was reused for a different request")
ErrExternalTaskIDReused = errors.New("external task id was reused")
ErrTaskExecutionLeaseUnavailable = errors.New("task execution lease is unavailable")
ErrTaskExecutionLeaseLost = errors.New("task execution lease was lost")
ErrTaskExecutionFinished = errors.New("task execution already finished")
@@ -67,7 +74,11 @@ var (
)
func Connect(ctx context.Context, databaseURL string) (*Store, error) {
pool, err := pgxpool.New(ctx, databaseURL)
config, err := postgresPoolConfig(databaseURL)
if err != nil {
return nil, err
}
pool, err := pgxpool.NewWithConfig(ctx, config)
if err != nil {
return nil, err
}
@@ -78,6 +89,38 @@ func Connect(ctx context.Context, databaseURL string) (*Store, error) {
return &Store{pool: pool}, nil
}
func postgresPoolConfig(databaseURL string) (*pgxpool.Config, error) {
config, err := pgxpool.ParseConfig(databaseURL)
if err != nil {
return nil, err
}
config.ConnConfig.ConnectTimeout = postgresConnectTimeout
config.ConnConfig.RuntimeParams["application_name"] = postgresApplicationName
return config, nil
}
func IsPostgresUnavailable(err error) bool {
if err == nil {
return false
}
if errors.Is(err, context.DeadlineExceeded) {
return true
}
var connectError *pgconn.ConnectError
if errors.As(err, &connectError) {
return true
}
var networkError net.Error
if errors.As(err, &networkError) {
return true
}
var postgresError *pgconn.PgError
if errors.As(err, &postgresError) {
return strings.HasPrefix(postgresError.Code, "08") || postgresError.Code == "53300" || strings.HasPrefix(postgresError.Code, "57P0")
}
return pgconn.SafeToRetry(err)
}
func (s *Store) Close() {
s.pool.Close()
}
@@ -174,33 +217,36 @@ type CreatedAPIKey struct {
}
type PlatformModel struct {
ID string `json:"id"`
PlatformID string `json:"platformId"`
BaseModelID string `json:"baseModelId,omitempty"`
Provider string `json:"provider,omitempty"`
PlatformName string `json:"platformName,omitempty"`
ModelName string `json:"modelName"`
ProviderModelName string `json:"providerModelName,omitempty"`
ModelAlias string `json:"modelAlias,omitempty"`
ModelType StringList `json:"modelType"`
DisplayName string `json:"displayName"`
CapabilityOverride map[string]any `json:"capabilityOverride,omitempty"`
Capabilities map[string]any `json:"capabilities,omitempty"`
BaseCapabilities map[string]any `json:"-"`
PricingMode string `json:"pricingMode"`
DiscountFactor float64 `json:"discountFactor,omitempty"`
PricingRuleSetID string `json:"pricingRuleSetId,omitempty"`
BillingConfigOverride map[string]any `json:"billingConfigOverride,omitempty"`
BillingConfig map[string]any `json:"billingConfig,omitempty"`
PermissionConfig map[string]any `json:"permissionConfig,omitempty"`
RetryPolicy map[string]any `json:"retryPolicy,omitempty"`
RateLimitPolicy map[string]any `json:"rateLimitPolicy,omitempty"`
RuntimePolicySetID string `json:"runtimePolicySetId,omitempty"`
RuntimePolicyOverride map[string]any `json:"runtimePolicyOverride,omitempty"`
CooldownUntil string `json:"cooldownUntil,omitempty"`
Enabled bool `json:"enabled"`
CreatedAt time.Time `json:"createdAt"`
UpdatedAt time.Time `json:"updatedAt"`
ID string `json:"id"`
PlatformID string `json:"platformId"`
BaseModelID string `json:"baseModelId,omitempty"`
Provider string `json:"provider,omitempty"`
PlatformName string `json:"platformName,omitempty"`
ModelName string `json:"modelName"`
ProviderModelName string `json:"providerModelName,omitempty"`
ModelAlias string `json:"modelAlias,omitempty"`
ModelType StringList `json:"modelType"`
DisplayName string `json:"displayName"`
CapabilityOverride map[string]any `json:"capabilityOverride,omitempty"`
Capabilities map[string]any `json:"capabilities,omitempty"`
BaseCapabilities map[string]any `json:"-"`
BaseBillingConfig map[string]any `json:"-"`
BasePricingRuleSetID string `json:"-"`
PlatformPricingRuleSetID string `json:"-"`
PricingMode string `json:"pricingMode"`
DiscountFactor float64 `json:"discountFactor,omitempty"`
PricingRuleSetID string `json:"pricingRuleSetId,omitempty"`
BillingConfigOverride map[string]any `json:"billingConfigOverride,omitempty"`
BillingConfig map[string]any `json:"billingConfig,omitempty"`
PermissionConfig map[string]any `json:"permissionConfig,omitempty"`
RetryPolicy map[string]any `json:"retryPolicy,omitempty"`
RateLimitPolicy map[string]any `json:"rateLimitPolicy,omitempty"`
RuntimePolicySetID string `json:"runtimePolicySetId,omitempty"`
RuntimePolicyOverride map[string]any `json:"runtimePolicyOverride,omitempty"`
CooldownUntil string `json:"cooldownUntil,omitempty"`
Enabled bool `json:"enabled"`
CreatedAt time.Time `json:"createdAt"`
UpdatedAt time.Time `json:"updatedAt"`
}
type AccessRule struct {
@@ -423,6 +469,7 @@ type RateLimitWindow struct {
type CreateTaskInput struct {
Kind string `json:"kind"`
Model string `json:"model"`
ExternalTaskID string `json:"externalTaskId,omitempty"`
RunMode string `json:"runMode"`
Async bool `json:"async"`
Request map[string]any `json:"request"`
@@ -440,6 +487,7 @@ type CreateTaskResult struct {
type GatewayTask struct {
ID string `json:"id"`
ExternalTaskID string `json:"externalTaskId,omitempty"`
Kind string `json:"kind"`
RunMode string `json:"runMode"`
UserID string `json:"userId"`
@@ -469,7 +517,7 @@ type GatewayTask struct {
Message string `json:"message,omitempty"`
AttemptCount int `json:"attemptCount"`
RemoteTaskID string `json:"remoteTaskId,omitempty"`
RemoteTaskPayload map[string]any `json:"remoteTaskPayload,omitempty"`
RemoteTaskPayload map[string]any `json:"-"`
Result map[string]any `json:"result,omitempty"`
Billings []any `json:"billings,omitempty"`
Usage map[string]any `json:"usage"`
@@ -499,7 +547,7 @@ type GatewayTask struct {
}
const gatewayTaskColumns = `
id::text, kind, run_mode, user_id, COALESCE(gateway_user_id::text, ''), user_source,
id::text, COALESCE(external_task_id, ''), kind, run_mode, user_id, COALESCE(gateway_user_id::text, ''), user_source,
COALESCE(gateway_tenant_id::text, ''), COALESCE(tenant_id, ''), COALESCE(tenant_key, ''),
COALESCE(api_key_id, ''), COALESCE(api_key_name, ''), COALESCE(api_key_prefix, ''),
COALESCE(user_group_id::text, ''), COALESCE(user_group_key, ''), model,
@@ -882,7 +930,9 @@ func (s *Store) listModels(ctx context.Context, platformID string) ([]PlatformMo
rows, err := s.pool.Query(ctx, `
SELECT m.id::text, m.platform_id::text, COALESCE(m.base_model_id::text, ''), p.provider, p.name,
m.model_name, COALESCE(NULLIF(m.provider_model_name, ''), m.model_name), COALESCE(m.model_alias, ''), m.model_type, m.display_name,
m.capability_override, m.capabilities, COALESCE(b.capabilities, '{}'::jsonb), m.pricing_mode, COALESCE(m.discount_factor, 0)::float8,
m.capability_override, m.capabilities, COALESCE(b.capabilities, '{}'::jsonb),
COALESCE(b.base_billing_config, '{}'::jsonb), COALESCE(b.pricing_rule_set_id::text, ''),
COALESCE(p.pricing_rule_set_id::text, ''), m.pricing_mode, COALESCE(m.discount_factor, 0)::float8,
COALESCE(m.pricing_rule_set_id::text, ''), m.billing_config_override, m.billing_config,
m.permission_config, m.retry_policy, m.rate_limit_policy, COALESCE(m.runtime_policy_set_id::text, ''), m.runtime_policy_override,
COALESCE(to_char(m.cooldown_until AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'), ''),
@@ -890,7 +940,7 @@ SELECT m.id::text, m.platform_id::text, COALESCE(m.base_model_id::text, ''), p.p
FROM platform_models m
JOIN integration_platforms p ON p.id = m.platform_id
LEFT JOIN LATERAL (
SELECT catalog.capabilities
SELECT catalog.capabilities, catalog.base_billing_config, catalog.pricing_rule_set_id
FROM base_model_catalog catalog
WHERE (m.base_model_id IS NOT NULL AND catalog.id = m.base_model_id)
OR (
@@ -917,6 +967,7 @@ ORDER BY m.model_type ASC, m.model_name ASC`, args...)
var capabilityOverride []byte
var capabilities []byte
var baseCapabilities []byte
var baseBillingConfig []byte
var billingConfigOverride []byte
var billingConfig []byte
var permissionConfig []byte
@@ -938,6 +989,9 @@ ORDER BY m.model_type ASC, m.model_name ASC`, args...)
&capabilityOverride,
&capabilities,
&baseCapabilities,
&baseBillingConfig,
&model.BasePricingRuleSetID,
&model.PlatformPricingRuleSetID,
&model.PricingMode,
&model.DiscountFactor,
&model.PricingRuleSetID,
@@ -958,6 +1012,7 @@ ORDER BY m.model_type ASC, m.model_name ASC`, args...)
model.CapabilityOverride = decodeObject(capabilityOverride)
model.Capabilities = decodeObject(capabilities)
model.BaseCapabilities = decodeObject(baseCapabilities)
model.BaseBillingConfig = decodeObject(baseBillingConfig)
model.ModelType = decodeStringArray(modelTypeBytes)
model.BillingConfigOverride = decodeObject(billingConfigOverride)
model.BillingConfig = decodeObject(billingConfig)
@@ -1519,11 +1574,35 @@ WHERE subject_type = 'api_key' AND subject_id = $1::uuid`, apiKeyID); err != nil
}
func (s *Store) VerifyLocalAPIKey(ctx context.Context, secret string) (*auth.User, error) {
return verifyLocalAPIKey(ctx, s.pool, secret)
}
type localAPIKeyDatabase interface {
Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error)
Exec(ctx context.Context, sql string, arguments ...any) (pgconn.CommandTag, error)
}
type localAPIKeyCandidate struct {
apiKeyID string
hash string
keyPrefix string
keyName string
scopesBytes []byte
userGroupID string
gatewayUserID string
username string
rolesBytes []byte
gatewayTenantID string
tenantID string
tenantKey string
}
func verifyLocalAPIKey(ctx context.Context, database localAPIKeyDatabase, secret string) (*auth.User, error) {
prefix := apiKeyPrefix(secret)
if prefix == "" {
return nil, auth.ErrUnauthorized
}
rows, err := s.pool.Query(ctx, `
rows, err := database.Query(ctx, `
SELECT k.id::text, k.key_hash, k.key_prefix, k.name, k.scopes, COALESCE(k.user_group_id::text, u.default_user_group_id::text, ''),
u.id::text, u.username, u.roles, COALESCE(u.gateway_tenant_id::text, ''),
COALESCE(u.tenant_id, ''), COALESCE(u.tenant_key, '')
@@ -1538,49 +1617,51 @@ WHERE k.key_prefix = $1
if err != nil {
return nil, err
}
defer rows.Close()
candidates, err := pgx.CollectRows(rows, func(row pgx.CollectableRow) (localAPIKeyCandidate, error) {
var candidate localAPIKeyCandidate
err := row.Scan(
&candidate.apiKeyID,
&candidate.hash,
&candidate.keyPrefix,
&candidate.keyName,
&candidate.scopesBytes,
&candidate.userGroupID,
&candidate.gatewayUserID,
&candidate.username,
&candidate.rolesBytes,
&candidate.gatewayTenantID,
&candidate.tenantID,
&candidate.tenantKey,
)
return candidate, err
})
if err != nil {
return nil, err
}
for rows.Next() {
var apiKeyID string
var hash string
var keyPrefix string
var keyName string
var scopesBytes []byte
var userGroupID string
var gatewayUserID string
var username string
var rolesBytes []byte
var gatewayTenantID string
var tenantID string
var tenantKey string
if err := rows.Scan(&apiKeyID, &hash, &keyPrefix, &keyName, &scopesBytes, &userGroupID, &gatewayUserID, &username, &rolesBytes, &gatewayTenantID, &tenantID, &tenantKey); err != nil {
return nil, err
}
if bcrypt.CompareHashAndPassword([]byte(hash), []byte(secret)) != nil {
for _, candidate := range candidates {
if bcrypt.CompareHashAndPassword([]byte(candidate.hash), []byte(secret)) != nil {
continue
}
if _, err := s.pool.Exec(ctx, `UPDATE gateway_api_keys SET last_used_at = now(), updated_at = now() WHERE id = $1::uuid`, apiKeyID); err != nil {
if _, err := database.Exec(ctx, `UPDATE gateway_api_keys SET last_used_at = now(), updated_at = now() WHERE id = $1::uuid`, candidate.apiKeyID); err != nil {
return nil, err
}
return &auth.User{
ID: gatewayUserID,
Username: username,
Roles: decodeStringArray(rolesBytes),
TenantID: tenantID,
GatewayTenantID: gatewayTenantID,
TenantKey: tenantKey,
ID: candidate.gatewayUserID,
Username: candidate.username,
Roles: decodeStringArray(candidate.rolesBytes),
TenantID: candidate.tenantID,
GatewayTenantID: candidate.gatewayTenantID,
TenantKey: candidate.tenantKey,
Source: "gateway",
GatewayUserID: gatewayUserID,
UserGroupID: userGroupID,
APIKeyID: apiKeyID,
APIKeyName: keyName,
APIKeyPrefix: keyPrefix,
APIKeyScopes: decodeStringArray(scopesBytes),
GatewayUserID: candidate.gatewayUserID,
UserGroupID: candidate.userGroupID,
APIKeyID: candidate.apiKeyID,
APIKeyName: candidate.keyName,
APIKeyPrefix: candidate.keyPrefix,
APIKeyScopes: decodeStringArray(candidate.scopesBytes),
}, nil
}
if err := rows.Err(); err != nil {
return nil, err
}
return nil, auth.ErrUnauthorized
}
@@ -1880,15 +1961,15 @@ func (s *Store) CreateTaskIdempotent(ctx context.Context, input CreateTaskInput,
task, err := scanGatewayTask(tx.QueryRow(ctx, `
INSERT INTO gateway_tasks (
kind, run_mode, user_id, gateway_user_id, user_source, gateway_tenant_id, tenant_id, tenant_key,
external_task_id, kind, run_mode, user_id, gateway_user_id, user_source, gateway_tenant_id, tenant_id, tenant_key,
api_key_id, api_key_name, api_key_prefix, user_group_id, user_group_key,
model, requested_model, request, async_mode, status, result, billings, conversation_id, new_message_count,
idempotency_key_hash, idempotency_request_hash, finished_at
)
VALUES ($1, $2, $3, NULLIF($4, '')::uuid, COALESCE(NULLIF($5, ''), 'gateway'), NULLIF($6, '')::uuid, NULLIF($7, ''), NULLIF($8, ''), NULLIF($9, ''), NULLIF($10, ''), NULLIF($11, ''), NULLIF($12, '')::uuid, NULLIF($13, ''), $14, $14, $15, $16, $17, $18::jsonb, $19::jsonb, NULLIF($20, '')::uuid, $21, NULLIF($22, ''), NULLIF($23, ''), NULL)
VALUES (NULLIF($1, ''), $2, $3, $4, NULLIF($5, '')::uuid, COALESCE(NULLIF($6, ''), 'gateway'), NULLIF($7, '')::uuid, NULLIF($8, ''), NULLIF($9, ''), NULLIF($10, ''), NULLIF($11, ''), NULLIF($12, ''), NULLIF($13, '')::uuid, NULLIF($14, ''), $15, $15, $16, $17, $18, $19::jsonb, $20::jsonb, NULLIF($21, '')::uuid, $22, NULLIF($23, ''), NULLIF($24, ''), NULL)
ON CONFLICT (user_id, idempotency_key_hash) WHERE idempotency_key_hash IS NOT NULL DO NOTHING
RETURNING `+gatewayTaskColumns,
input.Kind, runMode, user.ID, user.GatewayUserID, user.Source, user.GatewayTenantID, user.TenantID, user.TenantKey, user.APIKeyID, user.APIKeyName, user.APIKeyPrefix, user.UserGroupID, user.UserGroupKey, input.Model, requestBody, input.Async, status, resultBody, billingsBody, input.ConversationID, input.NewMessageCount, strings.TrimSpace(input.IdempotencyKeyHash), strings.TrimSpace(input.IdempotencyRequestHash),
strings.TrimSpace(input.ExternalTaskID), input.Kind, runMode, user.ID, user.GatewayUserID, user.Source, user.GatewayTenantID, user.TenantID, user.TenantKey, user.APIKeyID, user.APIKeyName, user.APIKeyPrefix, user.UserGroupID, user.UserGroupKey, input.Model, requestBody, input.Async, status, resultBody, billingsBody, input.ConversationID, input.NewMessageCount, strings.TrimSpace(input.IdempotencyKeyHash), strings.TrimSpace(input.IdempotencyRequestHash),
))
replayed := false
if errors.Is(err, pgx.ErrNoRows) && strings.TrimSpace(input.IdempotencyKeyHash) != "" {
@@ -1909,6 +1990,9 @@ FROM gateway_tasks
WHERE user_id = $1 AND idempotency_key_hash = $2`, user.ID, strings.TrimSpace(input.IdempotencyKeyHash)))
replayed = true
}
if isUniqueViolation(err) && strings.TrimSpace(input.ExternalTaskID) != "" {
return CreateTaskResult{}, ErrExternalTaskIDReused
}
if err != nil {
return CreateTaskResult{}, err
}
@@ -1973,6 +2057,7 @@ func scanGatewayTask(scanner taskScanner) (GatewayTask, error) {
var remoteTaskPayloadBytes []byte
if err := scanner.Scan(
&task.ID,
&task.ExternalTaskID,
&task.Kind,
&task.RunMode,
&task.UserID,
@@ -0,0 +1,49 @@
package store
import (
"context"
"testing"
"time"
"github.com/jackc/pgx/v5/pgconn"
)
func TestPostgresPoolConfigSetsDiagnosticAndConnectTimeout(t *testing.T) {
config, err := postgresPoolConfig("postgresql://gateway:password@localhost:5432/gateway?sslmode=disable")
if err != nil {
t.Fatalf("parse PostgreSQL pool config: %v", err)
}
if config.ConnConfig.ConnectTimeout != 5*time.Second {
t.Fatalf("connect timeout = %s, want 5s", config.ConnConfig.ConnectTimeout)
}
if applicationName := config.ConnConfig.RuntimeParams["application_name"]; applicationName != "easyai-ai-gateway" {
t.Fatalf("application_name = %q, want easyai-ai-gateway", applicationName)
}
}
func TestPostgresPoolConfigRejectsMalformedURL(t *testing.T) {
if _, err := postgresPoolConfig("://malformed"); err == nil {
t.Fatal("expected malformed PostgreSQL URL to fail")
}
}
func TestIsPostgresUnavailableClassifiesConnectivityFailures(t *testing.T) {
for _, testCase := range []struct {
name string
err error
}{
{name: "deadline", err: context.DeadlineExceeded},
{name: "connection exception", err: &pgconn.PgError{Code: "08006"}},
{name: "too many connections", err: &pgconn.PgError{Code: "53300"}},
{name: "cannot connect now", err: &pgconn.PgError{Code: "57P03"}},
} {
t.Run(testCase.name, func(t *testing.T) {
if !IsPostgresUnavailable(testCase.err) {
t.Fatalf("error %v was not classified as PostgreSQL unavailable", testCase.err)
}
})
}
if IsPostgresUnavailable(&pgconn.PgError{Code: "42601"}) {
t.Fatal("SQL syntax error was incorrectly classified as PostgreSQL unavailable")
}
}
+7 -1
View File
@@ -331,7 +331,7 @@ ORDER BY priority ASC, resource_type ASC, rule_key ASC`, id)
return EffectivePricingConfig{}, fmt.Errorf("pricing rule %s currency does not match its rule set", ruleKey)
}
switch calculatorType {
case "token_usage", "unit_weight", "duration_weight":
case "token_usage", "unit_weight", "duration_weight", "transition_matrix":
default:
return EffectivePricingConfig{}, fmt.Errorf("pricing rule %s uses unsupported calculator %s", ruleKey, calculatorType)
}
@@ -386,8 +386,12 @@ func ValidateEffectivePricingRuleShape(resourceType string, unit string, calcula
allowed = unit == "1k_tokens" && calculatorType == "token_usage"
case "image", "image_edit":
allowed = unit == "image" && calculatorType == "unit_weight"
case "image_vectorize":
allowed = unit == "conversion" && calculatorType == "unit_weight"
case "video":
allowed = unit == "5s" && calculatorType == "duration_weight"
case "video_enhance":
allowed = unit == "5s" && calculatorType == "transition_matrix"
case "music":
allowed = (unit == "song" || unit == "item") && calculatorType == "unit_weight"
case "audio":
@@ -415,6 +419,8 @@ func DefaultEffectivePricingCalculator(resourceType string) string {
return "token_usage"
case "video":
return "duration_weight"
case "video_enhance":
return "transition_matrix"
default:
return "unit_weight"
}
@@ -13,6 +13,8 @@ func TestValidateEffectivePricingRuleShape(t *testing.T) {
{name: "text", resource: "text_input", unit: "1k_tokens", calculator: "token_usage", valid: true},
{name: "image", resource: "image", unit: "image", calculator: "unit_weight", valid: true},
{name: "video", resource: "video", unit: "5s", calculator: "duration_weight", valid: true},
{name: "video enhance", resource: "video_enhance", unit: "5s", calculator: "transition_matrix", valid: true},
{name: "image vectorize", resource: "image_vectorize", unit: "conversion", calculator: "unit_weight", valid: true},
{name: "wrong video unit", resource: "video", unit: "second", calculator: "duration_weight"},
{name: "wrong image calculator", resource: "image", unit: "image", calculator: "token_usage"},
}
@@ -0,0 +1,25 @@
package store
import (
"encoding/json"
"strings"
"testing"
)
func TestGatewayTaskJSONOmitsPrivateRemoteTaskPayload(t *testing.T) {
raw, err := json.Marshal(GatewayTask{
ID: "task-private-state",
RemoteTaskID: "safe-upstream-id",
RemoteTaskPayload: map[string]any{"imageToken": "private-image-token", "receipt": "private-receipt"},
})
if err != nil {
t.Fatal(err)
}
serialized := strings.ToLower(string(raw))
if strings.Contains(serialized, "private-image-token") || strings.Contains(serialized, "private-receipt") || strings.Contains(serialized, "remotetaskpayload") {
t.Fatalf("private provider recovery state leaked into task JSON: %s", serialized)
}
if !strings.Contains(serialized, "safe-upstream-id") {
t.Fatalf("public upstream task id should remain available: %s", serialized)
}
}
@@ -0,0 +1,43 @@
package store
import (
"context"
"strings"
)
// GetRuntimeModelCandidateForRemoteTask restores the exact platform model used
// to submit an asynchronous provider task. It deliberately ignores enabled
// state so a task can still be cancelled after its platform is disabled.
func (s *Store) GetRuntimeModelCandidateForRemoteTask(ctx context.Context, platformModelID string, platformID string) (RuntimeModelCandidate, bool, error) {
platformModelID = strings.TrimSpace(platformModelID)
platformID = strings.TrimSpace(platformID)
if platformModelID == "" || platformID == "" {
return RuntimeModelCandidate{}, false, nil
}
var candidate RuntimeModelCandidate
var credentials, config []byte
err := s.pool.QueryRow(ctx, `
SELECT p.id::text, p.platform_key, p.name, p.provider,
COALESCE(NULLIF(p.config->>'specType', ''), p.provider), COALESCE(p.base_url, ''), p.auth_type,
p.credentials, p.config, m.id::text, COALESCE(NULLIF(m.provider_model_name, ''), m.model_name),
m.model_name, COALESCE(m.model_alias, ''),
COALESCE((m.model_type->>0), 'video_generate')
FROM platform_models m
JOIN integration_platforms p ON p.id = m.platform_id
WHERE m.id = $1::uuid AND p.id = $2::uuid AND p.deleted_at IS NULL`, platformModelID, platformID).Scan(
&candidate.PlatformID, &candidate.PlatformKey, &candidate.PlatformName, &candidate.Provider,
&candidate.SpecType, &candidate.BaseURL, &candidate.AuthType, &credentials, &config,
&candidate.PlatformModelID, &candidate.ProviderModelName, &candidate.ModelName, &candidate.ModelAlias, &candidate.ModelType,
)
if IsNotFound(err) {
return RuntimeModelCandidate{}, false, nil
}
if err != nil {
return RuntimeModelCandidate{}, false, err
}
candidate.Credentials = decodeObject(credentials)
candidate.PlatformConfig = decodeObject(config)
candidate.ClientID = candidate.PlatformKey + ":" + candidate.ModelType + ":" + firstNonEmpty(candidate.ProviderModelName, candidate.ModelName)
candidate.QueueKey = candidate.ClientID
return candidate, true, nil
}
+195 -4
View File
@@ -29,6 +29,16 @@ type TaskListResult struct {
PageSize int
}
type CompatTaskListFilter struct {
Provider string
Version string
Statuses []string
CreatedFrom *time.Time
CreatedTo *time.Time
Page int
PageSize int
}
func (s *Store) ListTasks(ctx context.Context, user *auth.User, filter TaskListFilter) (TaskListResult, error) {
page := filter.Page
if page <= 0 {
@@ -146,6 +156,122 @@ LIMIT $8 OFFSET $9`, queryArgs...)
}, nil
}
func (s *Store) GetCompatTask(ctx context.Context, user *auth.User, provider string, version string, identifier string) (GatewayTask, error) {
gatewayUserID := localGatewayUserID(user)
apiKeyID := ""
userID := ""
if user != nil {
apiKeyID = strings.TrimSpace(user.APIKeyID)
userID = strings.TrimSpace(user.ID)
}
if gatewayUserID == "" && userID == "" {
return GatewayTask{}, ErrLocalUserRequired
}
task, err := scanGatewayTask(s.pool.QueryRow(ctx, `
SELECT `+gatewayTaskColumns+`
FROM gateway_tasks
WHERE (
(NULLIF($1, '')::uuid IS NOT NULL AND gateway_user_id = NULLIF($1, '')::uuid)
OR (NULLIF($1, '')::uuid IS NULL AND NULLIF($2, '') IS NOT NULL AND user_id = $2)
)
AND (NULLIF($3, '') IS NULL OR api_key_id = $3)
AND request->>'_compat_provider' = $4
AND request->>'_kling_compat_version' = $5
AND (id::text = $6 OR external_task_id = $6)
ORDER BY created_at DESC
LIMIT 1`, gatewayUserID, userID, apiKeyID, strings.TrimSpace(provider), strings.TrimSpace(version), strings.TrimSpace(identifier)))
if err != nil {
return GatewayTask{}, err
}
attempts, err := s.ListTaskAttempts(ctx, task.ID)
if err != nil {
return GatewayTask{}, err
}
task.Attempts = attempts
return task, nil
}
func (s *Store) ListCompatTasks(ctx context.Context, user *auth.User, filter CompatTaskListFilter) (TaskListResult, error) {
page := filter.Page
if page <= 0 {
page = 1
}
pageSize := filter.PageSize
if pageSize <= 0 {
pageSize = 30
}
if pageSize > 500 {
pageSize = 500
}
offset := (page - 1) * pageSize
gatewayUserID := localGatewayUserID(user)
apiKeyID := ""
userID := ""
if user != nil {
apiKeyID = strings.TrimSpace(user.APIKeyID)
userID = strings.TrimSpace(user.ID)
}
if gatewayUserID == "" && userID == "" {
return TaskListResult{}, ErrLocalUserRequired
}
statuses := filter.Statuses
if len(statuses) == 0 {
statuses = nil
}
args := []any{
gatewayUserID,
userID,
apiKeyID,
strings.TrimSpace(filter.Provider),
strings.TrimSpace(filter.Version),
nullableTaskListTime(filter.CreatedFrom),
nullableTaskListTime(filter.CreatedTo),
statuses,
}
whereSQL := `
WHERE (
(NULLIF($1, '')::uuid IS NOT NULL AND gateway_user_id = NULLIF($1, '')::uuid)
OR (NULLIF($1, '')::uuid IS NULL AND NULLIF($2, '') IS NOT NULL AND user_id = $2)
)
AND (NULLIF($3, '') IS NULL OR api_key_id = $3)
AND request->>'_compat_provider' = $4
AND request->>'_kling_compat_version' = $5
AND ($6::timestamptz IS NULL OR created_at >= $6::timestamptz)
AND ($7::timestamptz IS NULL OR created_at <= $7::timestamptz)
AND ($8::text[] IS NULL OR status = ANY($8::text[]))`
var total int
if err := s.pool.QueryRow(ctx, `SELECT count(*) FROM gateway_tasks `+whereSQL, args...).Scan(&total); err != nil {
return TaskListResult{}, err
}
queryArgs := append(args, pageSize, offset)
rows, err := s.pool.Query(ctx, `
SELECT `+gatewayTaskColumns+`
FROM gateway_tasks
`+whereSQL+`
ORDER BY created_at DESC
LIMIT $9 OFFSET $10`, queryArgs...)
if err != nil {
return TaskListResult{}, err
}
defer rows.Close()
items := make([]GatewayTask, 0)
for rows.Next() {
task, err := scanGatewayTask(rows)
if err != nil {
return TaskListResult{}, err
}
items = append(items, task)
}
if err := rows.Err(); err != nil {
return TaskListResult{}, err
}
items, err = s.attachTaskAttempts(ctx, items)
if err != nil {
return TaskListResult{}, err
}
return TaskListResult{Items: items, Total: total, Page: page, PageSize: pageSize}, nil
}
func nullableTaskListTime(value *time.Time) any {
if value == nil {
return nil
@@ -402,12 +528,11 @@ WHERE id = $1::uuid
}
_, err = tx.Exec(ctx, `
UPDATE gateway_task_attempts
SET remote_task_id = NULLIF($2::text, ''),
response_snapshot = COALESCE(response_snapshot, '{}'::jsonb) || jsonb_build_object('remote_task_payload', $3::jsonb)
WHERE id = $1::uuid`,
SET remote_task_id = NULLIF($2::text, '')
WHERE id = $1::uuid
AND status = 'running'`,
attemptID,
remoteTaskID,
string(payloadJSON),
)
return err
})
@@ -459,6 +584,72 @@ WHERE id = $1::uuid
return task, true, nil
}
// CancelSubmittedTask records a confirmed upstream cancellation. Callers must
// first complete the provider-side DELETE so local status never claims a remote
// task was cancelled when the upstream request was not accepted.
func (s *Store) CancelSubmittedTask(ctx context.Context, taskID string, executionToken string, message string) (GatewayTask, bool, error) {
message = strings.TrimSpace(message)
if message == "" {
message = "任务已由上游取消"
}
var task GatewayTask
changed := false
err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
var err error
task, err = scanGatewayTask(tx.QueryRow(ctx, `
UPDATE gateway_tasks
SET status = 'cancelled',
error = NULLIF($2, ''),
error_code = 'task_cancelled',
error_message = NULLIF($2, ''),
billing_status = CASE
WHEN run_mode <> 'production' OR gateway_user_id IS NULL THEN 'not_required'
WHEN reservation_amount > 0 THEN 'pending'
ELSE 'released'
END,
billing_updated_at = now(),
locked_by = NULL,
locked_at = NULL,
heartbeat_at = NULL,
execution_token = NULL,
execution_lease_expires_at = NULL,
finished_at = now(),
updated_at = now()
WHERE id = $1::uuid
AND NULLIF(remote_task_id, '') IS NOT NULL
AND (
(status = 'running' AND execution_token = NULLIF($3, '')::uuid)
OR status = 'queued'
)
RETURNING `+gatewayTaskColumns, taskID, message, strings.TrimSpace(executionToken)))
if IsNotFound(err) {
return nil
}
if err != nil {
return err
}
changed = true
payloadJSON, _ := json.Marshal(map[string]any{"taskId": taskID, "reason": "upstream_cancelled"})
_, err = tx.Exec(ctx, `
INSERT INTO settlement_outbox (
task_id, event_type, action, amount, currency, pricing_snapshot, payload, status, next_attempt_at
)
SELECT id, 'task.billing.release', 'release', reservation_amount, billing_currency,
pricing_snapshot, $2::jsonb, 'pending', now()
FROM gateway_tasks
WHERE id = $1::uuid
AND run_mode = 'production'
AND gateway_user_id IS NOT NULL
AND reservation_amount > 0
ON CONFLICT (task_id, event_type) DO NOTHING`, taskID, string(payloadJSON))
return err
})
if err != nil {
return GatewayTask{}, false, err
}
return task, changed, nil
}
func (s *Store) ListRecoverableAsyncTasks(ctx context.Context, limit int) ([]AsyncTaskQueueItem, error) {
if limit <= 0 {
limit = 500
@@ -0,0 +1,111 @@
package store
import (
"context"
"strings"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
)
// VolcesCompatibleTaskListFilter mirrors the supported filters of Ark's
// ListContentsGenerationsTasks API. Task IDs are the gateway's public task
// IDs, which are the IDs returned by the compatibility create endpoint.
type VolcesCompatibleTaskListFilter struct {
CompatibilityMarker string
Status string
Model string
TaskIDs []string
Page int
PageSize int
}
// ListVolcesCompatibleTasks returns only video tasks created through a named
// compatibility surface. Keeping this query separate from ListTasks avoids
// broadening the ordinary task-list API's filtering semantics.
func (s *Store) ListVolcesCompatibleTasks(ctx context.Context, user *auth.User, filter VolcesCompatibleTaskListFilter) (TaskListResult, error) {
page := filter.Page
if page < 1 {
page = 1
}
if page > 500 {
page = 500
}
pageSize := filter.PageSize
if pageSize < 1 {
pageSize = 20
}
if pageSize > 500 {
pageSize = 500
}
gatewayUserID := localGatewayUserID(user)
userID, apiKeyID := "", ""
if user != nil {
userID = strings.TrimSpace(user.ID)
apiKeyID = strings.TrimSpace(user.APIKeyID)
}
if gatewayUserID == "" && userID == "" {
return TaskListResult{}, ErrLocalUserRequired
}
taskIDs := make([]string, 0, len(filter.TaskIDs))
seen := make(map[string]bool, len(filter.TaskIDs))
for _, taskID := range filter.TaskIDs {
taskID = strings.TrimSpace(taskID)
if taskID != "" && !seen[taskID] {
seen[taskID] = true
taskIDs = append(taskIDs, taskID)
}
}
args := []any{
gatewayUserID,
userID,
apiKeyID,
strings.TrimSpace(filter.CompatibilityMarker),
strings.ToLower(strings.TrimSpace(filter.Status)),
strings.TrimSpace(filter.Model),
taskIDs,
}
whereSQL := `
WHERE (
(
NULLIF($1, '')::uuid IS NOT NULL
AND gateway_user_id = NULLIF($1, '')::uuid
)
OR (
NULLIF($1, '')::uuid IS NULL
AND NULLIF($2, '') IS NOT NULL
AND user_id = $2
)
)
AND (NULLIF($3, '') IS NULL OR api_key_id = $3)
AND kind = 'videos.generations'
AND request->>'_gateway_compatibility' = $4
AND (NULLIF($5, '') IS NULL OR LOWER(status) = $5)
AND (NULLIF($6, '') IS NULL OR model = $6 OR resolved_model = $6)
AND (COALESCE(array_length($7::text[], 1), 0) = 0 OR id::text = ANY($7::text[]))`
var total int
if err := s.pool.QueryRow(ctx, `SELECT count(*) FROM gateway_tasks `+whereSQL, args...).Scan(&total); err != nil {
return TaskListResult{}, err
}
rows, err := s.pool.Query(ctx, `
SELECT `+gatewayTaskColumns+`
FROM gateway_tasks
`+whereSQL+`
ORDER BY created_at DESC
LIMIT $8 OFFSET $9`, append(args, pageSize, (page-1)*pageSize)...)
if err != nil {
return TaskListResult{}, err
}
defer rows.Close()
items := make([]GatewayTask, 0)
for rows.Next() {
task, err := scanGatewayTask(rows)
if err != nil {
return TaskListResult{}, err
}
items = append(items, task)
}
if err := rows.Err(); err != nil {
return TaskListResult{}, err
}
return TaskListResult{Items: items, Total: total, Page: page, PageSize: pageSize}, nil
}
@@ -0,0 +1,14 @@
UPDATE model_pricing_rules
SET formula_config = jsonb_set(
COALESCE(formula_config, '{}'::jsonb),
'{formula}',
to_jsonb(replace(
formula_config->>'formula',
'ceil(duration_seconds / 5)',
'(duration_seconds / 5)'
)),
true
),
updated_at = now()
WHERE resource_type = 'video'
AND strpos(formula_config->>'formula', 'ceil(duration_seconds / 5)') > 0;

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