Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
ee0256adb4 | ||
|
|
fce76a30ba | ||
|
|
d37fb3ea60 | ||
|
|
9872068596 | ||
|
|
a23b28c27d | ||
|
|
eb37b568ae | ||
|
|
a8d1c550ef | ||
|
|
5432760cf7 | ||
|
|
7c5a999e32 | ||
|
|
f7a5f2e808 | ||
|
|
152f9d1206 | ||
|
|
d95cecd0eb | ||
|
|
5954afef55 | ||
|
|
929a7a172c | ||
|
|
16149260c7 | ||
|
|
c72c43aaaa | ||
|
|
84aea01b5b | ||
|
|
293ef24bb7 | ||
|
|
0f1765b226 | ||
|
|
1fa58ba901 | ||
|
|
55595570c2 | ||
|
|
fba9759bc7 | ||
|
|
3d3460ce63 | ||
|
|
000ee1bbfd | ||
|
|
d0cfd0a385 | ||
|
|
0818f55235 | ||
|
|
fe83da56d2 | ||
|
|
002422b753 | ||
|
|
24b778b3ba | ||
|
|
d7951cfdd2 | ||
|
|
ddd68cfebd | ||
|
|
5a71643099 | ||
|
|
b04a7d9d3d | ||
|
|
6b675c406e | ||
|
|
56d4a3a6b7 | ||
|
|
276c0612d8 | ||
|
|
d818e7947a | ||
|
|
d5c2c58c67 | ||
|
|
9d4501bc42 | ||
|
|
e280c0875c | ||
|
|
142dcc7932 | ||
|
|
e3dfe8162b | ||
|
|
69b0c107d3 | ||
|
|
e533ec2367 | ||
|
|
bfa17a3aba | ||
|
|
86c374b5c2 | ||
|
|
505b074b47 | ||
|
|
8c38714296 | ||
|
|
1e55f7df8b | ||
|
|
0aa9b3e88f | ||
|
|
3561efa7da | ||
|
|
b7bb9ed8d5 | ||
|
|
8beb8501fa | ||
|
|
257ee09e58 | ||
|
|
1362970229 | ||
|
|
3c82c7b492 | ||
|
|
c879de18e2 | ||
|
|
62d25fcb11 | ||
|
|
5b2b94b1bd | ||
|
|
7cea21f765 | ||
|
|
5114686c35 | ||
|
|
01a013c809 | ||
|
|
dcf5c4f340 | ||
|
|
d956524690 | ||
|
|
c070cda22a |
@@ -23,6 +23,12 @@ CONFIG_JWT_SECRET=this is a very secret secret
|
||||
# - hybrid: both sources are accepted and separated by gateway_users.source.
|
||||
IDENTITY_MODE=hybrid
|
||||
|
||||
# Billing engine rollout mode:
|
||||
# - observe: keep legacy billing decisions and compare effective-pricing-v2 in logs.
|
||||
# - enforce: require v2 pricing, reserve the candidate maximum, then settle asynchronously.
|
||||
# - hold: reject new production generation before any upstream request; existing settlements continue.
|
||||
BILLING_ENGINE_MODE=observe
|
||||
|
||||
# Unified identity business settings are managed in System Settings > Unified
|
||||
# Identity. Deployment only supplies the SecretStore and infrastructure timing.
|
||||
AI_GATEWAY_PUBLIC_BASE_URL=http://localhost:8088
|
||||
@@ -43,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
|
||||
|
||||
@@ -9,8 +9,22 @@ on:
|
||||
jobs:
|
||||
verify:
|
||||
runs-on: easyai-gateway-ci-unprivileged-v2
|
||||
services:
|
||||
postgres:
|
||||
image: docker.io/library/postgres:16-alpine@sha256:57c72fd2a128e416c7fcc499958864df5301e940bca0a56f58fddf30ffc07777
|
||||
env:
|
||||
POSTGRES_USER: easyai_test
|
||||
POSTGRES_HOST_AUTH_METHOD: trust
|
||||
POSTGRES_DB: easyai_gateway_test
|
||||
options: >-
|
||||
--health-cmd "pg_isready -U easyai_test -d easyai_gateway_test"
|
||||
--health-interval 2s
|
||||
--health-timeout 5s
|
||||
--health-retries 30
|
||||
env:
|
||||
TRIVY_DB_REPOSITORY: ghcr.m.daocloud.io/aquasecurity/trivy-db:2
|
||||
AI_GATEWAY_DATABASE_URL: postgresql://easyai_test@postgres:5432/easyai_gateway_test?sslmode=disable
|
||||
AI_GATEWAY_TEST_DATABASE_URL: postgresql://easyai_test@postgres:5432/easyai_gateway_test?sslmode=disable
|
||||
steps:
|
||||
- name: Checkout without external Actions
|
||||
env:
|
||||
@@ -79,6 +93,9 @@ jobs:
|
||||
printf 'Go files require gofmt:\n%s\n' "$unformatted" >&2
|
||||
exit 1
|
||||
}
|
||||
- name: Migrate PostgreSQL 16 integration database
|
||||
working-directory: apps/api
|
||||
run: go run ./cmd/migrate
|
||||
- name: Verify Go code
|
||||
working-directory: apps/api
|
||||
env:
|
||||
|
||||
@@ -7,8 +7,22 @@ on:
|
||||
jobs:
|
||||
verify-tag:
|
||||
runs-on: easyai-gateway-ci-unprivileged-v2
|
||||
services:
|
||||
postgres:
|
||||
image: docker.io/library/postgres:16-alpine@sha256:57c72fd2a128e416c7fcc499958864df5301e940bca0a56f58fddf30ffc07777
|
||||
env:
|
||||
POSTGRES_USER: easyai_test
|
||||
POSTGRES_HOST_AUTH_METHOD: trust
|
||||
POSTGRES_DB: easyai_gateway_test
|
||||
options: >-
|
||||
--health-cmd "pg_isready -U easyai_test -d easyai_gateway_test"
|
||||
--health-interval 2s
|
||||
--health-timeout 5s
|
||||
--health-retries 30
|
||||
env:
|
||||
TRIVY_DB_REPOSITORY: ghcr.m.daocloud.io/aquasecurity/trivy-db:2
|
||||
AI_GATEWAY_DATABASE_URL: postgresql://easyai_test@postgres:5432/easyai_gateway_test?sslmode=disable
|
||||
AI_GATEWAY_TEST_DATABASE_URL: postgresql://easyai_test@postgres:5432/easyai_gateway_test?sslmode=disable
|
||||
steps:
|
||||
- name: Checkout without external Actions
|
||||
env:
|
||||
@@ -64,6 +78,9 @@ jobs:
|
||||
printf 'Go files require gofmt:\n%s\n' "$unformatted" >&2
|
||||
exit 1
|
||||
}
|
||||
- name: Migrate PostgreSQL 16 integration database
|
||||
working-directory: apps/api
|
||||
run: go run ./cmd/migrate
|
||||
- name: Verify Go code
|
||||
working-directory: apps/api
|
||||
env:
|
||||
|
||||
@@ -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` 在容器网络内的连接串是:
|
||||
|
||||
|
||||
+1737
-2532
File diff suppressed because it is too large
Load Diff
+1141
-1670
File diff suppressed because it is too large
Load Diff
+2
-2
@@ -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
@@ -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=
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -34,6 +34,11 @@ var gatewayOpenAIRequestExtensions = stringSet(
|
||||
"requestId", "request_id", "signal", "userMessage", "user_message", "platformId",
|
||||
"platform_id", "options", "enable_thinking", "thinking_budget_tokens", "enable_web_search",
|
||||
"modelType", "model_type", "capability", "capabilityType", "mode", "simulation", "testMode",
|
||||
"cacheAffinityKey", "cache_affinity_key", "simulationDurationMs", "simulationDurationSeconds",
|
||||
"simulationMinDurationMs", "simulationMaxDurationMs", "simulationMinDurationSeconds",
|
||||
"simulationMaxDurationSeconds", "simulationDurationMinMs", "simulationDurationMaxMs",
|
||||
"simulationDurationMinSeconds", "simulationDurationMaxSeconds", "simulationFailure",
|
||||
"simulationProfile", "simulationUsage",
|
||||
)
|
||||
|
||||
var gatewayResponsesRequestExtensions = stringSet("messages", "presence_penalty", "frequency_penalty")
|
||||
|
||||
@@ -60,6 +60,23 @@ func TestValidateOpenAIRequestParametersRejectsUnknownTopLevelField(t *testing.T
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateOpenAIRequestParametersAcceptsInternalSimulationAndAffinityFields(t *testing.T) {
|
||||
body := map[string]any{
|
||||
"model": "demo", "messages": []any{}, "simulation": true,
|
||||
"simulationDurationMs": 5, "simulationUsage": map[string]any{"inputTokens": 10},
|
||||
"cacheAffinityKey": "conversation-1",
|
||||
}
|
||||
if err := ValidateOpenAIRequestParameters("chat.completions", body); err != nil {
|
||||
t.Fatalf("expected documented gateway extensions to remain accepted, got %v", err)
|
||||
}
|
||||
filtered := FilterOpenAIChatRequestBody(body)
|
||||
for _, key := range []string{"simulation", "simulationDurationMs", "simulationUsage", "cacheAffinityKey"} {
|
||||
if _, ok := filtered[key]; ok {
|
||||
t.Fatalf("gateway-only field %q leaked upstream", key)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestResponsesFallbackMapsEquivalentCurrentParameters(t *testing.T) {
|
||||
body, err := ResponsesRequestToChat(map[string]any{
|
||||
"input": "hello", "store": false, "metadata": map[string]any{"trace": "1"},
|
||||
|
||||
@@ -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
|
||||
@@ -36,6 +37,7 @@ type ResponseTurn struct {
|
||||
}
|
||||
|
||||
type Response struct {
|
||||
AttemptID string
|
||||
Result map[string]any
|
||||
RequestID string
|
||||
Usage Usage
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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) != "" {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -47,6 +47,7 @@ type Config struct {
|
||||
GlobalHTTPProxy string
|
||||
GlobalHTTPProxySource string
|
||||
LogLevel slog.Level
|
||||
BillingEngineMode string
|
||||
}
|
||||
|
||||
func Load() Config {
|
||||
@@ -90,10 +91,16 @@ func Load() Config {
|
||||
GlobalHTTPProxy: globalProxy.HTTPProxy,
|
||||
GlobalHTTPProxySource: globalProxy.Source,
|
||||
LogLevel: logLevel(env("LOG_LEVEL", "info")),
|
||||
BillingEngineMode: strings.ToLower(env("BILLING_ENGINE_MODE", "observe")),
|
||||
}
|
||||
}
|
||||
|
||||
func (c Config) Validate() error {
|
||||
switch strings.ToLower(strings.TrimSpace(c.BillingEngineMode)) {
|
||||
case "", "observe", "enforce", "hold":
|
||||
default:
|
||||
return errors.New("BILLING_ENGINE_MODE must be observe, enforce, or hold")
|
||||
}
|
||||
switch strings.ToLower(strings.TrimSpace(c.IdentitySecretStore)) {
|
||||
case "":
|
||||
case "file":
|
||||
|
||||
@@ -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 管理端创建一条访问控制规则。
|
||||
|
||||
@@ -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, ®isterResponse)
|
||||
|
||||
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())
|
||||
}
|
||||
}
|
||||
@@ -13,7 +13,7 @@ import (
|
||||
|
||||
type walletBalanceRequest struct {
|
||||
Currency string `json:"currency" example:"USD"`
|
||||
Balance float64 `json:"balance" example:"100"`
|
||||
Balance json.Number `json:"balance" swaggertype:"number" example:"100"`
|
||||
Reason string `json:"reason" example:"manual recharge"`
|
||||
IdempotencyKey string `json:"idempotencyKey" example:"wallet-set-20260514-001"`
|
||||
Metadata map[string]any `json:"metadata"`
|
||||
@@ -21,7 +21,7 @@ type walletBalanceRequest struct {
|
||||
|
||||
type walletRechargeRequest struct {
|
||||
Currency string `json:"currency" example:"resource"`
|
||||
Amount float64 `json:"amount" example:"100"`
|
||||
Amount json.Number `json:"amount" swaggertype:"number" example:"100"`
|
||||
Reason string `json:"reason" example:"manual recharge"`
|
||||
IdempotencyKey string `json:"idempotencyKey" example:"wallet-recharge-20260514-001"`
|
||||
Metadata map[string]any `json:"metadata"`
|
||||
@@ -50,10 +50,6 @@ func (s *Server) setUserWalletBalance(w http.ResponseWriter, r *http.Request) {
|
||||
writeError(w, http.StatusBadRequest, "invalid json body")
|
||||
return
|
||||
}
|
||||
if input.Balance < 0 {
|
||||
writeError(w, http.StatusBadRequest, "wallet balance cannot be negative")
|
||||
return
|
||||
}
|
||||
gatewayUserID := strings.TrimSpace(r.PathValue("userID"))
|
||||
reason := strings.TrimSpace(input.Reason)
|
||||
if reason == "" {
|
||||
@@ -67,7 +63,7 @@ func (s *Server) setUserWalletBalance(w http.ResponseWriter, r *http.Request) {
|
||||
next, err := s.store.SetUserWalletBalanceTx(r.Context(), tx, store.WalletBalanceAdjustmentInput{
|
||||
GatewayUserID: gatewayUserID,
|
||||
Currency: input.Currency,
|
||||
Balance: input.Balance,
|
||||
BalanceText: input.Balance.String(),
|
||||
Reason: reason,
|
||||
IdempotencyKey: input.IdempotencyKey,
|
||||
Metadata: input.Metadata,
|
||||
@@ -89,6 +85,10 @@ func (s *Server) setUserWalletBalance(w http.ResponseWriter, r *http.Request) {
|
||||
writeError(w, http.StatusNotFound, "user not found")
|
||||
case errors.Is(err, store.ErrWalletBalanceUnchanged):
|
||||
writeError(w, http.StatusBadRequest, "wallet balance is unchanged")
|
||||
case errors.Is(err, store.ErrInvalidWalletAmount):
|
||||
writeError(w, http.StatusBadRequest, "wallet balance must be a non-negative decimal with at most nine fractional digits", "invalid_wallet_amount")
|
||||
case errors.Is(err, store.ErrBalanceBelowFrozen):
|
||||
writeError(w, http.StatusConflict, "wallet balance cannot be below frozen balance", "balance_below_frozen")
|
||||
default:
|
||||
s.logger.Error("set user wallet balance failed", "error", err)
|
||||
writeError(w, http.StatusInternalServerError, "set user wallet balance failed")
|
||||
@@ -126,10 +126,6 @@ func (s *Server) rechargeUserWalletBalance(w http.ResponseWriter, r *http.Reques
|
||||
writeError(w, http.StatusBadRequest, "invalid json body")
|
||||
return
|
||||
}
|
||||
if input.Amount <= 0 {
|
||||
writeError(w, http.StatusBadRequest, "wallet recharge amount must be positive")
|
||||
return
|
||||
}
|
||||
gatewayUserID := strings.TrimSpace(r.PathValue("userID"))
|
||||
reason := strings.TrimSpace(input.Reason)
|
||||
if reason == "" {
|
||||
@@ -143,7 +139,7 @@ func (s *Server) rechargeUserWalletBalance(w http.ResponseWriter, r *http.Reques
|
||||
next, err := s.store.RechargeUserWalletBalanceTx(r.Context(), tx, store.WalletRechargeInput{
|
||||
GatewayUserID: gatewayUserID,
|
||||
Currency: input.Currency,
|
||||
Amount: input.Amount,
|
||||
AmountText: input.Amount.String(),
|
||||
Reason: reason,
|
||||
IdempotencyKey: input.IdempotencyKey,
|
||||
Metadata: input.Metadata,
|
||||
@@ -163,6 +159,8 @@ func (s *Server) rechargeUserWalletBalance(w http.ResponseWriter, r *http.Reques
|
||||
switch {
|
||||
case store.IsNotFound(err):
|
||||
writeError(w, http.StatusNotFound, "user not found")
|
||||
case errors.Is(err, store.ErrInvalidWalletAmount):
|
||||
writeError(w, http.StatusBadRequest, "wallet recharge amount must be a positive decimal with at most nine fractional digits", "invalid_wallet_amount")
|
||||
default:
|
||||
s.logger.Error("recharge user wallet balance failed", "error", err)
|
||||
writeError(w, http.StatusInternalServerError, "recharge user wallet balance failed")
|
||||
|
||||
@@ -0,0 +1,145 @@
|
||||
package httpapi
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// listBillingSettlements godoc
|
||||
// @Summary 查询计费结算队列
|
||||
// @Description 管理端分页查询待结算、重试失败和人工复核记录。
|
||||
// @Tags billing
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param status query string false "结算状态"
|
||||
// @Param action query string false "动作:settle 或 release"
|
||||
// @Param page query int false "页码" default(1)
|
||||
// @Param pageSize query int false "每页数量" default(50)
|
||||
// @Success 200 {object} BillingSettlementListResponse
|
||||
// @Failure 400 {object} ErrorEnvelope
|
||||
// @Failure 401 {object} ErrorEnvelope
|
||||
// @Failure 403 {object} ErrorEnvelope
|
||||
// @Failure 500 {object} ErrorEnvelope
|
||||
// @Router /api/admin/runtime/billing-settlements [get]
|
||||
func (s *Server) listBillingSettlements(w http.ResponseWriter, r *http.Request) {
|
||||
query := r.URL.Query()
|
||||
page, err := positiveQueryInt(query.Get("page"), 1)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid page", "invalid_request")
|
||||
return
|
||||
}
|
||||
pageSize, err := positiveQueryInt(query.Get("pageSize"), 50)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid pageSize", "invalid_request")
|
||||
return
|
||||
}
|
||||
result, err := s.store.ListBillingSettlements(r.Context(), store.BillingSettlementListFilter{
|
||||
Status: query.Get("status"), Action: query.Get("action"), Page: page, PageSize: pageSize,
|
||||
})
|
||||
if err != nil {
|
||||
s.logger.Error("list billing settlements failed", "error_category", "billing_settlement_list_failed")
|
||||
writeError(w, http.StatusInternalServerError, "list billing settlements failed", "billing_settlement_list_failed")
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, result)
|
||||
}
|
||||
|
||||
// retryBillingSettlement godoc
|
||||
// @Summary 重试计费结算
|
||||
// @Description Manager 使用单值 Idempotency-Key 将重试失败或人工复核记录重新放入队列,并记录审计日志。
|
||||
// @Tags billing
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param settlementId path string true "结算记录 ID"
|
||||
// @Param Idempotency-Key header string true "幂等键"
|
||||
// @Success 200 {object} BillingSettlementRetryResponse
|
||||
// @Failure 400 {object} ErrorEnvelope
|
||||
// @Failure 401 {object} ErrorEnvelope
|
||||
// @Failure 403 {object} ErrorEnvelope
|
||||
// @Failure 404 {object} ErrorEnvelope
|
||||
// @Failure 409 {object} ErrorEnvelope
|
||||
// @Failure 500 {object} ErrorEnvelope
|
||||
// @Router /api/admin/runtime/billing-settlements/{settlementId}/retry [post]
|
||||
func (s *Server) retryBillingSettlement(w http.ResponseWriter, r *http.Request) {
|
||||
key, ok := singleIdempotencyKey(r)
|
||||
if !ok {
|
||||
writeError(w, http.StatusBadRequest, "a single Idempotency-Key is required", "idempotency_key_required")
|
||||
return
|
||||
}
|
||||
keyHash := sha256.Sum256([]byte(key))
|
||||
settlementID := strings.TrimSpace(r.PathValue("settlementId"))
|
||||
if _, err := uuid.Parse(settlementID); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid billing settlement ID", "invalid_request")
|
||||
return
|
||||
}
|
||||
actor, _ := auth.UserFromContext(r.Context())
|
||||
var item store.BillingSettlement
|
||||
var audit store.AuditLog
|
||||
var replayed bool
|
||||
err := s.store.InTx(r.Context(), func(tx store.Tx) error {
|
||||
var err error
|
||||
item, replayed, err = s.store.RetryBillingSettlementTx(r.Context(), tx, settlementID, hex.EncodeToString(keyHash[:]))
|
||||
if err != nil || replayed {
|
||||
return err
|
||||
}
|
||||
audit, err = s.store.RecordAuditLogTx(r.Context(), tx, billingSettlementRetryAuditInput(r, actor, item))
|
||||
return err
|
||||
})
|
||||
if err != nil {
|
||||
switch {
|
||||
case store.IsNotFound(err):
|
||||
writeError(w, http.StatusNotFound, "billing settlement not found", "billing_settlement_not_found")
|
||||
case errors.Is(err, store.ErrBillingSettlementNotRetryable):
|
||||
writeError(w, http.StatusConflict, "billing settlement is not retryable", "billing_settlement_not_retryable")
|
||||
default:
|
||||
s.logger.Error("retry billing settlement failed", "settlementID", settlementID, "error_category", "billing_settlement_retry_failed")
|
||||
writeError(w, http.StatusInternalServerError, "retry billing settlement failed", "billing_settlement_retry_failed")
|
||||
}
|
||||
return
|
||||
}
|
||||
if replayed {
|
||||
w.Header().Set("Idempotent-Replayed", "true")
|
||||
}
|
||||
response := map[string]any{"settlement": item}
|
||||
if !replayed {
|
||||
response["auditLog"] = audit
|
||||
}
|
||||
writeJSON(w, http.StatusOK, response)
|
||||
}
|
||||
|
||||
func singleIdempotencyKey(r *http.Request) (string, bool) {
|
||||
values := r.Header.Values("Idempotency-Key")
|
||||
if len(values) != 1 {
|
||||
return "", false
|
||||
}
|
||||
value := strings.TrimSpace(values[0])
|
||||
if value == "" || len(value) > 255 || strings.Contains(value, ",") {
|
||||
return "", false
|
||||
}
|
||||
return value, true
|
||||
}
|
||||
|
||||
func billingSettlementRetryAuditInput(r *http.Request, actor *auth.User, item store.BillingSettlement) store.AuditLogInput {
|
||||
input := store.AuditLogInput{
|
||||
Category: "billing", Action: "billing.settlement.retry",
|
||||
TargetType: "billing_settlement", TargetID: item.ID,
|
||||
RequestIP: requestIP(r), UserAgent: r.UserAgent(),
|
||||
AfterState: map[string]any{"status": item.Status, "attempts": item.Attempts},
|
||||
Metadata: map[string]any{"taskId": item.TaskID, "action": item.Action, "currency": item.Currency},
|
||||
}
|
||||
if actor != nil {
|
||||
input.ActorGatewayUserID = uuidText(firstNonEmptyText(actor.GatewayUserID, actor.ID))
|
||||
input.ActorUserID = actor.ID
|
||||
input.ActorUsername = actor.Username
|
||||
input.ActorSource = actor.Source
|
||||
input.ActorRoles = actor.Roles
|
||||
}
|
||||
return input
|
||||
}
|
||||
@@ -0,0 +1,21 @@
|
||||
package httpapi
|
||||
|
||||
import (
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestSingleIdempotencyKey(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
request := httptest.NewRequest("POST", "/", nil)
|
||||
request.Header.Add("Idempotency-Key", "retry-1")
|
||||
if got, ok := singleIdempotencyKey(request); !ok || got != "retry-1" {
|
||||
t.Fatalf("got key=%q ok=%v", got, ok)
|
||||
}
|
||||
|
||||
request.Header.Add("Idempotency-Key", "retry-2")
|
||||
if _, ok := singleIdempotencyKey(request); ok {
|
||||
t.Fatal("multiple Idempotency-Key values must be rejected")
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -48,6 +48,7 @@ func TestCoreLocalFlow(t *testing.T) {
|
||||
DatabaseURL: databaseURL,
|
||||
IdentityMode: "hybrid",
|
||||
JWTSecret: "test-secret",
|
||||
BillingEngineMode: "enforce",
|
||||
TaskProgressCallbackEnabled: true,
|
||||
TaskProgressCallbackURL: "http://callback.local/task-progress",
|
||||
CORSAllowedOrigin: "*",
|
||||
@@ -444,7 +445,6 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
|
||||
t.Fatalf("unexpected compatible chat response: %+v", compatChat)
|
||||
}
|
||||
|
||||
cancelMarker := "cancel-stream-" + suffixText
|
||||
cancelCtx, cancelRequest := context.WithCancel(context.Background())
|
||||
cancelPayload := map[string]any{
|
||||
"model": defaultTextModel,
|
||||
@@ -453,7 +453,6 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
|
||||
"stream": true,
|
||||
"simulation": true,
|
||||
"simulationDurationMs": 250,
|
||||
"cancelTestId": cancelMarker,
|
||||
}
|
||||
cancelRaw, err := json.Marshal(cancelPayload)
|
||||
if err != nil {
|
||||
@@ -466,15 +465,27 @@ VALUES ($1, 5, '{"purpose":"core-flow"}'::jsonb)`, inviteCode); err != nil {
|
||||
cancelReq.Header.Set("Authorization", "Bearer "+apiKeyResponse.Secret)
|
||||
cancelReq.Header.Set("Content-Type", "application/json")
|
||||
cancelErrCh := make(chan error, 1)
|
||||
cancelTaskIDCh := make(chan string, 1)
|
||||
go func() {
|
||||
resp, err := http.DefaultClient.Do(cancelReq)
|
||||
if resp != nil {
|
||||
cancelTaskIDCh <- strings.TrimSpace(resp.Header.Get("X-Gateway-Task-Id"))
|
||||
_, _ = io.ReadAll(resp.Body)
|
||||
_ = resp.Body.Close()
|
||||
} else {
|
||||
cancelTaskIDCh <- ""
|
||||
}
|
||||
cancelErrCh <- err
|
||||
}()
|
||||
cancelTaskID := waitForTaskIDByRequestMarker(t, ctx, testPool, cancelMarker, 2*time.Second)
|
||||
var cancelTaskID string
|
||||
select {
|
||||
case cancelTaskID = <-cancelTaskIDCh:
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("cancelled stream did not return response headers")
|
||||
}
|
||||
if cancelTaskID == "" {
|
||||
t.Fatal("cancelled stream response did not expose X-Gateway-Task-Id")
|
||||
}
|
||||
cancelRequest()
|
||||
select {
|
||||
case <-cancelErrCh:
|
||||
@@ -493,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",
|
||||
@@ -501,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)
|
||||
}
|
||||
@@ -513,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",
|
||||
@@ -521,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)
|
||||
}
|
||||
@@ -804,6 +819,7 @@ LIMIT 1`).Scan(&gptImageModelTypesRaw); err != nil {
|
||||
"currency": "resource",
|
||||
"rules": []map[string]any{
|
||||
{"ruleKey": "text_input", "displayName": "Text Input", "resourceType": "text_input", "unit": "1k_tokens", "basePrice": 1},
|
||||
{"ruleKey": "text_cached_input", "displayName": "Text Cached Input", "resourceType": "text_cached_input", "unit": "1k_tokens", "basePrice": 0.5},
|
||||
{"ruleKey": "text_output", "displayName": "Text Output", "resourceType": "text_output", "unit": "1k_tokens", "basePrice": 2},
|
||||
{"ruleKey": "image", "displayName": "Image", "resourceType": "image", "unit": "image", "basePrice": 7},
|
||||
{"ruleKey": "image_edit", "displayName": "Image Edit", "resourceType": "image_edit", "unit": "image", "basePrice": 11},
|
||||
@@ -835,6 +851,101 @@ WHERE gateway_user_id = $1::uuid
|
||||
AND currency = 'resource'`, smokeGatewayUserID).Scan(&walletBalanceBefore); err != nil {
|
||||
t.Fatalf("read wallet balance before pricing task: %v", err)
|
||||
}
|
||||
var estimateSideEffectsBefore struct {
|
||||
Tasks int
|
||||
Transactions int
|
||||
Frozen string
|
||||
}
|
||||
if err := testPool.QueryRow(ctx, `
|
||||
SELECT (SELECT count(*) FROM gateway_tasks),
|
||||
(SELECT count(*) FROM gateway_wallet_transactions WHERE gateway_user_id=$1::uuid),
|
||||
(SELECT frozen_balance::text FROM gateway_wallet_accounts WHERE gateway_user_id=$1::uuid AND currency='resource')`,
|
||||
smokeGatewayUserID,
|
||||
).Scan(&estimateSideEffectsBefore.Tasks, &estimateSideEffectsBefore.Transactions, &estimateSideEffectsBefore.Frozen); err != nil {
|
||||
t.Fatalf("read estimate side effects before request: %v", err)
|
||||
}
|
||||
var pricingEstimate struct {
|
||||
TotalAmount float64 `json:"totalAmount"`
|
||||
ReservationAmount float64 `json:"reservationAmount"`
|
||||
CandidateCount int `json:"candidateCount"`
|
||||
PricingVersion string `json:"pricingVersion"`
|
||||
RequestFingerprint string `json:"requestFingerprint"`
|
||||
}
|
||||
doJSON(t, server.URL, http.MethodPost, "/api/v1/pricing/estimate", apiKeyResponse.Secret, map[string]any{
|
||||
"kind": "chat.completions", "model": pricingModel, "max_completion_tokens": 32,
|
||||
"messages": []map[string]any{{"role": "user", "content": "estimate only"}},
|
||||
}, http.StatusOK, &pricingEstimate)
|
||||
if pricingEstimate.TotalAmount <= 0 || pricingEstimate.ReservationAmount < pricingEstimate.TotalAmount ||
|
||||
pricingEstimate.CandidateCount != 1 || pricingEstimate.PricingVersion != "effective-pricing-v2" || pricingEstimate.RequestFingerprint == "" {
|
||||
t.Fatalf("unexpected effective pricing estimate: %+v", pricingEstimate)
|
||||
}
|
||||
var estimateSideEffectsAfter struct {
|
||||
Tasks int
|
||||
Transactions int
|
||||
Frozen string
|
||||
}
|
||||
if err := testPool.QueryRow(ctx, `
|
||||
SELECT (SELECT count(*) FROM gateway_tasks),
|
||||
(SELECT count(*) FROM gateway_wallet_transactions WHERE gateway_user_id=$1::uuid),
|
||||
(SELECT frozen_balance::text FROM gateway_wallet_accounts WHERE gateway_user_id=$1::uuid AND currency='resource')`,
|
||||
smokeGatewayUserID,
|
||||
).Scan(&estimateSideEffectsAfter.Tasks, &estimateSideEffectsAfter.Transactions, &estimateSideEffectsAfter.Frozen); err != nil {
|
||||
t.Fatalf("read estimate side effects after request: %v", err)
|
||||
}
|
||||
if estimateSideEffectsAfter != estimateSideEffectsBefore {
|
||||
t.Fatalf("pricing estimate must be read-only, before=%+v after=%+v", estimateSideEffectsBefore, estimateSideEffectsAfter)
|
||||
}
|
||||
|
||||
idempotencyPayload := map[string]any{
|
||||
"model": pricingModel, "runMode": "simulation", "simulation": true, "simulationDurationMs": 5,
|
||||
"messages": []map[string]any{{"role": "user", "content": "idempotent replay"}},
|
||||
}
|
||||
idempotencyHeaders := map[string]string{"Idempotency-Key": "chat-replay-" + suffixText}
|
||||
var idempotentFirst map[string]any
|
||||
firstReplayHeaders := doJSONWithHeaders(t, server.URL, http.MethodPost, "/api/v1/chat/completions", apiKeyResponse.Secret, idempotencyPayload, idempotencyHeaders, http.StatusOK, &idempotentFirst)
|
||||
var idempotentSecond map[string]any
|
||||
secondReplayHeaders := doJSONWithHeaders(t, server.URL, http.MethodPost, "/api/v1/chat/completions", apiKeyResponse.Secret, idempotencyPayload, idempotencyHeaders, http.StatusOK, &idempotentSecond)
|
||||
idempotentTaskID := firstReplayHeaders.Get("X-Gateway-Task-Id")
|
||||
if idempotentTaskID == "" || secondReplayHeaders.Get("X-Gateway-Task-Id") != idempotentTaskID || secondReplayHeaders.Get("Idempotent-Replayed") != "true" {
|
||||
t.Fatalf("non-stream replay headers first=%v second=%v", firstReplayHeaders, secondReplayHeaders)
|
||||
}
|
||||
var idempotentAttempts int
|
||||
if err := testPool.QueryRow(ctx, `SELECT count(*) FROM gateway_task_attempts WHERE task_id=$1::uuid`, idempotentTaskID).Scan(&idempotentAttempts); err != nil {
|
||||
t.Fatalf("count idempotent task attempts: %v", err)
|
||||
}
|
||||
if idempotentAttempts != 1 {
|
||||
t.Fatalf("idempotent request called upstream %d times, want 1", idempotentAttempts)
|
||||
}
|
||||
idempotencyPayload["messages"] = []map[string]any{{"role": "user", "content": "different request"}}
|
||||
doJSONWithHeaders(t, server.URL, http.MethodPost, "/api/v1/chat/completions", apiKeyResponse.Secret, idempotencyPayload, idempotencyHeaders, http.StatusConflict, nil)
|
||||
|
||||
streamPayload := map[string]any{
|
||||
"model": pricingModel, "runMode": "simulation", "simulation": true, "simulationDurationMs": 5, "stream": true,
|
||||
"messages": []map[string]any{{"role": "user", "content": "stream replay"}},
|
||||
}
|
||||
streamRaw, err := json.Marshal(streamPayload)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
streamKey := "stream-replay-" + suffixText
|
||||
streamRequest, err := http.NewRequest(http.MethodPost, server.URL+"/api/v1/chat/completions", bytes.NewReader(streamRaw))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
streamRequest.Header.Set("Authorization", "Bearer "+apiKeyResponse.Secret)
|
||||
streamRequest.Header.Set("Content-Type", "application/json")
|
||||
streamRequest.Header.Set("Idempotency-Key", streamKey)
|
||||
streamResponse, err := http.DefaultClient.Do(streamRequest)
|
||||
if err != nil {
|
||||
t.Fatalf("execute first idempotent stream: %v", err)
|
||||
}
|
||||
_, _ = io.ReadAll(streamResponse.Body)
|
||||
_ = streamResponse.Body.Close()
|
||||
if streamResponse.StatusCode != http.StatusOK || streamResponse.Header.Get("X-Gateway-Task-Id") == "" {
|
||||
t.Fatalf("first idempotent stream status=%d headers=%v", streamResponse.StatusCode, streamResponse.Header)
|
||||
}
|
||||
doJSONWithHeaders(t, server.URL, http.MethodPost, "/api/v1/chat/completions", apiKeyResponse.Secret, streamPayload,
|
||||
map[string]string{"Idempotency-Key": streamKey}, http.StatusConflict, nil)
|
||||
var pricingTask struct {
|
||||
Task struct {
|
||||
ID string `json:"id"`
|
||||
@@ -845,7 +956,7 @@ WHERE gateway_user_id = $1::uuid
|
||||
}
|
||||
doAPIV1ChatCompletionAndLoadTask(t, ctx, testPool, server.URL, apiKeyResponse.Secret, map[string]any{
|
||||
"model": pricingModel,
|
||||
"runMode": "simulation",
|
||||
"runMode": "production",
|
||||
"simulation": true,
|
||||
"simulationDurationMs": 5,
|
||||
"messages": []map[string]any{{"role": "user", "content": "priced ping"}},
|
||||
@@ -853,6 +964,20 @@ WHERE gateway_user_id = $1::uuid
|
||||
if pricingTask.Task.Status != "succeeded" || !floatNear(pricingTask.Task.FinalChargeAmount, 0.028) {
|
||||
t.Fatalf("custom pricing rule set should drive text billing, got task=%+v", pricingTask.Task)
|
||||
}
|
||||
settlementDeadline := time.Now().Add(3 * time.Second)
|
||||
for {
|
||||
var billingStatus string
|
||||
if err := testPool.QueryRow(ctx, `SELECT billing_status FROM gateway_tasks WHERE id = $1::uuid`, pricingTask.Task.ID).Scan(&billingStatus); err != nil {
|
||||
t.Fatalf("read pricing task billing status: %v", err)
|
||||
}
|
||||
if billingStatus == "settled" {
|
||||
break
|
||||
}
|
||||
if time.Now().After(settlementDeadline) {
|
||||
t.Fatalf("pricing task billing did not settle before deadline, status=%s", billingStatus)
|
||||
}
|
||||
time.Sleep(25 * time.Millisecond)
|
||||
}
|
||||
var walletBalanceAfter float64
|
||||
var walletSpentAfter float64
|
||||
if err := testPool.QueryRow(ctx, `
|
||||
@@ -1075,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)
|
||||
}
|
||||
@@ -1097,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)
|
||||
}
|
||||
@@ -1484,7 +1614,7 @@ WHERE m.platform_id = $1::uuid
|
||||
t.Fatalf("workspace task list should include persisted task records, got %+v", workspaceTaskList.Items)
|
||||
}
|
||||
|
||||
req, err := http.NewRequest(http.MethodGet, server.URL+"/api/v1/tasks/"+taskResponse.Task.ID+"/events", nil)
|
||||
req, err := http.NewRequest(http.MethodGet, server.URL+"/api/v1/tasks/"+pricingTask.Task.ID+"/events", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("build events request: %v", err)
|
||||
}
|
||||
@@ -1533,6 +1663,17 @@ WHERE m.platform_id = $1::uuid
|
||||
t.Fatal("task progress callback outbox should receive events")
|
||||
}
|
||||
|
||||
restartModel := "worker-restart-" + suffixText
|
||||
createSimulationTextPlatformModel(
|
||||
t,
|
||||
server.URL,
|
||||
loginResponse.AccessToken,
|
||||
"openai-worker-restart-"+suffixText,
|
||||
"OpenAI Worker Restart",
|
||||
restartModel,
|
||||
1,
|
||||
nil,
|
||||
)
|
||||
var restartAsyncTask struct {
|
||||
TaskID string `json:"taskId"`
|
||||
Task struct {
|
||||
@@ -1542,7 +1683,7 @@ WHERE m.platform_id = $1::uuid
|
||||
} `json:"task"`
|
||||
}
|
||||
doJSONWithHeaders(t, server.URL, http.MethodPost, "/api/v1/responses", apiKeyResponse.Secret, map[string]any{
|
||||
"model": defaultTextModel,
|
||||
"model": restartModel,
|
||||
"runMode": "simulation",
|
||||
"simulation": true,
|
||||
"simulationDurationMs": 2000,
|
||||
@@ -1564,6 +1705,7 @@ WHERE m.platform_id = $1::uuid
|
||||
DatabaseURL: databaseURL,
|
||||
IdentityMode: "hybrid",
|
||||
JWTSecret: "test-secret",
|
||||
BillingEngineMode: "enforce",
|
||||
TaskProgressCallbackEnabled: true,
|
||||
TaskProgressCallbackURL: "http://callback.local/task-progress",
|
||||
CORSAllowedOrigin: "*",
|
||||
@@ -1864,7 +2006,7 @@ func doJSON(t *testing.T, baseURL string, method string, path string, token stri
|
||||
doJSONWithHeaders(t, baseURL, method, path, token, payload, nil, expectedStatus, out)
|
||||
}
|
||||
|
||||
func doJSONWithHeaders(t *testing.T, baseURL string, method string, path string, token string, payload any, headers map[string]string, expectedStatus int, out any) {
|
||||
func doJSONWithHeaders(t *testing.T, baseURL string, method string, path string, token string, payload any, headers map[string]string, expectedStatus int, out any) http.Header {
|
||||
t.Helper()
|
||||
var body io.Reader
|
||||
if payload != nil {
|
||||
@@ -1901,16 +2043,22 @@ func doJSONWithHeaders(t *testing.T, baseURL string, method string, path string,
|
||||
t.Fatalf("decode %s %s response: %v body=%s", method, path, err, string(raw))
|
||||
}
|
||||
}
|
||||
return resp.Header.Clone()
|
||||
}
|
||||
|
||||
func doAPIV1ChatCompletionAndLoadTask(t *testing.T, ctx context.Context, pool *pgxpool.Pool, baseURL string, token string, payload map[string]any, marker string, expectedStatus int, responseOut any, taskDetailOut any) string {
|
||||
t.Helper()
|
||||
payload["integrationTestMarker"] = marker
|
||||
_ = ctx
|
||||
_ = pool
|
||||
_ = marker
|
||||
if responseOut == nil {
|
||||
responseOut = &map[string]any{}
|
||||
}
|
||||
doJSON(t, baseURL, http.MethodPost, "/api/v1/chat/completions", token, payload, expectedStatus, responseOut)
|
||||
taskID := waitForTaskIDByRequestField(t, ctx, pool, "integrationTestMarker", marker, 2*time.Second)
|
||||
responseHeaders := doJSONWithHeaders(t, baseURL, http.MethodPost, "/api/v1/chat/completions", token, payload, nil, expectedStatus, responseOut)
|
||||
taskID := strings.TrimSpace(responseHeaders.Get("X-Gateway-Task-Id"))
|
||||
if taskID == "" {
|
||||
t.Fatal("chat completion response did not expose X-Gateway-Task-Id")
|
||||
}
|
||||
if taskDetailOut != nil {
|
||||
doJSON(t, baseURL, http.MethodGet, "/api/v1/tasks/"+taskID, token, nil, http.StatusOK, taskDetailOut)
|
||||
}
|
||||
@@ -2046,11 +2194,6 @@ func waitForTaskStatus(t *testing.T, baseURL string, token string, taskID string
|
||||
return detail
|
||||
}
|
||||
|
||||
func waitForTaskIDByRequestMarker(t *testing.T, ctx context.Context, pool *pgxpool.Pool, marker string, timeout time.Duration) string {
|
||||
t.Helper()
|
||||
return waitForTaskIDByRequestField(t, ctx, pool, "cancelTestId", marker, timeout)
|
||||
}
|
||||
|
||||
func waitForTaskIDByRequestField(t *testing.T, ctx context.Context, pool *pgxpool.Pool, key string, value string, timeout time.Duration) string {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(timeout)
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -4,6 +4,7 @@ import (
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"mime"
|
||||
@@ -39,9 +40,9 @@ type geminiUploadSession struct {
|
||||
}
|
||||
|
||||
var geminiGenerateContentRoutePrefixes = []string{
|
||||
"/api/v1/models/",
|
||||
"/v1beta/models/",
|
||||
"/v1/models/",
|
||||
"/models/",
|
||||
}
|
||||
|
||||
func (s *Server) registerGeminiGenerateContentRoutes(mux *http.ServeMux) {
|
||||
@@ -74,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 {
|
||||
@@ -94,6 +110,11 @@ func (s *Server) geminiGenerateContent(w http.ResponseWriter, r *http.Request) {
|
||||
writeError(w, http.StatusForbidden, "api key scope does not allow this capability")
|
||||
return
|
||||
}
|
||||
idempotencyKey, hasIdempotencyKey, err := optionalTaskIdempotencyKey(r)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, "Idempotency-Key must contain one non-empty value", "invalid_idempotency_key")
|
||||
return
|
||||
}
|
||||
prepared, err := s.prepareTaskRequest(r.Context(), r, user, mapping.Body)
|
||||
if err != nil {
|
||||
s.logger.Warn("prepare gemini task request failed", "kind", mapping.Kind, "error", err)
|
||||
@@ -104,7 +125,7 @@ func (s *Server) geminiGenerateContent(w http.ResponseWriter, r *http.Request) {
|
||||
writeError(w, status, err.Error(), clients.ErrorCode(err))
|
||||
return
|
||||
}
|
||||
task, err := s.store.CreateTask(r.Context(), store.CreateTaskInput{
|
||||
createInput := store.CreateTaskInput{
|
||||
Kind: mapping.Kind,
|
||||
Model: mapping.Model,
|
||||
RunMode: runModeFromRequest(prepared.Body),
|
||||
@@ -113,10 +134,33 @@ func (s *Server) geminiGenerateContent(w http.ResponseWriter, r *http.Request) {
|
||||
ConversationID: prepared.ConversationID,
|
||||
NewMessageCount: prepared.NewMessageCount,
|
||||
MessageRefs: prepared.MessageRefs,
|
||||
}, user)
|
||||
}
|
||||
if hasIdempotencyKey {
|
||||
createInput.IdempotencyKeyHash = taskIdempotencyKeyHash(idempotencyKey)
|
||||
createInput.IdempotencyRequestHash = taskIdempotencyRequestHash(mapping.Kind, false, false, prepared.Body)
|
||||
}
|
||||
created, err := s.store.CreateTaskIdempotent(r.Context(), createInput, user)
|
||||
if err != nil {
|
||||
s.logger.Error("create gemini task failed", "kind", mapping.Kind, "error", err)
|
||||
writeError(w, http.StatusInternalServerError, "create task failed")
|
||||
if errors.Is(err, store.ErrIdempotencyKeyReused) {
|
||||
writeError(w, http.StatusConflict, "Idempotency-Key was reused for a different request", "idempotency_key_reused")
|
||||
return
|
||||
}
|
||||
s.logger.Error("create gemini task failed", "kind", mapping.Kind, "error_category", "task_create_failed")
|
||||
writeError(w, http.StatusInternalServerError, "create task failed", "task_create_failed")
|
||||
return
|
||||
}
|
||||
task := created.Task
|
||||
if created.Replayed {
|
||||
if s.billingMetrics != nil {
|
||||
s.billingMetrics.ObserveBillingEvent("idempotent_replay")
|
||||
}
|
||||
w.Header().Set("Idempotent-Replayed", "true")
|
||||
w.Header().Set("X-Gateway-Task-Id", task.ID)
|
||||
if task.Status == "succeeded" {
|
||||
writeJSON(w, http.StatusOK, geminiGenerateContentResponse(task.Result, mapping.Model))
|
||||
return
|
||||
}
|
||||
writeIdempotentTaskReplay(w, task, true)
|
||||
return
|
||||
}
|
||||
runCtx, cancelRun := s.requestExecutionContext(r)
|
||||
@@ -366,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))
|
||||
@@ -383,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))
|
||||
@@ -411,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{
|
||||
|
||||
@@ -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"
|
||||
@@ -892,6 +927,7 @@ func (s *Server) deleteAPIKey(w http.ResponseWriter, r *http.Request) {
|
||||
// @Failure 403 {object} ErrorEnvelope
|
||||
// @Failure 404 {object} ErrorEnvelope
|
||||
// @Failure 429 {object} ErrorEnvelope
|
||||
// @Failure 503 {object} ErrorEnvelope
|
||||
// @Failure 500 {object} ErrorEnvelope
|
||||
// @Router /api/v1/pricing/estimate [post]
|
||||
func (s *Server) estimatePricing(w http.ResponseWriter, r *http.Request) {
|
||||
@@ -916,10 +952,24 @@ func (s *Server) estimatePricing(w http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
estimate, err := s.runner.Estimate(r.Context(), kind, model, body, user)
|
||||
if err != nil {
|
||||
if s.billingMetrics != nil {
|
||||
s.billingMetrics.ObserveBillingEvent("estimate_failed")
|
||||
}
|
||||
if runner.IsPricingUnavailable(err) {
|
||||
if s.billingMetrics != nil {
|
||||
s.billingMetrics.ObserveBillingEvent("pricing_unavailable")
|
||||
}
|
||||
writeErrorWithDetails(w, http.StatusServiceUnavailable, runErrorMessage(err), runErrorDetails(err), "pricing_unavailable")
|
||||
return
|
||||
}
|
||||
if errors.Is(err, store.ErrNoModelCandidate) {
|
||||
writeErrorWithDetails(w, statusFromRunError(err), runErrorMessage(err), runErrorDetails(err), store.ModelCandidateErrorCode(err))
|
||||
return
|
||||
}
|
||||
if code := clients.ErrorCode(err); code == "bad_request" || code == "invalid_parameter" {
|
||||
writeErrorWithDetails(w, http.StatusBadRequest, runErrorMessage(err), runErrorDetails(err), code)
|
||||
return
|
||||
}
|
||||
s.logger.Error("estimate pricing failed", "error", err)
|
||||
writeError(w, http.StatusInternalServerError, "estimate pricing failed")
|
||||
return
|
||||
@@ -971,12 +1021,13 @@ 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
|
||||
// @Security BearerAuth
|
||||
// @Param X-Async header bool false "true 时异步创建任务并返回 202"
|
||||
// @Param Idempotency-Key header string false "可选请求幂等键;同一用户范围内唯一"
|
||||
// @Param input body TaskRequest true "AI 任务请求,字段随任务类型变化"
|
||||
// @Success 200 {object} CompatibleResponse
|
||||
// @Success 202 {object} TaskAcceptedResponse
|
||||
@@ -996,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())
|
||||
@@ -1035,16 +1070,25 @@ 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
|
||||
}
|
||||
responsePlan := planTaskResponse(kind, compatible, body, r)
|
||||
idempotencyKey, hasIdempotencyKey, err := optionalTaskIdempotencyKey(r)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, "Idempotency-Key must contain one non-empty value", "invalid_idempotency_key")
|
||||
return
|
||||
}
|
||||
prepared, err := s.prepareTaskRequest(r.Context(), r, user, body)
|
||||
if err != nil {
|
||||
s.logger.Warn("prepare task request failed", "kind", kind, "error", err)
|
||||
@@ -1056,7 +1100,7 @@ func (s *Server) createTask(kind string, compatible bool) http.Handler {
|
||||
return
|
||||
}
|
||||
|
||||
task, err := s.store.CreateTask(r.Context(), store.CreateTaskInput{
|
||||
createInput := store.CreateTaskInput{
|
||||
Kind: kind,
|
||||
Model: model,
|
||||
RunMode: runModeFromRequest(prepared.Body),
|
||||
@@ -1065,10 +1109,32 @@ func (s *Server) createTask(kind string, compatible bool) http.Handler {
|
||||
ConversationID: prepared.ConversationID,
|
||||
NewMessageCount: prepared.NewMessageCount,
|
||||
MessageRefs: prepared.MessageRefs,
|
||||
}, user)
|
||||
}
|
||||
if hasIdempotencyKey {
|
||||
createInput.IdempotencyKeyHash = taskIdempotencyKeyHash(idempotencyKey)
|
||||
createInput.IdempotencyRequestHash = taskIdempotencyRequestHash(kind, responsePlan.asyncMode, responsePlan.streamMode, prepared.Body)
|
||||
}
|
||||
created, err := s.store.CreateTaskIdempotent(r.Context(), createInput, user)
|
||||
if err != nil {
|
||||
s.logger.Error("create task failed", "kind", kind, "error", err)
|
||||
writeError(w, http.StatusInternalServerError, "create task failed")
|
||||
if errors.Is(err, store.ErrIdempotencyKeyReused) {
|
||||
writeError(w, http.StatusConflict, "Idempotency-Key was reused for a different request", "idempotency_key_reused")
|
||||
return
|
||||
}
|
||||
s.logger.Error("create task failed", "kind", kind, "error_category", "task_create_failed")
|
||||
writeError(w, http.StatusInternalServerError, "create task failed", "task_create_failed")
|
||||
return
|
||||
}
|
||||
task := created.Task
|
||||
if created.Replayed {
|
||||
if s.billingMetrics != nil {
|
||||
s.billingMetrics.ObserveBillingEvent("idempotent_replay")
|
||||
}
|
||||
w.Header().Set("Idempotent-Replayed", "true")
|
||||
if responsePlan.streamMode {
|
||||
writeError(w, http.StatusConflict, "streaming idempotent replay is not supported", "idempotency_stream_replay_unsupported")
|
||||
return
|
||||
}
|
||||
writeIdempotentTaskReplay(w, task, responsePlan.compatibleMode)
|
||||
return
|
||||
}
|
||||
if responsePlan.asyncMode {
|
||||
@@ -1136,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
|
||||
@@ -1156,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() {}
|
||||
|
||||
@@ -1344,6 +1406,17 @@ func requestModelName(body map[string]any) string {
|
||||
return modelNameFromValue(body["model"])
|
||||
}
|
||||
|
||||
func canonicalTaskModelName(kind string, model string) string {
|
||||
model = strings.TrimSpace(model)
|
||||
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:
|
||||
@@ -1383,6 +1456,10 @@ func scopeForTaskKind(kind string) string {
|
||||
|
||||
func statusFromRunError(err error) int {
|
||||
switch {
|
||||
case clients.ErrorCode(err) == "billing_hold":
|
||||
return http.StatusServiceUnavailable
|
||||
case runner.IsPricingUnavailable(err):
|
||||
return http.StatusServiceUnavailable
|
||||
case clients.ErrorCode(err) == "invalid_previous_response_id" || clients.ErrorCode(err) == "response_chain_too_deep" || clients.ErrorCode(err) == "unsupported_model_protocol" || clients.ErrorCode(err) == "unsupported_response_tool" || clients.ErrorCode(err) == "unsupported_response_parameter":
|
||||
return http.StatusBadRequest
|
||||
case clients.ErrorCode(err) == "response_chain_unavailable":
|
||||
@@ -1411,6 +1488,9 @@ func statusFromRunError(err error) int {
|
||||
}
|
||||
|
||||
func runErrorCode(err error) string {
|
||||
if runner.IsPricingUnavailable(err) {
|
||||
return "pricing_unavailable"
|
||||
}
|
||||
if errors.Is(err, store.ErrNoModelCandidate) {
|
||||
return store.ModelCandidateErrorCode(err)
|
||||
}
|
||||
@@ -1428,6 +1508,9 @@ func runErrorMessage(err error) string {
|
||||
}
|
||||
|
||||
func runErrorDetails(err error) map[string]any {
|
||||
if detail := runner.PricingUnavailableDetails(err); len(detail) > 0 {
|
||||
return map[string]any{"pricing": detail}
|
||||
}
|
||||
if detail := rateLimitErrorDetail(err); len(detail) > 0 {
|
||||
return map[string]any{"rateLimit": detail}
|
||||
}
|
||||
@@ -1572,7 +1655,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 {
|
||||
@@ -1681,7 +1763,6 @@ 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) {
|
||||
task, err := s.store.GetTask(r.Context(), r.PathValue("taskID"))
|
||||
if err == nil {
|
||||
@@ -1714,8 +1795,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 {
|
||||
@@ -1752,7 +1831,6 @@ 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) {
|
||||
task, err := s.store.GetTask(r.Context(), r.PathValue("taskID"))
|
||||
if err != nil {
|
||||
@@ -1786,7 +1864,6 @@ 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) {
|
||||
task, err := s.store.GetTask(r.Context(), r.PathValue("taskID"))
|
||||
if err != nil {
|
||||
|
||||
@@ -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-video;Bearer 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, ®istered)
|
||||
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 == "" ||
|
||||
|
||||
@@ -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, ®isterResponse)
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
@@ -43,7 +43,7 @@ func TestOIDCJITFullGatewayUserFlowAndSecurityBoundary(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Fatalf("connect store: %v", err)
|
||||
}
|
||||
defer db.Close()
|
||||
t.Cleanup(db.Close)
|
||||
|
||||
key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
|
||||
if err != nil {
|
||||
@@ -247,6 +247,9 @@ DELETE FROM gateway_users WHERE source = 'oidc' AND external_user_id = ANY($1::t
|
||||
t.Fatalf("rejected OIDC tokens created %d local users", rejectedWrites)
|
||||
}
|
||||
|
||||
if _, err := db.DisableActiveIdentityRevision(ctx, activeRevision.Version, "oidc-jit-disabled", "oidc-jit-disabled"); err != nil {
|
||||
t.Fatalf("disable initial OIDC JIT revision: %v", err)
|
||||
}
|
||||
disabledJITRevision := prepareOIDCJITRevision(t, ctx, db, baseConfig, issuer, "default", false)
|
||||
testRevisionIDs = append(testRevisionIDs, disabledJITRevision.ID)
|
||||
disabledJITRevision, _, err = db.ActivateIdentityRevision(ctx, disabledJITRevision.ID, disabledJITRevision.Version, "oidc-jit-disabled", "oidc-jit-disabled")
|
||||
@@ -298,6 +301,9 @@ func prepareOIDCJITRevision(t *testing.T, ctx context.Context, db *store.Store,
|
||||
if err := secrets.Put(ctx, sessionReference, bytes.Repeat([]byte{8}, 32)); err != nil {
|
||||
t.Fatalf("store OIDC JIT session key: %v", err)
|
||||
}
|
||||
if err := db.QueueIdentitySecretCleanup(ctx, sessionReference, time.Now().Add(10*time.Minute)); err != nil {
|
||||
t.Fatalf("stage OIDC JIT session key for adoption: %v", err)
|
||||
}
|
||||
draft, err = db.ApplyIdentityManifest(ctx, draft.ID, draft.Version, identity.ManifestApplication{
|
||||
Manifest: identity.ManifestV1{
|
||||
SchemaVersion: 1, Issuer: issuer, TenantID: oidcJITTenantID, ApplicationID: uuid.NewString(),
|
||||
@@ -360,6 +366,10 @@ var currentOIDCTestNonce string
|
||||
|
||||
func createOIDCBFFSessionCookie(t *testing.T, baseURL string) *http.Cookie {
|
||||
t.Helper()
|
||||
parsedBaseURL, err := url.Parse(baseURL)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
jar, err := cookiejar.New(nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
@@ -371,6 +381,11 @@ func createOIDCBFFSessionCookie(t *testing.T, baseURL string) *http.Cookie {
|
||||
if len(via) > 10 {
|
||||
return errors.New("too many redirects")
|
||||
}
|
||||
// Follow the synthetic issuer's authorization redirect, then stop before
|
||||
// the callback redirects the browser to the separately hosted web app.
|
||||
if request.URL.Host != parsedBaseURL.Host && request.URL.Path != "/authorize" {
|
||||
return http.ErrUseLastResponse
|
||||
}
|
||||
return nil
|
||||
}}
|
||||
response, err := client.Get(baseURL + "/api/v1/auth/oidc/login?returnTo=%2Fapi%2Fv1%2Fme")
|
||||
@@ -378,11 +393,10 @@ func createOIDCBFFSessionCookie(t *testing.T, baseURL string) *http.Cookie {
|
||||
t.Fatalf("complete OIDC BFF login: %v", err)
|
||||
}
|
||||
defer response.Body.Close()
|
||||
if response.StatusCode != http.StatusOK {
|
||||
if response.StatusCode != http.StatusOK && response.StatusCode != http.StatusSeeOther {
|
||||
body, _ := io.ReadAll(response.Body)
|
||||
t.Fatalf("OIDC BFF login status = %d, want 200: %s", response.StatusCode, body)
|
||||
}
|
||||
parsedBaseURL, _ := url.Parse(baseURL)
|
||||
for _, cookie := range jar.Cookies(parsedBaseURL) {
|
||||
if cookie.Name == auth.OIDCSessionCookieName {
|
||||
if strings.Count(cookie.Value, ".") == 2 {
|
||||
|
||||
@@ -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 {
|
||||
@@ -114,6 +114,18 @@ type AuditLogListResponse struct {
|
||||
Items []store.AuditLog `json:"items"`
|
||||
}
|
||||
|
||||
type BillingSettlementListResponse struct {
|
||||
Items []store.BillingSettlement `json:"items"`
|
||||
Total int `json:"total" example:"42"`
|
||||
Page int `json:"page" example:"1"`
|
||||
PageSize int `json:"pageSize" example:"50"`
|
||||
}
|
||||
|
||||
type BillingSettlementRetryResponse struct {
|
||||
Settlement store.BillingSettlement `json:"settlement"`
|
||||
AuditLog store.AuditLog `json:"auditLog"`
|
||||
}
|
||||
|
||||
type WalletTransactionListResponse struct {
|
||||
Items []store.GatewayWalletTransaction `json:"items"`
|
||||
Total int `json:"total" example:"42"`
|
||||
@@ -188,8 +200,14 @@ type PricingEstimateRequest struct {
|
||||
}
|
||||
|
||||
type PricingEstimateResponse struct {
|
||||
Items []map[string]interface{} `json:"items"`
|
||||
Resolver string `json:"resolver" example:"effective-pricing-v1"`
|
||||
Items []map[string]interface{} `json:"items"`
|
||||
Resolver string `json:"resolver" example:"effective-pricing-v2"`
|
||||
TotalAmount float64 `json:"totalAmount" example:"1.25"`
|
||||
ReservationAmount float64 `json:"reservationAmount" example:"2.75"`
|
||||
Currency string `json:"currency" example:"resource"`
|
||||
CandidateCount int `json:"candidateCount" example:"2"`
|
||||
PricingVersion string `json:"pricingVersion" example:"effective-pricing-v2"`
|
||||
RequestFingerprint string `json:"requestFingerprint" example:"76ef6a537de8e71bd1ca93acadc078dbdbfa9f17e45224e4f9df59f535d2886f"`
|
||||
}
|
||||
|
||||
type TaskRequest struct {
|
||||
@@ -211,6 +229,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"`
|
||||
|
||||
@@ -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 JSON,category 必须是 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 ""
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
package httpapi
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"testing"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/runner"
|
||||
)
|
||||
|
||||
func TestPricingUnavailableUsesStructuredServiceUnavailableError(t *testing.T) {
|
||||
err := &runner.PricingUnavailableError{
|
||||
Reason: "missing, invalid, or not explicitly free",
|
||||
ResourceType: "image",
|
||||
RuleSetID: "rule-set-id",
|
||||
}
|
||||
if got := statusFromRunError(err); got != http.StatusServiceUnavailable {
|
||||
t.Fatalf("statusFromRunError()=%d, want %d", got, http.StatusServiceUnavailable)
|
||||
}
|
||||
if got := runErrorCode(err); got != "pricing_unavailable" {
|
||||
t.Fatalf("runErrorCode()=%q, want pricing_unavailable", got)
|
||||
}
|
||||
details := runErrorDetails(err)
|
||||
pricing, _ := details["pricing"].(map[string]any)
|
||||
if pricing["resourceType"] != "image" || pricing["ruleSetId"] != "rule-set-id" {
|
||||
t.Fatalf("unexpected pricing details: %+v", details)
|
||||
}
|
||||
}
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
@@ -146,10 +147,45 @@ func validPricingRuleSetInput(input store.PricingRuleSetInput) bool {
|
||||
if strings.TrimSpace(input.RuleSetKey) == "" || strings.TrimSpace(input.Name) == "" || len(input.Rules) == 0 {
|
||||
return false
|
||||
}
|
||||
if currency := strings.TrimSpace(input.Currency); currency != "" && currency != "resource" {
|
||||
return false
|
||||
}
|
||||
for _, rule := range input.Rules {
|
||||
if strings.TrimSpace(rule.ResourceType) == "" || strings.TrimSpace(rule.Unit) == "" {
|
||||
return false
|
||||
}
|
||||
if rule.BasePrice < 0 || (rule.BasePrice == 0 && !rule.IsFree) {
|
||||
return false
|
||||
}
|
||||
if currency := strings.TrimSpace(rule.Currency); currency != "" && currency != "resource" {
|
||||
return false
|
||||
}
|
||||
switch calculator := strings.TrimSpace(rule.CalculatorType); calculator {
|
||||
case "", "token_usage", "unit_weight", "duration_weight":
|
||||
default:
|
||||
return false
|
||||
}
|
||||
calculator := strings.TrimSpace(rule.CalculatorType)
|
||||
if calculator == "" {
|
||||
calculator = store.DefaultEffectivePricingCalculator(rule.ResourceType)
|
||||
}
|
||||
if store.ValidateEffectivePricingRuleShape(strings.TrimSpace(rule.ResourceType), store.NormalizeEffectivePricingRuleUnit(rule.ResourceType, rule.Unit), calculator) != nil {
|
||||
return false
|
||||
}
|
||||
effectiveFrom, fromOK := pricingEffectiveTime(rule.EffectiveFrom)
|
||||
effectiveTo, toOK := pricingEffectiveTime(rule.EffectiveTo)
|
||||
if !fromOK || !toOK || (!effectiveFrom.IsZero() && !effectiveTo.IsZero() && !effectiveFrom.Before(effectiveTo)) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func pricingEffectiveTime(value string) (time.Time, bool) {
|
||||
value = strings.TrimSpace(value)
|
||||
if value == "" {
|
||||
return time.Time{}, true
|
||||
}
|
||||
parsed, err := time.Parse(time.RFC3339, value)
|
||||
return parsed, err == nil
|
||||
}
|
||||
|
||||
@@ -0,0 +1,28 @@
|
||||
package httpapi
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func TestPricingRuleSetRejectsImplicitZeroAndUnsupportedCalculator(t *testing.T) {
|
||||
base := store.PricingRuleSetInput{
|
||||
RuleSetKey: "test", Name: "test",
|
||||
Rules: []store.PricingRuleInput{{
|
||||
ResourceType: "image", Unit: "image", BasePrice: 0,
|
||||
Currency: "resource", CalculatorType: "unit_weight",
|
||||
}},
|
||||
}
|
||||
if validPricingRuleSetInput(base) {
|
||||
t.Fatal("implicit zero price must be rejected")
|
||||
}
|
||||
base.Rules[0].IsFree = true
|
||||
if !validPricingRuleSetInput(base) {
|
||||
t.Fatal("explicit free price should be accepted")
|
||||
}
|
||||
base.Rules[0].CalculatorType = "formula"
|
||||
if validPricingRuleSetInput(base) {
|
||||
t.Fatal("arbitrary formula calculator must be rejected")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,72 @@
|
||||
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/videos/generations",
|
||||
"/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)
|
||||
|
||||
@@ -43,6 +43,7 @@ type Server struct {
|
||||
identityTestRevision identity.Revision
|
||||
identityTestCookieSecure bool
|
||||
identityTestBrowserEnabled bool
|
||||
billingMetrics *ssfreceiver.Metrics
|
||||
}
|
||||
|
||||
type oidcPublicClient interface {
|
||||
@@ -66,18 +67,19 @@ func NewServer(cfg config.Config, db *store.Store, logger *slog.Logger) http.Han
|
||||
}
|
||||
|
||||
func NewServerWithContext(ctx context.Context, cfg config.Config, db *store.Store, logger *slog.Logger) http.Handler {
|
||||
securityEventMetrics := &ssfreceiver.Metrics{}
|
||||
server := &Server{
|
||||
ctx: ctx,
|
||||
cfg: cfg,
|
||||
store: db,
|
||||
oidcUserResolver: db,
|
||||
auth: auth.New(cfg.JWTSecret, cfg.ServerMainBaseURL, cfg.ServerMainInternalToken),
|
||||
runner: runner.New(cfg, db, logger),
|
||||
runner: runner.New(cfg, db, logger, securityEventMetrics),
|
||||
logger: logger,
|
||||
billingMetrics: securityEventMetrics,
|
||||
}
|
||||
server.auth.ServerMainInternalKey = cfg.ServerMainInternalKey
|
||||
server.auth.ServerMainInternalSecret = cfg.ServerMainInternalSecret
|
||||
securityEventMetrics := &ssfreceiver.Metrics{}
|
||||
secretStore, err := identitySecretStore(cfg)
|
||||
if err != nil {
|
||||
panic("invalid identity SecretStore: " + err.Error())
|
||||
@@ -118,18 +120,23 @@ func NewServerWithContext(ctx context.Context, cfg config.Config, db *store.Stor
|
||||
}
|
||||
server.auth.LocalAPIKeyVerifier = db.VerifyLocalAPIKey
|
||||
server.runner.StartAsyncQueueWorker(ctx)
|
||||
server.runner.StartBillingSettlementWorker(ctx)
|
||||
server.startLocalTempAssetCleanup(ctx)
|
||||
server.startOIDCSessionCleanup(ctx)
|
||||
|
||||
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)
|
||||
|
||||
@@ -179,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)))
|
||||
@@ -204,6 +212,8 @@ func NewServerWithContext(ctx context.Context, cfg config.Config, db *store.Stor
|
||||
mux.Handle("DELETE /api/admin/runtime/policy-sets/{policySetID}", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.deleteRuntimePolicySet)))
|
||||
mux.Handle("GET /api/admin/runtime/runner-policy", server.requireAdmin(auth.PermissionPower, http.HandlerFunc(server.getRunnerPolicy)))
|
||||
mux.Handle("PATCH /api/admin/runtime/runner-policy", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.updateRunnerPolicy)))
|
||||
mux.Handle("GET /api/admin/runtime/billing-settlements", server.requireAdmin(auth.PermissionPower, http.HandlerFunc(server.listBillingSettlements)))
|
||||
mux.Handle("POST /api/admin/runtime/billing-settlements/{settlementId}/retry", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.retryBillingSettlement)))
|
||||
mux.Handle("POST /api/admin/runtime/model-rate-limits/{platformModelID}/restore", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.restorePlatformModelRuntimeStatus)))
|
||||
mux.Handle("GET /api/admin/config/network-proxy", server.requireAdmin(auth.PermissionPower, http.HandlerFunc(server.getNetworkProxyConfig)))
|
||||
mux.Handle("GET /api/admin/system/file-storage/settings", server.requireAdmin(auth.PermissionPower, http.HandlerFunc(server.getFileStorageSettings)))
|
||||
@@ -246,12 +256,14 @@ 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/videos/generations", server.requireUser(auth.PermissionBasic, server.createTask("videos.generations", true)))
|
||||
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)))
|
||||
@@ -259,9 +271,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)))
|
||||
@@ -284,6 +311,10 @@ 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 /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)))
|
||||
|
||||
@@ -0,0 +1,85 @@
|
||||
package httpapi
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
var errInvalidTaskIdempotencyKey = errors.New("invalid Idempotency-Key header")
|
||||
|
||||
func optionalTaskIdempotencyKey(r *http.Request) (string, bool, error) {
|
||||
values := r.Header.Values("Idempotency-Key")
|
||||
if len(values) == 0 {
|
||||
return "", false, nil
|
||||
}
|
||||
if len(values) != 1 {
|
||||
return "", false, errInvalidTaskIdempotencyKey
|
||||
}
|
||||
value := strings.TrimSpace(values[0])
|
||||
if value == "" || len(value) > 255 || strings.Contains(value, ",") {
|
||||
return "", false, errInvalidTaskIdempotencyKey
|
||||
}
|
||||
return value, true, nil
|
||||
}
|
||||
|
||||
func taskIdempotencyKeyHash(key string) string {
|
||||
digest := sha256.Sum256([]byte(key))
|
||||
return hex.EncodeToString(digest[:])
|
||||
}
|
||||
|
||||
func taskIdempotencyRequestHash(kind string, async bool, stream bool, body map[string]any) string {
|
||||
payload, _ := json.Marshal(map[string]any{
|
||||
"kind": kind, "async": async, "stream": stream, "request": body,
|
||||
})
|
||||
digest := sha256.Sum256(payload)
|
||||
return hex.EncodeToString(digest[:])
|
||||
}
|
||||
|
||||
func writeIdempotentTaskReplay(w http.ResponseWriter, task store.GatewayTask, compatible bool) {
|
||||
w.Header().Set("Idempotent-Replayed", "true")
|
||||
w.Header().Set("X-Gateway-Task-Id", task.ID)
|
||||
if !compatible || task.AsyncMode || (task.Status != "succeeded" && task.Status != "failed" && task.Status != "cancelled") {
|
||||
writeTaskAccepted(w, task)
|
||||
return
|
||||
}
|
||||
if task.Status == "succeeded" {
|
||||
writeJSON(w, http.StatusOK, task.Result)
|
||||
return
|
||||
}
|
||||
status := storedTaskErrorStatus(task.ErrorCode)
|
||||
message := strings.TrimSpace(task.ErrorMessage)
|
||||
if message == "" {
|
||||
message = strings.TrimSpace(task.Error)
|
||||
}
|
||||
if message == "" {
|
||||
message = "task failed"
|
||||
}
|
||||
code := strings.TrimSpace(task.ErrorCode)
|
||||
if code == "" {
|
||||
code = "task_failed"
|
||||
}
|
||||
writeError(w, status, message, code)
|
||||
}
|
||||
|
||||
func storedTaskErrorStatus(code string) int {
|
||||
switch strings.TrimSpace(code) {
|
||||
case "pricing_unavailable", "response_chain_unavailable", "billing_hold":
|
||||
return http.StatusServiceUnavailable
|
||||
case "insufficient_balance":
|
||||
return http.StatusPaymentRequired
|
||||
case "bad_request", "invalid_parameter", "invalid_previous_response_id", "unsupported_operation":
|
||||
return http.StatusBadRequest
|
||||
case "no_model_candidate", "cloned_voice_not_found":
|
||||
return http.StatusNotFound
|
||||
case "rate_limit", "platform_cooling_down", "model_cooling_down":
|
||||
return http.StatusTooManyRequests
|
||||
default:
|
||||
return http.StatusBadGateway
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
package httpapi
|
||||
|
||||
import (
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestOptionalTaskIdempotencyKeyRejectsMultipleValues(t *testing.T) {
|
||||
t.Parallel()
|
||||
request := httptest.NewRequest("POST", "/", nil)
|
||||
if _, present, err := optionalTaskIdempotencyKey(request); err != nil || present {
|
||||
t.Fatalf("missing key present=%v err=%v", present, err)
|
||||
}
|
||||
request.Header.Add("Idempotency-Key", "one")
|
||||
request.Header.Add("Idempotency-Key", "two")
|
||||
if _, _, err := optionalTaskIdempotencyKey(request); err == nil {
|
||||
t.Fatal("multiple keys must be rejected")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskIdempotencyRequestHashIsCanonical(t *testing.T) {
|
||||
t.Parallel()
|
||||
first := map[string]any{"model": "m", "n": float64(1), "nested": map[string]any{"b": true, "a": "x"}}
|
||||
second := map[string]any{"nested": map[string]any{"a": "x", "b": true}, "n": float64(1), "model": "m"}
|
||||
if taskIdempotencyRequestHash("images.generations", false, false, first) != taskIdempotencyRequestHash("images.generations", false, false, second) {
|
||||
t.Fatal("equivalent JSON objects must have the same request hash")
|
||||
}
|
||||
if taskIdempotencyRequestHash("images.generations", true, false, first) == taskIdempotencyRequestHash("images.generations", false, false, first) {
|
||||
t.Fatal("async response semantics must be part of the request hash")
|
||||
}
|
||||
}
|
||||
@@ -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,330 @@
|
||||
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
|
||||
// @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
|
||||
// @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
|
||||
// @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)
|
||||
}
|
||||
}
|
||||
@@ -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,27 @@
|
||||
package runner
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestNormalizedBillingEngineMode(t *testing.T) {
|
||||
t.Parallel()
|
||||
if got := normalizedBillingEngineMode(""); got != "observe" {
|
||||
t.Fatalf("empty mode = %q", got)
|
||||
}
|
||||
if got := normalizedBillingEngineMode("ENFORCE"); got != "enforce" {
|
||||
t.Fatalf("enforce mode = %q", got)
|
||||
}
|
||||
if got := normalizedBillingEngineMode("hold"); got != "hold" {
|
||||
t.Fatalf("hold mode = %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBillingItemsFixedTotalKeepsNineDecimalPlaces(t *testing.T) {
|
||||
t.Parallel()
|
||||
items := []any{
|
||||
map[string]any{"amount": "0.000000001"},
|
||||
map[string]any{"amount": float64(0.000000002)},
|
||||
}
|
||||
if got := billingItemsFixedTotal(items).String(); got != "0.000000003" {
|
||||
t.Fatalf("total = %s", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,96 @@
|
||||
package runner
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
const billingSettlementPollInterval = time.Second
|
||||
|
||||
func billingSettlementRetryDelay(attempt int) time.Duration {
|
||||
if attempt < 1 {
|
||||
attempt = 1
|
||||
}
|
||||
delay := time.Second
|
||||
for current := 1; current < attempt && delay < 15*time.Minute; current++ {
|
||||
delay *= 2
|
||||
}
|
||||
if delay > 15*time.Minute {
|
||||
return 15 * time.Minute
|
||||
}
|
||||
return delay
|
||||
}
|
||||
|
||||
func billingSettlementErrorCode(err error) string {
|
||||
if errors.Is(err, store.ErrInsufficientWalletBalance) {
|
||||
return "insufficient_balance"
|
||||
}
|
||||
return "settlement_failed"
|
||||
}
|
||||
|
||||
func billingSettlementErrorMessage(code string) string {
|
||||
if code == "insufficient_balance" {
|
||||
return "wallet balance is insufficient for settlement"
|
||||
}
|
||||
return "billing settlement processing failed"
|
||||
}
|
||||
|
||||
func (s *Service) StartBillingSettlementWorker(ctx context.Context) {
|
||||
workerID := "billing-" + uuid.NewString()
|
||||
go s.runBillingSettlementWorker(ctx, workerID)
|
||||
}
|
||||
|
||||
func (s *Service) runBillingSettlementWorker(ctx context.Context, workerID string) {
|
||||
ticker := time.NewTicker(billingSettlementPollInterval)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
s.processBillingSettlementBatch(ctx, workerID)
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) processBillingSettlementBatch(ctx context.Context, workerID string) {
|
||||
items, err := s.store.ClaimBillingSettlements(ctx, workerID, store.BillingSettlementBatchSize, store.BillingSettlementLockTimeout)
|
||||
if err != nil {
|
||||
if ctx.Err() == nil {
|
||||
s.logger.Error("claim billing settlements failed", "error_category", "billing_settlement_claim_failed")
|
||||
}
|
||||
return
|
||||
}
|
||||
for _, item := range items {
|
||||
if ctx.Err() != nil {
|
||||
return
|
||||
}
|
||||
if err := s.store.ProcessBillingSettlement(ctx, item); err != nil {
|
||||
code := billingSettlementErrorCode(err)
|
||||
markErr := s.store.MarkBillingSettlementFailed(
|
||||
context.WithoutCancel(ctx),
|
||||
item,
|
||||
code,
|
||||
billingSettlementErrorMessage(code),
|
||||
billingSettlementRetryDelay(item.Attempts),
|
||||
)
|
||||
if markErr != nil {
|
||||
s.logger.Error("mark billing settlement failed", "settlementID", item.ID, "taskID", item.TaskID, "error_category", "billing_settlement_state_failed")
|
||||
continue
|
||||
}
|
||||
s.observeBillingEvent("settlement_retry")
|
||||
if item.Attempts >= store.BillingSettlementMaxAttempts {
|
||||
s.observeBillingEvent("manual_review")
|
||||
}
|
||||
s.logger.Warn("billing settlement scheduled for retry", "settlementID", item.ID, "taskID", item.TaskID, "action", item.Action, "error_category", code, "attempt", item.Attempts)
|
||||
continue
|
||||
}
|
||||
s.observeBillingEvent("settlement_completed")
|
||||
s.logger.Debug("billing settlement completed", "settlementID", item.ID, "taskID", item.TaskID, "action", item.Action)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,40 @@
|
||||
package runner
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func TestBillingSettlementRetryDelay(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
attempt int
|
||||
want time.Duration
|
||||
}{
|
||||
{attempt: 1, want: time.Second},
|
||||
{attempt: 2, want: 2 * time.Second},
|
||||
{attempt: 10, want: 512 * time.Second},
|
||||
{attempt: 11, want: 15 * time.Minute},
|
||||
{attempt: 20, want: 15 * time.Minute},
|
||||
}
|
||||
for _, test := range tests {
|
||||
if got := billingSettlementRetryDelay(test.attempt); got != test.want {
|
||||
t.Fatalf("attempt %d: got %s, want %s", test.attempt, got, test.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestBillingSettlementErrorCode(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
if got := billingSettlementErrorCode(store.ErrInsufficientWalletBalance); got != "insufficient_balance" {
|
||||
t.Fatalf("got %q", got)
|
||||
}
|
||||
if got := billingSettlementErrorCode(errors.New("boom")); got != "settlement_failed" {
|
||||
t.Fatalf("got %q", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
package runner
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
const (
|
||||
taskExecutionLeaseTTL = 5 * time.Minute
|
||||
taskExecutionRenewInterval = 30 * time.Second
|
||||
)
|
||||
|
||||
func (s *Service) renewTaskExecutionLease(ctx context.Context, cancel context.CancelFunc, taskID string, executionToken string) {
|
||||
ticker := time.NewTicker(taskExecutionRenewInterval)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
if err := s.store.RenewTaskExecutionLease(ctx, taskID, executionToken, taskExecutionLeaseTTL); err != nil {
|
||||
if errors.Is(err, store.ErrTaskExecutionFinished) {
|
||||
return
|
||||
}
|
||||
s.logger.Warn("task execution lease lost", "taskID", taskID, "error_category", "task_execution_lease_lost")
|
||||
cancel()
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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" }
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -10,11 +10,17 @@ 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"`
|
||||
TotalAmount float64 `json:"totalAmount"`
|
||||
Currency string `json:"currency"`
|
||||
Items []any `json:"items"`
|
||||
Resolver string `json:"resolver"`
|
||||
TotalAmount float64 `json:"totalAmount"`
|
||||
ReservationAmount float64 `json:"reservationAmount"`
|
||||
Currency string `json:"currency"`
|
||||
CandidateCount int `json:"candidateCount"`
|
||||
PricingVersion string `json:"pricingVersion"`
|
||||
RequestFingerprint string `json:"requestFingerprint"`
|
||||
}
|
||||
|
||||
func (s *Service) Estimate(ctx context.Context, kind string, model string, body map[string]any, user *auth.User) (EstimateResult, error) {
|
||||
@@ -32,15 +38,19 @@ func (s *Service) Estimate(ctx context.Context, kind string, model string, body
|
||||
if err != nil {
|
||||
return EstimateResult{}, err
|
||||
}
|
||||
candidate := candidates[0]
|
||||
body = preprocessRequest(kind, body, candidate)
|
||||
items := s.estimatedBillings(ctx, user, kind, body, candidate)
|
||||
return EstimateResult{
|
||||
Items: items,
|
||||
Resolver: "effective-pricing-v1",
|
||||
TotalAmount: totalBillingAmount(items),
|
||||
Currency: billingCurrency(items),
|
||||
}, nil
|
||||
estimates := make([]candidateEstimate, 0, len(candidates))
|
||||
for _, candidate := range candidates {
|
||||
preprocessing := s.preprocessRequestWithScripts(ctx, kind, body, candidate)
|
||||
if preprocessing.Err != nil {
|
||||
return EstimateResult{}, parameterPreprocessClientError(preprocessing.Err)
|
||||
}
|
||||
estimate, candidateErr := s.estimateCandidateV2(ctx, user, kind, preprocessing.Body, candidate)
|
||||
if candidateErr != nil {
|
||||
return EstimateResult{}, candidateErr
|
||||
}
|
||||
estimates = append(estimates, estimate)
|
||||
}
|
||||
return buildEstimateResult(estimates, pricingRequestFingerprint(kind, model, body))
|
||||
}
|
||||
|
||||
func (s *Service) estimatedBillings(ctx context.Context, user *auth.User, kind string, body map[string]any, candidate store.RuntimeModelCandidate) []any {
|
||||
@@ -127,7 +137,7 @@ func (s *Service) billings(ctx context.Context, user *auth.User, kind string, bo
|
||||
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") *
|
||||
@@ -136,7 +146,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,
|
||||
@@ -184,24 +194,25 @@ func (s *Service) billings(ctx context.Context, user *auth.User, kind string, bo
|
||||
}
|
||||
|
||||
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 {
|
||||
@@ -408,6 +419,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 {
|
||||
@@ -468,11 +489,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
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
@@ -172,13 +172,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 {
|
||||
|
||||
@@ -0,0 +1,784 @@
|
||||
package runner
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"math"
|
||||
"math/big"
|
||||
"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/store"
|
||||
)
|
||||
|
||||
const (
|
||||
pricingVersionV2 = "effective-pricing-v2"
|
||||
fixedScale = int64(1_000_000_000)
|
||||
maxFixedAmount = fixedAmount(math.MaxInt64)
|
||||
)
|
||||
|
||||
var ErrPricingUnavailable = errors.New("pricing unavailable")
|
||||
var errFixedAmountOverflow = errors.New("fixed amount overflow")
|
||||
|
||||
type PricingUnavailableError struct {
|
||||
Reason string
|
||||
ResourceType string
|
||||
RuleSetID string
|
||||
}
|
||||
|
||||
func (e *PricingUnavailableError) Error() string {
|
||||
message := "pricing unavailable"
|
||||
if e.ResourceType != "" {
|
||||
message += " for " + e.ResourceType
|
||||
}
|
||||
if e.Reason != "" {
|
||||
message += ": " + e.Reason
|
||||
}
|
||||
return message
|
||||
}
|
||||
|
||||
func (e *PricingUnavailableError) Unwrap() error { return ErrPricingUnavailable }
|
||||
|
||||
func isPricingUnavailable(err error) bool { return errors.Is(err, ErrPricingUnavailable) }
|
||||
|
||||
func IsPricingUnavailable(err error) bool { return errors.Is(err, ErrPricingUnavailable) }
|
||||
|
||||
func PricingUnavailableDetails(err error) map[string]any {
|
||||
var pricingErr *PricingUnavailableError
|
||||
if !errors.As(err, &pricingErr) {
|
||||
return nil
|
||||
}
|
||||
details := map[string]any{"reason": pricingErr.Reason}
|
||||
if pricingErr.ResourceType != "" {
|
||||
details["resourceType"] = pricingErr.ResourceType
|
||||
}
|
||||
if pricingErr.RuleSetID != "" {
|
||||
details["ruleSetId"] = pricingErr.RuleSetID
|
||||
}
|
||||
return details
|
||||
}
|
||||
|
||||
// fixedAmount is a signed decimal with exactly nine fractional digits.
|
||||
// Monetary decisions use this type; float64 is only produced at JSON-compatible boundaries.
|
||||
type fixedAmount int64
|
||||
|
||||
func parseFixedAmount(value string) (fixedAmount, error) {
|
||||
value = strings.TrimSpace(value)
|
||||
if value == "" {
|
||||
return 0, fmt.Errorf("empty decimal")
|
||||
}
|
||||
sign := int64(1)
|
||||
if value[0] == '-' || value[0] == '+' {
|
||||
if value[0] == '-' {
|
||||
sign = -1
|
||||
}
|
||||
value = value[1:]
|
||||
}
|
||||
parts := strings.Split(value, ".")
|
||||
if len(parts) > 2 || value == "" {
|
||||
return 0, fmt.Errorf("invalid decimal %q", value)
|
||||
}
|
||||
whole := parts[0]
|
||||
if whole == "" {
|
||||
whole = "0"
|
||||
}
|
||||
fraction := ""
|
||||
if len(parts) == 2 {
|
||||
fraction = parts[1]
|
||||
}
|
||||
for _, digits := range []string{whole, fraction} {
|
||||
for _, char := range digits {
|
||||
if char < '0' || char > '9' {
|
||||
return 0, fmt.Errorf("invalid decimal %q", value)
|
||||
}
|
||||
}
|
||||
}
|
||||
roundUp := false
|
||||
if len(fraction) > 9 {
|
||||
roundUp = fraction[9] >= '5'
|
||||
fraction = fraction[:9]
|
||||
}
|
||||
fraction += strings.Repeat("0", 9-len(fraction))
|
||||
combined := strings.TrimLeft(whole+fraction, "0")
|
||||
if combined == "" {
|
||||
combined = "0"
|
||||
}
|
||||
number := new(big.Int)
|
||||
if _, ok := number.SetString(combined, 10); !ok {
|
||||
return 0, fmt.Errorf("invalid decimal %q", value)
|
||||
}
|
||||
if roundUp {
|
||||
number.Add(number, big.NewInt(1))
|
||||
}
|
||||
if !number.IsInt64() {
|
||||
return 0, fmt.Errorf("decimal %q exceeds fixed amount range", value)
|
||||
}
|
||||
return fixedAmount(sign * number.Int64()), nil
|
||||
}
|
||||
|
||||
func fixedAmountFromAny(value any) (fixedAmount, error) {
|
||||
switch typed := value.(type) {
|
||||
case fixedAmount:
|
||||
return typed, nil
|
||||
case string:
|
||||
return parseFixedAmount(typed)
|
||||
case json.Number:
|
||||
return parseFixedAmount(typed.String())
|
||||
case int:
|
||||
return parseFixedAmount(strconv.FormatInt(int64(typed), 10))
|
||||
case int64:
|
||||
return parseFixedAmount(strconv.FormatInt(typed, 10))
|
||||
case int32:
|
||||
return parseFixedAmount(strconv.FormatInt(int64(typed), 10))
|
||||
case float64:
|
||||
if math.IsNaN(typed) || math.IsInf(typed, 0) {
|
||||
return 0, fmt.Errorf("non-finite decimal")
|
||||
}
|
||||
return parseFixedAmount(strconv.FormatFloat(typed, 'f', 9, 64))
|
||||
case float32:
|
||||
return fixedAmountFromAny(float64(typed))
|
||||
default:
|
||||
return 0, fmt.Errorf("unsupported decimal type %T", value)
|
||||
}
|
||||
}
|
||||
|
||||
func (a fixedAmount) String() string {
|
||||
value := int64(a)
|
||||
sign := ""
|
||||
if value < 0 {
|
||||
sign = "-"
|
||||
value = -value
|
||||
}
|
||||
return fmt.Sprintf("%s%d.%09d", sign, value/fixedScale, value%fixedScale)
|
||||
}
|
||||
|
||||
func (a fixedAmount) Float64() float64 { return float64(a) / float64(fixedScale) }
|
||||
func (a fixedAmount) IsZero() bool { return a == 0 }
|
||||
|
||||
func addFixedAmounts(left fixedAmount, right fixedAmount) (fixedAmount, error) {
|
||||
value := new(big.Int).Add(big.NewInt(int64(left)), big.NewInt(int64(right)))
|
||||
return fixedAmountFromBigInt(value)
|
||||
}
|
||||
|
||||
func multiplyFixedAmountByInt(amount fixedAmount, multiplier int) (fixedAmount, error) {
|
||||
value := new(big.Int).Mul(big.NewInt(int64(amount)), big.NewInt(int64(multiplier)))
|
||||
return fixedAmountFromBigInt(value)
|
||||
}
|
||||
|
||||
func multiplyFixedAmounts(left fixedAmount, right fixedAmount) (fixedAmount, error) {
|
||||
product := new(big.Int).Mul(big.NewInt(int64(left)), big.NewInt(int64(right)))
|
||||
return fixedAmountFromBigInt(roundBigIntRatio(product, big.NewInt(fixedScale)))
|
||||
}
|
||||
|
||||
func multiplyFixedAmountRatio(amount fixedAmount, numerator int, denominator int) (fixedAmount, error) {
|
||||
if denominator == 0 {
|
||||
return 0, fmt.Errorf("division by zero")
|
||||
}
|
||||
product := new(big.Int).Mul(big.NewInt(int64(amount)), big.NewInt(int64(numerator)))
|
||||
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
|
||||
}
|
||||
return fixedAmount(value.Int64()), nil
|
||||
}
|
||||
|
||||
func roundBigIntRatio(numerator *big.Int, denominator *big.Int) *big.Int {
|
||||
negative := numerator.Sign() < 0
|
||||
absolute := new(big.Int).Abs(new(big.Int).Set(numerator))
|
||||
quotient, remainder := new(big.Int), new(big.Int)
|
||||
quotient.QuoRem(absolute, denominator, remainder)
|
||||
if new(big.Int).Mul(remainder, big.NewInt(2)).Cmp(denominator) >= 0 {
|
||||
quotient.Add(quotient, big.NewInt(1))
|
||||
}
|
||||
if negative {
|
||||
quotient.Neg(quotient)
|
||||
}
|
||||
return quotient
|
||||
}
|
||||
|
||||
type resolvedPricing struct {
|
||||
Config map[string]any
|
||||
Currency string
|
||||
RuleSetID string
|
||||
RuleSetKey string
|
||||
Source string
|
||||
FreeResource map[string]bool
|
||||
Snapshot map[string]any
|
||||
}
|
||||
|
||||
func (pricing resolvedPricing) requiredPrice(resource string, keys ...string) (fixedAmount, error) {
|
||||
if price, found, err := pricing.price(resource, keys...); err != nil {
|
||||
return 0, &PricingUnavailableError{Reason: err.Error(), ResourceType: resource, RuleSetID: pricing.RuleSetID}
|
||||
} else if found && price > 0 {
|
||||
return price, nil
|
||||
} else if found && price == 0 && pricing.FreeResource[resource] {
|
||||
return 0, nil
|
||||
}
|
||||
return 0, &PricingUnavailableError{Reason: "missing, invalid, or not explicitly free", ResourceType: resource, RuleSetID: pricing.RuleSetID}
|
||||
}
|
||||
|
||||
func (pricing resolvedPricing) price(resource string, keys ...string) (fixedAmount, bool, error) {
|
||||
for _, key := range keys {
|
||||
if value, ok := pricing.Config[key]; ok {
|
||||
amount, err := fixedAmountFromAny(value)
|
||||
return amount, true, err
|
||||
}
|
||||
}
|
||||
resourceConfig, _ := pricing.Config[resource].(map[string]any)
|
||||
if len(resourceConfig) == 0 && resource == "image_edit" {
|
||||
resourceConfig, _ = pricing.Config["image"].(map[string]any)
|
||||
}
|
||||
for _, key := range keys {
|
||||
if value, ok := resourceConfig[key]; ok {
|
||||
amount, err := fixedAmountFromAny(value)
|
||||
return amount, true, err
|
||||
}
|
||||
}
|
||||
if formula, ok := resourceConfig["formulaConfig"].(map[string]any); ok {
|
||||
for _, key := range keys {
|
||||
if value, exists := formula[key]; exists {
|
||||
amount, err := fixedAmountFromAny(value)
|
||||
return amount, true, err
|
||||
}
|
||||
}
|
||||
}
|
||||
if value, ok := resourceConfig["basePrice"]; ok {
|
||||
amount, err := fixedAmountFromAny(value)
|
||||
return amount, true, err
|
||||
}
|
||||
if resource == "text" {
|
||||
return resolvedPricing{Config: pricing.Config}.price("text_total", keys...)
|
||||
}
|
||||
return 0, false, nil
|
||||
}
|
||||
|
||||
func (pricing resolvedPricing) weight(resource string, key string, name string) (fixedAmount, error) {
|
||||
if strings.TrimSpace(name) == "" {
|
||||
return fixedAmount(fixedScale), nil
|
||||
}
|
||||
keys := weightKeyAliases(key)
|
||||
names := weightValueAliases(key, name)
|
||||
value, found := pricingWeightValue(pricing.Config, resource, keys, names)
|
||||
if !found {
|
||||
return fixedAmount(fixedScale), nil
|
||||
}
|
||||
weight, err := fixedAmountFromAny(value)
|
||||
if err != nil || weight <= 0 {
|
||||
reason := fmt.Sprintf("invalid %s weight for %s", key, name)
|
||||
if err != nil {
|
||||
reason += ": " + err.Error()
|
||||
}
|
||||
return 0, &PricingUnavailableError{Reason: reason, ResourceType: resource, RuleSetID: pricing.RuleSetID}
|
||||
}
|
||||
return weight, nil
|
||||
}
|
||||
|
||||
func (pricing resolvedPricing) calculate(resource string, base fixedAmount, integerFactors []int, fixedFactors ...fixedAmount) (fixedAmount, error) {
|
||||
amount := base
|
||||
var err error
|
||||
for _, factor := range integerFactors {
|
||||
amount, err = multiplyFixedAmountByInt(amount, factor)
|
||||
if err != nil {
|
||||
return 0, pricing.calculationError(resource, err)
|
||||
}
|
||||
}
|
||||
for _, factor := range fixedFactors {
|
||||
amount, err = multiplyFixedAmounts(amount, factor)
|
||||
if err != nil {
|
||||
return 0, pricing.calculationError(resource, err)
|
||||
}
|
||||
}
|
||||
return amount, nil
|
||||
}
|
||||
|
||||
func (pricing resolvedPricing) calculateRatio(resource string, base fixedAmount, numerator int, denominator int, fixedFactors ...fixedAmount) (fixedAmount, error) {
|
||||
amount, err := multiplyFixedAmountRatio(base, numerator, denominator)
|
||||
if err != nil {
|
||||
return 0, pricing.calculationError(resource, err)
|
||||
}
|
||||
return pricing.calculate(resource, amount, nil, fixedFactors...)
|
||||
}
|
||||
|
||||
func (pricing resolvedPricing) add(resource string, left fixedAmount, right fixedAmount) (fixedAmount, error) {
|
||||
amount, err := addFixedAmounts(left, right)
|
||||
if err != nil {
|
||||
return 0, pricing.calculationError(resource, err)
|
||||
}
|
||||
return amount, nil
|
||||
}
|
||||
|
||||
func (pricing resolvedPricing) calculationError(resource string, err error) error {
|
||||
return &PricingUnavailableError{Reason: "pricing calculation failed: " + err.Error(), ResourceType: resource, RuleSetID: pricing.RuleSetID}
|
||||
}
|
||||
|
||||
func pricingWeightValue(config map[string]any, resource string, keys []string, names []string) (any, bool) {
|
||||
if value, ok := pricingWeightValueFromConfig(config, keys, names); ok {
|
||||
return value, true
|
||||
}
|
||||
resourceConfig, _ := config[resource].(map[string]any)
|
||||
if len(resourceConfig) == 0 && resource == "image_edit" {
|
||||
resourceConfig, _ = config["image"].(map[string]any)
|
||||
}
|
||||
return pricingWeightValueFromConfig(resourceConfig, keys, names)
|
||||
}
|
||||
|
||||
func pricingWeightValueFromConfig(config map[string]any, keys []string, names []string) (any, bool) {
|
||||
if len(config) == 0 {
|
||||
return nil, false
|
||||
}
|
||||
for _, key := range keys {
|
||||
weights, _ := config[key].(map[string]any)
|
||||
for _, name := range names {
|
||||
if value, ok := weights[name]; ok {
|
||||
return value, true
|
||||
}
|
||||
}
|
||||
}
|
||||
dynamic, _ := config["dynamicWeight"].(map[string]any)
|
||||
for _, name := range names {
|
||||
if value, ok := dynamic[name]; ok {
|
||||
return value, true
|
||||
}
|
||||
}
|
||||
for _, key := range keys {
|
||||
weights, _ := dynamic[key].(map[string]any)
|
||||
for _, name := range names {
|
||||
if value, ok := weights[name]; ok {
|
||||
return value, true
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil, false
|
||||
}
|
||||
|
||||
type candidateEstimate struct {
|
||||
Items []any
|
||||
Amount fixedAmount
|
||||
Currency string
|
||||
Snapshot map[string]any
|
||||
Pricing resolvedPricing
|
||||
}
|
||||
|
||||
func buildEstimateResult(estimates []candidateEstimate, fingerprint string) (EstimateResult, error) {
|
||||
if len(estimates) == 0 {
|
||||
return EstimateResult{}, &PricingUnavailableError{Reason: "no candidate has effective pricing"}
|
||||
}
|
||||
preferred := estimates[0]
|
||||
reservation := preferred.Amount
|
||||
currency := preferred.Currency
|
||||
if currency == "" {
|
||||
currency = "resource"
|
||||
}
|
||||
for _, estimate := range estimates[1:] {
|
||||
candidateCurrency := estimate.Currency
|
||||
if candidateCurrency == "" {
|
||||
candidateCurrency = "resource"
|
||||
}
|
||||
if candidateCurrency != currency {
|
||||
return EstimateResult{}, &PricingUnavailableError{Reason: "candidate currencies do not match"}
|
||||
}
|
||||
if estimate.Amount > reservation {
|
||||
reservation = estimate.Amount
|
||||
}
|
||||
}
|
||||
return EstimateResult{
|
||||
Items: preferred.Items,
|
||||
Resolver: pricingVersionV2,
|
||||
TotalAmount: preferred.Amount.Float64(),
|
||||
ReservationAmount: reservation.Float64(),
|
||||
Currency: currency,
|
||||
CandidateCount: len(estimates),
|
||||
PricingVersion: pricingVersionV2,
|
||||
RequestFingerprint: fingerprint,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func estimatedOutputTokens(body map[string]any, candidate store.RuntimeModelCandidate) int {
|
||||
for _, key := range []string{"max_completion_tokens", "max_output_tokens", "max_tokens"} {
|
||||
if value, ok := positiveInteger(body[key]); ok {
|
||||
return value
|
||||
}
|
||||
}
|
||||
if limit, _, _, ok := candidateMaxOutputTokens(candidate, candidate.ModelType); ok {
|
||||
return limit
|
||||
}
|
||||
return 4096
|
||||
}
|
||||
|
||||
func pricingRequestFingerprint(kind string, model string, body map[string]any) string {
|
||||
payload, _ := json.Marshal(map[string]any{"kind": kind, "model": model, "request": body})
|
||||
digest := sha256.Sum256(payload)
|
||||
return hex.EncodeToString(digest[:])
|
||||
}
|
||||
|
||||
func (s *Service) resolveEffectivePricing(ctx context.Context, candidate store.RuntimeModelCandidate) (resolvedPricing, error) {
|
||||
ruleSetID := firstNonEmptyString(
|
||||
candidate.ModelPricingRuleSetID,
|
||||
candidate.PlatformPricingRuleSetID,
|
||||
candidate.BasePricingRuleSetID,
|
||||
)
|
||||
if ruleSetID != "" {
|
||||
if s.store == nil {
|
||||
return resolvedPricing{}, &PricingUnavailableError{Reason: "pricing store is unavailable", RuleSetID: ruleSetID}
|
||||
}
|
||||
config, err := s.store.PricingRuleSetBillingConfigV2(ctx, ruleSetID)
|
||||
if err != nil {
|
||||
return resolvedPricing{}, &PricingUnavailableError{Reason: err.Error(), RuleSetID: ruleSetID}
|
||||
}
|
||||
if len(candidate.BillingConfigOverride) > 0 {
|
||||
config.Config = mergeMap(config.Config, candidate.BillingConfigOverride)
|
||||
}
|
||||
return resolvedPricing{
|
||||
Config: config.Config, Currency: config.Currency, RuleSetID: config.RuleSetID,
|
||||
RuleSetKey: config.RuleSetKey, Source: pricingRuleSource(candidate, ruleSetID),
|
||||
FreeResource: config.FreeResource, Snapshot: config.Snapshot,
|
||||
}, nil
|
||||
}
|
||||
|
||||
config := candidate.BaseBillingConfig
|
||||
source := "base_model_config"
|
||||
if len(candidate.BillingConfig) > 0 {
|
||||
config = candidate.BillingConfig
|
||||
source = "platform_model_config"
|
||||
}
|
||||
if len(candidate.BillingConfigOverride) > 0 {
|
||||
config = mergeMap(config, candidate.BillingConfigOverride)
|
||||
source += "_override"
|
||||
}
|
||||
if len(config) == 0 {
|
||||
return resolvedPricing{}, &PricingUnavailableError{Reason: "candidate has no pricing configuration"}
|
||||
}
|
||||
freeResource := explicitFreeResources(config)
|
||||
snapshot := map[string]any{
|
||||
"pricingVersion": pricingVersionV2,
|
||||
"source": source,
|
||||
"currency": "resource",
|
||||
"config": config,
|
||||
}
|
||||
return resolvedPricing{
|
||||
Config: config, Currency: "resource", Source: source,
|
||||
FreeResource: freeResource, Snapshot: snapshot,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func pricingRuleSource(candidate store.RuntimeModelCandidate, ruleSetID string) string {
|
||||
switch ruleSetID {
|
||||
case candidate.ModelPricingRuleSetID:
|
||||
return "platform_model_rule_set"
|
||||
case candidate.PlatformPricingRuleSetID:
|
||||
return "platform_rule_set"
|
||||
default:
|
||||
return "base_model_rule_set"
|
||||
}
|
||||
}
|
||||
|
||||
func explicitFreeResources(config map[string]any) map[string]bool {
|
||||
resources := map[string]bool{}
|
||||
if isFree, _ := config["isFree"].(bool); isFree {
|
||||
for _, resource := range []string{"text_input", "text_cached_input", "text_output", "text_total", "image", "image_edit", "video", "music", "audio"} {
|
||||
resources[resource] = true
|
||||
}
|
||||
}
|
||||
for _, resource := range []string{"text", "text_total", "image", "image_edit", "video", "music", "audio"} {
|
||||
resourceConfig, _ := config[resource].(map[string]any)
|
||||
if isFree, _ := resourceConfig["isFree"].(bool); isFree {
|
||||
resources[resource] = true
|
||||
}
|
||||
}
|
||||
return resources
|
||||
}
|
||||
|
||||
func (s *Service) estimateCandidateV2(ctx context.Context, user *auth.User, kind string, body map[string]any, candidate store.RuntimeModelCandidate) (candidateEstimate, error) {
|
||||
usage := clients.Usage{InputTokens: estimateRequestTokens(body)}
|
||||
if isTextGenerationKind(kind) {
|
||||
usage.OutputTokens = estimatedOutputTokens(body, candidate)
|
||||
}
|
||||
usage.TotalTokens = usage.InputTokens + usage.OutputTokens
|
||||
response := clients.Response{Usage: usage}
|
||||
items, total, pricing, err := s.billingsV2(ctx, user, kind, body, candidate, response, true)
|
||||
if err != nil {
|
||||
return candidateEstimate{}, err
|
||||
}
|
||||
if isTextBillingKind(kind) {
|
||||
if _, err := pricing.requiredTextPrice("text_cached_input", "textCachedInputPer1k", "cachedInputTokenPrice", "cacheInputTokenPrice", "textInputCacheHitPer1k", "inputCacheHitTokenPrice"); err != nil {
|
||||
return candidateEstimate{}, err
|
||||
}
|
||||
}
|
||||
return candidateEstimate{
|
||||
Items: items, Amount: total, Currency: pricing.Currency, Snapshot: pricing.Snapshot, Pricing: pricing,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *Service) billingsV2(
|
||||
ctx context.Context,
|
||||
user *auth.User,
|
||||
kind string,
|
||||
body map[string]any,
|
||||
candidate store.RuntimeModelCandidate,
|
||||
response clients.Response,
|
||||
simulated bool,
|
||||
) ([]any, fixedAmount, resolvedPricing, error) {
|
||||
pricing, err := s.resolveEffectivePricing(ctx, candidate)
|
||||
if err != nil {
|
||||
return nil, 0, resolvedPricing{}, err
|
||||
}
|
||||
return s.billingsWithResolvedPricingV2(ctx, user, kind, body, candidate, response, simulated, pricing)
|
||||
}
|
||||
|
||||
func (s *Service) billingsWithResolvedPricingV2(
|
||||
ctx context.Context,
|
||||
user *auth.User,
|
||||
kind string,
|
||||
body map[string]any,
|
||||
candidate store.RuntimeModelCandidate,
|
||||
response clients.Response,
|
||||
simulated bool,
|
||||
pricing resolvedPricing,
|
||||
) ([]any, fixedAmount, resolvedPricing, error) {
|
||||
discount, err := fixedAmountFromAny(effectiveDiscount(ctx, s.store, user, candidate))
|
||||
if err != nil || discount <= 0 {
|
||||
return nil, 0, resolvedPricing{}, &PricingUnavailableError{Reason: "invalid discount factor", RuleSetID: pricing.RuleSetID}
|
||||
}
|
||||
buildLine := func(resourceType string, unit string, quantity any, amount fixedAmount, details map[string]any) map[string]any {
|
||||
line := billingLineWithDetails(candidate, resourceType, unit, quantity, amount.Float64(), discount.Float64(), simulated, details)
|
||||
line["pricingVersion"] = pricingVersionV2
|
||||
line["pricingSource"] = pricing.Source
|
||||
if pricing.RuleSetID != "" {
|
||||
line["pricingRuleSetId"] = pricing.RuleSetID
|
||||
}
|
||||
line["isFree"] = amount.IsZero()
|
||||
return line
|
||||
}
|
||||
|
||||
if isTextBillingKind(kind) {
|
||||
inputTokens := response.Usage.InputTokens
|
||||
outputTokens := response.Usage.OutputTokens
|
||||
cachedInputTokens := response.Usage.CachedInputTokens
|
||||
if isTextInputOnlyKind(kind) && inputTokens == 0 && response.Usage.TotalTokens > 0 {
|
||||
inputTokens = response.Usage.TotalTokens
|
||||
}
|
||||
if inputTokens == 0 && outputTokens == 0 {
|
||||
inputTokens = estimateRequestTokens(body)
|
||||
if isTextGenerationKind(kind) {
|
||||
outputTokens = 1
|
||||
}
|
||||
}
|
||||
if cachedInputTokens > inputTokens && inputTokens > 0 {
|
||||
cachedInputTokens = inputTokens
|
||||
}
|
||||
uncachedInputTokens := inputTokens - cachedInputTokens
|
||||
if uncachedInputTokens < 0 {
|
||||
uncachedInputTokens = 0
|
||||
}
|
||||
inputPrice, err := pricing.requiredTextPrice("text_input", "textInputPer1k", "inputTokenPrice", "basePrice")
|
||||
if err != nil {
|
||||
return nil, 0, resolvedPricing{}, err
|
||||
}
|
||||
items := make([]any, 0, 3)
|
||||
total := fixedAmount(0)
|
||||
if uncachedInputTokens > 0 || cachedInputTokens == 0 {
|
||||
amount, calculationErr := pricing.calculateRatio("text_input", inputPrice, uncachedInputTokens, 1000, discount)
|
||||
if calculationErr != nil {
|
||||
return nil, 0, resolvedPricing{}, calculationErr
|
||||
}
|
||||
total, calculationErr = pricing.add("text", total, amount)
|
||||
if calculationErr != nil {
|
||||
return nil, 0, resolvedPricing{}, calculationErr
|
||||
}
|
||||
items = append(items, buildLine("text_input", "1k_tokens", uncachedInputTokens, amount, map[string]any{
|
||||
"inputTokens": inputTokens, "uncachedInputTokens": uncachedInputTokens,
|
||||
"cachedInputTokens": cachedInputTokens, "pricePer1k": inputPrice.Float64(),
|
||||
}))
|
||||
}
|
||||
if cachedInputTokens > 0 {
|
||||
cachedPrice, priceErr := pricing.requiredTextPrice("text_cached_input", "textCachedInputPer1k", "cachedInputTokenPrice", "cacheInputTokenPrice", "textInputCacheHitPer1k", "inputCacheHitTokenPrice")
|
||||
if priceErr != nil {
|
||||
return nil, 0, resolvedPricing{}, priceErr
|
||||
}
|
||||
amount, calculationErr := pricing.calculateRatio("text_cached_input", cachedPrice, cachedInputTokens, 1000, discount)
|
||||
if calculationErr != nil {
|
||||
return nil, 0, resolvedPricing{}, calculationErr
|
||||
}
|
||||
total, calculationErr = pricing.add("text", total, amount)
|
||||
if calculationErr != nil {
|
||||
return nil, 0, resolvedPricing{}, calculationErr
|
||||
}
|
||||
items = append(items, buildLine("text_cached_input", "1k_tokens", cachedInputTokens, amount, map[string]any{
|
||||
"inputTokens": inputTokens, "uncachedInputTokens": uncachedInputTokens,
|
||||
"cachedInputTokens": cachedInputTokens, "pricePer1k": cachedPrice.Float64(),
|
||||
}))
|
||||
}
|
||||
if isTextGenerationKind(kind) {
|
||||
outputPrice, priceErr := pricing.requiredTextPrice("text_output", "textOutputPer1k", "outputTokenPrice", "basePrice")
|
||||
if priceErr != nil {
|
||||
return nil, 0, resolvedPricing{}, priceErr
|
||||
}
|
||||
amount, calculationErr := pricing.calculateRatio("text_output", outputPrice, outputTokens, 1000, discount)
|
||||
if calculationErr != nil {
|
||||
return nil, 0, resolvedPricing{}, calculationErr
|
||||
}
|
||||
total, calculationErr = pricing.add("text", total, amount)
|
||||
if calculationErr != nil {
|
||||
return nil, 0, resolvedPricing{}, calculationErr
|
||||
}
|
||||
items = append(items, buildLine("text_output", "1k_tokens", outputTokens, amount, map[string]any{"pricePer1k": outputPrice.Float64()}))
|
||||
}
|
||||
return items, total, pricing, nil
|
||||
}
|
||||
|
||||
count := requestOutputCount(body)
|
||||
resource := "image"
|
||||
unit := "image"
|
||||
baseKey := "imageBase"
|
||||
if kind == "images.edits" {
|
||||
resource = "image_edit"
|
||||
baseKey = "editBase"
|
||||
}
|
||||
if kind == "videos.generations" {
|
||||
resource = "video"
|
||||
unit = "5s_video"
|
||||
baseKey = "videoBase"
|
||||
duration, durationSource := billingDurationSeconds(body, response)
|
||||
audioEnabled, audioSource := billingAudioEnabled(body, response)
|
||||
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
|
||||
}
|
||||
resolutionWeight, weightErr := pricing.weight(resource, "resolutionWeights", firstNonEmptyString(stringFromMap(body, "resolution"), stringFromMap(body, "size")))
|
||||
if weightErr != nil {
|
||||
return nil, 0, resolvedPricing{}, weightErr
|
||||
}
|
||||
audioWeight, weightErr := pricing.weight(resource, "audioWeights", boolWeightKey(audioEnabled))
|
||||
if weightErr != nil {
|
||||
return nil, 0, resolvedPricing{}, weightErr
|
||||
}
|
||||
referenceVideoWeight, weightErr := pricing.weight(resource, "referenceVideoWeights", boolWeightKey(requestHasReferenceVideo(body)))
|
||||
if weightErr != nil {
|
||||
return nil, 0, resolvedPricing{}, weightErr
|
||||
}
|
||||
voiceWeight, weightErr := pricing.weight(resource, "voiceSpecifiedWeights", boolWeightKey(requestHasVoiceID(body, audioEnabled)))
|
||||
if weightErr != nil {
|
||||
return nil, 0, resolvedPricing{}, weightErr
|
||||
}
|
||||
amount, calculationErr := multiplyFixedProductRatio(
|
||||
price,
|
||||
[]int{count},
|
||||
[]fixedAmount{durationFixed, resolutionWeight, audioWeight, referenceVideoWeight, voiceWeight, discount},
|
||||
videoBillingUnitSeconds,
|
||||
)
|
||||
if calculationErr != nil {
|
||||
return nil, 0, resolvedPricing{}, pricing.calculationError(resource, calculationErr)
|
||||
}
|
||||
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,
|
||||
})
|
||||
return []any{item}, amount, pricing, nil
|
||||
}
|
||||
if kind == "song.generations" || kind == "music.generations" {
|
||||
resource = "music"
|
||||
unit = "song"
|
||||
baseKey = "musicBase"
|
||||
}
|
||||
if kind == "speech.generations" || kind == "voice.clone" {
|
||||
resource = "audio"
|
||||
unit = "character"
|
||||
baseKey = "audioBase"
|
||||
count = len([]rune(stringFromMap(body, "text")))
|
||||
if count <= 0 {
|
||||
count = 1
|
||||
}
|
||||
}
|
||||
price, err := pricing.requiredPrice(resource, baseKey, "basePrice")
|
||||
if err != nil {
|
||||
return nil, 0, resolvedPricing{}, err
|
||||
}
|
||||
amount, calculationErr := pricing.calculate(resource, price, []int{count})
|
||||
if calculationErr != nil {
|
||||
return nil, 0, resolvedPricing{}, calculationErr
|
||||
}
|
||||
if resource == "image" || resource == "image_edit" {
|
||||
qualityWeight, weightErr := pricing.weight(resource, "qualityWeights", stringFromMap(body, "quality"))
|
||||
if weightErr != nil {
|
||||
return nil, 0, resolvedPricing{}, weightErr
|
||||
}
|
||||
sizeWeight, weightErr := pricing.weight(resource, "sizeWeights", stringFromMap(body, "size"))
|
||||
if weightErr != nil {
|
||||
return nil, 0, resolvedPricing{}, weightErr
|
||||
}
|
||||
resolutionWeight, weightErr := pricing.weight(resource, "resolutionWeights", firstNonEmptyString(stringFromMap(body, "resolution"), stringFromMap(body, "size")))
|
||||
if weightErr != nil {
|
||||
return nil, 0, resolvedPricing{}, weightErr
|
||||
}
|
||||
amount, calculationErr = pricing.calculate(resource, amount, nil, qualityWeight, sizeWeight, resolutionWeight)
|
||||
if calculationErr != nil {
|
||||
return nil, 0, resolvedPricing{}, calculationErr
|
||||
}
|
||||
}
|
||||
amount, calculationErr = pricing.calculate(resource, amount, nil, discount)
|
||||
if calculationErr != nil {
|
||||
return nil, 0, resolvedPricing{}, calculationErr
|
||||
}
|
||||
return []any{buildLine(resource, unit, count, amount, nil)}, amount, pricing, nil
|
||||
}
|
||||
|
||||
func (pricing resolvedPricing) requiredTextPrice(resource string, keys ...string) (fixedAmount, error) {
|
||||
price, found, err := pricing.price("text", keys...)
|
||||
if err != nil {
|
||||
return 0, &PricingUnavailableError{Reason: err.Error(), ResourceType: resource, RuleSetID: pricing.RuleSetID}
|
||||
}
|
||||
if found && price > 0 {
|
||||
return price, nil
|
||||
}
|
||||
if found && price == 0 && (pricing.FreeResource[resource] || pricing.FreeResource["text_total"] || pricing.FreeResource["text"]) {
|
||||
return 0, nil
|
||||
}
|
||||
return 0, &PricingUnavailableError{Reason: "missing, invalid, or not explicitly free", ResourceType: resource, RuleSetID: pricing.RuleSetID}
|
||||
}
|
||||
|
||||
func maximumCandidateEstimate(estimates []candidateEstimate) (candidateEstimate, error) {
|
||||
if len(estimates) == 0 {
|
||||
return candidateEstimate{}, &PricingUnavailableError{Reason: "no candidate has effective pricing"}
|
||||
}
|
||||
maximum := estimates[0]
|
||||
for _, estimate := range estimates[1:] {
|
||||
if estimate.Currency != maximum.Currency {
|
||||
return candidateEstimate{}, &PricingUnavailableError{Reason: "candidate currencies do not match"}
|
||||
}
|
||||
if estimate.Amount > maximum.Amount {
|
||||
maximum = estimate
|
||||
}
|
||||
}
|
||||
return maximum, nil
|
||||
}
|
||||
@@ -0,0 +1,255 @@
|
||||
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"
|
||||
)
|
||||
|
||||
func TestFixedAmountPreservesNineDecimalPlaces(t *testing.T) {
|
||||
amount, err := parseFixedAmount("0.123456789")
|
||||
if err != nil {
|
||||
t.Fatalf("parse fixed amount: %v", err)
|
||||
}
|
||||
if got, want := amount.String(), "0.123456789"; got != want {
|
||||
t.Fatalf("amount.String()=%q, want %q", got, want)
|
||||
}
|
||||
|
||||
price := mustFixedAmount(t, "0.000000001")
|
||||
product, err := multiplyFixedAmountByInt(price, 3)
|
||||
if err != nil {
|
||||
t.Fatalf("multiply fixed amount: %v", err)
|
||||
}
|
||||
if got, want := product.String(), "0.000000003"; got != want {
|
||||
t.Fatalf("nano amount multiplication=%q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFixedAmountOperationsRejectOverflow(t *testing.T) {
|
||||
maximum := fixedAmount(math.MaxInt64)
|
||||
if _, err := multiplyFixedAmountByInt(maximum, 2); !errors.Is(err, errFixedAmountOverflow) {
|
||||
t.Fatalf("multiply overflow error=%v", err)
|
||||
}
|
||||
if _, err := addFixedAmounts(maximum, 1); !errors.Is(err, errFixedAmountOverflow) {
|
||||
t.Fatalf("add overflow error=%v", err)
|
||||
}
|
||||
pricing := resolvedPricing{RuleSetID: "overflow-rule"}
|
||||
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) {
|
||||
candidate := store.RuntimeModelCandidate{
|
||||
ModelType: "text_generate",
|
||||
Capabilities: map[string]any{
|
||||
"text_generate": map[string]any{"max_output_tokens": 8192},
|
||||
},
|
||||
}
|
||||
tests := []struct {
|
||||
name string
|
||||
body map[string]any
|
||||
want int
|
||||
}{
|
||||
{name: "max completion has highest priority", body: map[string]any{"max_completion_tokens": 700, "max_output_tokens": 600, "max_tokens": 500}, want: 700},
|
||||
{name: "max output precedes legacy max tokens", body: map[string]any{"max_output_tokens": 600, "max_tokens": 500}, want: 600},
|
||||
{name: "legacy max tokens", body: map[string]any{"max_tokens": 500}, want: 500},
|
||||
{name: "model capability fallback", body: map[string]any{}, want: 8192},
|
||||
}
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
if got := estimatedOutputTokens(test.body, candidate); got != test.want {
|
||||
t.Fatalf("estimatedOutputTokens()=%d, want %d", got, test.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
if got := estimatedOutputTokens(nil, store.RuntimeModelCandidate{}); got != 4096 {
|
||||
t.Fatalf("missing capability fallback=%d, want 4096", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildEstimateResultUsesPreferredTotalAndMaximumReservation(t *testing.T) {
|
||||
result, err := buildEstimateResult([]candidateEstimate{
|
||||
{Items: []any{map[string]any{"amount": 1.25, "currency": "resource"}}, Amount: mustFixedAmount(t, "1.25")},
|
||||
{Items: []any{map[string]any{"amount": 2.75, "currency": "resource"}}, Amount: mustFixedAmount(t, "2.75")},
|
||||
}, "fingerprint")
|
||||
if err != nil {
|
||||
t.Fatalf("build estimate result: %v", err)
|
||||
}
|
||||
if result.TotalAmount != 1.25 {
|
||||
t.Fatalf("preferred total=%v, want 1.25", result.TotalAmount)
|
||||
}
|
||||
if result.ReservationAmount != 2.75 {
|
||||
t.Fatalf("reservation=%v, want 2.75", result.ReservationAmount)
|
||||
}
|
||||
if result.CandidateCount != 2 || result.PricingVersion != pricingVersionV2 || result.RequestFingerprint != "fingerprint" {
|
||||
t.Fatalf("unexpected estimate metadata: %+v", result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPricingAvailabilityRequiresExplicitFree(t *testing.T) {
|
||||
paid := resolvedPricing{Config: map[string]any{"imageBase": 1.5}, Currency: "resource"}
|
||||
if _, err := paid.requiredPrice("image", "imageBase", "basePrice"); err != nil {
|
||||
t.Fatalf("positive price should be available: %v", err)
|
||||
}
|
||||
|
||||
missing := resolvedPricing{Config: map[string]any{}, Currency: "resource"}
|
||||
if _, err := missing.requiredPrice("image", "imageBase", "basePrice"); !isPricingUnavailable(err) {
|
||||
t.Fatalf("missing price should be unavailable, got %v", err)
|
||||
}
|
||||
|
||||
ordinaryZero := resolvedPricing{Config: map[string]any{"imageBase": 0}, Currency: "resource"}
|
||||
if _, err := ordinaryZero.requiredPrice("image", "imageBase", "basePrice"); !isPricingUnavailable(err) {
|
||||
t.Fatalf("ordinary zero should be unavailable, got %v", err)
|
||||
}
|
||||
|
||||
explicitFree := resolvedPricing{
|
||||
Config: map[string]any{"imageBase": 0},
|
||||
Currency: "resource",
|
||||
FreeResource: map[string]bool{"image": true},
|
||||
}
|
||||
price, err := explicitFree.requiredPrice("image", "imageBase", "basePrice")
|
||||
if err != nil || !price.IsZero() {
|
||||
t.Fatalf("explicit free should resolve to zero: price=%s err=%v", price.String(), err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPricingWeightsUseFixedPrecisionAndRejectInvalidValues(t *testing.T) {
|
||||
pricing := resolvedPricing{Config: map[string]any{
|
||||
"image": map[string]any{
|
||||
"dynamicWeight": map[string]any{
|
||||
"qualityFactors": map[string]any{"high": "1.123456789", "broken": 0},
|
||||
},
|
||||
},
|
||||
}}
|
||||
weight, err := pricing.weight("image", "qualityWeights", "high")
|
||||
if err != nil || weight.String() != "1.123456789" {
|
||||
t.Fatalf("exact weight=%s err=%v", weight.String(), err)
|
||||
}
|
||||
if _, err := pricing.weight("image", "qualityWeights", "broken"); !isPricingUnavailable(err) {
|
||||
t.Fatalf("invalid configured weight should make pricing unavailable: %v", err)
|
||||
}
|
||||
defaultWeight, err := pricing.weight("image", "qualityWeights", "unconfigured")
|
||||
if err != nil || defaultWeight.String() != "1.000000000" {
|
||||
t.Fatalf("default weight=%s err=%v", defaultWeight.String(), err)
|
||||
}
|
||||
}
|
||||
|
||||
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)
|
||||
if err != nil {
|
||||
t.Fatalf("parse %q: %v", value, err)
|
||||
}
|
||||
return amount
|
||||
}
|
||||
@@ -10,6 +10,7 @@ import (
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
"github.com/google/uuid"
|
||||
"github.com/riverqueue/river"
|
||||
"github.com/riverqueue/river/riverdriver/riverpgxv5"
|
||||
"github.com/riverqueue/river/rivermigrate"
|
||||
@@ -38,16 +39,26 @@ func (w *asyncTaskWorker) Work(ctx context.Context, job *river.Job[asyncTaskArgs
|
||||
if task.Status == "succeeded" || task.Status == "failed" || task.Status == "cancelled" {
|
||||
return nil
|
||||
}
|
||||
result, runErr := w.service.Execute(ctx, task, authUserFromTask(task))
|
||||
executionToken := uuid.NewString()
|
||||
result, runErr := w.service.executeWithToken(ctx, task, authUserFromTask(task), nil, executionToken)
|
||||
if runErr == nil {
|
||||
w.service.logger.Debug("river async task completed", "taskID", task.ID, "status", result.Task.Status, "riverJobID", job.ID)
|
||||
return nil
|
||||
}
|
||||
if errors.Is(runErr, store.ErrTaskExecutionLeaseUnavailable) {
|
||||
w.service.logger.Debug("river async task execution lease already held", "taskID", task.ID, "riverJobID", job.ID)
|
||||
return nil
|
||||
}
|
||||
if errors.Is(runErr, store.ErrTaskExecutionManualReview) {
|
||||
w.service.logger.Warn("river async task moved to manual review after ambiguous upstream submission", "taskID", task.ID, "riverJobID", job.ID)
|
||||
return nil
|
||||
}
|
||||
var queuedErr *TaskQueuedError
|
||||
if errors.As(runErr, &queuedErr) {
|
||||
return river.JobSnooze(queuedErr.Delay)
|
||||
}
|
||||
if ctx.Err() != nil {
|
||||
task.ExecutionToken = executionToken
|
||||
queued, queueErr := w.service.requeueInterruptedAsyncTask(context.WithoutCancel(ctx), task)
|
||||
if queueErr != nil {
|
||||
return queueErr
|
||||
@@ -145,8 +156,7 @@ func (s *Service) recoverAsyncRiverJobs(ctx context.Context) error {
|
||||
return err
|
||||
}
|
||||
for _, item := range items {
|
||||
task := store.GatewayTask{ID: item.ID}
|
||||
result, err := s.riverClient.Insert(ctx, asyncTaskArgs{TaskID: item.ID}, asyncTaskInsertOpts(task))
|
||||
result, err := s.riverClient.Insert(ctx, asyncTaskArgs{TaskID: item.ID}, asyncTaskRecoveryInsertOpts(item, time.Now()))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
@@ -186,6 +196,18 @@ func asyncTaskInsertOpts(task store.GatewayTask) *river.InsertOpts {
|
||||
}
|
||||
}
|
||||
|
||||
func asyncTaskRecoveryInsertOpts(item store.AsyncTaskQueueItem, now time.Time) *river.InsertOpts {
|
||||
opts := asyncTaskInsertOpts(store.GatewayTask{ID: item.ID})
|
||||
if item.NextRunAt.After(now) {
|
||||
opts.ScheduledAt = item.NextRunAt
|
||||
}
|
||||
// A replacement process must not be blocked by a River row that the dead
|
||||
// process left in running state. PostgreSQL execution leases still ensure
|
||||
// that only one recovery job can call the upstream provider.
|
||||
opts.UniqueOpts = river.UniqueOpts{}
|
||||
return opts
|
||||
}
|
||||
|
||||
func authUserFromTask(task store.GatewayTask) *auth.User {
|
||||
roles := []string{"user"}
|
||||
if strings.TrimSpace(task.UserID) == "" {
|
||||
|
||||
@@ -0,0 +1,33 @@
|
||||
package runner
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func TestAsyncTaskRecoveryInsertOptsMakesDueTaskImmediatelyAvailable(t *testing.T) {
|
||||
now := time.Date(2026, time.July, 21, 3, 30, 0, 0, time.UTC)
|
||||
opts := asyncTaskRecoveryInsertOpts(store.AsyncTaskQueueItem{
|
||||
ID: "due-task",
|
||||
NextRunAt: now.Add(-time.Second),
|
||||
}, now)
|
||||
|
||||
if !opts.ScheduledAt.IsZero() {
|
||||
t.Fatalf("due recovery task should be immediately available, scheduled at %s", opts.ScheduledAt)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAsyncTaskRecoveryInsertOptsPreservesFutureDelay(t *testing.T) {
|
||||
now := time.Date(2026, time.July, 21, 3, 30, 0, 0, time.UTC)
|
||||
nextRunAt := now.Add(30 * time.Second)
|
||||
opts := asyncTaskRecoveryInsertOpts(store.AsyncTaskQueueItem{
|
||||
ID: "delayed-task",
|
||||
NextRunAt: nextRunAt,
|
||||
}, now)
|
||||
|
||||
if !opts.ScheduledAt.Equal(nextRunAt) {
|
||||
t.Fatalf("future recovery task should stay delayed until %s, got %s", nextRunAt, opts.ScheduledAt)
|
||||
}
|
||||
}
|
||||
@@ -14,6 +14,7 @@ import (
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
|
||||
scriptengine "github.com/easyai/easyai-ai-gateway/apps/api/internal/script"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
"github.com/google/uuid"
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/riverqueue/river"
|
||||
)
|
||||
@@ -26,6 +27,11 @@ type Service struct {
|
||||
scriptExecutor *scriptengine.Executor
|
||||
httpClients *httpClientCache
|
||||
riverClient *river.Client[pgx.Tx]
|
||||
billingMetrics billingMetricsObserver
|
||||
}
|
||||
|
||||
type billingMetricsObserver interface {
|
||||
ObserveBillingEvent(string)
|
||||
}
|
||||
|
||||
type Result struct {
|
||||
@@ -39,6 +45,19 @@ type TaskQueuedError struct {
|
||||
Delay time.Duration
|
||||
}
|
||||
|
||||
type upstreamSubmissionUnknownError struct {
|
||||
AttemptID string
|
||||
Cause error
|
||||
}
|
||||
|
||||
func (e *upstreamSubmissionUnknownError) Error() string {
|
||||
return "upstream submission result is unknown"
|
||||
}
|
||||
|
||||
func (e *upstreamSubmissionUnknownError) Unwrap() error {
|
||||
return e.Cause
|
||||
}
|
||||
|
||||
func (e *TaskQueuedError) Error() string {
|
||||
return ErrTaskQueued.Error()
|
||||
}
|
||||
@@ -47,10 +66,10 @@ func (e *TaskQueuedError) Is(target error) bool {
|
||||
return target == ErrTaskQueued
|
||||
}
|
||||
|
||||
func New(cfg config.Config, db *store.Store, logger *slog.Logger) *Service {
|
||||
func New(cfg config.Config, db *store.Store, logger *slog.Logger, observers ...billingMetricsObserver) *Service {
|
||||
httpClients := newHTTPClientCache()
|
||||
scriptExecutor := &scriptengine.Executor{Logger: logger}
|
||||
return &Service{
|
||||
service := &Service{
|
||||
cfg: cfg,
|
||||
store: db,
|
||||
logger: logger,
|
||||
@@ -76,6 +95,16 @@ func New(cfg config.Config, db *store.Store, logger *slog.Logger) *Service {
|
||||
},
|
||||
httpClients: httpClients,
|
||||
}
|
||||
if len(observers) > 0 {
|
||||
service.billingMetrics = observers[0]
|
||||
}
|
||||
return service
|
||||
}
|
||||
|
||||
func (s *Service) observeBillingEvent(event string) {
|
||||
if s.billingMetrics != nil {
|
||||
s.billingMetrics.ObserveBillingEvent(event)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) Execute(ctx context.Context, task store.GatewayTask, user *auth.User) (Result, error) {
|
||||
@@ -87,6 +116,20 @@ func (s *Service) ExecuteStream(ctx context.Context, task store.GatewayTask, use
|
||||
}
|
||||
|
||||
func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *auth.User, onDelta clients.StreamDelta) (Result, error) {
|
||||
return s.executeWithToken(ctx, task, user, onDelta, uuid.NewString())
|
||||
}
|
||||
|
||||
func (s *Service) executeWithToken(ctx context.Context, task store.GatewayTask, user *auth.User, onDelta clients.StreamDelta, executionToken string) (Result, error) {
|
||||
wasRunning := task.Status == "running"
|
||||
claimed, err := s.store.ClaimTaskExecution(ctx, task.ID, executionToken, taskExecutionLeaseTTL)
|
||||
if err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
task = claimed
|
||||
executionCtx, stopExecution := context.WithCancel(ctx)
|
||||
defer stopExecution()
|
||||
go s.renewTaskExecutionLease(executionCtx, stopExecution, task.ID, task.ExecutionToken)
|
||||
ctx = executionCtx
|
||||
executeStartedAt := time.Now()
|
||||
restoredRequest, err := s.restoreTaskRequestReferences(ctx, task)
|
||||
if err != nil {
|
||||
@@ -95,10 +138,10 @@ func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *aut
|
||||
body := normalizeRequest(task.Kind, restoredRequest)
|
||||
responseExecution := responseExecutionContext{}
|
||||
modelType := modelTypeFromKind(task.Kind, body)
|
||||
if err := s.store.MarkTaskRunning(ctx, task.ID, modelType, s.slimTaskRequestSnapshot(task, body)); err != nil {
|
||||
if err := s.store.MarkTaskRunning(ctx, task.ID, task.ExecutionToken, modelType, s.slimTaskRequestSnapshot(task, body)); err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
if task.Status != "running" {
|
||||
if !wasRunning {
|
||||
if err := s.emit(ctx, task.ID, "task.running", "running", "starting", 0.12, "task pulled from queue and started", map[string]any{"modelType": modelType}, task.RunMode == "simulation"); err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
@@ -115,7 +158,7 @@ func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *aut
|
||||
Reason: "request_validation_failed",
|
||||
ModelType: modelType,
|
||||
})
|
||||
failed, finishErr := s.failTask(ctx, task.ID, "bad_request", err.Error(), task.RunMode == "simulation", err)
|
||||
failed, finishErr := s.failTask(ctx, task.ID, task.ExecutionToken, "bad_request", err.Error(), task.RunMode == "simulation", err)
|
||||
if finishErr != nil {
|
||||
return Result{}, finishErr
|
||||
}
|
||||
@@ -135,14 +178,14 @@ func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *aut
|
||||
Reason: "cloned_voice_binding_failed",
|
||||
ModelType: modelType,
|
||||
})
|
||||
failed, finishErr := s.failTask(ctx, task.ID, clients.ErrorCode(err), err.Error(), task.RunMode == "simulation", err)
|
||||
failed, finishErr := s.failTask(ctx, task.ID, task.ExecutionToken, clients.ErrorCode(err), err.Error(), task.RunMode == "simulation", err)
|
||||
if finishErr != nil {
|
||||
return Result{}, finishErr
|
||||
}
|
||||
return Result{Task: failed, Output: failed.Result}, err
|
||||
}
|
||||
if clonedVoice.Found {
|
||||
if err := s.store.MarkTaskRunning(ctx, task.ID, modelType, s.slimTaskRequestSnapshot(task, body)); err != nil {
|
||||
if err := s.store.MarkTaskRunning(ctx, task.ID, task.ExecutionToken, modelType, s.slimTaskRequestSnapshot(task, body)); err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
}
|
||||
@@ -151,7 +194,7 @@ func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *aut
|
||||
if err != nil {
|
||||
code, message := responseExecutionFailure(err)
|
||||
s.recordFailedAttempt(ctx, failedAttemptRecord{Task: task, Body: body, AttemptNo: task.AttemptCount + 1, Code: code, Cause: err, Simulated: task.RunMode == "simulation", Scope: "response_chain", Reason: code, ModelType: modelType})
|
||||
failed, finishErr := s.failTask(ctx, task.ID, code, message, task.RunMode == "simulation", err)
|
||||
failed, finishErr := s.failTask(ctx, task.ID, task.ExecutionToken, code, message, task.RunMode == "simulation", err)
|
||||
if finishErr != nil {
|
||||
return Result{}, finishErr
|
||||
}
|
||||
@@ -183,7 +226,7 @@ func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *aut
|
||||
Reason: "candidate_selection_failed",
|
||||
ModelType: modelType,
|
||||
})
|
||||
failed, finishErr := s.failTask(ctx, task.ID, store.ModelCandidateErrorCode(err), err.Error(), task.RunMode == "simulation", err)
|
||||
failed, finishErr := s.failTask(ctx, task.ID, task.ExecutionToken, store.ModelCandidateErrorCode(err), err.Error(), task.RunMode == "simulation", err)
|
||||
if finishErr != nil {
|
||||
return Result{}, finishErr
|
||||
}
|
||||
@@ -202,7 +245,7 @@ func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *aut
|
||||
Reason: store.ModelCandidateErrorCode(err),
|
||||
ModelType: modelType,
|
||||
})
|
||||
failed, finishErr := s.failTask(ctx, task.ID, store.ModelCandidateErrorCode(err), err.Error(), task.RunMode == "simulation", err)
|
||||
failed, finishErr := s.failTask(ctx, task.ID, task.ExecutionToken, store.ModelCandidateErrorCode(err), err.Error(), task.RunMode == "simulation", err)
|
||||
if finishErr != nil {
|
||||
return Result{}, finishErr
|
||||
}
|
||||
@@ -224,7 +267,7 @@ func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *aut
|
||||
ExtraMetrics: []map[string]any{candidateFilterMetrics},
|
||||
ModelType: modelType,
|
||||
})
|
||||
failed, finishErr := s.failTask(ctx, task.ID, store.ModelCandidateErrorCode(err), err.Error(), task.RunMode == "simulation", err, candidateFilterMetrics)
|
||||
failed, finishErr := s.failTask(ctx, task.ID, task.ExecutionToken, store.ModelCandidateErrorCode(err), err.Error(), task.RunMode == "simulation", err, candidateFilterMetrics)
|
||||
if finishErr != nil {
|
||||
return Result{}, finishErr
|
||||
}
|
||||
@@ -240,7 +283,7 @@ func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *aut
|
||||
Simulated: task.RunMode == "simulation", Scope: "candidate_output_token_filter", Reason: store.ModelCandidateErrorCode(err),
|
||||
ExtraMetrics: []map[string]any{candidateFilterMetrics}, ModelType: modelType,
|
||||
})
|
||||
failed, finishErr := s.failTask(ctx, task.ID, store.ModelCandidateErrorCode(err), err.Error(), task.RunMode == "simulation", err, candidateFilterMetrics)
|
||||
failed, finishErr := s.failTask(ctx, task.ID, task.ExecutionToken, store.ModelCandidateErrorCode(err), err.Error(), task.RunMode == "simulation", err, candidateFilterMetrics)
|
||||
if finishErr != nil {
|
||||
return Result{}, finishErr
|
||||
}
|
||||
@@ -251,13 +294,120 @@ func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *aut
|
||||
if err != nil {
|
||||
code, message := responseExecutionFailure(err)
|
||||
s.recordFailedAttempt(ctx, failedAttemptRecord{Task: task, Body: body, AttemptNo: task.AttemptCount + 1, Code: code, Cause: err, Simulated: task.RunMode == "simulation", Scope: "response_protocol", Reason: code, ModelType: modelType})
|
||||
failed, finishErr := s.failTask(ctx, task.ID, code, message, task.RunMode == "simulation", err)
|
||||
failed, finishErr := s.failTask(ctx, task.ID, task.ExecutionToken, code, message, task.RunMode == "simulation", err)
|
||||
if finishErr != nil {
|
||||
return Result{}, finishErr
|
||||
}
|
||||
return Result{Task: failed, Output: failed.Result}, err
|
||||
}
|
||||
}
|
||||
pricingByCandidate := map[string]resolvedPricing{}
|
||||
preprocessingByCandidate := map[string]parameterPreprocessResult{}
|
||||
reservationBillings := []any(nil)
|
||||
reservationPricingSnapshot := map[string]any(nil)
|
||||
if task.RunMode == "production" {
|
||||
billingMode := normalizedBillingEngineMode(s.cfg.BillingEngineMode)
|
||||
if billingMode == "hold" {
|
||||
holdErr := &clients.ClientError{Code: "billing_hold", Message: "production billing is temporarily on hold", Retryable: true}
|
||||
failed, finishErr := s.failTask(ctx, task.ID, task.ExecutionToken, "billing_hold", holdErr.Error(), false, holdErr)
|
||||
if finishErr != nil {
|
||||
return Result{}, finishErr
|
||||
}
|
||||
return Result{Task: failed, Output: failed.Result}, holdErr
|
||||
}
|
||||
estimates := make([]candidateEstimate, 0, len(candidates))
|
||||
var pricingErr error
|
||||
var preprocessingErr error
|
||||
var preprocessingCandidate store.RuntimeModelCandidate
|
||||
for _, candidate := range candidates {
|
||||
preprocessing := s.preprocessRequestWithScripts(ctx, task.Kind, body, candidate)
|
||||
preprocessingByCandidate[pricingCandidateKey(candidate)] = preprocessing
|
||||
if preprocessing.Err != nil {
|
||||
preprocessingErr = parameterPreprocessClientError(preprocessing.Err)
|
||||
preprocessingCandidate = candidate
|
||||
break
|
||||
}
|
||||
estimate, estimateErr := s.estimateCandidateV2(ctx, user, task.Kind, preprocessing.Body, candidate)
|
||||
if estimateErr != nil {
|
||||
pricingErr = estimateErr
|
||||
if billingMode == "enforce" {
|
||||
break
|
||||
}
|
||||
continue
|
||||
}
|
||||
estimates = append(estimates, estimate)
|
||||
pricingByCandidate[pricingCandidateKey(candidate)] = estimate.Pricing
|
||||
}
|
||||
if preprocessingErr != nil {
|
||||
preprocessing := preprocessingByCandidate[pricingCandidateKey(preprocessingCandidate)]
|
||||
s.recordFailedAttempt(ctx, failedAttemptRecord{
|
||||
Task: task, Body: preprocessing.Body, Candidate: &preprocessingCandidate,
|
||||
AttemptNo: task.AttemptCount + 1, Code: clients.ErrorCode(preprocessingErr), Cause: preprocessingErr,
|
||||
Simulated: false, Scope: "parameter_preprocessing", Reason: "parameter_preprocessing_failed",
|
||||
ExtraMetrics: []map[string]any{parameterPreprocessingMetrics(preprocessing.Log)}, Preprocessing: &preprocessing.Log,
|
||||
ModelType: preprocessingCandidate.ModelType,
|
||||
})
|
||||
failed, finishErr := s.failTask(ctx, task.ID, task.ExecutionToken, clients.ErrorCode(preprocessingErr), preprocessingErr.Error(), false, preprocessingErr, parameterPreprocessingMetrics(preprocessing.Log))
|
||||
if finishErr != nil {
|
||||
return Result{}, finishErr
|
||||
}
|
||||
return Result{Task: failed, Output: failed.Result}, preprocessingErr
|
||||
}
|
||||
if billingMode == "observe" {
|
||||
legacyItems, legacyAmount := s.maximumLegacyCandidateEstimate(ctx, user, task.Kind, body, candidates, preprocessingByCandidate)
|
||||
reservationBillings = legacyItems
|
||||
candidateSnapshots := make([]any, 0, len(estimates))
|
||||
for _, estimate := range estimates {
|
||||
candidateSnapshots = append(candidateSnapshots, estimate.Snapshot)
|
||||
}
|
||||
observedAmount := ""
|
||||
if maximumEstimate, estimateErr := maximumCandidateEstimate(estimates); estimateErr == nil {
|
||||
observedAmount = maximumEstimate.Amount.String()
|
||||
}
|
||||
reservationPricingSnapshot = map[string]any{
|
||||
"pricingVersion": "legacy-observe", "observedPricingVersion": pricingVersionV2,
|
||||
"requestFingerprint": pricingRequestFingerprint(task.Kind, task.Model, body),
|
||||
"reservationAmount": legacyAmount.String(), "observedReservationAmount": observedAmount,
|
||||
"candidateCount": len(candidates), "observedCandidates": candidateSnapshots,
|
||||
}
|
||||
s.logger.Info("billing observe comparison", "taskID", task.ID, "legacyAmount", legacyAmount.String(), "v2Amount", observedAmount, "pricedCandidates", len(estimates), "candidateCount", len(candidates))
|
||||
} else if pricingErr != nil || len(estimates) != len(candidates) {
|
||||
if pricingErr == nil {
|
||||
pricingErr = &PricingUnavailableError{Reason: "not every candidate has effective pricing"}
|
||||
}
|
||||
s.observeBillingEvent("pricing_unavailable")
|
||||
failed, finishErr := s.failTask(ctx, task.ID, task.ExecutionToken, "pricing_unavailable", pricingErr.Error(), false, pricingErr)
|
||||
if finishErr != nil {
|
||||
return Result{}, finishErr
|
||||
}
|
||||
return Result{Task: failed, Output: failed.Result}, pricingErr
|
||||
} else {
|
||||
maximumEstimate, estimateErr := maximumCandidateEstimate(estimates)
|
||||
if estimateErr != nil {
|
||||
failed, finishErr := s.failTask(ctx, task.ID, task.ExecutionToken, "pricing_unavailable", estimateErr.Error(), false, estimateErr)
|
||||
if finishErr != nil {
|
||||
return Result{}, finishErr
|
||||
}
|
||||
return Result{Task: failed, Output: failed.Result}, estimateErr
|
||||
}
|
||||
pricedCandidates := make([]store.RuntimeModelCandidate, 0, len(estimates))
|
||||
candidateSnapshots := make([]any, 0, len(estimates))
|
||||
for _, candidate := range candidates {
|
||||
if _, ok := pricingByCandidate[pricingCandidateKey(candidate)]; ok {
|
||||
pricedCandidates = append(pricedCandidates, candidate)
|
||||
}
|
||||
}
|
||||
for _, estimate := range estimates {
|
||||
candidateSnapshots = append(candidateSnapshots, estimate.Snapshot)
|
||||
}
|
||||
candidates = pricedCandidates
|
||||
reservationBillings = maximumEstimate.Items
|
||||
reservationPricingSnapshot = map[string]any{
|
||||
"pricingVersion": pricingVersionV2, "requestFingerprint": pricingRequestFingerprint(task.Kind, task.Model, body),
|
||||
"reservationAmount": maximumEstimate.Amount.String(), "candidateCount": len(estimates), "candidates": candidateSnapshots,
|
||||
}
|
||||
}
|
||||
}
|
||||
firstCandidateBody := body
|
||||
normalizedModelType := modelType
|
||||
attemptNo := task.AttemptCount
|
||||
@@ -270,7 +420,10 @@ func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *aut
|
||||
}
|
||||
}()
|
||||
if len(candidates) > 0 {
|
||||
preprocessing := s.preprocessRequestWithScripts(ctx, task.Kind, body, candidates[0])
|
||||
preprocessing, ok := preprocessingByCandidate[pricingCandidateKey(candidates[0])]
|
||||
if !ok {
|
||||
preprocessing = s.preprocessRequestWithScripts(ctx, task.Kind, body, candidates[0])
|
||||
}
|
||||
firstCandidateBody = preprocessing.Body
|
||||
firstPreprocessing = preprocessing.Log
|
||||
normalizedModelType = candidates[0].ModelType
|
||||
@@ -290,18 +443,17 @@ func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *aut
|
||||
Preprocessing: &firstPreprocessing,
|
||||
ModelType: normalizedModelType,
|
||||
})
|
||||
failed, finishErr := s.failTask(ctx, task.ID, clients.ErrorCode(clientErr), clientErr.Error(), task.RunMode == "simulation", clientErr, parameterPreprocessingMetrics(firstPreprocessing))
|
||||
failed, finishErr := s.failTask(ctx, task.ID, task.ExecutionToken, clients.ErrorCode(clientErr), clientErr.Error(), task.RunMode == "simulation", clientErr, parameterPreprocessingMetrics(firstPreprocessing))
|
||||
if finishErr != nil {
|
||||
return Result{}, finishErr
|
||||
}
|
||||
return Result{Task: failed, Output: failed.Result}, clientErr
|
||||
}
|
||||
if err := s.store.MarkTaskRunning(ctx, task.ID, candidates[0].ModelType, s.slimTaskRequestSnapshot(task, firstCandidateBody)); err != nil {
|
||||
if err := s.store.MarkTaskRunning(ctx, task.ID, task.ExecutionToken, candidates[0].ModelType, s.slimTaskRequestSnapshot(task, firstCandidateBody)); err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
estimatedBillings := s.estimatedBillings(ctx, user, task.Kind, firstCandidateBody, candidates[0])
|
||||
var reserveErr error
|
||||
walletReservations, reserveErr = s.store.ReserveTaskBilling(ctx, task, user, estimatedBillings)
|
||||
walletReservations, reserveErr = s.store.ReserveTaskBilling(ctx, task, user, reservationBillings, reservationPricingSnapshot)
|
||||
if reserveErr != nil {
|
||||
if errors.Is(reserveErr, store.ErrInsufficientWalletBalance) {
|
||||
attemptNo = s.recordFailedAttempt(ctx, failedAttemptRecord{
|
||||
@@ -318,7 +470,7 @@ func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *aut
|
||||
Preprocessing: &firstPreprocessing,
|
||||
ModelType: normalizedModelType,
|
||||
})
|
||||
failed, finishErr := s.failTask(ctx, task.ID, "insufficient_balance", reserveErr.Error(), task.RunMode == "simulation", reserveErr, parameterPreprocessingMetrics(firstPreprocessing))
|
||||
failed, finishErr := s.failTask(ctx, task.ID, task.ExecutionToken, "insufficient_balance", reserveErr.Error(), task.RunMode == "simulation", reserveErr, parameterPreprocessingMetrics(firstPreprocessing))
|
||||
if finishErr != nil {
|
||||
return Result{}, finishErr
|
||||
}
|
||||
@@ -347,7 +499,10 @@ candidatesLoop:
|
||||
var candidateErr error
|
||||
for clientAttempt := 1; clientAttempt <= clientAttempts; clientAttempt++ {
|
||||
nextAttemptNo := attemptNo + 1
|
||||
preprocessing := s.preprocessRequestWithScripts(ctx, task.Kind, body, candidate)
|
||||
preprocessing, ok := preprocessingByCandidate[pricingCandidateKey(candidate)]
|
||||
if !ok {
|
||||
preprocessing = s.preprocessRequestWithScripts(ctx, task.Kind, body, candidate)
|
||||
}
|
||||
preprocessingLog := preprocessing.Log
|
||||
lastPreprocessing = &preprocessingLog
|
||||
if preprocessing.Err != nil {
|
||||
@@ -369,47 +524,101 @@ candidatesLoop:
|
||||
break candidatesLoop
|
||||
}
|
||||
candidateBody := preprocessing.Body
|
||||
response, err := s.runCandidate(ctx, task, user, candidateBody, preprocessing.Log, candidate, nextAttemptNo, onDelta, responseExecution, singleSourceProtected, runnerPolicy.CacheAffinityPolicy, cacheAffinityKeys.Record)
|
||||
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
|
||||
billings := s.billings(ctx, user, task.Kind, candidateBody, candidate, response, isSimulation(task, candidate))
|
||||
var billings []any
|
||||
finalAmount := fixedAmount(0)
|
||||
finalAmountText := ""
|
||||
pricingSnapshot := map[string]any{"pricingVersion": pricingVersionV2, "source": "simulation"}
|
||||
if task.RunMode == "production" {
|
||||
billingMode := normalizedBillingEngineMode(s.cfg.BillingEngineMode)
|
||||
var billingErr error
|
||||
if billingMode == "observe" {
|
||||
billings = s.billings(ctx, user, task.Kind, candidateBody, candidate, response, false)
|
||||
finalAmount = billingItemsFixedTotal(billings)
|
||||
pricingSnapshot = map[string]any{
|
||||
"pricingVersion": "legacy-observe", "observedPricingVersion": pricingVersionV2,
|
||||
"observedPricing": candidatePricing.Snapshot,
|
||||
}
|
||||
} else {
|
||||
billings, finalAmount, _, billingErr = s.billingsWithResolvedPricingV2(ctx, user, task.Kind, candidateBody, candidate, response, false, candidatePricing)
|
||||
pricingSnapshot = candidatePricing.Snapshot
|
||||
}
|
||||
if billingErr != nil {
|
||||
review, reviewErr := s.store.FinishTaskManualReview(context.WithoutCancel(ctx), store.FinishTaskManualReviewInput{
|
||||
TaskID: task.ID, ExecutionToken: task.ExecutionToken, AttemptID: response.AttemptID, TaskStatus: "succeeded",
|
||||
Code: "billing_calculation_failed", Message: "billing calculation requires manual review",
|
||||
Result: response.Result, RequestID: response.RequestID,
|
||||
PricingSnapshot: candidatePricing.Snapshot,
|
||||
RequestFingerprint: pricingRequestFingerprint(task.Kind, task.Model, candidateBody),
|
||||
ResponseStartedAt: response.ResponseStartedAt, ResponseFinishedAt: response.ResponseFinishedAt,
|
||||
ResponseDurationMS: response.ResponseDurationMS,
|
||||
})
|
||||
if reviewErr != nil {
|
||||
return Result{}, reviewErr
|
||||
}
|
||||
walletReservationFinalized = true
|
||||
s.logger.Warn("task succeeded but billing requires manual review", "taskID", task.ID, "error_category", "billing_calculation_failed")
|
||||
return Result{Task: review, Output: response.Result}, nil
|
||||
}
|
||||
finalAmountText = finalAmount.String()
|
||||
} else {
|
||||
billings = s.billings(ctx, user, task.Kind, candidateBody, candidate, response, true)
|
||||
}
|
||||
record := buildSuccessRecord(task, user, candidateBody, candidate, response, billings, isSimulation(task, candidate))
|
||||
record.Metrics = mergeMetrics(record.Metrics, candidateCapabilityFilterMetrics(candidateFilterSummary))
|
||||
record.Metrics = mergeMetrics(record.Metrics, parameterPreprocessingMetrics(preprocessing.Log))
|
||||
record.Metrics = s.withAttemptHistory(ctx, task.ID, record.Metrics)
|
||||
finished, finishErr := s.store.FinishTaskSuccess(ctx, store.FinishTaskSuccessInput{
|
||||
TaskID: task.ID,
|
||||
Result: response.Result,
|
||||
Billings: billings,
|
||||
RequestID: record.RequestID,
|
||||
ResolvedModel: record.ResolvedModel,
|
||||
Usage: record.Usage,
|
||||
Metrics: record.Metrics,
|
||||
BillingSummary: record.BillingSummary,
|
||||
FinalChargeAmount: record.FinalChargeAmount,
|
||||
ResponseStartedAt: record.ResponseStartedAt,
|
||||
ResponseFinishedAt: record.ResponseFinishedAt,
|
||||
ResponseDurationMS: record.ResponseDurationMS,
|
||||
TaskID: task.ID,
|
||||
ExecutionToken: task.ExecutionToken,
|
||||
AttemptID: response.AttemptID,
|
||||
Result: response.Result,
|
||||
Billings: billings,
|
||||
RequestID: record.RequestID,
|
||||
ResolvedModel: record.ResolvedModel,
|
||||
Usage: record.Usage,
|
||||
Metrics: record.Metrics,
|
||||
BillingSummary: record.BillingSummary,
|
||||
FinalChargeAmount: record.FinalChargeAmount,
|
||||
FinalChargeAmountText: finalAmountText,
|
||||
BillingCurrency: stringFromAny(record.BillingSummary["currency"]),
|
||||
PricingSnapshot: pricingSnapshot,
|
||||
RequestFingerprint: pricingRequestFingerprint(task.Kind, task.Model, candidateBody),
|
||||
ResponseStartedAt: record.ResponseStartedAt,
|
||||
ResponseFinishedAt: record.ResponseFinishedAt,
|
||||
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
|
||||
}
|
||||
if finished.FinalChargeAmount > 0 {
|
||||
walletReservationFinalized = true
|
||||
if settleErr := s.store.SettleTaskBilling(ctx, finished); settleErr != nil {
|
||||
return Result{}, settleErr
|
||||
}
|
||||
} else if len(walletReservations) > 0 {
|
||||
if releaseErr := s.store.ReleaseTaskBillingReservations(ctx, walletReservations, "task_billing_zero"); releaseErr != nil {
|
||||
return Result{}, releaseErr
|
||||
}
|
||||
walletReservationFinalized = true
|
||||
}
|
||||
walletReservationFinalized = true
|
||||
if finished.FinalChargeAmount > 0 {
|
||||
if err := s.emit(ctx, task.ID, "task.billing.settled", "succeeded", "billing", 0.98, "task billing settled", map[string]any{
|
||||
"amount": finished.FinalChargeAmount,
|
||||
"currency": stringFromAny(record.BillingSummary["currency"]),
|
||||
if finished.BillingStatus == "pending" {
|
||||
if err := s.emit(ctx, task.ID, "task.billing.pending", "succeeded", "billing", 0.98, "task billing queued", map[string]any{
|
||||
"amount": finished.FinalChargeAmount, "currency": finished.BillingCurrency,
|
||||
}, isSimulation(task, candidate)); err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
@@ -426,6 +635,21 @@ candidatesLoop:
|
||||
}
|
||||
return Result{Task: finished, Output: response.Result}, nil
|
||||
}
|
||||
var submissionUnknown *upstreamSubmissionUnknownError
|
||||
if errors.As(err, &submissionUnknown) {
|
||||
review, reviewErr := s.store.FinishTaskManualReview(context.WithoutCancel(ctx), store.FinishTaskManualReviewInput{
|
||||
TaskID: task.ID, ExecutionToken: task.ExecutionToken, AttemptID: submissionUnknown.AttemptID, TaskStatus: "failed",
|
||||
Code: "upstream_submission_unknown", Message: submissionUnknown.Error(),
|
||||
PricingSnapshot: candidatePricing.Snapshot,
|
||||
RequestFingerprint: pricingRequestFingerprint(task.Kind, task.Model, candidateBody),
|
||||
})
|
||||
if reviewErr != nil {
|
||||
return Result{}, reviewErr
|
||||
}
|
||||
walletReservationFinalized = true
|
||||
s.logger.Warn("upstream submission requires manual review", "taskID", task.ID, "attemptID", submissionUnknown.AttemptID, "error_category", "upstream_submission_unknown")
|
||||
return Result{Task: review, Output: review.Result}, submissionUnknown
|
||||
}
|
||||
if isLocalRateLimitError(err) {
|
||||
lastErr = err
|
||||
candidateErr = err
|
||||
@@ -573,14 +797,66 @@ candidatesLoop:
|
||||
if lastPreprocessing != nil {
|
||||
extraMetrics = append(extraMetrics, parameterPreprocessingMetrics(*lastPreprocessing))
|
||||
}
|
||||
failed, err := s.failTask(ctx, task.ID, code, message, task.RunMode == "simulation", lastErr, extraMetrics...)
|
||||
failed, err := s.failTask(ctx, task.ID, task.ExecutionToken, code, message, task.RunMode == "simulation", lastErr, extraMetrics...)
|
||||
if err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
// FinishTaskFailure atomically transfers ownership of any reservation to the release Outbox.
|
||||
walletReservationFinalized = true
|
||||
return Result{Task: failed, Output: failed.Result}, lastErr
|
||||
}
|
||||
|
||||
func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user *auth.User, body map[string]any, preprocessing parameterPreprocessingLog, candidate store.RuntimeModelCandidate, attemptNo int, onDelta clients.StreamDelta, responseExecution responseExecutionContext, singleSourceProtected bool, cacheAffinityPolicy map[string]any, cacheAffinityRecordKeys []string) (clients.Response, error) {
|
||||
func pricingCandidateKey(candidate store.RuntimeModelCandidate) string {
|
||||
return firstNonEmptyString(candidate.PlatformModelID, candidate.PlatformID+":"+candidate.ModelName)
|
||||
}
|
||||
|
||||
func normalizedBillingEngineMode(value string) string {
|
||||
switch strings.ToLower(strings.TrimSpace(value)) {
|
||||
case "enforce", "hold":
|
||||
return strings.ToLower(strings.TrimSpace(value))
|
||||
default:
|
||||
return "observe"
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) maximumLegacyCandidateEstimate(ctx context.Context, user *auth.User, kind string, body map[string]any, candidates []store.RuntimeModelCandidate, preprocessingByCandidate map[string]parameterPreprocessResult) ([]any, fixedAmount) {
|
||||
var maximumItems []any
|
||||
maximumAmount := fixedAmount(0)
|
||||
for index, candidate := range candidates {
|
||||
candidateBody := preprocessRequest(kind, cloneMap(body), candidate)
|
||||
if preprocessing, ok := preprocessingByCandidate[pricingCandidateKey(candidate)]; ok && preprocessing.Err == nil {
|
||||
candidateBody = preprocessing.Body
|
||||
}
|
||||
items := s.estimatedBillings(ctx, user, kind, candidateBody, candidate)
|
||||
amount := billingItemsFixedTotal(items)
|
||||
if index == 0 || amount > maximumAmount {
|
||||
maximumItems = items
|
||||
maximumAmount = amount
|
||||
}
|
||||
}
|
||||
return maximumItems, maximumAmount
|
||||
}
|
||||
|
||||
func billingItemsFixedTotal(items []any) fixedAmount {
|
||||
total := fixedAmount(0)
|
||||
for _, raw := range items {
|
||||
line, _ := raw.(map[string]any)
|
||||
if line == nil {
|
||||
continue
|
||||
}
|
||||
amount, err := fixedAmountFromAny(line["amount"])
|
||||
if err == nil && amount > 0 {
|
||||
next, addErr := addFixedAmounts(total, amount)
|
||||
if addErr != nil {
|
||||
return maxFixedAmount
|
||||
}
|
||||
total = next
|
||||
}
|
||||
}
|
||||
return total
|
||||
}
|
||||
|
||||
func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user *auth.User, body map[string]any, preprocessing parameterPreprocessingLog, candidate store.RuntimeModelCandidate, pricing resolvedPricing, attemptNo int, onDelta clients.StreamDelta, responseExecution responseExecutionContext, singleSourceProtected bool, cacheAffinityPolicy map[string]any, cacheAffinityRecordKeys []string) (clients.Response, error) {
|
||||
simulated := isSimulation(task, candidate)
|
||||
baseAttemptMetrics := mergeMetrics(attemptMetrics(candidate, attemptNo, simulated), parameterPreprocessingMetrics(preprocessing))
|
||||
reservations := s.rateLimitReservations(ctx, user, candidate, body)
|
||||
@@ -599,16 +875,18 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
|
||||
defer s.store.ReleaseConcurrencyLeases(context.WithoutCancel(ctx), limitResult.LeaseIDs)
|
||||
|
||||
attemptID, err := s.store.CreateTaskAttempt(ctx, store.CreateTaskAttemptInput{
|
||||
TaskID: task.ID,
|
||||
AttemptNo: attemptNo,
|
||||
PlatformID: candidate.PlatformID,
|
||||
PlatformModelID: candidate.PlatformModelID,
|
||||
ClientID: candidate.ClientID,
|
||||
QueueKey: candidate.QueueKey,
|
||||
Status: "running",
|
||||
Simulated: simulated,
|
||||
RequestSnapshot: s.slimTaskRequestSnapshot(task, body),
|
||||
Metrics: baseAttemptMetrics,
|
||||
TaskID: task.ID,
|
||||
AttemptNo: attemptNo,
|
||||
PlatformID: candidate.PlatformID,
|
||||
PlatformModelID: candidate.PlatformModelID,
|
||||
ClientID: candidate.ClientID,
|
||||
QueueKey: candidate.QueueKey,
|
||||
Status: "running",
|
||||
Simulated: simulated,
|
||||
RequestSnapshot: s.slimTaskRequestSnapshot(task, body),
|
||||
Metrics: baseAttemptMetrics,
|
||||
PricingSnapshot: pricing.Snapshot,
|
||||
RequestFingerprint: pricingRequestFingerprint(task.Kind, task.Model, body),
|
||||
})
|
||||
if err != nil {
|
||||
return clients.Response{}, fmt.Errorf("create task attempt: %w", err)
|
||||
@@ -659,7 +937,19 @@ 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.hydrateProviderRequestAssets(ctx, providerBody, candidate)
|
||||
if err != nil {
|
||||
_ = s.store.FinishTaskAttempt(ctx, store.FinishTaskAttemptInput{
|
||||
AttemptID: attemptID,
|
||||
@@ -678,6 +968,9 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
|
||||
publicResponseID = responseExecution.PublicResponseID
|
||||
publicPreviousResponseID = responseExecution.PublicPreviousResponseID
|
||||
}
|
||||
if err := s.store.SetAttemptUpstreamSubmissionStatus(ctx, attemptID, "submitting"); err != nil {
|
||||
return clients.Response{}, fmt.Errorf("mark upstream submission: %w", err)
|
||||
}
|
||||
response, err := client.Run(ctx, clients.Request{
|
||||
Kind: task.Kind,
|
||||
ModelType: candidate.ModelType,
|
||||
@@ -691,7 +984,13 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
|
||||
if strings.TrimSpace(remoteTaskID) == "" {
|
||||
return nil
|
||||
}
|
||||
return s.store.SetTaskRemoteTask(context.WithoutCancel(ctx), task.ID, attemptID, remoteTaskID, payload)
|
||||
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,
|
||||
@@ -702,6 +1001,11 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
|
||||
PreviousResponseTurns: responseExecution.PreviousTurns,
|
||||
})
|
||||
callFinishedAt := time.Now()
|
||||
if err == nil {
|
||||
if markErr := s.store.SetAttemptUpstreamSubmissionStatus(context.WithoutCancel(ctx), attemptID, "response_received"); markErr != nil {
|
||||
return clients.Response{}, &upstreamSubmissionUnknownError{AttemptID: attemptID, Cause: markErr}
|
||||
}
|
||||
}
|
||||
if response.ResponseStartedAt.IsZero() {
|
||||
response.ResponseStartedAt = callStartedAt
|
||||
}
|
||||
@@ -715,6 +1019,11 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
if clients.ErrorResponseMetadata(err).StatusCode > 0 {
|
||||
if markErr := s.store.SetAttemptUpstreamSubmissionStatus(context.WithoutCancel(ctx), attemptID, "response_received"); markErr != nil {
|
||||
return clients.Response{}, &upstreamSubmissionUnknownError{AttemptID: attemptID, Cause: markErr}
|
||||
}
|
||||
}
|
||||
retryable := clients.IsRetryable(err)
|
||||
requestID, metrics, responseStartedAt, responseFinishedAt, responseDurationMS := failureMetrics(err, simulated)
|
||||
if responseStartedAt.IsZero() {
|
||||
@@ -743,6 +1052,9 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
|
||||
ErrorMessage: err.Error(),
|
||||
})
|
||||
_ = s.emit(ctx, task.ID, "task.attempt.failed", "running", "attempt_failed", 0.45, err.Error(), map[string]any{"attempt": attemptNo, "retryable": retryable, "requestId": requestID, "statusCode": clients.ErrorResponseMetadata(err).StatusCode, "metrics": metrics}, simulated)
|
||||
if !simulated && clients.ErrorResponseMetadata(err).StatusCode == 0 {
|
||||
return clients.Response{}, &upstreamSubmissionUnknownError{AttemptID: attemptID, Cause: err}
|
||||
}
|
||||
s.applyCandidateFailurePolicies(ctx, task.ID, candidate, err, simulated, singleSourceProtected)
|
||||
return clients.Response{}, err
|
||||
}
|
||||
@@ -824,19 +1136,7 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
|
||||
return clients.Response{}, fmt.Errorf("commit rate limit reservations: %w", err)
|
||||
}
|
||||
rateReservationsFinalized = true
|
||||
if err := s.store.FinishTaskAttempt(ctx, store.FinishTaskAttemptInput{
|
||||
AttemptID: attemptID,
|
||||
Status: "succeeded",
|
||||
RequestID: response.RequestID,
|
||||
Usage: usageToMap(response.Usage),
|
||||
Metrics: mergeMetrics(taskMetrics(task, user, body, candidate, response, simulated), parameterPreprocessingMetrics(preprocessing)),
|
||||
ResponseSnapshot: response.Result,
|
||||
ResponseStartedAt: response.ResponseStartedAt,
|
||||
ResponseFinishedAt: response.ResponseFinishedAt,
|
||||
ResponseDurationMS: response.ResponseDurationMS,
|
||||
}); err != nil {
|
||||
return clients.Response{}, fmt.Errorf("finish task attempt: %w", err)
|
||||
}
|
||||
response.AttemptID = attemptID
|
||||
if err := s.store.RecordCacheAffinityObservation(context.WithoutCancel(ctx), store.CacheAffinityObservationInput{
|
||||
CacheAffinityKey: candidate.CacheAffinity.Key,
|
||||
CacheAffinityKeys: cacheAffinityRecordKeys,
|
||||
@@ -904,7 +1204,7 @@ func (s *Service) clientFor(candidate store.RuntimeModelCandidate, simulated boo
|
||||
return s.clients["openai"]
|
||||
}
|
||||
|
||||
func (s *Service) failTask(ctx context.Context, taskID string, code string, message string, simulated bool, cause error, extraMetrics ...map[string]any) (store.GatewayTask, error) {
|
||||
func (s *Service) failTask(ctx context.Context, taskID string, executionToken string, code string, message string, simulated bool, cause error, extraMetrics ...map[string]any) (store.GatewayTask, error) {
|
||||
requestID, metrics, responseStartedAt, responseFinishedAt, responseDurationMS := failureMetrics(cause, simulated)
|
||||
if len(extraMetrics) > 0 {
|
||||
values := append([]map[string]any{metrics}, extraMetrics...)
|
||||
@@ -913,6 +1213,7 @@ func (s *Service) failTask(ctx context.Context, taskID string, code string, mess
|
||||
metrics = s.withAttemptHistory(ctx, taskID, metrics)
|
||||
failed, err := s.store.FinishTaskFailure(ctx, store.FinishTaskFailureInput{
|
||||
TaskID: taskID,
|
||||
ExecutionToken: executionToken,
|
||||
Code: code,
|
||||
Message: message,
|
||||
Result: buildFailureResult(code, message, requestID, cause),
|
||||
@@ -925,12 +1226,19 @@ func (s *Service) failTask(ctx context.Context, taskID string, code string, mess
|
||||
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
|
||||
@@ -1041,7 +1349,7 @@ func (s *Service) requeueRateLimitedTask(ctx context.Context, task store.Gateway
|
||||
if delay <= 0 {
|
||||
delay = 5 * time.Second
|
||||
}
|
||||
queued, err := s.store.RequeueTask(ctx, task.ID, delay, candidate.QueueKey)
|
||||
queued, err := s.store.RequeueTask(ctx, task.ID, task.ExecutionToken, delay, candidate.QueueKey)
|
||||
if err != nil {
|
||||
return store.GatewayTask{}, 0, err
|
||||
}
|
||||
@@ -1057,7 +1365,7 @@ func (s *Service) requeueRateLimitedTask(ctx context.Context, task store.Gateway
|
||||
}
|
||||
|
||||
func (s *Service) requeueInterruptedAsyncTask(ctx context.Context, task store.GatewayTask) (store.GatewayTask, error) {
|
||||
queued, err := s.store.RequeueTask(ctx, task.ID, 0, "")
|
||||
queued, err := s.store.RequeueTask(ctx, task.ID, task.ExecutionToken, 0, "")
|
||||
if err != nil {
|
||||
return store.GatewayTask{}, err
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -182,6 +182,20 @@ func TestExecuteWithMockClientRejectsConcurrentTasksBeyondWalletBalance(t *testi
|
||||
if got := mockClient.calls.Load(); got != 1 {
|
||||
t.Fatalf("mock client calls = %d, want 1", got)
|
||||
}
|
||||
settlements, err := db.ClaimBillingSettlements(ctx, "wallet-execute-test", store.BillingSettlementBatchSize, store.BillingSettlementLockTimeout)
|
||||
if err != nil {
|
||||
t.Fatalf("claim billing settlements: %v", err)
|
||||
}
|
||||
processed := 0
|
||||
for _, settlement := range settlements {
|
||||
if err := db.ProcessBillingSettlement(ctx, settlement); err != nil {
|
||||
t.Fatalf("process billing settlement %s: %v", settlement.ID, err)
|
||||
}
|
||||
processed++
|
||||
}
|
||||
if processed != 1 {
|
||||
t.Fatalf("processed billing settlements = %d, want 1", processed)
|
||||
}
|
||||
|
||||
summary, err := db.GetWalletSummary(ctx, user, "resource")
|
||||
if err != nil {
|
||||
|
||||
@@ -18,25 +18,32 @@ type MetricsSnapshotProvider interface {
|
||||
type DynamicMetricsSnapshotProvider interface {
|
||||
MetricsSnapshotProvider
|
||||
SecurityEventConnection(context.Context) (store.SecurityEventConnection, error)
|
||||
BillingMetrics(context.Context) (store.BillingMetricsSnapshot, error)
|
||||
}
|
||||
|
||||
type Metrics struct {
|
||||
accepted atomic.Uint64
|
||||
rejected atomic.Uint64
|
||||
duplicate atomic.Uint64
|
||||
sessionsDeleted atomic.Uint64
|
||||
watermarkRejected atomic.Uint64
|
||||
verificationAccepted atomic.Uint64
|
||||
heartbeatAccepted atomic.Uint64
|
||||
heartbeatFailed atomic.Uint64
|
||||
introspectionActive atomic.Uint64
|
||||
introspectionInactive atomic.Uint64
|
||||
introspectionFailed atomic.Uint64
|
||||
jwksSSFFailed atomic.Uint64
|
||||
jwksOIDCFailed atomic.Uint64
|
||||
processingCount atomic.Uint64
|
||||
processingNanos atomic.Uint64
|
||||
processingBuckets [6]atomic.Uint64
|
||||
accepted atomic.Uint64
|
||||
rejected atomic.Uint64
|
||||
duplicate atomic.Uint64
|
||||
sessionsDeleted atomic.Uint64
|
||||
watermarkRejected atomic.Uint64
|
||||
verificationAccepted atomic.Uint64
|
||||
heartbeatAccepted atomic.Uint64
|
||||
heartbeatFailed atomic.Uint64
|
||||
introspectionActive atomic.Uint64
|
||||
introspectionInactive atomic.Uint64
|
||||
introspectionFailed atomic.Uint64
|
||||
jwksSSFFailed atomic.Uint64
|
||||
jwksOIDCFailed atomic.Uint64
|
||||
processingCount atomic.Uint64
|
||||
processingNanos atomic.Uint64
|
||||
processingBuckets [6]atomic.Uint64
|
||||
billingSettlementCompleted atomic.Uint64
|
||||
billingSettlementRetry atomic.Uint64
|
||||
billingManualReview atomic.Uint64
|
||||
billingEstimateFailed atomic.Uint64
|
||||
billingIdempotentReplay atomic.Uint64
|
||||
billingPricingUnavailable atomic.Uint64
|
||||
}
|
||||
|
||||
var processingDurationBounds = [...]time.Duration{
|
||||
@@ -99,6 +106,23 @@ func (m *Metrics) ObserveJWKSRefreshFailure(source string) {
|
||||
m.jwksOIDCFailed.Add(1)
|
||||
}
|
||||
|
||||
func (m *Metrics) ObserveBillingEvent(event string) {
|
||||
switch event {
|
||||
case "settlement_completed":
|
||||
m.billingSettlementCompleted.Add(1)
|
||||
case "settlement_retry":
|
||||
m.billingSettlementRetry.Add(1)
|
||||
case "manual_review":
|
||||
m.billingManualReview.Add(1)
|
||||
case "estimate_failed":
|
||||
m.billingEstimateFailed.Add(1)
|
||||
case "idempotent_replay":
|
||||
m.billingIdempotentReplay.Add(1)
|
||||
case "pricing_unavailable":
|
||||
m.billingPricingUnavailable.Add(1)
|
||||
}
|
||||
}
|
||||
|
||||
func (m *Metrics) Handler(provider MetricsSnapshotProvider, issuer, audience string, enabled bool) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
mode := "disabled"
|
||||
@@ -113,6 +137,17 @@ func (m *Metrics) Handler(provider MetricsSnapshotProvider, issuer, audience str
|
||||
return
|
||||
}
|
||||
}
|
||||
billing := store.BillingMetricsSnapshot{}
|
||||
if billingProvider, ok := provider.(interface {
|
||||
BillingMetrics(context.Context) (store.BillingMetricsSnapshot, error)
|
||||
}); ok {
|
||||
var err error
|
||||
billing, err = billingProvider.BillingMetrics(r.Context())
|
||||
if err != nil {
|
||||
http.Error(w, "metrics unavailable", http.StatusServiceUnavailable)
|
||||
return
|
||||
}
|
||||
}
|
||||
w.Header().Set("Content-Type", "text/plain; version=0.0.4; charset=utf-8")
|
||||
outcomeCounters(w, "easyai_gateway_ssf_receipts_total", "Received SETs by bounded outcome.", []outcomeValue{
|
||||
{"accepted", m.accepted.Load()}, {"rejected", m.rejected.Load()}, {"duplicate", m.duplicate.Load()},
|
||||
@@ -156,6 +191,27 @@ func (m *Metrics) Handler(provider MetricsSnapshotProvider, issuer, audience str
|
||||
}
|
||||
fmt.Fprintf(w, "easyai_gateway_ssf_mode{mode=\"%s\"} %d\n", modeName, value)
|
||||
}
|
||||
fmt.Fprintln(w, "# HELP easyai_gateway_billing_settlement_backlog Current unsettled Outbox records.")
|
||||
fmt.Fprintln(w, "# TYPE easyai_gateway_billing_settlement_backlog gauge")
|
||||
fmt.Fprintf(w, "easyai_gateway_billing_settlement_backlog %d\n", billing.SettlementBacklog)
|
||||
fmt.Fprintln(w, "# HELP easyai_gateway_billing_settlement_delay_seconds Age of the oldest unsettled Outbox record.")
|
||||
fmt.Fprintln(w, "# TYPE easyai_gateway_billing_settlement_delay_seconds gauge")
|
||||
fmt.Fprintf(w, "easyai_gateway_billing_settlement_delay_seconds %.6f\n", billing.SettlementDelaySecs)
|
||||
fmt.Fprintln(w, "# HELP easyai_gateway_billing_manual_review Current billing records requiring manual review.")
|
||||
fmt.Fprintln(w, "# TYPE easyai_gateway_billing_manual_review gauge")
|
||||
fmt.Fprintf(w, "easyai_gateway_billing_manual_review %d\n", billing.ManualReview)
|
||||
fmt.Fprintln(w, "# HELP easyai_gateway_billing_orphan_frozen Wallets whose frozen amount differs from active reservations.")
|
||||
fmt.Fprintln(w, "# TYPE easyai_gateway_billing_orphan_frozen gauge")
|
||||
fmt.Fprintf(w, "easyai_gateway_billing_orphan_frozen %d\n", billing.OrphanFrozen)
|
||||
fmt.Fprintln(w, "# HELP easyai_gateway_billing_pricing_unavailable_tasks Current tasks rejected because no effective price was available.")
|
||||
fmt.Fprintln(w, "# TYPE easyai_gateway_billing_pricing_unavailable_tasks gauge")
|
||||
fmt.Fprintf(w, "easyai_gateway_billing_pricing_unavailable_tasks %d\n", billing.PricingUnavailable)
|
||||
plainCounter(w, "easyai_gateway_billing_settlements_completed_total", "Billing settlements completed by this process.", m.billingSettlementCompleted.Load())
|
||||
plainCounter(w, "easyai_gateway_billing_settlement_retries_total", "Billing settlement retries scheduled by this process.", m.billingSettlementRetry.Load())
|
||||
plainCounter(w, "easyai_gateway_billing_manual_review_transitions_total", "Billing records transitioned to manual review by this process.", m.billingManualReview.Load())
|
||||
plainCounter(w, "easyai_gateway_billing_estimate_failures_total", "Pricing estimate requests that failed.", m.billingEstimateFailed.Load())
|
||||
plainCounter(w, "easyai_gateway_billing_idempotent_replays_total", "Generation requests replayed idempotently.", m.billingIdempotentReplay.Load())
|
||||
plainCounter(w, "easyai_gateway_billing_pricing_unavailable_total", "Pricing requests rejected because no effective price was available.", m.billingPricingUnavailable.Load())
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
+8
-7
@@ -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.
|
||||
|
||||
+3
-3
@@ -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.
|
||||
|
||||
+4
-4
@@ -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
|
||||
```
|
||||
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -488,6 +488,8 @@ func modelTypeAliases(value string) []string {
|
||||
return []string{"image_edit"}
|
||||
case "video", "videos.generations":
|
||||
return []string{"video_generate"}
|
||||
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":
|
||||
|
||||
@@ -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},
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,637 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
)
|
||||
|
||||
const (
|
||||
BillingSettlementBatchSize = 50
|
||||
BillingSettlementLockTimeout = 120 * time.Second
|
||||
BillingSettlementMaxAttempts = 20
|
||||
)
|
||||
|
||||
var ErrBillingSettlementNotRetryable = errors.New("billing settlement is not retryable")
|
||||
|
||||
type BillingSettlement struct {
|
||||
ID string `json:"id"`
|
||||
TaskID string `json:"taskId"`
|
||||
Action string `json:"action"`
|
||||
Amount float64 `json:"amount"`
|
||||
Currency string `json:"currency"`
|
||||
Status string `json:"status"`
|
||||
Attempts int `json:"attempts"`
|
||||
NextAttemptAt time.Time `json:"nextAttemptAt"`
|
||||
LockedBy string `json:"lockedBy,omitempty"`
|
||||
LockedAt string `json:"lockedAt,omitempty"`
|
||||
LastErrorCode string `json:"lastErrorCode,omitempty"`
|
||||
LastErrorMessage string `json:"lastErrorMessage,omitempty"`
|
||||
ManualReviewReason string `json:"manualReviewReason,omitempty"`
|
||||
PricingSnapshot map[string]any `json:"pricingSnapshot,omitempty"`
|
||||
Payload map[string]any `json:"payload,omitempty"`
|
||||
CompletedAt string `json:"completedAt,omitempty"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
LockToken string `json:"-"`
|
||||
amountExact string
|
||||
}
|
||||
|
||||
type BillingSettlementListFilter struct {
|
||||
Status string
|
||||
Action string
|
||||
Page int
|
||||
PageSize int
|
||||
}
|
||||
|
||||
type BillingSettlementListResult struct {
|
||||
Items []BillingSettlement `json:"items"`
|
||||
Total int `json:"total"`
|
||||
Page int `json:"page"`
|
||||
PageSize int `json:"pageSize"`
|
||||
}
|
||||
|
||||
type BillingMetricsSnapshot struct {
|
||||
SettlementBacklog int64
|
||||
SettlementDelaySecs float64
|
||||
ManualReview int64
|
||||
OrphanFrozen int64
|
||||
PricingUnavailable int64
|
||||
}
|
||||
|
||||
func (s *Store) BillingMetrics(ctx context.Context) (BillingMetricsSnapshot, error) {
|
||||
var snapshot BillingMetricsSnapshot
|
||||
err := s.pool.QueryRow(ctx, `
|
||||
WITH open_outbox AS (
|
||||
SELECT created_at
|
||||
FROM settlement_outbox
|
||||
WHERE status IN ('pending', 'processing', 'retryable_failed')
|
||||
), active_reservations AS (
|
||||
SELECT reserve.account_id, COALESCE(SUM(reserve.amount), 0) AS amount
|
||||
FROM gateway_wallet_transactions reserve
|
||||
WHERE reserve.transaction_type = 'reserve'
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM gateway_wallet_transactions release
|
||||
WHERE release.account_id = reserve.account_id
|
||||
AND release.transaction_type = 'release'
|
||||
AND release.idempotency_key = reserve.idempotency_key || ':release'
|
||||
)
|
||||
GROUP BY reserve.account_id
|
||||
)
|
||||
SELECT
|
||||
(SELECT COUNT(*) FROM open_outbox),
|
||||
COALESCE((SELECT EXTRACT(EPOCH FROM now() - MIN(created_at)) FROM open_outbox), 0),
|
||||
(SELECT COUNT(*) FROM gateway_tasks WHERE billing_status = 'manual_review'),
|
||||
(SELECT COUNT(*) FROM gateway_wallet_accounts account
|
||||
LEFT JOIN active_reservations reservation ON reservation.account_id = account.id
|
||||
WHERE account.frozen_balance <> COALESCE(reservation.amount, 0)),
|
||||
(SELECT COUNT(*) FROM gateway_tasks WHERE error_code = 'pricing_unavailable')`).Scan(
|
||||
&snapshot.SettlementBacklog,
|
||||
&snapshot.SettlementDelaySecs,
|
||||
&snapshot.ManualReview,
|
||||
&snapshot.OrphanFrozen,
|
||||
&snapshot.PricingUnavailable,
|
||||
)
|
||||
return snapshot, err
|
||||
}
|
||||
|
||||
func (s *Store) ClaimBillingSettlements(ctx context.Context, workerID string, limit int, staleAfter time.Duration) ([]BillingSettlement, error) {
|
||||
if limit <= 0 || limit > BillingSettlementBatchSize {
|
||||
limit = BillingSettlementBatchSize
|
||||
}
|
||||
if staleAfter <= 0 {
|
||||
staleAfter = BillingSettlementLockTimeout
|
||||
}
|
||||
rows, err := s.pool.Query(ctx, `
|
||||
WITH picked AS (
|
||||
SELECT id
|
||||
FROM settlement_outbox
|
||||
WHERE (
|
||||
status IN ('pending', 'retryable_failed')
|
||||
AND next_attempt_at <= now()
|
||||
)
|
||||
OR (
|
||||
status = 'processing'
|
||||
AND locked_at < now() - ($3::int * interval '1 second')
|
||||
)
|
||||
ORDER BY next_attempt_at ASC, created_at ASC
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
)
|
||||
UPDATE settlement_outbox outbox
|
||||
SET status = 'processing',
|
||||
attempts = outbox.attempts + 1,
|
||||
locked_by = $1,
|
||||
lock_token = gen_random_uuid(),
|
||||
locked_at = now(),
|
||||
updated_at = now()
|
||||
FROM picked
|
||||
WHERE outbox.id = picked.id
|
||||
RETURNING outbox.id::text, outbox.task_id::text, outbox.action, outbox.amount::text,
|
||||
outbox.currency, outbox.status, outbox.attempts, outbox.next_attempt_at,
|
||||
COALESCE(outbox.locked_by, ''), COALESCE(outbox.lock_token::text, ''),
|
||||
COALESCE(outbox.locked_at::text, ''), COALESCE(outbox.last_error_code, ''),
|
||||
COALESCE(outbox.last_error_message, ''), COALESCE(outbox.manual_review_reason, ''),
|
||||
outbox.pricing_snapshot, outbox.payload, COALESCE(outbox.completed_at::text, ''),
|
||||
outbox.created_at, outbox.updated_at`,
|
||||
workerID, limit, int(staleAfter/time.Second))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := make([]BillingSettlement, 0)
|
||||
for rows.Next() {
|
||||
item, err := scanBillingSettlement(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
|
||||
func (s *Store) ProcessBillingSettlement(ctx context.Context, settlement BillingSettlement) error {
|
||||
return pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
|
||||
var action string
|
||||
var amount string
|
||||
var currency string
|
||||
var taskID string
|
||||
var taskStatus string
|
||||
var gatewayUserID string
|
||||
var gatewayTenantID string
|
||||
var userID string
|
||||
var tenantID string
|
||||
var tenantKey string
|
||||
err := tx.QueryRow(ctx, `
|
||||
SELECT outbox.action, outbox.amount::text, outbox.currency, outbox.task_id::text,
|
||||
task.status, COALESCE(task.gateway_user_id::text, ''),
|
||||
COALESCE(task.gateway_tenant_id::text, ''), task.user_id,
|
||||
COALESCE(task.tenant_id, ''), COALESCE(task.tenant_key, '')
|
||||
FROM settlement_outbox outbox
|
||||
JOIN gateway_tasks task ON task.id = outbox.task_id
|
||||
WHERE outbox.id = $1::uuid
|
||||
AND outbox.status = 'processing'
|
||||
AND outbox.lock_token = $2::uuid
|
||||
FOR UPDATE OF outbox, task`, settlement.ID, settlement.LockToken).Scan(
|
||||
&action, &amount, ¤cy, &taskID, &taskStatus, &gatewayUserID,
|
||||
&gatewayTenantID, &userID, &tenantID, &tenantKey,
|
||||
)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := tx.Exec(ctx, `
|
||||
UPDATE gateway_tasks
|
||||
SET billing_status = 'processing', billing_updated_at = now(), updated_at = now()
|
||||
WHERE id = $1::uuid`, taskID); err != nil {
|
||||
return err
|
||||
}
|
||||
switch action {
|
||||
case "settle":
|
||||
if taskStatus != "succeeded" {
|
||||
return fmt.Errorf("settle action requires succeeded task")
|
||||
}
|
||||
if err := settleBillingOutboxTx(ctx, tx, taskID, gatewayUserID, gatewayTenantID, userID, tenantID, tenantKey, currency, amount); err != nil {
|
||||
return err
|
||||
}
|
||||
return completeBillingOutboxTx(ctx, tx, settlement.ID, settlement.LockToken, taskID, "settled")
|
||||
case "release":
|
||||
if err := releaseBillingOutboxTx(ctx, tx, taskID, gatewayUserID, gatewayTenantID, currency); err != nil {
|
||||
return err
|
||||
}
|
||||
return completeBillingOutboxTx(ctx, tx, settlement.ID, settlement.LockToken, taskID, "released")
|
||||
default:
|
||||
return fmt.Errorf("unsupported billing settlement action %q", action)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func settleBillingOutboxTx(
|
||||
ctx context.Context,
|
||||
tx pgx.Tx,
|
||||
taskID string,
|
||||
gatewayUserID string,
|
||||
gatewayTenantID string,
|
||||
userID string,
|
||||
tenantID string,
|
||||
tenantKey string,
|
||||
currency string,
|
||||
amount string,
|
||||
) error {
|
||||
if gatewayUserID == "" {
|
||||
return fmt.Errorf("task %s has no gateway wallet user", taskID)
|
||||
}
|
||||
if _, err := tx.Exec(ctx, `
|
||||
INSERT INTO gateway_wallet_accounts (
|
||||
gateway_tenant_id, gateway_user_id, tenant_id, tenant_key, user_id, currency
|
||||
)
|
||||
VALUES (NULLIF($1, '')::uuid, $2::uuid, NULLIF($3, ''), NULLIF($4, ''), NULLIF($5, ''), $6)
|
||||
ON CONFLICT (gateway_user_id, currency) DO NOTHING`,
|
||||
gatewayTenantID, gatewayUserID, tenantID, tenantKey, userID, currency); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := ensureWalletAccountAuditGuard(ctx, tx, gatewayUserID, currency); err != nil {
|
||||
return err
|
||||
}
|
||||
var accountID string
|
||||
var balanceBefore string
|
||||
var frozenBefore string
|
||||
if err := tx.QueryRow(ctx, `
|
||||
SELECT id::text, balance::text, frozen_balance::text
|
||||
FROM gateway_wallet_accounts
|
||||
WHERE gateway_user_id = $1::uuid AND currency = $2
|
||||
FOR UPDATE`, gatewayUserID, currency).Scan(&accountID, &balanceBefore, &frozenBefore); err != nil {
|
||||
return err
|
||||
}
|
||||
var alreadySettled bool
|
||||
if err := tx.QueryRow(ctx, `
|
||||
SELECT EXISTS (
|
||||
SELECT 1 FROM gateway_wallet_transactions
|
||||
WHERE account_id = $1::uuid AND idempotency_key = $2
|
||||
)`, accountID, billingIdempotencyKey(taskID)).Scan(&alreadySettled); err != nil {
|
||||
return err
|
||||
}
|
||||
if alreadySettled {
|
||||
return nil
|
||||
}
|
||||
|
||||
var reservationKey string
|
||||
var reservedAmount string
|
||||
if err := tx.QueryRow(ctx, `
|
||||
SELECT COALESCE(reserve.idempotency_key, ''), reserve.amount::text
|
||||
FROM gateway_wallet_transactions reserve
|
||||
WHERE reserve.account_id = $1::uuid
|
||||
AND reserve.reference_type = 'gateway_task'
|
||||
AND reserve.reference_id = $2
|
||||
AND reserve.transaction_type = 'reserve'
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM gateway_wallet_transactions release
|
||||
WHERE release.account_id = reserve.account_id
|
||||
AND release.transaction_type = 'release'
|
||||
AND release.idempotency_key = reserve.idempotency_key || ':release'
|
||||
)
|
||||
ORDER BY reserve.created_at DESC
|
||||
LIMIT 1`, accountID, taskID).Scan(&reservationKey, &reservedAmount); errors.Is(err, pgx.ErrNoRows) {
|
||||
reservationKey = ""
|
||||
reservedAmount = "0"
|
||||
} else if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
var balanceAfter string
|
||||
var frozenAfter string
|
||||
err := tx.QueryRow(ctx, `
|
||||
UPDATE gateway_wallet_accounts
|
||||
SET balance = balance - $2::numeric,
|
||||
total_spent = total_spent + $2::numeric,
|
||||
frozen_balance = GREATEST(0, frozen_balance - $3::numeric),
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid
|
||||
AND balance - frozen_balance + $3::numeric >= $2::numeric
|
||||
RETURNING balance::text, frozen_balance::text`,
|
||||
accountID, amount, reservedAmount).Scan(&balanceAfter, &frozenAfter)
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return fmt.Errorf("%w: task %s settlement amount exceeds spendable balance", ErrInsufficientWalletBalance, taskID)
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
var reservedPositive bool
|
||||
if err := tx.QueryRow(ctx, `SELECT $1::numeric > 0`, reservedAmount).Scan(&reservedPositive); err != nil {
|
||||
return err
|
||||
}
|
||||
if reservedPositive {
|
||||
releaseMetadata, _ := json.Marshal(map[string]any{
|
||||
"taskId": taskID, "reason": "task_billing_settled",
|
||||
"reserved": reservedAmount, "frozenBefore": frozenBefore, "frozenAfter": frozenAfter,
|
||||
})
|
||||
if _, err := tx.Exec(ctx, `
|
||||
INSERT INTO gateway_wallet_transactions (
|
||||
account_id, gateway_tenant_id, gateway_user_id, direction, transaction_type,
|
||||
amount, balance_before, balance_after, idempotency_key, reference_type, reference_id, metadata
|
||||
)
|
||||
VALUES (
|
||||
$1::uuid, NULLIF($2, '')::uuid, $3::uuid, 'credit', 'release',
|
||||
$4::numeric, $5::numeric, $5::numeric, $6, 'gateway_task', $7, $8::jsonb
|
||||
)
|
||||
ON CONFLICT (account_id, idempotency_key) WHERE idempotency_key IS NOT NULL DO NOTHING`,
|
||||
accountID, gatewayTenantID, gatewayUserID, reservedAmount, balanceBefore,
|
||||
billingReservationReleaseIdempotencyKey(reservationKey), taskID, string(releaseMetadata)); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
metadata, _ := json.Marshal(map[string]any{
|
||||
"taskId": taskID, "reservedAmount": reservedAmount,
|
||||
"frozenBefore": frozenBefore, "frozenAfter": frozenAfter,
|
||||
})
|
||||
_, err = tx.Exec(ctx, `
|
||||
INSERT INTO gateway_wallet_transactions (
|
||||
account_id, gateway_tenant_id, gateway_user_id, direction, transaction_type,
|
||||
amount, balance_before, balance_after, idempotency_key, reference_type, reference_id, metadata
|
||||
)
|
||||
VALUES (
|
||||
$1::uuid, NULLIF($2, '')::uuid, $3::uuid, 'debit', 'task_billing',
|
||||
$4::numeric, $5::numeric, $6::numeric, $7, 'gateway_task', $8, $9::jsonb
|
||||
)
|
||||
ON CONFLICT (account_id, idempotency_key) WHERE idempotency_key IS NOT NULL DO NOTHING`,
|
||||
accountID, gatewayTenantID, gatewayUserID, amount, balanceBefore, balanceAfter,
|
||||
billingIdempotencyKey(taskID), taskID, string(metadata))
|
||||
return err
|
||||
}
|
||||
|
||||
func releaseBillingOutboxTx(ctx context.Context, tx pgx.Tx, taskID string, gatewayUserID string, gatewayTenantID string, currency string) error {
|
||||
if gatewayUserID == "" {
|
||||
return nil
|
||||
}
|
||||
var accountID string
|
||||
var balance string
|
||||
var frozenBefore string
|
||||
err := tx.QueryRow(ctx, `
|
||||
SELECT id::text, balance::text, frozen_balance::text
|
||||
FROM gateway_wallet_accounts
|
||||
WHERE gateway_user_id = $1::uuid AND currency = $2
|
||||
FOR UPDATE`, gatewayUserID, currency).Scan(&accountID, &balance, &frozenBefore)
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return nil
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
var reservationKey string
|
||||
var reservedAmount string
|
||||
err = tx.QueryRow(ctx, `
|
||||
SELECT COALESCE(reserve.idempotency_key, ''), reserve.amount::text
|
||||
FROM gateway_wallet_transactions reserve
|
||||
WHERE reserve.account_id = $1::uuid
|
||||
AND reserve.reference_type = 'gateway_task'
|
||||
AND reserve.reference_id = $2
|
||||
AND reserve.transaction_type = 'reserve'
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM gateway_wallet_transactions release
|
||||
WHERE release.account_id = reserve.account_id
|
||||
AND release.transaction_type = 'release'
|
||||
AND release.idempotency_key = reserve.idempotency_key || ':release'
|
||||
)
|
||||
ORDER BY reserve.created_at DESC
|
||||
LIMIT 1`, accountID, taskID).Scan(&reservationKey, &reservedAmount)
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return nil
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
var frozenAfter string
|
||||
if err := tx.QueryRow(ctx, `
|
||||
UPDATE gateway_wallet_accounts
|
||||
SET frozen_balance = frozen_balance - $2::numeric,
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid
|
||||
AND frozen_balance >= $2::numeric
|
||||
RETURNING frozen_balance::text`, accountID, reservedAmount).Scan(&frozenAfter); err != nil {
|
||||
return err
|
||||
}
|
||||
metadata, _ := json.Marshal(map[string]any{
|
||||
"taskId": taskID, "reason": "task_terminal_release",
|
||||
"reserved": reservedAmount, "frozenBefore": frozenBefore, "frozenAfter": frozenAfter,
|
||||
})
|
||||
_, err = tx.Exec(ctx, `
|
||||
INSERT INTO gateway_wallet_transactions (
|
||||
account_id, gateway_tenant_id, gateway_user_id, direction, transaction_type,
|
||||
amount, balance_before, balance_after, idempotency_key, reference_type, reference_id, metadata
|
||||
)
|
||||
VALUES (
|
||||
$1::uuid, NULLIF($2, '')::uuid, $3::uuid, 'credit', 'release',
|
||||
$4::numeric, $5::numeric, $5::numeric, $6, 'gateway_task', $7, $8::jsonb
|
||||
)
|
||||
ON CONFLICT (account_id, idempotency_key) WHERE idempotency_key IS NOT NULL DO NOTHING`,
|
||||
accountID, gatewayTenantID, gatewayUserID, reservedAmount, balance,
|
||||
billingReservationReleaseIdempotencyKey(reservationKey), taskID, string(metadata))
|
||||
return err
|
||||
}
|
||||
|
||||
func completeBillingOutboxTx(ctx context.Context, tx pgx.Tx, settlementID string, lockToken string, taskID string, billingStatus string) error {
|
||||
if _, err := tx.Exec(ctx, `
|
||||
UPDATE gateway_tasks
|
||||
SET billing_status = $2,
|
||||
reservation_amount = 0,
|
||||
billing_updated_at = now(),
|
||||
billing_settled_at = CASE WHEN $2 = 'settled' THEN now() ELSE billing_settled_at END,
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid`, taskID, billingStatus); err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := tx.Exec(ctx, `
|
||||
INSERT INTO gateway_task_events (task_id, seq, event_type, status, phase, progress, message, payload, simulated)
|
||||
SELECT task.id,
|
||||
COALESCE((SELECT MAX(event.seq) + 1 FROM gateway_task_events event WHERE event.task_id = task.id), 1),
|
||||
CASE WHEN $2 = 'settled' THEN 'task.billing.settled' ELSE 'task.billing.released' END,
|
||||
task.status,
|
||||
'billing',
|
||||
1,
|
||||
CASE WHEN $2 = 'settled' THEN 'task billing settled' ELSE 'task billing reservation released' END,
|
||||
jsonb_build_object('settlementId', $3::text, 'billingStatus', $2::text),
|
||||
task.run_mode = 'simulation'
|
||||
FROM gateway_tasks task
|
||||
WHERE task.id = $1::uuid`, taskID, billingStatus, settlementID); err != nil {
|
||||
return err
|
||||
}
|
||||
tag, err := tx.Exec(ctx, `
|
||||
UPDATE settlement_outbox
|
||||
SET status = 'completed',
|
||||
completed_at = now(),
|
||||
locked_by = NULL,
|
||||
lock_token = NULL,
|
||||
locked_at = NULL,
|
||||
last_error_code = NULL,
|
||||
last_error_message = NULL,
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid
|
||||
AND status = 'processing'
|
||||
AND lock_token = $2::uuid`, settlementID, lockToken)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() != 1 {
|
||||
return pgx.ErrNoRows
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Store) MarkBillingSettlementFailed(ctx context.Context, settlement BillingSettlement, code string, message string, delay time.Duration) error {
|
||||
if delay < time.Second {
|
||||
delay = time.Second
|
||||
}
|
||||
if delay > 15*time.Minute {
|
||||
delay = 15 * time.Minute
|
||||
}
|
||||
manualReview := settlement.Attempts >= BillingSettlementMaxAttempts
|
||||
status := "retryable_failed"
|
||||
taskStatus := "retryable_failed"
|
||||
manualReason := ""
|
||||
if manualReview {
|
||||
status = "manual_review"
|
||||
taskStatus = "manual_review"
|
||||
manualReason = "maximum settlement attempts exceeded"
|
||||
}
|
||||
return pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
|
||||
tag, err := tx.Exec(ctx, `
|
||||
UPDATE settlement_outbox
|
||||
SET status = $3,
|
||||
next_attempt_at = now() + ($4::int * interval '1 second'),
|
||||
locked_by = NULL,
|
||||
lock_token = NULL,
|
||||
locked_at = NULL,
|
||||
last_error_code = NULLIF($5, ''),
|
||||
last_error_message = NULLIF($6, ''),
|
||||
manual_review_reason = NULLIF($7, ''),
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid
|
||||
AND lock_token = $2::uuid
|
||||
AND status = 'processing'`,
|
||||
settlement.ID, settlement.LockToken, status, int(delay/time.Second), code, message, manualReason)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() != 1 {
|
||||
return pgx.ErrNoRows
|
||||
}
|
||||
_, err = tx.Exec(ctx, `
|
||||
UPDATE gateway_tasks
|
||||
SET billing_status = $2,
|
||||
billing_updated_at = now(),
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid`, settlement.TaskID, taskStatus)
|
||||
return err
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Store) ListBillingSettlements(ctx context.Context, filter BillingSettlementListFilter) (BillingSettlementListResult, error) {
|
||||
page := filter.Page
|
||||
if page <= 0 {
|
||||
page = 1
|
||||
}
|
||||
pageSize := filter.PageSize
|
||||
if pageSize <= 0 {
|
||||
pageSize = 50
|
||||
}
|
||||
if pageSize > 100 {
|
||||
pageSize = 100
|
||||
}
|
||||
status := strings.TrimSpace(filter.Status)
|
||||
action := strings.TrimSpace(filter.Action)
|
||||
var total int
|
||||
if err := s.pool.QueryRow(ctx, `
|
||||
SELECT count(*)
|
||||
FROM settlement_outbox
|
||||
WHERE (NULLIF($1, '') IS NULL OR status = $1)
|
||||
AND (NULLIF($2, '') IS NULL OR action = $2)`, status, action).Scan(&total); err != nil {
|
||||
return BillingSettlementListResult{}, err
|
||||
}
|
||||
rows, err := s.pool.Query(ctx, `
|
||||
SELECT id::text, task_id::text, action, amount::text, currency, status, attempts,
|
||||
next_attempt_at, COALESCE(locked_by, ''), COALESCE(lock_token::text, ''),
|
||||
COALESCE(locked_at::text, ''), COALESCE(last_error_code, ''),
|
||||
COALESCE(last_error_message, ''), COALESCE(manual_review_reason, ''),
|
||||
pricing_snapshot, payload, COALESCE(completed_at::text, ''), created_at, updated_at
|
||||
FROM settlement_outbox
|
||||
WHERE (NULLIF($1, '') IS NULL OR status = $1)
|
||||
AND (NULLIF($2, '') IS NULL OR action = $2)
|
||||
ORDER BY created_at DESC
|
||||
LIMIT $3 OFFSET $4`, status, action, pageSize, (page-1)*pageSize)
|
||||
if err != nil {
|
||||
return BillingSettlementListResult{}, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := make([]BillingSettlement, 0)
|
||||
for rows.Next() {
|
||||
item, err := scanBillingSettlement(rows)
|
||||
if err != nil {
|
||||
return BillingSettlementListResult{}, err
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return BillingSettlementListResult{}, err
|
||||
}
|
||||
return BillingSettlementListResult{Items: items, Total: total, Page: page, PageSize: pageSize}, nil
|
||||
}
|
||||
|
||||
func (s *Store) RetryBillingSettlementTx(ctx context.Context, tx Tx, id string, idempotencyKeyHash string) (BillingSettlement, bool, error) {
|
||||
current, err := scanBillingSettlement(tx.QueryRow(ctx, `
|
||||
SELECT id::text, task_id::text, action, amount::text, currency, status, attempts,
|
||||
next_attempt_at, COALESCE(locked_by, ''), COALESCE(lock_token::text, ''),
|
||||
COALESCE(locked_at::text, ''), COALESCE(last_error_code, ''),
|
||||
COALESCE(last_error_message, ''), COALESCE(manual_review_reason, ''),
|
||||
pricing_snapshot, payload, COALESCE(completed_at::text, ''), created_at, updated_at
|
||||
FROM settlement_outbox
|
||||
WHERE id = $1::uuid
|
||||
FOR UPDATE`, id))
|
||||
if err != nil {
|
||||
return BillingSettlement{}, false, err
|
||||
}
|
||||
var recordedHash string
|
||||
if err := tx.QueryRow(ctx, `
|
||||
SELECT COALESCE(retry_idempotency_key_hash, '')
|
||||
FROM settlement_outbox
|
||||
WHERE id = $1::uuid`, id).Scan(&recordedHash); err != nil {
|
||||
return BillingSettlement{}, false, err
|
||||
}
|
||||
if recordedHash != "" && recordedHash == idempotencyKeyHash {
|
||||
return current, true, nil
|
||||
}
|
||||
if current.Status != "retryable_failed" && current.Status != "manual_review" {
|
||||
return BillingSettlement{}, false, ErrBillingSettlementNotRetryable
|
||||
}
|
||||
item, err := scanBillingSettlement(tx.QueryRow(ctx, `
|
||||
UPDATE settlement_outbox
|
||||
SET status = 'pending',
|
||||
next_attempt_at = now(),
|
||||
locked_by = NULL,
|
||||
lock_token = NULL,
|
||||
locked_at = NULL,
|
||||
manual_review_reason = NULL,
|
||||
retry_idempotency_key_hash = $2,
|
||||
retry_requested_at = now(),
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid
|
||||
AND status IN ('retryable_failed', 'manual_review')
|
||||
RETURNING id::text, task_id::text, action, amount::text, currency, status, attempts,
|
||||
next_attempt_at, COALESCE(locked_by, ''), COALESCE(lock_token::text, ''),
|
||||
COALESCE(locked_at::text, ''), COALESCE(last_error_code, ''),
|
||||
COALESCE(last_error_message, ''), COALESCE(manual_review_reason, ''),
|
||||
pricing_snapshot, payload, COALESCE(completed_at::text, ''), created_at, updated_at`, id, idempotencyKeyHash))
|
||||
if err != nil {
|
||||
return BillingSettlement{}, false, err
|
||||
}
|
||||
if _, err := tx.Exec(ctx, `
|
||||
UPDATE gateway_tasks
|
||||
SET billing_status = 'pending', billing_updated_at = now(), updated_at = now()
|
||||
WHERE id = $1::uuid`, item.TaskID); err != nil {
|
||||
return BillingSettlement{}, false, err
|
||||
}
|
||||
return item, false, nil
|
||||
}
|
||||
|
||||
func scanBillingSettlement(scanner taskScanner) (BillingSettlement, error) {
|
||||
var item BillingSettlement
|
||||
var pricingSnapshotBytes []byte
|
||||
var payloadBytes []byte
|
||||
if err := scanner.Scan(
|
||||
&item.ID, &item.TaskID, &item.Action, &item.amountExact, &item.Currency,
|
||||
&item.Status, &item.Attempts, &item.NextAttemptAt, &item.LockedBy, &item.LockToken,
|
||||
&item.LockedAt, &item.LastErrorCode, &item.LastErrorMessage, &item.ManualReviewReason,
|
||||
&pricingSnapshotBytes, &payloadBytes, &item.CompletedAt, &item.CreatedAt, &item.UpdatedAt,
|
||||
); err != nil {
|
||||
return BillingSettlement{}, err
|
||||
}
|
||||
item.Amount, _ = strconv.ParseFloat(item.amountExact, 64)
|
||||
item.PricingSnapshot = decodeObject(pricingSnapshotBytes)
|
||||
item.Payload = decodeObject(payloadBytes)
|
||||
return item, nil
|
||||
}
|
||||
@@ -0,0 +1,408 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
||||
"github.com/google/uuid"
|
||||
"github.com/jackc/pgx/v5"
|
||||
)
|
||||
|
||||
func TestTaskIdempotencyAndExecutionLease(t *testing.T) {
|
||||
db := billingV2IntegrationStore(t)
|
||||
ctx := context.Background()
|
||||
tenantID, gatewayUserID := seedWalletReservationUser(t, ctx, db)
|
||||
user := &auth.User{ID: "billing-v2-" + uuid.NewString(), GatewayUserID: gatewayUserID, GatewayTenantID: tenantID}
|
||||
input := CreateTaskInput{
|
||||
Kind: "images.generations", Model: "billing-v2-model", RunMode: "simulation",
|
||||
Request: map[string]any{"model": "billing-v2-model", "prompt": "lease"},
|
||||
IdempotencyKeyHash: "key-hash-" + uuid.NewString(), IdempotencyRequestHash: "request-a",
|
||||
}
|
||||
created, err := db.CreateTaskIdempotent(ctx, input, user)
|
||||
if err != nil || created.Replayed {
|
||||
t.Fatalf("create task result=%+v err=%v", created, err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_tasks WHERE id=$1::uuid`, created.Task.ID)
|
||||
})
|
||||
|
||||
replayed, err := db.CreateTaskIdempotent(ctx, input, user)
|
||||
if err != nil || !replayed.Replayed || replayed.Task.ID != created.Task.ID {
|
||||
t.Fatalf("replay result=%+v err=%v", replayed, err)
|
||||
}
|
||||
input.IdempotencyRequestHash = "request-b"
|
||||
if _, err := db.CreateTaskIdempotent(ctx, input, user); !errors.Is(err, ErrIdempotencyKeyReused) {
|
||||
t.Fatalf("different request error=%v", err)
|
||||
}
|
||||
|
||||
firstToken := uuid.NewString()
|
||||
claimed, err := db.ClaimTaskExecution(ctx, created.Task.ID, firstToken, 5*time.Minute)
|
||||
if err != nil || claimed.ExecutionToken != firstToken {
|
||||
t.Fatalf("first claim task=%+v err=%v", claimed, err)
|
||||
}
|
||||
if err := db.RenewTaskExecutionLease(ctx, created.Task.ID, firstToken, 5*time.Minute); err != nil {
|
||||
t.Fatalf("renew first lease: %v", err)
|
||||
}
|
||||
if _, err := db.pool.Exec(ctx, `UPDATE gateway_tasks SET execution_lease_expires_at=now()-interval '1 second' WHERE id=$1::uuid`, created.Task.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
secondToken := uuid.NewString()
|
||||
if _, err := db.ClaimTaskExecution(ctx, created.Task.ID, secondToken, 5*time.Minute); err != nil {
|
||||
t.Fatalf("take over expired lease: %v", err)
|
||||
}
|
||||
if _, err := db.FinishTaskFailure(ctx, FinishTaskFailureInput{TaskID: created.Task.ID, ExecutionToken: firstToken, Code: "old_worker", Message: "old"}); !errors.Is(err, ErrTaskExecutionLeaseLost) {
|
||||
t.Fatalf("old worker terminal error=%v", err)
|
||||
}
|
||||
finished, err := db.FinishTaskFailure(ctx, FinishTaskFailureInput{TaskID: created.Task.ID, ExecutionToken: secondToken, Code: "new_worker", Message: "new"})
|
||||
if err != nil || finished.ErrorCode != "new_worker" {
|
||||
t.Fatalf("new worker terminal task=%+v err=%v", finished, err)
|
||||
}
|
||||
if err := db.RenewTaskExecutionLease(ctx, created.Task.ID, secondToken, 5*time.Minute); !errors.Is(err, ErrTaskExecutionFinished) {
|
||||
t.Fatalf("terminal task renewal error=%v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExpiredExecutionLeaseDoesNotReplayAmbiguousUpstreamSubmission(t *testing.T) {
|
||||
db := billingV2IntegrationStore(t)
|
||||
ctx := context.Background()
|
||||
tenantID, gatewayUserID := seedWalletReservationUser(t, ctx, db)
|
||||
user := &auth.User{ID: "billing-review-" + uuid.NewString(), GatewayUserID: gatewayUserID, GatewayTenantID: tenantID}
|
||||
created, err := db.CreateTask(ctx, CreateTaskInput{
|
||||
Kind: "images.generations", Model: "billing-v2-model", RunMode: "production",
|
||||
Request: map[string]any{"model": "billing-v2-model"},
|
||||
}, user)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_tasks WHERE id=$1::uuid`, created.ID)
|
||||
})
|
||||
|
||||
firstToken := uuid.NewString()
|
||||
if _, err := db.ClaimTaskExecution(ctx, created.ID, firstToken, 5*time.Minute); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
attemptID, err := db.CreateTaskAttempt(ctx, CreateTaskAttemptInput{
|
||||
TaskID: created.ID, AttemptNo: 1, Status: "running", RequestSnapshot: map[string]any{"model": created.Model},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.SetAttemptUpstreamSubmissionStatus(ctx, attemptID, "submitting"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := db.pool.Exec(ctx, `UPDATE gateway_tasks SET execution_lease_expires_at=now()-interval '1 second' WHERE id=$1::uuid`, created.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := db.ClaimTaskExecution(ctx, created.ID, uuid.NewString(), 5*time.Minute); !errors.Is(err, ErrTaskExecutionManualReview) {
|
||||
t.Fatalf("ambiguous submission takeover error=%v", err)
|
||||
}
|
||||
|
||||
review, err := db.GetTask(ctx, created.ID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if review.Status != "failed" || review.BillingStatus != "manual_review" || review.ErrorCode != "upstream_submission_unknown" {
|
||||
t.Fatalf("manual review task=%+v", review)
|
||||
}
|
||||
var outboxStatus string
|
||||
var outboxAction string
|
||||
var reviewReason string
|
||||
if err := db.pool.QueryRow(ctx, `
|
||||
SELECT status, action, COALESCE(manual_review_reason, '')
|
||||
FROM settlement_outbox
|
||||
WHERE task_id=$1::uuid AND event_type='task.billing.review'`, created.ID).Scan(&outboxStatus, &outboxAction, &reviewReason); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if outboxStatus != "manual_review" || outboxAction != "release" || reviewReason != "upstream_submission_unknown" {
|
||||
t.Fatalf("review outbox status=%s action=%s reason=%s", outboxStatus, outboxAction, reviewReason)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExpiredExecutionLeaseCanResumeAfterKnownRejectedResponse(t *testing.T) {
|
||||
db := billingV2IntegrationStore(t)
|
||||
ctx := context.Background()
|
||||
_, gatewayUserID := seedWalletReservationUser(t, ctx, db)
|
||||
user := &auth.User{ID: "billing-known-response-" + uuid.NewString(), GatewayUserID: gatewayUserID}
|
||||
created, err := db.CreateTask(ctx, CreateTaskInput{
|
||||
Kind: "images.generations", Model: "billing-v2-model", RunMode: "production",
|
||||
Request: map[string]any{"model": "billing-v2-model"},
|
||||
}, user)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_tasks WHERE id=$1::uuid`, created.ID)
|
||||
})
|
||||
|
||||
if _, err := db.ClaimTaskExecution(ctx, created.ID, uuid.NewString(), 5*time.Minute); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
attemptID, err := db.CreateTaskAttempt(ctx, CreateTaskAttemptInput{TaskID: created.ID, AttemptNo: 1, Status: "running"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.SetAttemptUpstreamSubmissionStatus(ctx, attemptID, "response_received"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.FinishTaskAttempt(ctx, FinishTaskAttemptInput{AttemptID: attemptID, Status: "failed", ErrorCode: "upstream_rejected"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := db.pool.Exec(ctx, `UPDATE gateway_tasks SET execution_lease_expires_at=now()-interval '1 second' WHERE id=$1::uuid`, created.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := db.ClaimTaskExecution(ctx, created.ID, uuid.NewString(), 5*time.Minute); err != nil {
|
||||
t.Fatalf("known rejected response should remain retryable: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFinishTaskManualReviewCreatesVisibleBillingRecord(t *testing.T) {
|
||||
db := billingV2IntegrationStore(t)
|
||||
ctx := context.Background()
|
||||
_, gatewayUserID := seedWalletReservationUser(t, ctx, db)
|
||||
user := &auth.User{ID: "billing-direct-review-" + uuid.NewString(), GatewayUserID: gatewayUserID}
|
||||
created, err := db.CreateTask(ctx, CreateTaskInput{
|
||||
Kind: "images.generations", Model: "billing-v2-model", RunMode: "production",
|
||||
Request: map[string]any{"model": "billing-v2-model"},
|
||||
}, user)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_tasks WHERE id=$1::uuid`, created.ID)
|
||||
})
|
||||
token := uuid.NewString()
|
||||
if _, err := db.ClaimTaskExecution(ctx, created.ID, token, 5*time.Minute); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
attemptID, err := db.CreateTaskAttempt(ctx, CreateTaskAttemptInput{TaskID: created.ID, AttemptNo: 1, Status: "running"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.SetAttemptUpstreamSubmissionStatus(ctx, attemptID, "submitting"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := db.FinishTaskManualReview(ctx, FinishTaskManualReviewInput{
|
||||
TaskID: created.ID, ExecutionToken: token, AttemptID: attemptID, TaskStatus: "failed",
|
||||
Code: "upstream_submission_unknown", Message: "upstream submission result is unknown",
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var visible bool
|
||||
if err := db.pool.QueryRow(ctx, `
|
||||
SELECT EXISTS (
|
||||
SELECT 1 FROM settlement_outbox
|
||||
WHERE task_id=$1::uuid AND status='manual_review'
|
||||
AND action='release' AND manual_review_reason='upstream_submission_unknown'
|
||||
)`, created.ID).Scan(&visible); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !visible {
|
||||
t.Fatal("manual review billing record is not visible in settlement outbox")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBillingSettlementStaleTakeoverDebitsExactlyOnce(t *testing.T) {
|
||||
db := billingV2IntegrationStore(t)
|
||||
ctx := context.Background()
|
||||
tenantID, gatewayUserID := seedWalletReservationUser(t, ctx, db)
|
||||
user := &auth.User{ID: "billing-settle-" + uuid.NewString(), GatewayUserID: gatewayUserID, GatewayTenantID: tenantID}
|
||||
if _, err := db.SetUserWalletBalance(ctx, WalletBalanceAdjustmentInput{GatewayUserID: gatewayUserID, Currency: "resource", Balance: 10, Reason: "billing v2 test"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
created, err := db.CreateTask(ctx, CreateTaskInput{Kind: "images.generations", Model: "billing-v2-model", RunMode: "production", Request: map[string]any{"model": "billing-v2-model"}}, user)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
cleanupCtx := context.Background()
|
||||
_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_wallet_transactions WHERE gateway_user_id=$1::uuid`, gatewayUserID)
|
||||
_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_tasks WHERE id=$1::uuid`, created.ID)
|
||||
_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_wallet_accounts WHERE gateway_user_id=$1::uuid`, gatewayUserID)
|
||||
})
|
||||
token := uuid.NewString()
|
||||
claimed, err := db.ClaimTaskExecution(ctx, created.ID, token, 5*time.Minute)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
amount := "1.000000001"
|
||||
reservations, err := db.ReserveTaskBilling(ctx, claimed, user, []any{map[string]any{"currency": "resource", "amount": amount}}, map[string]any{
|
||||
"pricingVersion": "effective-pricing-v2", "reservationAmount": amount, "currency": "resource", "requestFingerprint": "billing-v2-test",
|
||||
})
|
||||
if err != nil || len(reservations) != 1 {
|
||||
t.Fatalf("reserve=%+v err=%v", reservations, err)
|
||||
}
|
||||
if _, err := db.FinishTaskSuccess(ctx, FinishTaskSuccessInput{
|
||||
TaskID: created.ID, ExecutionToken: token, Result: map[string]any{"ok": true},
|
||||
FinalChargeAmountText: amount, BillingCurrency: "resource",
|
||||
PricingSnapshot: map[string]any{"pricingVersion": "effective-pricing-v2"},
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
firstClaims, err := db.ClaimBillingSettlements(ctx, "worker-one", BillingSettlementBatchSize, BillingSettlementLockTimeout)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
first := settlementForTask(t, firstClaims, created.ID)
|
||||
if _, err := db.pool.Exec(ctx, `UPDATE settlement_outbox SET locked_at=now()-interval '3 minutes' WHERE id=$1::uuid`, first.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
secondClaims, err := db.ClaimBillingSettlements(ctx, "worker-two", BillingSettlementBatchSize, BillingSettlementLockTimeout)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
second := settlementForTask(t, secondClaims, created.ID)
|
||||
if err := db.ProcessBillingSettlement(ctx, first); !errors.Is(err, pgx.ErrNoRows) {
|
||||
t.Fatalf("stale settlement lock error=%v", err)
|
||||
}
|
||||
if err := db.ProcessBillingSettlement(ctx, second); err != nil {
|
||||
t.Fatalf("process takeover: %v", err)
|
||||
}
|
||||
|
||||
var walletExact bool
|
||||
if err := db.pool.QueryRow(ctx, `
|
||||
SELECT balance = 8.999999999::numeric
|
||||
AND frozen_balance = 0::numeric
|
||||
AND total_spent = 1.000000001::numeric
|
||||
FROM gateway_wallet_accounts
|
||||
WHERE gateway_user_id=$1::uuid AND currency='resource'`, gatewayUserID).Scan(&walletExact); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !walletExact {
|
||||
t.Fatal("wallet amounts did not preserve nine-decimal settlement")
|
||||
}
|
||||
var billingTransactions int
|
||||
if err := db.pool.QueryRow(ctx, `SELECT count(*) FROM gateway_wallet_transactions WHERE reference_id=$1 AND transaction_type='task_billing'`, created.ID).Scan(&billingTransactions); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if billingTransactions != 1 {
|
||||
t.Fatalf("task billing transactions=%d", billingTransactions)
|
||||
}
|
||||
if _, err := db.pool.Exec(ctx, `
|
||||
DELETE FROM gateway_wallet_accounts
|
||||
WHERE gateway_user_id=$1::uuid AND currency='resource'`, gatewayUserID); err == nil {
|
||||
t.Fatal("wallet account with audit transactions must not be deletable")
|
||||
}
|
||||
if err := db.pool.QueryRow(ctx, `
|
||||
SELECT count(*) FROM gateway_wallet_transactions
|
||||
WHERE gateway_user_id=$1::uuid AND reference_id=$2`, gatewayUserID, created.ID).Scan(&billingTransactions); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if billingTransactions == 0 {
|
||||
t.Fatal("wallet audit transactions were lost after rejected account deletion")
|
||||
}
|
||||
settled, err := db.GetTask(ctx, created.ID)
|
||||
if err != nil || settled.BillingStatus != "settled" {
|
||||
t.Fatalf("settled task=%+v err=%v", settled, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReleaseTaskBillingReservationsPreservesNineDecimalPlaces(t *testing.T) {
|
||||
db := billingV2IntegrationStore(t)
|
||||
ctx := context.Background()
|
||||
tenantID, gatewayUserID := seedWalletReservationUser(t, ctx, db)
|
||||
user := &auth.User{ID: "billing-release-" + uuid.NewString(), GatewayUserID: gatewayUserID, GatewayTenantID: tenantID}
|
||||
if _, err := db.SetUserWalletBalance(ctx, WalletBalanceAdjustmentInput{
|
||||
GatewayUserID: gatewayUserID, Currency: "resource", Balance: 1, Reason: "billing v2 release test",
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
created, err := db.CreateTask(ctx, CreateTaskInput{
|
||||
Kind: "images.generations", Model: "billing-v2-model", RunMode: "production",
|
||||
Request: map[string]any{"model": "billing-v2-model"},
|
||||
}, user)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
cleanupCtx := context.Background()
|
||||
_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_wallet_transactions WHERE gateway_user_id=$1::uuid`, gatewayUserID)
|
||||
_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_tasks WHERE id=$1::uuid`, created.ID)
|
||||
_, _ = db.pool.Exec(cleanupCtx, `DELETE FROM gateway_wallet_accounts WHERE gateway_user_id=$1::uuid`, gatewayUserID)
|
||||
})
|
||||
amount := "0.000000001"
|
||||
reservations, err := db.ReserveTaskBilling(ctx, created, user, nil, map[string]any{
|
||||
"pricingVersion": "effective-pricing-v2", "reservationAmount": amount, "currency": "resource",
|
||||
})
|
||||
if err != nil || len(reservations) != 1 {
|
||||
t.Fatalf("reserve=%+v err=%v", reservations, err)
|
||||
}
|
||||
if _, err := db.SetUserWalletBalance(ctx, WalletBalanceAdjustmentInput{
|
||||
GatewayUserID: gatewayUserID, Currency: "resource", BalanceText: "0", Reason: "must not cross frozen balance",
|
||||
}); !errors.Is(err, ErrBalanceBelowFrozen) {
|
||||
t.Fatalf("balance below frozen error=%v", err)
|
||||
}
|
||||
if err := db.ReleaseTaskBillingReservations(ctx, reservations, "integration_test"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := db.SetUserWalletBalance(ctx, WalletBalanceAdjustmentInput{
|
||||
GatewayUserID: gatewayUserID, Currency: "resource", BalanceText: "0.123456789", Reason: "exact adjustment",
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := db.RechargeUserWalletBalance(ctx, WalletRechargeInput{
|
||||
GatewayUserID: gatewayUserID, Currency: "resource", AmountText: "0.000000001", Reason: "exact recharge",
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var exact bool
|
||||
if err := db.pool.QueryRow(ctx, `
|
||||
SELECT account.balance = 0.123456790::numeric
|
||||
AND account.frozen_balance = 0::numeric
|
||||
AND task.reservation_amount = 0::numeric
|
||||
AND task.billing_status = 'not_started'
|
||||
AND EXISTS (
|
||||
SELECT 1 FROM gateway_wallet_transactions transaction
|
||||
WHERE transaction.reference_id = task.id::text
|
||||
AND transaction.transaction_type = 'release'
|
||||
AND transaction.amount = 0.000000001::numeric
|
||||
)
|
||||
FROM gateway_wallet_accounts account
|
||||
JOIN gateway_tasks task ON task.gateway_user_id = account.gateway_user_id
|
||||
WHERE task.id=$1::uuid AND account.currency='resource'`, created.ID).Scan(&exact); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !exact {
|
||||
t.Fatal("nine-decimal reservation was not released exactly")
|
||||
}
|
||||
}
|
||||
|
||||
func billingV2IntegrationStore(t *testing.T) *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 billing v2 PostgreSQL integration tests")
|
||||
}
|
||||
db, err := Connect(context.Background(), databaseURL)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var databaseName string
|
||||
if err := db.pool.QueryRow(context.Background(), `SELECT current_database()`).Scan(&databaseName); err != nil {
|
||||
db.Close()
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(strings.ToLower(databaseName), "test") {
|
||||
db.Close()
|
||||
t.Fatalf("refusing to use non-test database %q", databaseName)
|
||||
}
|
||||
t.Cleanup(db.Close)
|
||||
return db
|
||||
}
|
||||
|
||||
func settlementForTask(t *testing.T, items []BillingSettlement, taskID string) BillingSettlement {
|
||||
t.Helper()
|
||||
for _, item := range items {
|
||||
if item.TaskID == taskID {
|
||||
return item
|
||||
}
|
||||
}
|
||||
t.Fatalf("settlement for task %s not claimed", taskID)
|
||||
return BillingSettlement{}
|
||||
}
|
||||
@@ -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 图像编辑",
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
+348
-179
@@ -7,6 +7,7 @@ import (
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net"
|
||||
"strings"
|
||||
"time"
|
||||
"unicode"
|
||||
@@ -22,6 +23,11 @@ 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"}
|
||||
}
|
||||
@@ -47,20 +53,32 @@ func normalizeAPIKeyScopes(scopes []string) []string {
|
||||
}
|
||||
|
||||
var (
|
||||
ErrInvalidCredentials = errors.New("invalid account or password")
|
||||
ErrInvalidInvitation = errors.New("invalid or expired invitation code")
|
||||
ErrInvalidAPIKeyScopes = errors.New("api key scopes must not be empty")
|
||||
ErrAccessRuleResourceDenied = errors.New("access rule resource is not available")
|
||||
ErrInsufficientWalletBalance = errors.New("insufficient wallet balance")
|
||||
ErrLocalUserRequired = errors.New("local gateway user is required")
|
||||
ErrWalletBalanceUnchanged = errors.New("wallet balance unchanged")
|
||||
ErrProtectedDefault = errors.New("protected default resource cannot be deleted")
|
||||
ErrUserAlreadyExists = errors.New("user already exists")
|
||||
ErrWeakPassword = errors.New("password must be at least 8 characters")
|
||||
ErrInvalidCredentials = errors.New("invalid account or password")
|
||||
ErrInvalidInvitation = errors.New("invalid or expired invitation code")
|
||||
ErrInvalidAPIKeyScopes = errors.New("api key scopes must not be empty")
|
||||
ErrAccessRuleResourceDenied = errors.New("access rule resource is not available")
|
||||
ErrInsufficientWalletBalance = errors.New("insufficient wallet balance")
|
||||
ErrLocalUserRequired = errors.New("local gateway user is required")
|
||||
ErrWalletBalanceUnchanged = errors.New("wallet balance unchanged")
|
||||
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")
|
||||
ErrTaskExecutionManualReview = errors.New("task execution requires manual review")
|
||||
ErrProtectedDefault = errors.New("protected default resource cannot be deleted")
|
||||
ErrUserAlreadyExists = errors.New("user already exists")
|
||||
ErrWeakPassword = errors.New("password must be at least 8 characters")
|
||||
)
|
||||
|
||||
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
|
||||
}
|
||||
@@ -71,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()
|
||||
}
|
||||
@@ -167,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 {
|
||||
@@ -295,6 +348,7 @@ type PricingRule struct {
|
||||
ResourceType string `json:"resourceType"`
|
||||
Unit string `json:"unit"`
|
||||
BasePrice float64 `json:"basePrice"`
|
||||
IsFree bool `json:"isFree"`
|
||||
Currency string `json:"currency"`
|
||||
BaseWeight map[string]any `json:"baseWeight,omitempty"`
|
||||
DynamicWeight map[string]any `json:"dynamicWeight,omitempty"`
|
||||
@@ -304,6 +358,8 @@ type PricingRule struct {
|
||||
Priority int `json:"priority"`
|
||||
Status string `json:"status"`
|
||||
Metadata map[string]any `json:"metadata,omitempty"`
|
||||
EffectiveFrom string `json:"effectiveFrom,omitempty"`
|
||||
EffectiveTo string `json:"effectiveTo,omitempty"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
}
|
||||
@@ -411,68 +467,87 @@ type RateLimitWindow struct {
|
||||
}
|
||||
|
||||
type CreateTaskInput struct {
|
||||
Kind string `json:"kind"`
|
||||
Model string `json:"model"`
|
||||
RunMode string `json:"runMode"`
|
||||
Async bool `json:"async"`
|
||||
Request map[string]any `json:"request"`
|
||||
ConversationID string `json:"conversationId"`
|
||||
NewMessageCount int `json:"newMessageCount"`
|
||||
MessageRefs []TaskMessageRefInput `json:"messageRefs"`
|
||||
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"`
|
||||
ConversationID string `json:"conversationId"`
|
||||
NewMessageCount int `json:"newMessageCount"`
|
||||
MessageRefs []TaskMessageRefInput `json:"messageRefs"`
|
||||
IdempotencyKeyHash string `json:"-"`
|
||||
IdempotencyRequestHash string `json:"-"`
|
||||
}
|
||||
|
||||
type CreateTaskResult struct {
|
||||
Task GatewayTask
|
||||
Replayed bool
|
||||
}
|
||||
|
||||
type GatewayTask struct {
|
||||
ID string `json:"id"`
|
||||
Kind string `json:"kind"`
|
||||
RunMode string `json:"runMode"`
|
||||
UserID string `json:"userId"`
|
||||
GatewayUserID string `json:"gatewayUserId,omitempty"`
|
||||
UserSource string `json:"userSource,omitempty"`
|
||||
GatewayTenantID string `json:"gatewayTenantId,omitempty"`
|
||||
TenantID string `json:"tenantId,omitempty"`
|
||||
TenantKey string `json:"tenantKey,omitempty"`
|
||||
APIKeyID string `json:"apiKeyId,omitempty"`
|
||||
APIKeyName string `json:"apiKeyName,omitempty"`
|
||||
APIKeyPrefix string `json:"apiKeyPrefix,omitempty"`
|
||||
UserGroupID string `json:"userGroupId,omitempty"`
|
||||
UserGroupKey string `json:"userGroupKey,omitempty"`
|
||||
Model string `json:"model"`
|
||||
ModelType string `json:"modelType,omitempty"`
|
||||
RequestedModel string `json:"requestedModel,omitempty"`
|
||||
ResolvedModel string `json:"resolvedModel,omitempty"`
|
||||
RequestID string `json:"requestId,omitempty"`
|
||||
ConversationID string `json:"conversationId,omitempty"`
|
||||
NewMessageCount int `json:"newMessageCount,omitempty"`
|
||||
Request map[string]any `json:"request,omitempty"`
|
||||
AsyncMode bool `json:"asyncMode"`
|
||||
RiverJobID int64 `json:"riverJobId,omitempty"`
|
||||
Status string `json:"status"`
|
||||
Cancellable *bool `json:"cancellable,omitempty"`
|
||||
Submitted *bool `json:"submitted,omitempty"`
|
||||
Message string `json:"message,omitempty"`
|
||||
AttemptCount int `json:"attemptCount"`
|
||||
RemoteTaskID string `json:"remoteTaskId,omitempty"`
|
||||
RemoteTaskPayload map[string]any `json:"remoteTaskPayload,omitempty"`
|
||||
Result map[string]any `json:"result,omitempty"`
|
||||
Billings []any `json:"billings,omitempty"`
|
||||
Usage map[string]any `json:"usage"`
|
||||
Metrics map[string]any `json:"metrics"`
|
||||
BillingSummary map[string]any `json:"billingSummary"`
|
||||
FinalChargeAmount float64 `json:"finalChargeAmount"`
|
||||
ResponseStartedAt string `json:"responseStartedAt,omitempty"`
|
||||
ResponseFinishedAt string `json:"responseFinishedAt,omitempty"`
|
||||
ResponseDurationMS int64 `json:"responseDurationMs"`
|
||||
FinishedAt string `json:"finishedAt,omitempty"`
|
||||
Error string `json:"error,omitempty"`
|
||||
ErrorCode string `json:"errorCode,omitempty"`
|
||||
ErrorMessage string `json:"errorMessage,omitempty"`
|
||||
Attempts []TaskAttempt `json:"attempts,omitempty"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
ID string `json:"id"`
|
||||
ExternalTaskID string `json:"externalTaskId,omitempty"`
|
||||
Kind string `json:"kind"`
|
||||
RunMode string `json:"runMode"`
|
||||
UserID string `json:"userId"`
|
||||
GatewayUserID string `json:"gatewayUserId,omitempty"`
|
||||
UserSource string `json:"userSource,omitempty"`
|
||||
GatewayTenantID string `json:"gatewayTenantId,omitempty"`
|
||||
TenantID string `json:"tenantId,omitempty"`
|
||||
TenantKey string `json:"tenantKey,omitempty"`
|
||||
APIKeyID string `json:"apiKeyId,omitempty"`
|
||||
APIKeyName string `json:"apiKeyName,omitempty"`
|
||||
APIKeyPrefix string `json:"apiKeyPrefix,omitempty"`
|
||||
UserGroupID string `json:"userGroupId,omitempty"`
|
||||
UserGroupKey string `json:"userGroupKey,omitempty"`
|
||||
Model string `json:"model"`
|
||||
ModelType string `json:"modelType,omitempty"`
|
||||
RequestedModel string `json:"requestedModel,omitempty"`
|
||||
ResolvedModel string `json:"resolvedModel,omitempty"`
|
||||
RequestID string `json:"requestId,omitempty"`
|
||||
ConversationID string `json:"conversationId,omitempty"`
|
||||
NewMessageCount int `json:"newMessageCount,omitempty"`
|
||||
Request map[string]any `json:"request,omitempty"`
|
||||
AsyncMode bool `json:"asyncMode"`
|
||||
RiverJobID int64 `json:"riverJobId,omitempty"`
|
||||
Status string `json:"status"`
|
||||
Cancellable *bool `json:"cancellable,omitempty"`
|
||||
Submitted *bool `json:"submitted,omitempty"`
|
||||
Message string `json:"message,omitempty"`
|
||||
AttemptCount int `json:"attemptCount"`
|
||||
RemoteTaskID string `json:"remoteTaskId,omitempty"`
|
||||
RemoteTaskPayload map[string]any `json:"remoteTaskPayload,omitempty"`
|
||||
Result map[string]any `json:"result,omitempty"`
|
||||
Billings []any `json:"billings,omitempty"`
|
||||
Usage map[string]any `json:"usage"`
|
||||
Metrics map[string]any `json:"metrics"`
|
||||
BillingSummary map[string]any `json:"billingSummary"`
|
||||
FinalChargeAmount float64 `json:"finalChargeAmount"`
|
||||
BillingVersion string `json:"billingVersion"`
|
||||
BillingStatus string `json:"billingStatus"`
|
||||
BillingCurrency string `json:"billingCurrency"`
|
||||
PricingSnapshot map[string]any `json:"pricingSnapshot,omitempty"`
|
||||
RequestFingerprint string `json:"requestFingerprint,omitempty"`
|
||||
ReservationAmount float64 `json:"reservationAmount"`
|
||||
ExecutionToken string `json:"-"`
|
||||
ExecutionLeaseUntil string `json:"executionLeaseExpiresAt,omitempty"`
|
||||
BillingUpdatedAt string `json:"billingUpdatedAt,omitempty"`
|
||||
BillingSettledAt string `json:"billingSettledAt,omitempty"`
|
||||
ResponseStartedAt string `json:"responseStartedAt,omitempty"`
|
||||
ResponseFinishedAt string `json:"responseFinishedAt,omitempty"`
|
||||
ResponseDurationMS int64 `json:"responseDurationMs"`
|
||||
FinishedAt string `json:"finishedAt,omitempty"`
|
||||
Error string `json:"error,omitempty"`
|
||||
ErrorCode string `json:"errorCode,omitempty"`
|
||||
ErrorMessage string `json:"errorMessage,omitempty"`
|
||||
Attempts []TaskAttempt `json:"attempts,omitempty"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
}
|
||||
|
||||
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,
|
||||
@@ -482,7 +557,11 @@ request, COALESCE(async_mode, false), COALESCE(river_job_id, 0), status, COALESC
|
||||
COALESCE(remote_task_id, ''), COALESCE(remote_task_payload, '{}'::jsonb),
|
||||
COALESCE(result, '{}'::jsonb), COALESCE(billings, '[]'::jsonb),
|
||||
COALESCE(usage, '{}'::jsonb), COALESCE(metrics, '{}'::jsonb), COALESCE(billing_summary, '{}'::jsonb),
|
||||
COALESCE(final_charge_amount, 0)::float8, COALESCE(response_started_at::text, ''),
|
||||
COALESCE(final_charge_amount, 0)::float8, billing_version, billing_status, billing_currency,
|
||||
COALESCE(pricing_snapshot, '{}'::jsonb), COALESCE(request_fingerprint, ''),
|
||||
COALESCE(reservation_amount, 0)::float8, COALESCE(execution_token::text, ''),
|
||||
COALESCE(execution_lease_expires_at::text, ''), COALESCE(billing_updated_at::text, ''),
|
||||
COALESCE(billing_settled_at::text, ''), COALESCE(response_started_at::text, ''),
|
||||
COALESCE(response_finished_at::text, ''), COALESCE(response_duration_ms, 0), COALESCE(error, ''),
|
||||
COALESCE(error_code, ''), COALESCE(error_message, ''),
|
||||
created_at, updated_at, COALESCE(finished_at::text, '')`
|
||||
@@ -502,35 +581,39 @@ type TaskEvent struct {
|
||||
}
|
||||
|
||||
type TaskAttempt struct {
|
||||
ID string `json:"id"`
|
||||
TaskID string `json:"taskId"`
|
||||
AttemptNo int `json:"attemptNo"`
|
||||
PlatformID string `json:"platformId,omitempty"`
|
||||
PlatformName string `json:"platformName,omitempty"`
|
||||
Provider string `json:"provider,omitempty"`
|
||||
PlatformModelID string `json:"platformModelId,omitempty"`
|
||||
ModelName string `json:"modelName,omitempty"`
|
||||
ProviderModelName string `json:"providerModelName,omitempty"`
|
||||
ModelAlias string `json:"modelAlias,omitempty"`
|
||||
ModelType string `json:"modelType,omitempty"`
|
||||
ClientID string `json:"clientId,omitempty"`
|
||||
QueueKey string `json:"queueKey"`
|
||||
Status string `json:"status"`
|
||||
Retryable bool `json:"retryable"`
|
||||
Simulated bool `json:"simulated"`
|
||||
RequestID string `json:"requestId,omitempty"`
|
||||
StatusCode int `json:"statusCode,omitempty"`
|
||||
Usage map[string]any `json:"usage,omitempty"`
|
||||
Metrics map[string]any `json:"metrics,omitempty"`
|
||||
RequestSnapshot map[string]any `json:"requestSnapshot,omitempty"`
|
||||
ResponseSnapshot map[string]any `json:"responseSnapshot,omitempty"`
|
||||
ResponseStartedAt string `json:"responseStartedAt,omitempty"`
|
||||
ResponseFinishedAt string `json:"responseFinishedAt,omitempty"`
|
||||
ResponseDurationMS int64 `json:"responseDurationMs"`
|
||||
ErrorCode string `json:"errorCode,omitempty"`
|
||||
ErrorMessage string `json:"errorMessage,omitempty"`
|
||||
StartedAt time.Time `json:"startedAt"`
|
||||
FinishedAt string `json:"finishedAt,omitempty"`
|
||||
ID string `json:"id"`
|
||||
TaskID string `json:"taskId"`
|
||||
AttemptNo int `json:"attemptNo"`
|
||||
PlatformID string `json:"platformId,omitempty"`
|
||||
PlatformName string `json:"platformName,omitempty"`
|
||||
Provider string `json:"provider,omitempty"`
|
||||
PlatformModelID string `json:"platformModelId,omitempty"`
|
||||
ModelName string `json:"modelName,omitempty"`
|
||||
ProviderModelName string `json:"providerModelName,omitempty"`
|
||||
ModelAlias string `json:"modelAlias,omitempty"`
|
||||
ModelType string `json:"modelType,omitempty"`
|
||||
ClientID string `json:"clientId,omitempty"`
|
||||
QueueKey string `json:"queueKey"`
|
||||
Status string `json:"status"`
|
||||
Retryable bool `json:"retryable"`
|
||||
Simulated bool `json:"simulated"`
|
||||
RequestID string `json:"requestId,omitempty"`
|
||||
StatusCode int `json:"statusCode,omitempty"`
|
||||
Usage map[string]any `json:"usage,omitempty"`
|
||||
Metrics map[string]any `json:"metrics,omitempty"`
|
||||
RequestSnapshot map[string]any `json:"requestSnapshot,omitempty"`
|
||||
ResponseSnapshot map[string]any `json:"responseSnapshot,omitempty"`
|
||||
ResponseStartedAt string `json:"responseStartedAt,omitempty"`
|
||||
ResponseFinishedAt string `json:"responseFinishedAt,omitempty"`
|
||||
ResponseDurationMS int64 `json:"responseDurationMs"`
|
||||
PricingSnapshot map[string]any `json:"pricingSnapshot,omitempty"`
|
||||
RequestFingerprint string `json:"requestFingerprint,omitempty"`
|
||||
UpstreamSubmissionStatus string `json:"upstreamSubmissionStatus"`
|
||||
UpstreamSubmissionUpdatedAt string `json:"upstreamSubmissionUpdatedAt,omitempty"`
|
||||
ErrorCode string `json:"errorCode,omitempty"`
|
||||
ErrorMessage string `json:"errorMessage,omitempty"`
|
||||
StartedAt time.Time `json:"startedAt"`
|
||||
FinishedAt string `json:"finishedAt,omitempty"`
|
||||
}
|
||||
|
||||
type TaskParamPreprocessingLog struct {
|
||||
@@ -847,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"'), ''),
|
||||
@@ -855,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 (
|
||||
@@ -882,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
|
||||
@@ -903,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,
|
||||
@@ -923,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)
|
||||
@@ -1484,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, '')
|
||||
@@ -1503,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
|
||||
}
|
||||
|
||||
@@ -1713,6 +1829,9 @@ ON CONFLICT (gateway_user_id, currency) DO NOTHING`,
|
||||
); err != nil {
|
||||
return GatewayUser{}, err
|
||||
}
|
||||
if err := ensureWalletAccountAuditGuard(ctx, tx, user.ID, "resource"); err != nil {
|
||||
return GatewayUser{}, err
|
||||
}
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
return GatewayUser{}, err
|
||||
}
|
||||
@@ -1823,6 +1942,11 @@ ORDER BY window_start DESC, scope_type ASC, scope_key ASC, metric ASC`)
|
||||
}
|
||||
|
||||
func (s *Store) CreateTask(ctx context.Context, input CreateTaskInput, user *auth.User) (GatewayTask, error) {
|
||||
result, err := s.CreateTaskIdempotent(ctx, input, user)
|
||||
return result.Task, err
|
||||
}
|
||||
|
||||
func (s *Store) CreateTaskIdempotent(ctx context.Context, input CreateTaskInput, user *auth.User) (CreateTaskResult, error) {
|
||||
requestBody, _ := json.Marshal(input.Request)
|
||||
runMode := normalizeRunMode(input.RunMode, input.Request)
|
||||
status := "queued"
|
||||
@@ -1831,41 +1955,73 @@ func (s *Store) CreateTask(ctx context.Context, input CreateTaskInput, user *aut
|
||||
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return GatewayTask{}, err
|
||||
return CreateTaskResult{}, err
|
||||
}
|
||||
defer tx.Rollback(ctx)
|
||||
|
||||
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, finished_at
|
||||
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, CASE WHEN $22 THEN now() ELSE NULL END)
|
||||
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, false,
|
||||
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) != "" {
|
||||
var existingRequestHash string
|
||||
if err := tx.QueryRow(ctx, `
|
||||
SELECT COALESCE(idempotency_request_hash, '')
|
||||
FROM gateway_tasks
|
||||
WHERE user_id = $1 AND idempotency_key_hash = $2
|
||||
FOR UPDATE`, user.ID, strings.TrimSpace(input.IdempotencyKeyHash)).Scan(&existingRequestHash); err != nil {
|
||||
return CreateTaskResult{}, err
|
||||
}
|
||||
if existingRequestHash != strings.TrimSpace(input.IdempotencyRequestHash) {
|
||||
return CreateTaskResult{}, ErrIdempotencyKeyReused
|
||||
}
|
||||
task, err = scanGatewayTask(tx.QueryRow(ctx, `
|
||||
SELECT `+gatewayTaskColumns+`
|
||||
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 GatewayTask{}, err
|
||||
return CreateTaskResult{}, err
|
||||
}
|
||||
if err := insertTaskMessageRefs(ctx, tx, task.ID, input.MessageRefs); err != nil {
|
||||
return GatewayTask{}, err
|
||||
}
|
||||
events := taskEventsForCreate(task.ID, runMode, status, nil)
|
||||
for _, event := range events {
|
||||
payload, _ := json.Marshal(event.Payload)
|
||||
if _, err := tx.Exec(ctx, `
|
||||
if !replayed {
|
||||
if err := insertTaskMessageRefs(ctx, tx, task.ID, input.MessageRefs); err != nil {
|
||||
return CreateTaskResult{}, err
|
||||
}
|
||||
events := taskEventsForCreate(task.ID, runMode, status, nil)
|
||||
for _, event := range events {
|
||||
payload, _ := json.Marshal(event.Payload)
|
||||
if _, err := tx.Exec(ctx, `
|
||||
INSERT INTO gateway_task_events (task_id, seq, event_type, status, phase, progress, message, payload, simulated)
|
||||
VALUES ($1::uuid, $2, $3::text, NULLIF($4::text, ''), NULLIF($5::text, ''), $6, NULLIF($7::text, ''), $8::jsonb, $9)`,
|
||||
task.ID, event.Seq, event.EventType, event.Status, event.Phase, event.Progress, event.Message, string(payload), event.Simulated,
|
||||
); err != nil {
|
||||
return GatewayTask{}, err
|
||||
task.ID, event.Seq, event.EventType, event.Status, event.Phase, event.Progress, event.Message, string(payload), event.Simulated,
|
||||
); err != nil {
|
||||
return CreateTaskResult{}, err
|
||||
}
|
||||
}
|
||||
}
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
return GatewayTask{}, err
|
||||
return CreateTaskResult{}, err
|
||||
}
|
||||
return task, nil
|
||||
if replayed {
|
||||
task, err = s.GetTask(ctx, task.ID)
|
||||
if err != nil {
|
||||
return CreateTaskResult{}, err
|
||||
}
|
||||
}
|
||||
return CreateTaskResult{Task: task, Replayed: replayed}, nil
|
||||
}
|
||||
|
||||
func (s *Store) GetTask(ctx context.Context, taskID string) (GatewayTask, error) {
|
||||
@@ -1897,9 +2053,11 @@ func scanGatewayTask(scanner taskScanner) (GatewayTask, error) {
|
||||
var usageBytes []byte
|
||||
var metricsBytes []byte
|
||||
var billingSummaryBytes []byte
|
||||
var pricingSnapshotBytes []byte
|
||||
var remoteTaskPayloadBytes []byte
|
||||
if err := scanner.Scan(
|
||||
&task.ID,
|
||||
&task.ExternalTaskID,
|
||||
&task.Kind,
|
||||
&task.RunMode,
|
||||
&task.UserID,
|
||||
@@ -1933,6 +2091,16 @@ func scanGatewayTask(scanner taskScanner) (GatewayTask, error) {
|
||||
&metricsBytes,
|
||||
&billingSummaryBytes,
|
||||
&task.FinalChargeAmount,
|
||||
&task.BillingVersion,
|
||||
&task.BillingStatus,
|
||||
&task.BillingCurrency,
|
||||
&pricingSnapshotBytes,
|
||||
&task.RequestFingerprint,
|
||||
&task.ReservationAmount,
|
||||
&task.ExecutionToken,
|
||||
&task.ExecutionLeaseUntil,
|
||||
&task.BillingUpdatedAt,
|
||||
&task.BillingSettledAt,
|
||||
&task.ResponseStartedAt,
|
||||
&task.ResponseFinishedAt,
|
||||
&task.ResponseDurationMS,
|
||||
@@ -1952,6 +2120,7 @@ func scanGatewayTask(scanner taskScanner) (GatewayTask, error) {
|
||||
task.Usage = decodeObject(usageBytes)
|
||||
task.Metrics = decodeObject(metricsBytes)
|
||||
task.BillingSummary = decodeObject(billingSummaryBytes)
|
||||
task.PricingSnapshot = decodeObject(pricingSnapshotBytes)
|
||||
return task, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user