Compare commits
274
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
cfa44b5f28 | ||
|
|
335b21d1f8 | ||
|
|
21f72da7a8 | ||
|
|
7786692d32 | ||
|
|
03c0873649 | ||
|
|
b13392ef50 | ||
|
|
f9b945e4aa | ||
|
|
dd10a807df | ||
|
|
9ce9053d7b | ||
|
|
ebdb96e7d7 | ||
|
|
c79c2a7b44 | ||
|
|
8ef9de0b06 | ||
|
|
4b639df30a | ||
|
|
60d16b11ba | ||
|
|
a5dd7f36f5 | ||
|
|
0b2bd1d512 | ||
|
|
fe8dcb40ca | ||
|
|
b2c9b4f6d9 | ||
|
|
484b05d005 | ||
|
|
9bcfec890b | ||
|
|
f1402309c8 | ||
|
|
0d9d773ce8 | ||
|
|
f4214dd489 | ||
|
|
44d5cf2b9d | ||
|
|
5c7d6ac9aa | ||
|
|
394cc6288f | ||
|
|
6a64d3936a | ||
|
|
fe56aa46b9 | ||
|
|
0f0998cbcf | ||
|
|
d129bcccbd | ||
|
|
3af6655e86 | ||
|
|
eccfe568ad | ||
|
|
c8d04ca731 | ||
|
|
13faf1d072 | ||
|
|
8362d6d27a | ||
|
|
b125599354 | ||
|
|
3c8d9839a4 | ||
|
|
b56769512f | ||
|
|
ad8cdd525b | ||
|
|
7376d6fab6 | ||
|
|
c9393af43a | ||
|
|
76bb730b13 | ||
|
|
b933d59783 | ||
|
|
d740e9f676 | ||
|
|
8485f14bc4 | ||
|
|
aa4f8532bb | ||
|
|
0d33400e09 | ||
|
|
75b3ed4748 | ||
|
|
2636e93353 | ||
|
|
d8b86bb88c | ||
|
|
0e292bb1d8 | ||
|
|
699d004c5b | ||
|
|
4ff85d96e0 | ||
|
|
4fa1981bb8 | ||
|
|
61a74445bc | ||
|
|
5e9a7888ae | ||
|
|
dc2605fb39 | ||
|
|
cc97e6649c | ||
|
|
c28bf74230 | ||
|
|
9a01fd4657 | ||
|
|
e20688bbe2 | ||
|
|
c55e52af11 | ||
|
|
ac8d6dcf45 | ||
|
|
a4046c2b81 | ||
|
|
9c093974a1 | ||
|
|
50a13de70b | ||
|
|
56a5176c92 | ||
|
|
c89c56ca65 | ||
|
|
92e328a575 | ||
|
|
89a0ff7e99 | ||
|
|
edcbbd6a87 | ||
|
|
215b6b607f | ||
|
|
7623449e33 | ||
|
|
8ca3e2dbc2 | ||
|
|
3514c13b0d | ||
|
|
53fc79691a | ||
|
|
8749f38eb3 | ||
|
|
4d84210344 | ||
|
|
80801b7bbb | ||
|
|
773d310214 | ||
|
|
fcc3dc3cb9 | ||
|
|
c9fbf6e6f6 | ||
|
|
ed3b970178 | ||
|
|
e77bd17754 | ||
|
|
53589fe3ef | ||
|
|
d3b36cf63d | ||
|
|
c63709d785 | ||
|
|
70b0ffb9ae | ||
|
|
132cda35d8 | ||
|
|
be6ce7f78a | ||
|
|
9fd9c267d6 | ||
|
|
3a01a5c93e | ||
|
|
a612c3cb4a | ||
|
|
b6ef0435ab | ||
|
|
993fb7d8be | ||
|
|
744e4acd1f | ||
|
|
709434a256 | ||
|
|
0b681275ed | ||
|
|
a95184b5b6 | ||
|
|
015ff8ea6c | ||
|
|
add99e5421 | ||
|
|
29a9b8c89f | ||
|
|
29533537ec | ||
|
|
f190af00be | ||
|
|
bfdabd3853 | ||
|
|
2c7ed905a9 | ||
|
|
26f3b3fb0a | ||
|
|
805b677e86 | ||
|
|
3ba6026e12 | ||
|
|
38207494d7 | ||
|
|
0bc4c00581 | ||
|
|
ed3bae4ba6 | ||
|
|
aec61eecb1 | ||
|
|
0b9634b74c | ||
|
|
e05922b0f4 | ||
|
|
93d36d0e55 | ||
|
|
59f0817938 | ||
|
|
c0296dbf06 | ||
|
|
88c971564a | ||
|
|
cae97b6f77 | ||
|
|
447e7ed701 | ||
|
|
580f84f569 | ||
|
|
387bf5842d | ||
|
|
11ec057e6a | ||
|
|
b0dac9f8ec | ||
|
|
1e12522f0c | ||
|
|
a34a508140 | ||
|
|
17433bf2da | ||
|
|
d5b5d4b4cd | ||
|
|
5dd7765ae3 | ||
|
|
14d13ad3a7 | ||
|
|
b625edcd71 | ||
|
|
d1b482c2f3 | ||
|
|
bfd1257db0 | ||
|
|
238798e47c | ||
|
|
3d9ae74b87 | ||
|
|
4fa44f4c41 | ||
|
|
4af86b22ee | ||
|
|
352e23e099 | ||
|
|
7b2ab786a2 | ||
|
|
1bb4b2ab12 | ||
|
|
61b4417695 | ||
|
|
262090db7b | ||
|
|
8ce120631f | ||
|
|
add80f781e | ||
|
|
36546bd5c4 | ||
|
|
75a56b21bc | ||
|
|
2529679b67 | ||
|
|
2606cc8223 | ||
|
|
1bd9e618ca | ||
|
|
78005aaf85 | ||
|
|
489208cef9 | ||
|
|
7a193934ed | ||
|
|
e0f841e8fb | ||
|
|
3976cfb64d | ||
|
|
66db98ec1c | ||
|
|
a2dea335e0 | ||
|
|
2bfeb3e179 | ||
|
|
0c9960d2e6 | ||
|
|
cd9c8d56ab | ||
|
|
3b67038660 | ||
|
|
0e22ccb0b8 | ||
|
|
02f9071a37 | ||
|
|
cec3d429af | ||
|
|
3053ba4925 | ||
|
|
f33d6d64e0 | ||
|
|
9e06ed6162 | ||
|
|
e1e18dbe58 | ||
|
|
38f87b6970 | ||
|
|
fa818e9ebf | ||
|
|
e8893e2b8f | ||
|
|
6bab0f0749 | ||
|
|
bc44af751e | ||
|
|
3aa78d3e5c | ||
|
|
ea58d21d03 | ||
|
|
fd3b6bf042 | ||
|
|
f3dd7cd262 | ||
|
|
98820378b7 | ||
|
|
91451b3c86 | ||
|
|
3886048e0f | ||
|
|
012823adff | ||
|
|
d810881901 | ||
|
|
4371e57212 | ||
|
|
8ce1f857a4 | ||
|
|
ba01cbef53 | ||
|
|
4eb394bfbd | ||
|
|
de1ce2274b | ||
|
|
a073c840ac | ||
|
|
39b1541fa1 | ||
|
|
9e4fc7362d | ||
|
|
2e5a90731b | ||
|
|
fd9bbbb508 | ||
|
|
135fb3d5b8 | ||
|
|
ea33be5d69 | ||
|
|
a642f43cf2 | ||
|
|
5d5284e068 | ||
|
|
9a3cc582ce | ||
|
|
0bef13fbd7 | ||
|
|
6897fb3b66 | ||
|
|
006609569d | ||
|
|
8e387ca553 | ||
|
|
39a8fd2ce0 | ||
|
|
f12f352bf0 | ||
|
|
945fc83c60 | ||
|
|
5d2796337b | ||
|
|
b7c03d348e | ||
|
|
74cf0ed5c8 | ||
|
|
021ee9ab8f | ||
|
|
206e367dbf | ||
|
|
5c679ff13f | ||
|
|
0b02e62c72 | ||
|
|
9dfd1aafa5 | ||
|
|
7c142f5960 | ||
|
|
4a464cacca | ||
|
|
f5b6ff72f2 | ||
|
|
5d8ed1f293 | ||
|
|
b51ecc5eed | ||
|
|
1c0c298267 | ||
|
|
6d7b2e88ba | ||
|
|
95ddfb2a04 | ||
|
|
44915c1343 | ||
|
|
19ce2d4272 | ||
|
|
d342462a3d | ||
|
|
4e60bf1c64 | ||
|
|
1b48a91af0 | ||
|
|
85e14f2cb8 | ||
|
|
971540a2a4 | ||
|
|
edf7e66941 | ||
|
|
2c7f8c14eb | ||
|
|
9e8722dc9f | ||
|
|
31565af07a | ||
|
|
046c16fc69 | ||
|
|
95776c84f6 | ||
|
|
039220c835 | ||
|
|
31c32690b2 | ||
|
|
de10875439 | ||
|
|
b29d539d49 | ||
|
|
152885bbb6 | ||
|
|
d7a0ec56f5 | ||
|
|
1fc80ffe23 | ||
|
|
a2ed32447a | ||
|
|
4f163ea6d7 | ||
|
|
2457de6a56 | ||
|
|
810dcfeee6 | ||
|
|
09375bfae7 | ||
|
|
20945dbbcb | ||
|
|
6c5daf29ca | ||
|
|
290b8c1854 | ||
|
|
d80b2e5ebf | ||
|
|
36d1b18e10 | ||
|
|
8df87875de | ||
|
|
db640c1211 | ||
|
|
1da712854f | ||
|
|
8b8e09cc22 | ||
|
|
fb7e08fe5c | ||
|
|
22ee5bb624 | ||
|
|
76a7702925 | ||
|
|
8b7d3e9c9a | ||
|
|
46cdb1a288 | ||
|
|
8d86c4b3b3 | ||
|
|
e3f6c0fa3e | ||
|
|
fac6ec95da | ||
|
|
c305603589 | ||
|
|
3de26d5157 | ||
|
|
fbe8d1d3ec | ||
|
|
2812cadd4f | ||
|
|
0b7c46cca9 | ||
|
|
e00851d7e6 | ||
|
|
e07a997aa9 | ||
|
|
42e8b517fd | ||
|
|
762d61c9cf | ||
|
|
dae5d16a58 | ||
|
|
3056cf8fca | ||
|
|
cbebfd7baa |
@@ -8,6 +8,7 @@
|
||||
|
||||
.env
|
||||
*.log
|
||||
.local-secrets
|
||||
|
||||
node_modules
|
||||
**/node_modules
|
||||
|
||||
@@ -29,6 +29,11 @@ IDENTITY_MODE=hybrid
|
||||
# - hold: reject new production generation before any upstream request; existing settlements continue.
|
||||
BILLING_ENGINE_MODE=observe
|
||||
|
||||
# River 执行 worker 会按平台模型和活跃用户组中更严格的 concurrent 总量每 5 秒重算。
|
||||
# hard limit 是单进程安全边界;平台模型/用户组 concurrency lease 才是业务并发真值。
|
||||
AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT=2048
|
||||
AI_GATEWAY_ASYNC_WORKER_REFRESH_INTERVAL_SECONDS=5
|
||||
|
||||
# Unified identity business settings are managed in System Settings > Unified
|
||||
# Identity. Deployment only supplies the SecretStore and infrastructure timing.
|
||||
AI_GATEWAY_PUBLIC_BASE_URL=http://localhost:8088
|
||||
@@ -68,6 +73,18 @@ TASK_PROGRESS_CALLBACK_ENABLED=true
|
||||
TASK_PROGRESS_CALLBACK_URL=http://localhost:3000/internal/platform/task-progress-callbacks
|
||||
TASK_PROGRESS_CALLBACK_TIMEOUT_MS=5000
|
||||
TASK_PROGRESS_CALLBACK_MAX_ATTEMPTS=10
|
||||
# First deploy with cleanup disabled; enable only after compatibility and
|
||||
# asynchronous-resume verification has passed.
|
||||
AI_GATEWAY_TASK_CLEANUP_ENABLED=false
|
||||
AI_GATEWAY_TASK_RETENTION_DAYS=30
|
||||
AI_GATEWAY_TASK_ANALYSIS_RETENTION_DAYS=7
|
||||
AI_GATEWAY_TASK_CLEANUP_INTERVAL_SECONDS=300
|
||||
AI_GATEWAY_TASK_CLEANUP_BATCH_SIZE=1000
|
||||
AI_GATEWAY_LOCAL_RESULT_TTL_HOURS=24
|
||||
AI_GATEWAY_LOCAL_RESULT_MIN_FREE_BYTES=10737418240
|
||||
AI_GATEWAY_LOCAL_RESULT_MAX_BYTES=268435456
|
||||
AI_GATEWAY_LOCAL_RESULT_MAX_TASK_BYTES=536870912
|
||||
AI_GATEWAY_ASYNC_QUEUE_WORKER_ENABLED=true
|
||||
|
||||
CORS_ALLOWED_ORIGIN=http://localhost:5178,http://127.0.0.1:5178
|
||||
VITE_GATEWAY_API_BASE_URL=http://localhost:8088
|
||||
|
||||
@@ -1,129 +0,0 @@
|
||||
name: ci
|
||||
|
||||
on:
|
||||
push:
|
||||
branches: [main]
|
||||
pull_request:
|
||||
branches: [main]
|
||||
|
||||
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:
|
||||
CI_REPOSITORY: ${{ github.repository }}
|
||||
CI_SERVER_URL: ${{ github.server_url }}
|
||||
CI_SHA: ${{ github.sha }}
|
||||
CI_JOB_TOKEN: ${{ github.token }}
|
||||
CI_EVENT_BEFORE: ${{ github.event.before }}
|
||||
CI_PR_BASE_SHA: ${{ github.event.pull_request.base.sha }}
|
||||
run: |
|
||||
set -eu
|
||||
test -n "$CI_JOB_TOKEN"
|
||||
authorization=$(printf 'x-access-token:%s' "$CI_JOB_TOKEN" | base64 | tr -d '\n')
|
||||
git init .
|
||||
git -c "http.extraHeader=AUTHORIZATION: basic $authorization" \
|
||||
fetch --no-tags "$CI_SERVER_URL/$CI_REPOSITORY.git" "$CI_SHA"
|
||||
for comparison_sha in "$CI_EVENT_BEFORE" "$CI_PR_BASE_SHA"; do
|
||||
case "$comparison_sha" in
|
||||
[0-9a-f][0-9a-f][0-9a-f][0-9a-f][0-9a-f][0-9a-f][0-9a-f][0-9a-f]*)
|
||||
if test "${#comparison_sha}" -eq 40 && \
|
||||
test "$comparison_sha" != 0000000000000000000000000000000000000000; then
|
||||
git -c "http.extraHeader=AUTHORIZATION: basic $authorization" \
|
||||
fetch --no-tags "$CI_SERVER_URL/$CI_REPOSITORY.git" "$comparison_sha"
|
||||
fi
|
||||
;;
|
||||
esac
|
||||
done
|
||||
unset authorization CI_JOB_TOKEN
|
||||
test ! -f .git/shallow
|
||||
git checkout --detach "$CI_SHA"
|
||||
test "$(git rev-parse HEAD)" = "$CI_SHA"
|
||||
- name: Verify pinned host toolchains
|
||||
run: |
|
||||
go version
|
||||
node --version
|
||||
pnpm --version
|
||||
docker-compose version
|
||||
shellcheck --version
|
||||
trivy --version
|
||||
govulncheck -version
|
||||
- name: Verify production migration safety
|
||||
env:
|
||||
CI_EVENT_NAME: ${{ github.event_name }}
|
||||
CI_EVENT_BEFORE: ${{ github.event.before }}
|
||||
CI_PR_BASE_SHA: ${{ github.event.pull_request.base.sha }}
|
||||
run: |
|
||||
production_base=$(cat deploy/ci/production-migration-base)
|
||||
immutable_base=$CI_EVENT_BEFORE
|
||||
if test "$CI_EVENT_NAME" = pull_request; then
|
||||
immutable_base=$CI_PR_BASE_SHA
|
||||
fi
|
||||
case "$immutable_base" in
|
||||
[0-9a-f][0-9a-f][0-9a-f][0-9a-f][0-9a-f][0-9a-f][0-9a-f][0-9a-f]*)
|
||||
test "${#immutable_base}" -eq 40
|
||||
test "$immutable_base" != 0000000000000000000000000000000000000000
|
||||
;;
|
||||
*) exit 1 ;;
|
||||
esac
|
||||
git merge-base --is-ancestor "$immutable_base" HEAD
|
||||
node ./scripts/ci-validate-migrations.mjs \
|
||||
"$production_base" "$immutable_base"
|
||||
- name: Verify Go formatting
|
||||
run: |
|
||||
unformatted=$(gofmt -l apps/api)
|
||||
test -z "$unformatted" || {
|
||||
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:
|
||||
GOFLAGS: "-p=1"
|
||||
GOMAXPROCS: "1"
|
||||
run: |
|
||||
go vet ./...
|
||||
go test ./...
|
||||
govulncheck ./...
|
||||
- run: pnpm install --frozen-lockfile
|
||||
- run: pnpm lint
|
||||
- run: pnpm test
|
||||
- run: pnpm build
|
||||
- name: Audit JavaScript dependencies
|
||||
run: pnpm audit --audit-level high
|
||||
- name: Validate deployment configuration
|
||||
run: |
|
||||
docker-compose -f docker-compose.yml config --quiet
|
||||
shellcheck scripts/ci-build-images.sh scripts/ci-validate-semver.sh \
|
||||
scripts/provision-ci-runner.sh tests/ci/ci-build-images-test.sh \
|
||||
tests/ci/migrations-test.sh tests/ci/pipeline-test.sh \
|
||||
tests/ci/semver-test.sh
|
||||
./tests/ci/ci-build-images-test.sh
|
||||
./tests/ci/migrations-test.sh
|
||||
./tests/ci/pipeline-test.sh
|
||||
./tests/ci/semver-test.sh
|
||||
- name: Scan repository
|
||||
run: |
|
||||
trivy fs --scanners vuln,secret,misconfig --severity HIGH,CRITICAL \
|
||||
--ignore-unfixed --exit-code 1 --timeout 15m --skip-dirs .git \
|
||||
--skip-dirs node_modules .
|
||||
@@ -1,114 +0,0 @@
|
||||
name: release-ci
|
||||
|
||||
on:
|
||||
push:
|
||||
tags: ['v*']
|
||||
|
||||
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:
|
||||
CI_REPOSITORY: ${{ github.repository }}
|
||||
CI_SERVER_URL: ${{ github.server_url }}
|
||||
CI_SHA: ${{ github.sha }}
|
||||
CI_JOB_TOKEN: ${{ github.token }}
|
||||
run: |
|
||||
set -eu
|
||||
test -n "$CI_JOB_TOKEN"
|
||||
authorization=$(printf 'x-access-token:%s' "$CI_JOB_TOKEN" | base64 | tr -d '\n')
|
||||
git init .
|
||||
git -c "http.extraHeader=AUTHORIZATION: basic $authorization" \
|
||||
fetch --no-tags "$CI_SERVER_URL/$CI_REPOSITORY.git" "$CI_SHA"
|
||||
unset authorization CI_JOB_TOKEN
|
||||
test ! -f .git/shallow
|
||||
git checkout --detach FETCH_HEAD
|
||||
- name: Verify pinned host toolchains
|
||||
run: |
|
||||
go version
|
||||
node --version
|
||||
pnpm --version
|
||||
docker-compose version
|
||||
shellcheck --version
|
||||
trivy --version
|
||||
govulncheck -version
|
||||
- name: Verify release tag ancestry
|
||||
env:
|
||||
CI_REPOSITORY: ${{ github.repository }}
|
||||
CI_SERVER_URL: ${{ github.server_url }}
|
||||
CI_SHA: ${{ github.sha }}
|
||||
CI_JOB_TOKEN: ${{ github.token }}
|
||||
run: |
|
||||
set -eu
|
||||
tag_name=${GITHUB_REF#refs/tags/}
|
||||
./scripts/ci-validate-semver.sh "$tag_name"
|
||||
authorization=$(printf 'x-access-token:%s' "$CI_JOB_TOKEN" | base64 | tr -d '\n')
|
||||
git -c "http.extraHeader=AUTHORIZATION: basic $authorization" \
|
||||
fetch --no-tags "$CI_SERVER_URL/$CI_REPOSITORY.git" \
|
||||
+refs/heads/main:refs/remotes/origin/main
|
||||
unset authorization CI_JOB_TOKEN
|
||||
test ! -f .git/shallow
|
||||
test "$(git rev-parse HEAD)" = "$CI_SHA"
|
||||
git merge-base --is-ancestor "$CI_SHA" refs/remotes/origin/main
|
||||
- name: Verify production migration safety
|
||||
run: |
|
||||
production_base=$(cat deploy/ci/production-migration-base)
|
||||
node ./scripts/ci-validate-migrations.mjs "$production_base"
|
||||
- name: Verify Go formatting
|
||||
run: |
|
||||
unformatted=$(gofmt -l apps/api)
|
||||
test -z "$unformatted" || {
|
||||
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:
|
||||
GOFLAGS: "-p=1"
|
||||
GOMAXPROCS: "1"
|
||||
run: |
|
||||
go vet ./...
|
||||
go test ./...
|
||||
govulncheck ./...
|
||||
- run: pnpm install --frozen-lockfile
|
||||
- run: pnpm lint
|
||||
- run: pnpm test
|
||||
- run: pnpm build
|
||||
- name: Audit JavaScript dependencies
|
||||
run: pnpm audit --audit-level high
|
||||
- name: Validate deployment configuration
|
||||
run: |
|
||||
docker-compose -f docker-compose.yml config --quiet
|
||||
shellcheck scripts/ci-build-images.sh scripts/ci-validate-semver.sh \
|
||||
scripts/provision-ci-runner.sh tests/ci/ci-build-images-test.sh \
|
||||
tests/ci/migrations-test.sh tests/ci/pipeline-test.sh \
|
||||
tests/ci/semver-test.sh
|
||||
./tests/ci/ci-build-images-test.sh
|
||||
./tests/ci/migrations-test.sh
|
||||
./tests/ci/pipeline-test.sh
|
||||
./tests/ci/semver-test.sh
|
||||
- name: Scan repository
|
||||
run: |
|
||||
trivy fs --scanners vuln,secret,misconfig --severity HIGH,CRITICAL \
|
||||
--ignore-unfixed --exit-code 1 --timeout 15m --skip-dirs .git \
|
||||
--skip-dirs node_modules .
|
||||
@@ -9,19 +9,28 @@
|
||||
3. `type` 使用小写英文,可选值为 `feat`、`fix`、`docs`、`test`、`refactor`、`perf`、`build`、`ci`、`chore`、`revert`。
|
||||
4. 提交摘要和正文必须使用中文;专有名词、协议名称、命令、路径、代码标识符和第三方原始错误可保留原文。
|
||||
5. 摘要应简洁明确,末尾不加句号,不得使用“更新代码”“修复问题”等无法说明意图的模糊描述。
|
||||
6. 非简单变更应在提交正文或 PR 描述中用中文说明原因、影响、风险和验证结果。
|
||||
6. 非简单变更应在提交正文中用中文说明原因、影响、风险和验证结果。
|
||||
|
||||
## 仓库边界
|
||||
|
||||
1. 后端位于 `apps/api`,前端位于 `apps/web`,共享 TypeScript 契约位于 `packages/contracts`。
|
||||
2. 修改 HTTP 接口、请求或响应类型后,必须执行 `pnpm openapi` 并提交匹配的 OpenAPI 产物。
|
||||
3. 数据库迁移只能新增,禁止修改已经进入生产基线的历史迁移;迁移必须通过生产迁移安全检查。
|
||||
4. 不得把 `.env`、密码、Secret、Token、授权码、私钥或生产凭据提交到 Git、日志、测试输出或验收证据中。
|
||||
3. 数据库迁移只能新增,禁止修改或删除已经进入 Git 历史的迁移;发布前必须相对当前线上 SHA 执行迁移安全检查。
|
||||
4. 不得把 `.env`、密码、Secret、Token、授权码、私钥或生产凭据提交到 Git、日志、测试输出、release manifest 或验收证据中。
|
||||
5. 当前工作区存在用户改动时必须保留,不得覆盖、清理、重置或混入当前任务提交;需要隔离时使用独立分支、克隆或 worktree。
|
||||
|
||||
## 验证要求
|
||||
|
||||
根据改动范围执行最小充分验证;准备合并或发布时执行完整门禁:
|
||||
根据改动范围执行最小充分验证。修改 Shell 脚本后必须执行 `bash -n` 和 ShellCheck;修改 Go 文件后必须确认 `gofmt -l` 没有输出。
|
||||
|
||||
本地发布的快速强制门禁由 `scripts/publish-release-images.sh` 固定执行:
|
||||
|
||||
```bash
|
||||
cd apps/api && env -u AI_GATEWAY_TEST_DATABASE_URL go test ./... -count=1
|
||||
node scripts/ci-validate-migrations.mjs <当前线上完整 Git SHA>
|
||||
```
|
||||
|
||||
完整测试、审计和镜像扫描不再由 Gitea 自动运行。需要合并大范围重构、准备安全审计或人工要求完整验证时执行:
|
||||
|
||||
```bash
|
||||
cd apps/api && go vet ./... && go test ./... && govulncheck ./...
|
||||
@@ -31,18 +40,30 @@ pnpm test
|
||||
pnpm build
|
||||
pnpm audit --audit-level high
|
||||
docker compose -f docker-compose.yml config --quiet
|
||||
./tests/ci/ci-build-images-test.sh
|
||||
./tests/ci/migrations-test.sh
|
||||
./tests/ci/pipeline-test.sh
|
||||
./tests/ci/semver-test.sh
|
||||
./tests/release/manual-release-test.sh
|
||||
```
|
||||
|
||||
修改 Shell 脚本后还必须执行 `bash -n` 和 ShellCheck。修改 Go 文件后必须确认 `gofmt -l` 没有输出。
|
||||
仓库没有自动 CI 兜底,合并或推送前必须以本地命令的实际结果为验收依据;报告中应明确区分已运行、命中缓存、未运行和无法运行的项目,不得把未执行的检查表述为通过。
|
||||
|
||||
## Git 与 CI/CD
|
||||
## Git 与人工发布
|
||||
|
||||
1. 一个提交只包含一个逻辑变更,提交前检查暂存差异并确认不含敏感信息。
|
||||
2. `main` 只能通过短生命周期分支和 PR 合并,禁止直接推送或强制推送。
|
||||
3. PR 必须通过精确的 `ci / verify (pull_request)` 状态后才能合并;合并后还要确认同一 SHA 的 `ci / verify (push)` 成功。
|
||||
4. 生产版本仅使用稳定 SemVer `vMAJOR.MINOR.PATCH` 标签,并同时要求 `release-ci / verify-tag (push)` 成功。
|
||||
5. 未验证 protected tag、发布账本、数据库备份、健康检查和回滚路径时,不得声称 CI/CD 或生产发布已经完成。
|
||||
1. 仓库不使用 Gitea Actions,不存在 Push、PR、Tag、Webhook、轮询或定时触发的自动构建和自动部署。
|
||||
2. `main` 不设置受保护分支或 required status;Agent 在获得用户明确授权后,可以合并已验证的 PR,也可以直接提交并推送 `main`。操作前必须获取最新 `origin/main`、确认工作区边界、检查提交差异和敏感信息,并执行与风险匹配的本地验证。
|
||||
3. 不得等待、伪造或手工补写已经取消的 CI status,也不得为了满足旧流程擅自恢复 Gitea Actions。Git 提交、合并、Push 和 Tag 授权只覆盖源码历史变更,永远不等于 publish 或 deploy 授权。
|
||||
4. 发布只能使用工作区干净、已经提交并属于 `origin/main` 历史的完整 SHA。镜像 Tag 使用完整 SHA,线上只能使用 Registry digest,禁止使用 `latest`。
|
||||
5. 本地镜像发布必须由用户明确要求后执行:
|
||||
|
||||
```bash
|
||||
./scripts/publish-release-images.sh --components auto
|
||||
```
|
||||
|
||||
该命令只能构建、冒烟、推送镜像和生成 `dist/releases/<SHA>.json`,不得修改生产。用户只要求“构建镜像”“推送镜像”或明确限定为 publish 时,执行到此为止。
|
||||
6. 用户明确要求“发布”“上线”或“部署到线上”时,该次授权同时覆盖 publish 和 deploy。publish 成功并核验 manifest、组件、digest 和验证结果后,应直接继续执行下列部署命令,无需再次请求确认:
|
||||
|
||||
```bash
|
||||
./scripts/deploy-production-release.sh dist/releases/<SHA>.json
|
||||
```
|
||||
|
||||
7. deploy 前后必须验证线上基线、镜像架构/revision、数据库备份条件、内部及公网健康检查和自动应用回滚。未实际完成这些验证时不得声称生产发布成功。
|
||||
8. 回滚也必须由用户明确要求,只能选择服务器已有的历史 manifest;应用回滚不自动恢复数据库。
|
||||
|
||||
+28
-6
@@ -6,8 +6,8 @@ ARG NODE_BUILD_IMAGE=node:${NODE_VERSION}-alpine
|
||||
ARG WEB_RUNTIME_IMAGE=nginx:1.27-alpine
|
||||
|
||||
FROM --platform=$BUILDPLATFORM ${GO_BUILD_IMAGE} AS api-builder
|
||||
ARG TARGETOS=linux
|
||||
ARG TARGETARCH=amd64
|
||||
ARG TARGETOS
|
||||
ARG TARGETARCH
|
||||
ARG GOPROXY=https://goproxy.cn,direct
|
||||
|
||||
ENV GOPROXY=$GOPROXY
|
||||
@@ -32,16 +32,28 @@ RUN --mount=type=cache,target=/go/pkg/mod \
|
||||
--mount=type=cache,target=/root/.cache/go-build \
|
||||
cd apps/api && \
|
||||
CGO_ENABLED=0 GOOS=$TARGETOS GOARCH=$TARGETARCH go build -trimpath -ldflags="-s -w" -o /out/easyai-ai-gateway ./cmd/gateway && \
|
||||
CGO_ENABLED=0 GOOS=$TARGETOS GOARCH=$TARGETARCH go build -trimpath -ldflags="-s -w" -o /out/easyai-ai-gateway-migrate ./cmd/migrate
|
||||
CGO_ENABLED=0 GOOS=$TARGETOS GOARCH=$TARGETARCH go build -trimpath -ldflags="-s -w" -o /out/easyai-ai-gateway-migrate ./cmd/migrate && \
|
||||
CGO_ENABLED=0 GOOS=$TARGETOS GOARCH=$TARGETARCH go build -trimpath -ldflags="-s -w" -o /out/easyai-ai-gateway-backfill-binary-results ./cmd/backfill-binary-results && \
|
||||
CGO_ENABLED=0 GOOS=$TARGETOS GOARCH=$TARGETARCH go build -trimpath -ldflags="-s -w" -o /out/easyai-ai-gateway-acceptance-emulator ./cmd/acceptance-emulator && \
|
||||
CGO_ENABLED=0 GOOS=$TARGETOS GOARCH=$TARGETARCH go build -trimpath -ldflags="-s -w" -o /out/easyai-ai-gateway-acceptance-callback-collector ./cmd/acceptance-callback-collector && \
|
||||
CGO_ENABLED=0 GOOS=$TARGETOS GOARCH=$TARGETARCH go build -trimpath -ldflags="-s -w" -o /out/easyai-ai-gateway-acceptance-load ./cmd/acceptance-load && \
|
||||
CGO_ENABLED=0 GOOS=$TARGETOS GOARCH=$TARGETARCH go build -trimpath -ldflags="-s -w" -o /out/easyai-ai-gateway-acceptance-snapshot ./cmd/acceptance-snapshot && \
|
||||
CGO_ENABLED=0 GOOS=$TARGETOS GOARCH=$TARGETARCH go build -trimpath -ldflags="-s -w" -o /out/easyai-ai-gateway-acceptance-bootstrap ./cmd/acceptance-bootstrap
|
||||
|
||||
FROM ${API_RUNTIME_IMAGE} AS api
|
||||
|
||||
RUN apk add --no-cache ca-certificates tzdata wget && \
|
||||
RUN apk add --no-cache ca-certificates tzdata wget ffmpeg && \
|
||||
adduser -D -H -u 10001 appuser
|
||||
WORKDIR /app
|
||||
|
||||
COPY --from=api-builder /out/easyai-ai-gateway /app/easyai-ai-gateway
|
||||
COPY --from=api-builder /out/easyai-ai-gateway-migrate /app/easyai-ai-gateway-migrate
|
||||
COPY --from=api-builder /out/easyai-ai-gateway-backfill-binary-results /app/easyai-ai-gateway-backfill-binary-results
|
||||
COPY --from=api-builder /out/easyai-ai-gateway-acceptance-emulator /app/easyai-ai-gateway-acceptance-emulator
|
||||
COPY --from=api-builder /out/easyai-ai-gateway-acceptance-callback-collector /app/easyai-ai-gateway-acceptance-callback-collector
|
||||
COPY --from=api-builder /out/easyai-ai-gateway-acceptance-load /app/easyai-ai-gateway-acceptance-load
|
||||
COPY --from=api-builder /out/easyai-ai-gateway-acceptance-snapshot /app/easyai-ai-gateway-acceptance-snapshot
|
||||
COPY --from=api-builder /out/easyai-ai-gateway-acceptance-bootstrap /app/easyai-ai-gateway-acceptance-bootstrap
|
||||
COPY apps/api/migrations /app/migrations
|
||||
|
||||
RUN mkdir -p /app/data/static/generated /app/data/static/uploaded && \
|
||||
@@ -57,15 +69,16 @@ ENV APP_ENV=production \
|
||||
CMD ["/app/easyai-ai-gateway"]
|
||||
|
||||
FROM --platform=$BUILDPLATFORM ${NODE_BUILD_IMAGE} AS web-builder
|
||||
ARG NPM_CONFIG_REGISTRY=https://registry.npmjs.org
|
||||
|
||||
WORKDIR /src
|
||||
RUN npm install -g pnpm@10.18.1
|
||||
RUN npm install --registry "$NPM_CONFIG_REGISTRY" -g pnpm@10.18.1
|
||||
|
||||
COPY package.json pnpm-lock.yaml pnpm-workspace.yaml nx.json ./
|
||||
COPY apps/web/package.json apps/web/package.json
|
||||
COPY packages/contracts/package.json packages/contracts/package.json
|
||||
RUN --mount=type=cache,id=pnpm-store,target=/root/.local/share/pnpm/store \
|
||||
pnpm install --frozen-lockfile
|
||||
pnpm install --frozen-lockfile --registry "$NPM_CONFIG_REGISTRY"
|
||||
|
||||
COPY packages packages
|
||||
COPY apps/web apps/web
|
||||
@@ -82,3 +95,12 @@ COPY docker/nginx.conf /etc/nginx/conf.d/default.conf
|
||||
COPY --from=web-builder /src/apps/web/dist /usr/share/nginx/html
|
||||
|
||||
EXPOSE 80
|
||||
|
||||
FROM alpine:3.22 AS acceptance-netem
|
||||
|
||||
RUN apk add --no-cache ca-certificates curl iproute2 socat
|
||||
COPY scripts/acceptance/netem-entrypoint.sh /usr/local/bin/netem-entrypoint
|
||||
RUN chmod 0555 /usr/local/bin/netem-entrypoint
|
||||
|
||||
USER 65534:65534
|
||||
ENTRYPOINT ["/usr/local/bin/netem-entrypoint"]
|
||||
|
||||
@@ -59,9 +59,9 @@ OIDC 用户通过签名、Issuer、Audience、`tid`、Scope 和应用角色校
|
||||
|
||||
后端热更新可通过 `GO_WATCH_SHUTDOWN_GRACE_MS` 和 `GO_WATCH_RESTART_DELAY_MS` 调整旧进程退出等待时间与重启间隔。
|
||||
|
||||
## Docker Compose 一键部署
|
||||
## Docker Compose 本地运行
|
||||
|
||||
仓库内提供了面向 `linux/amd64` 的 Docker Compose 构建和部署脚本,会自动构建 API/Web 镜像、启动 PostgreSQL、执行数据库迁移,并验证 API 与 Web 是否可访问:
|
||||
仓库内的 Compose 脚本只用于本地构建和运行:启动 PostgreSQL、执行数据库迁移,并验证 API 与 Web 是否可访问。它不再包含 Registry 推送或生产部署能力:
|
||||
|
||||
```bash
|
||||
scripts/deploy-compose.sh
|
||||
@@ -78,28 +78,21 @@ scripts/deploy-compose.sh
|
||||
常用覆盖项:
|
||||
|
||||
```bash
|
||||
AI_GATEWAY_IMAGE_TAG=2026.05.23-1 scripts/deploy-compose.sh
|
||||
AI_GATEWAY_IMAGE_TAG=2026.05.23-1 AI_GATEWAY_PUSH=1 scripts/deploy-compose.sh
|
||||
AI_GATEWAY_IMAGE_TAG=2026.05.23-1 scripts/deploy-compose.sh push
|
||||
AI_GATEWAY_IMAGE_TAG=local-test scripts/deploy-compose.sh
|
||||
AI_GATEWAY_WEB_PORT=8080 AI_GATEWAY_API_PORT=18088 scripts/deploy-compose.sh
|
||||
AI_GATEWAY_GO_PROXY='https://proxy.golang.org,direct' scripts/deploy-compose.sh
|
||||
AI_GATEWAY_NPM_REGISTRY='https://registry.npmmirror.com' scripts/deploy-compose.sh
|
||||
AI_GATEWAY_SKIP_BUILD=1 scripts/deploy-compose.sh
|
||||
scripts/deploy-compose.sh down
|
||||
scripts/deploy-compose.sh clean
|
||||
```
|
||||
|
||||
默认镜像地址为:
|
||||
默认本地镜像地址为:
|
||||
|
||||
- API: `registry.cn-shanghai.aliyuncs.com/easyaigc/ai-gateway:latest`
|
||||
- Web: `registry.cn-shanghai.aliyuncs.com/easyaigc/ai-gateway-web:latest`
|
||||
- API: `registry.cn-shanghai.aliyuncs.com/easyaigc/ai-gateway:local`
|
||||
- Web: `registry.cn-shanghai.aliyuncs.com/easyaigc/ai-gateway-web:local`
|
||||
|
||||
执行 `scripts/deploy-compose.sh push` 或设置 `AI_GATEWAY_PUSH=1` 时,会同时推送当前版本 tag 和 `latest`。当前版本 tag 优先使用 `AI_GATEWAY_IMAGE_TAG`;如果没有设置,则使用根 `package.json` 里的 `version`。
|
||||
|
||||
推送前需要先登录阿里云镜像仓库:
|
||||
|
||||
```bash
|
||||
docker login --username=<your-aliyun-account> registry.cn-shanghai.aliyuncs.com
|
||||
```
|
||||
生产镜像禁止使用 `latest`,只能通过后文的人工 publish 命令推送完整 Git SHA Tag。
|
||||
|
||||
Web 容器的 Nginx 配置通过 bind mount 挂载自仓库文件 [docker/nginx.conf](docker/nginx.conf),可直接修改该文件调整静态资源、规范 `/api/v1` 公开入口和旧 `/gateway-api` 兼容反向代理。修改后执行以下命令使配置生效:
|
||||
|
||||
@@ -107,9 +100,40 @@ Web 容器的 Nginx 配置通过 bind mount 挂载自仓库文件 [docker/nginx.
|
||||
docker compose -f docker-compose.yml restart web
|
||||
```
|
||||
|
||||
## 生产 CI/CD
|
||||
## 生产人工发布
|
||||
|
||||
Gitea Actions 会在隔离的 rootless DinD Runner 中对 Pull Request、`main` Push 和版本 Tag 执行完整质量门禁;Tag 使用独立的 `release-ci / verify-tag (push)` context,不能复用旧的 `main` 成功状态。源码 Job 没有宿主 Docker、`sudo` 或生产部署权限。部署仓的 root-owned dispatcher 只在相同 SHA 的 `main` 与 Tag context 都成功后,用固定命令构建并扫描镜像,再以 Registry digest 发布 `ai.51easyai.com`。安装 Runner、Fork PR 审批、发布验证和回滚步骤见 [生产 CI/CD 运行手册](docs/runbooks/production-ci-cd.md),信任边界见 [ADR-001](docs/decisions/001-production-cicd.md)。
|
||||
本仓库没有 Gitea Actions、Webhook、Tag、`main` Push、轮询或定时发布。`main` 不使用受保护分支限制;commit、push 和 Tag 都不会构建镜像或更新生产。
|
||||
|
||||
生产发布由 Agent 在用户明确指令下按两个阶段连续执行。第一步在本机构建 `linux/amd64` 镜像、运行临时 PostgreSQL + simulation API 冒烟、推送完整 SHA Tag,并生成带内容完整性校验的 digest-pinned manifest;该命令不会修改生产:
|
||||
|
||||
```bash
|
||||
docker login --username=<your-aliyun-account> registry.cn-shanghai.aliyuncs.com
|
||||
./scripts/publish-release-images.sh --components auto
|
||||
```
|
||||
|
||||
用户要求“发布”“上线”或“部署到线上”时,Agent 核验并报告 `dist/releases/<SHA>.json`、组件和 digest 后直接执行第二步,无需再次确认;用户仅要求构建或推送镜像时才在第一步后停止:
|
||||
|
||||
```bash
|
||||
./scripts/deploy-production-release.sh dist/releases/<SHA>.json
|
||||
```
|
||||
|
||||
查看生产版本和显式回滚:
|
||||
|
||||
```bash
|
||||
./scripts/deploy-production-release.sh --status
|
||||
./scripts/deploy-production-release.sh --rollback <历史完整 Git SHA>
|
||||
```
|
||||
|
||||
完整安装、验证、失败处理和停用旧自动化步骤见[人工生产发布运行手册](docs/runbooks/production-ci-cd.md);
|
||||
三节点 K3s、CloudNativePG、双 NGINX、备份恢复和跨节点文件验收见
|
||||
[K3s 高可用运行手册](docs/operations/k3s-ha-runbook.md);正式流量接入前的 Gemini
|
||||
带图、多参考图视频、Worker 强杀和容量阶梯验收见
|
||||
[生产同构验收手册](docs/operations/production-acceptance.md)。决策背景见
|
||||
[ADR-003](docs/decisions/003-manual-agent-release.md)。
|
||||
|
||||
生产验收前先运行[本地三节点同构验收](docs/operations/local-isomorphic-acceptance.md);
|
||||
本地原生高并发、精确 amd64 制品冒烟、线上模拟、真实金丝雀和人工开闸使用同一
|
||||
`acceptance-report/v1` 身份链路。
|
||||
|
||||
Compose 默认使用独立容器数据库 `postgres:18-alpine`,数据卷会保留在 `postgres_data` 和 `api_data`。为避免本地开发 `.env` 中的 `localhost` 数据库地址污染容器部署,compose 使用 `AI_GATEWAY_COMPOSE_*` 变量作为容器部署专用覆盖,例如:
|
||||
|
||||
@@ -144,6 +168,8 @@ AI_GATEWAY_DATABASE_URL=postgresql://easyai:easyai2025@localhost:5432/easyai_ai_
|
||||
|
||||
如果现有 `easyai-pgvector` 没有把 `5432` 映射到宿主机,就需要补端口映射,或者把 AI Gateway 后端容器化后接入同一个 `easyai` Docker network。
|
||||
|
||||
异步队列 worker 不使用固定业务并发。服务分别汇总启用平台模型和活跃用户组的有效 `concurrent` 策略,采用两者中更严格的集群容量,默认每 5 秒在线调整 River 执行容量。策略解析兼容历史 `platformLimits/modelLimits.max_concurrent_requests`,运行时统一转换为 `rules`。`AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT` 默认 `2048`,限制策略推导出的集群目标;`AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT` 默认 `32`,按单 Worker 的内存安全容量限制实例分配,即使其他实例失活也不会突破。Worker 的 `AI_GATEWAY_DATABASE_MAX_CONNS` 必须高于实例执行容量,为心跳、选主、租约续期和健康检查保留连接;生产环境按每实例执行容量 `24` 配置连接池上限 `32`。`AI_GATEWAY_ASYNC_ADMISSION_MICROBATCH_SIZE` 默认 `8`,把短任务准入、容量租约和唯一 River job 合并为有界原子微批次,摊薄同步提交,同时禁止重新形成整窗大事务。异步任务首次排队时会持久化已经鉴权的候选和 admission scope 快照;启用 `AI_GATEWAY_ASYNC_ADMISSION_DISPATCHER_ENABLED=true` 的进程批量读取快照并仅刷新动态执行容量,只有显式重新选路才重新水化媒体和计算候选。生产 K3s 在两地 API 启用 dispatcher、在 Worker 禁用,使当前 PostgreSQL 主库同站点 API 负责低延迟准入,远端 Worker 只执行 River job;数据库任务锁和 scope 锁保证主库切换时另一地 API 可安全接管。未显式配置该变量时保持兼容行为,由执行 Worker 同时运行 dispatcher。`AI_GATEWAY_DATABASE_MIN_IDLE_CONNS` 控制启动预热连接数,默认 `0`,生产 K3s 配置为 `4`;`AI_GATEWAY_DATABASE_MAX_CONN_IDLE_SECONDS` 生产配置为 `300` 秒。平台模型和用户组的 PostgreSQL concurrency lease 仍是业务并发真值。可通过 `AI_GATEWAY_ASYNC_WORKER_REFRESH_INTERVAL_SECONDS` 调整刷新周期。
|
||||
|
||||
## 迁移原则
|
||||
|
||||
1. 新服务先并行运行,不直接删除 `easyai-server-main` 内现有模块。
|
||||
|
||||
@@ -0,0 +1,628 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"flag"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/acceptancesnapshot"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
var (
|
||||
fullSHAPattern = regexp.MustCompile(`^[0-9a-f]{40}$`)
|
||||
digestPattern = regexp.MustCompile(`^sha256:[0-9a-f]{64}$`)
|
||||
)
|
||||
|
||||
type options struct {
|
||||
clusterID string
|
||||
releaseSHA string
|
||||
apiImageDigest string
|
||||
workerDigest string
|
||||
emulatorBaseURL string
|
||||
callbackURL string
|
||||
output string
|
||||
identityShards int
|
||||
}
|
||||
|
||||
type runtimeFile struct {
|
||||
SchemaVersion string `json:"schemaVersion"`
|
||||
LocalClusterID string `json:"localClusterId"`
|
||||
RunID string `json:"runId"`
|
||||
RunToken string `json:"runToken"`
|
||||
APIKeys []string `json:"apiKeys"`
|
||||
Participants []runtimeParticipant `json:"participants"`
|
||||
GeminiModel string `json:"geminiModel"`
|
||||
VideoModel string `json:"videoModel"`
|
||||
EmulatorBaseURL string `json:"emulatorBaseUrl"`
|
||||
CallbackURL string `json:"callbackUrl"`
|
||||
ReleaseSHA string `json:"releaseSha"`
|
||||
APIImageDigest string `json:"apiImageDigest"`
|
||||
WorkerImageDigest string `json:"workerImageDigest"`
|
||||
SnapshotConfigHash string `json:"snapshotConfigHash"`
|
||||
SnapshotSHA256 string `json:"snapshotSha256"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
}
|
||||
|
||||
type runtimeParticipant struct {
|
||||
APIKeyID string `json:"apiKeyId"`
|
||||
UserID string `json:"userId"`
|
||||
}
|
||||
|
||||
func main() {
|
||||
opts, err := parseOptions()
|
||||
if err != nil {
|
||||
fmt.Fprintln(os.Stderr, "acceptance bootstrap:", err)
|
||||
os.Exit(64)
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
|
||||
defer cancel()
|
||||
if err := run(ctx, opts); err != nil {
|
||||
fmt.Fprintln(os.Stderr, "acceptance bootstrap:", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
}
|
||||
|
||||
func parseOptions() (options, error) {
|
||||
var opts options
|
||||
flag.StringVar(&opts.clusterID, "local-cluster-id", "", "required local cluster marker")
|
||||
flag.StringVar(&opts.releaseSHA, "release-sha", "", "full source Git SHA")
|
||||
flag.StringVar(&opts.apiImageDigest, "api-image-digest", "", "immutable API image digest")
|
||||
flag.StringVar(&opts.workerDigest, "worker-image-digest", "", "immutable Worker image digest")
|
||||
flag.StringVar(&opts.emulatorBaseURL, "emulator-base-url", "", "in-cluster protocol emulator URL")
|
||||
flag.StringVar(&opts.callbackURL, "callback-url", "", "in-cluster callback collector URL")
|
||||
flag.StringVar(&opts.output, "output", "", "private runtime output file")
|
||||
flag.IntVar(&opts.identityShards, "identity-shards", 32, "isolated acceptance identities")
|
||||
flag.Parse()
|
||||
opts.clusterID = strings.TrimSpace(opts.clusterID)
|
||||
opts.releaseSHA = strings.ToLower(strings.TrimSpace(opts.releaseSHA))
|
||||
opts.apiImageDigest = strings.ToLower(strings.TrimSpace(opts.apiImageDigest))
|
||||
opts.workerDigest = strings.ToLower(strings.TrimSpace(opts.workerDigest))
|
||||
opts.emulatorBaseURL = strings.TrimRight(strings.TrimSpace(opts.emulatorBaseURL), "/")
|
||||
opts.callbackURL = strings.TrimSpace(opts.callbackURL)
|
||||
if opts.clusterID == "" || !fullSHAPattern.MatchString(opts.releaseSHA) ||
|
||||
!digestPattern.MatchString(opts.apiImageDigest) || !digestPattern.MatchString(opts.workerDigest) {
|
||||
return options{}, errors.New("local cluster ID, full release SHA, and immutable image digests are required")
|
||||
}
|
||||
if opts.emulatorBaseURL == "" || opts.callbackURL == "" {
|
||||
return options{}, errors.New("emulator and callback URLs are required")
|
||||
}
|
||||
if opts.identityShards < 1 || opts.identityShards > 128 {
|
||||
return options{}, errors.New("identity shards must be between 1 and 128")
|
||||
}
|
||||
if strings.TrimSpace(opts.output) == "" {
|
||||
return options{}, errors.New("private runtime output path is required")
|
||||
}
|
||||
if flag.NArg() != 0 {
|
||||
return options{}, errors.New("unexpected positional arguments")
|
||||
}
|
||||
return opts, nil
|
||||
}
|
||||
|
||||
func run(ctx context.Context, opts options) error {
|
||||
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_DATABASE_URL"))
|
||||
if databaseURL == "" {
|
||||
return errors.New("AI_GATEWAY_DATABASE_URL is required")
|
||||
}
|
||||
database, err := store.ConnectWithMaxConns(ctx, databaseURL, 4)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer database.Close()
|
||||
if err := verifyLocalClusterMarker(ctx, database, opts.clusterID); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := resetPreviousLocalRun(ctx, database, opts); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := verifyLocalDatabase(ctx, database, opts.clusterID); err != nil {
|
||||
return err
|
||||
}
|
||||
snapshotConfigHash, snapshotSHA, err := importedSnapshot(ctx, database)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
groupID, tenantID, err := ensureAcceptanceIdentityDomain(ctx, database)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
participants := make([]runtimeParticipant, 0, opts.identityShards)
|
||||
apiKeys := make([]string, 0, opts.identityShards)
|
||||
for ordinal := 0; ordinal < opts.identityShards; ordinal++ {
|
||||
userID, err := ensureAcceptanceUser(ctx, database, groupID, tenantID, ordinal)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
keyID, secret, err := ensureAcceptanceAPIKey(ctx, database, userID, groupID, tenantID, ordinal)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := database.SetUserWalletBalance(ctx, store.WalletBalanceAdjustmentInput{
|
||||
GatewayUserID: userID,
|
||||
Currency: "resource",
|
||||
BalanceText: "1000000000",
|
||||
Reason: "local acceptance isolated wallet",
|
||||
}); err != nil && !errors.Is(err, store.ErrWalletBalanceUnchanged) {
|
||||
return err
|
||||
}
|
||||
participants = append(participants, runtimeParticipant{APIKeyID: keyID, UserID: userID})
|
||||
apiKeys = append(apiKeys, secret)
|
||||
}
|
||||
if err := ensureAcceptanceAccessRules(ctx, database, groupID); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := ensureAcceptanceObjectStorage(ctx, database, opts.emulatorBaseURL); err != nil {
|
||||
return err
|
||||
}
|
||||
geminiModel, videoModel, err := selectedModels(ctx, database)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
runToken, err := randomToken()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
run, err := database.CreateAcceptanceRun(ctx, store.CreateAcceptanceRunInput{
|
||||
ReleaseSHA: opts.releaseSHA,
|
||||
APIImageDigest: opts.apiImageDigest,
|
||||
WorkerImageDigest: opts.workerDigest,
|
||||
APIKeyID: participants[0].APIKeyID,
|
||||
UserID: participants[0].UserID,
|
||||
Token: runToken,
|
||||
EmulatorBaseURL: opts.emulatorBaseURL,
|
||||
CallbackURL: opts.callbackURL,
|
||||
CapacityProfile: "P24",
|
||||
Config: map[string]any{
|
||||
"workloads": []any{"gemini_image_edit", "multi_reference_video"},
|
||||
"participants": participants,
|
||||
"localClusterId": opts.clusterID,
|
||||
"snapshotConfigHash": snapshotConfigHash,
|
||||
"snapshotSha256": snapshotSHA,
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := database.ActivateAcceptanceRun(ctx, run.ID); err != nil {
|
||||
return err
|
||||
}
|
||||
outputPath, err := privateOutputPath(opts.output)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
payload, err := json.MarshalIndent(runtimeFile{
|
||||
SchemaVersion: "acceptance-runtime/v1",
|
||||
LocalClusterID: opts.clusterID,
|
||||
RunID: run.ID,
|
||||
RunToken: runToken,
|
||||
APIKeys: apiKeys,
|
||||
Participants: participants,
|
||||
GeminiModel: geminiModel,
|
||||
VideoModel: videoModel,
|
||||
EmulatorBaseURL: opts.emulatorBaseURL,
|
||||
CallbackURL: opts.callbackURL,
|
||||
ReleaseSHA: opts.releaseSHA,
|
||||
APIImageDigest: opts.apiImageDigest,
|
||||
WorkerImageDigest: opts.workerDigest,
|
||||
SnapshotConfigHash: snapshotConfigHash,
|
||||
SnapshotSHA256: snapshotSHA,
|
||||
CreatedAt: time.Now().UTC(),
|
||||
}, "", " ")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
outputFile, err := os.OpenFile(outputPath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := outputFile.Write(append(payload, '\n')); err != nil {
|
||||
_ = outputFile.Close()
|
||||
return err
|
||||
}
|
||||
if err := outputFile.Close(); err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Printf(
|
||||
"acceptance_bootstrap=PASS cluster_id=%s run_id=%s identities=%d config_hash=%s runtime_file=%s\n",
|
||||
opts.clusterID,
|
||||
run.ID,
|
||||
len(participants),
|
||||
snapshotConfigHash,
|
||||
outputPath,
|
||||
)
|
||||
return nil
|
||||
}
|
||||
|
||||
func ensureAcceptanceObjectStorage(ctx context.Context, database *store.Store, emulatorBaseURL string) error {
|
||||
baseURL := strings.TrimRight(strings.TrimSpace(emulatorBaseURL), "/")
|
||||
if baseURL == "" {
|
||||
return errors.New("acceptance object storage requires the emulator URL")
|
||||
}
|
||||
channels := []struct {
|
||||
key string
|
||||
name string
|
||||
provider string
|
||||
endpoint string
|
||||
priority int
|
||||
pathStyle bool
|
||||
}{
|
||||
{
|
||||
key: "local-acceptance-oss", name: "Local Acceptance Aliyun OSS", provider: "aliyun_oss",
|
||||
endpoint: baseURL + "/storage/oss/bucket", priority: 10, pathStyle: true,
|
||||
},
|
||||
{
|
||||
key: "local-acceptance-s3", name: "Local Acceptance S3", provider: "s3",
|
||||
endpoint: baseURL + "/storage/s3", priority: 20, pathStyle: true,
|
||||
},
|
||||
}
|
||||
for _, channel := range channels {
|
||||
_, err := database.Pool().Exec(ctx, `
|
||||
INSERT INTO file_storage_channels (
|
||||
channel_key, name, provider, credentials, config, retry_policy, priority, status
|
||||
)
|
||||
VALUES (
|
||||
$1, $2, $3,
|
||||
'{"accessKeyId":"local-acceptance","accessKeySecret":"local-acceptance-secret"}'::jsonb,
|
||||
jsonb_build_object(
|
||||
'endpoint', $4::text,
|
||||
'region', 'local-acceptance-1',
|
||||
'bucket', 'bucket',
|
||||
'objectPrefix', 'acceptance/media',
|
||||
'accessScope', 'private',
|
||||
'forcePathStyle', $5::boolean,
|
||||
'scenes', jsonb_build_array('upload', 'image_result', 'request_asset'),
|
||||
'acceptanceEmulatorOnly', true
|
||||
),
|
||||
'{"enabled":true,"maxRetries":2,"backoffSeconds":[0.25,1],"strategy":"exponential"}'::jsonb,
|
||||
$6, 'enabled'
|
||||
)
|
||||
ON CONFLICT (channel_key) DO UPDATE
|
||||
SET name = EXCLUDED.name,
|
||||
provider = EXCLUDED.provider,
|
||||
credentials = EXCLUDED.credentials,
|
||||
config = EXCLUDED.config,
|
||||
retry_policy = EXCLUDED.retry_policy,
|
||||
priority = EXCLUDED.priority,
|
||||
status = 'enabled',
|
||||
deleted_at = NULL,
|
||||
last_error = NULL,
|
||||
updated_at = now()`, channel.key, channel.name, channel.provider, channel.endpoint, channel.pathStyle, channel.priority)
|
||||
if err != nil {
|
||||
return fmt.Errorf("configure %s object storage channel: %w", channel.provider, err)
|
||||
}
|
||||
}
|
||||
_, err := database.Pool().Exec(ctx, `
|
||||
INSERT INTO system_settings (setting_key, value)
|
||||
VALUES ('file_storage', '{"resultUploadPolicy":"default"}'::jsonb)
|
||||
ON CONFLICT (setting_key) DO UPDATE
|
||||
SET value = jsonb_set(COALESCE(system_settings.value, '{}'::jsonb), '{resultUploadPolicy}', '"default"'::jsonb, true),
|
||||
updated_at = now()`)
|
||||
return err
|
||||
}
|
||||
|
||||
func resetPreviousLocalRun(ctx context.Context, database *store.Store, opts options) error {
|
||||
mode, err := database.GetGatewayTrafficMode(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if mode.Mode == "live" {
|
||||
return nil
|
||||
}
|
||||
if mode.Mode != "validation" || strings.TrimSpace(mode.RunID) == "" {
|
||||
return fmt.Errorf("refusing bootstrap while traffic mode is %s", mode.Mode)
|
||||
}
|
||||
run, err := database.GetAcceptanceRun(ctx, mode.RunID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if run.Config["localClusterId"] != opts.clusterID {
|
||||
return errors.New("refusing to replace an acceptance run owned by another cluster")
|
||||
}
|
||||
var outstanding int
|
||||
if err := database.Pool().QueryRow(ctx, `
|
||||
SELECT count(*)
|
||||
FROM gateway_tasks
|
||||
WHERE acceptance_run_id = $1::uuid
|
||||
AND status IN ('queued','running')`, run.ID).Scan(&outstanding); err != nil {
|
||||
return err
|
||||
}
|
||||
if outstanding != 0 {
|
||||
return fmt.Errorf("refusing to replace local acceptance run with %d outstanding tasks", outstanding)
|
||||
}
|
||||
_, err = database.AbortAcceptanceRun(ctx, store.PromoteAcceptanceRunInput{
|
||||
RunID: mode.RunID,
|
||||
Revision: mode.Revision,
|
||||
ReleaseSHA: mode.ReleaseSHA,
|
||||
APIImageDigest: mode.APIImageDigest,
|
||||
WorkerImageDigest: mode.WorkerImageDigest,
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
func verifyLocalDatabase(ctx context.Context, database *store.Store, clusterID string) error {
|
||||
if err := verifyLocalClusterMarker(ctx, database, clusterID); err != nil {
|
||||
return err
|
||||
}
|
||||
mode, err := database.GetGatewayTrafficMode(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if mode.Mode != "live" {
|
||||
return fmt.Errorf("refusing bootstrap while traffic mode is %s", mode.Mode)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func verifyLocalClusterMarker(ctx context.Context, database *store.Store, clusterID string) error {
|
||||
var marker string
|
||||
err := database.Pool().QueryRow(ctx, `
|
||||
SELECT COALESCE(value->>'clusterId', '')
|
||||
FROM system_settings
|
||||
WHERE setting_key = $1`, acceptancesnapshot.LocalClusterSettingKey).Scan(&marker)
|
||||
if err != nil {
|
||||
return fmt.Errorf("read local acceptance cluster marker: %w", err)
|
||||
}
|
||||
if marker != clusterID {
|
||||
return errors.New("refusing bootstrap: local cluster marker mismatch")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func importedSnapshot(ctx context.Context, database *store.Store) (string, string, error) {
|
||||
var configHash, snapshotSHA string
|
||||
err := database.Pool().QueryRow(ctx, `
|
||||
SELECT COALESCE(value->>'configHash', ''),
|
||||
COALESCE(value->>'snapshotSha256', '')
|
||||
FROM system_settings
|
||||
WHERE setting_key = 'acceptance_snapshot'`).Scan(&configHash, &snapshotSHA)
|
||||
if err != nil {
|
||||
return "", "", fmt.Errorf("read imported acceptance snapshot: %w", err)
|
||||
}
|
||||
if len(configHash) != 64 || len(snapshotSHA) != 64 {
|
||||
return "", "", errors.New("imported acceptance snapshot hashes are invalid")
|
||||
}
|
||||
return configHash, snapshotSHA, nil
|
||||
}
|
||||
|
||||
func ensureAcceptanceIdentityDomain(ctx context.Context, database *store.Store) (string, string, error) {
|
||||
var groupID string
|
||||
err := database.Pool().QueryRow(ctx, `
|
||||
INSERT INTO gateway_user_groups (
|
||||
group_key, name, description, source, priority,
|
||||
recharge_discount_policy, billing_discount_policy,
|
||||
rate_limit_policy, quota_policy, metadata, status
|
||||
)
|
||||
VALUES (
|
||||
'local-acceptance', 'Local Acceptance', 'Isolated local homologous acceptance identities',
|
||||
'gateway', 1, '{"discountFactor":1}'::jsonb, '{"discountFactor":1}'::jsonb,
|
||||
'{"rules":[]}'::jsonb, '{}'::jsonb,
|
||||
'{"purpose":"local_acceptance","isolated":true}'::jsonb, 'active'
|
||||
)
|
||||
ON CONFLICT (group_key) DO UPDATE
|
||||
SET name = EXCLUDED.name,
|
||||
description = EXCLUDED.description,
|
||||
priority = EXCLUDED.priority,
|
||||
rate_limit_policy = EXCLUDED.rate_limit_policy,
|
||||
quota_policy = EXCLUDED.quota_policy,
|
||||
metadata = EXCLUDED.metadata,
|
||||
status = 'active',
|
||||
updated_at = now()
|
||||
RETURNING id::text`).Scan(&groupID)
|
||||
if err != nil {
|
||||
return "", "", err
|
||||
}
|
||||
var tenantID string
|
||||
err = database.Pool().QueryRow(ctx, `
|
||||
INSERT INTO gateway_tenants (
|
||||
tenant_key, source, external_tenant_id, name, default_user_group_id,
|
||||
billing_profile, rate_limit_policy, auth_policy, metadata, status
|
||||
)
|
||||
VALUES (
|
||||
'local-acceptance', 'gateway', 'local-acceptance', 'Local Acceptance',
|
||||
$1::uuid, '{}'::jsonb, '{"rules":[]}'::jsonb, '{}'::jsonb,
|
||||
'{"purpose":"local_acceptance","isolated":true}'::jsonb, 'active'
|
||||
)
|
||||
ON CONFLICT (tenant_key) DO UPDATE
|
||||
SET default_user_group_id = EXCLUDED.default_user_group_id,
|
||||
metadata = EXCLUDED.metadata,
|
||||
status = 'active',
|
||||
updated_at = now()
|
||||
RETURNING id::text`, groupID).Scan(&tenantID)
|
||||
return groupID, tenantID, err
|
||||
}
|
||||
|
||||
func ensureAcceptanceUser(
|
||||
ctx context.Context,
|
||||
database *store.Store,
|
||||
groupID string,
|
||||
tenantID string,
|
||||
ordinal int,
|
||||
) (string, error) {
|
||||
userKey := fmt.Sprintf("local-acceptance-%03d", ordinal)
|
||||
roles := `["user"]`
|
||||
if ordinal == 0 {
|
||||
roles = `["manager"]`
|
||||
}
|
||||
var userID string
|
||||
err := database.Pool().QueryRow(ctx, `
|
||||
INSERT INTO gateway_users (
|
||||
user_key, source, external_user_id, username, display_name,
|
||||
gateway_tenant_id, tenant_id, tenant_key, default_user_group_id,
|
||||
roles, auth_profile, metadata, status
|
||||
)
|
||||
VALUES (
|
||||
$1, 'gateway', $1, $1, $2,
|
||||
$3::uuid, 'local-acceptance', 'local-acceptance', $4::uuid,
|
||||
$5::jsonb, '{}'::jsonb,
|
||||
jsonb_build_object('purpose','local_acceptance','ordinal',$6::integer,'isolated',true),
|
||||
'active'
|
||||
)
|
||||
ON CONFLICT (user_key) DO UPDATE
|
||||
SET gateway_tenant_id = EXCLUDED.gateway_tenant_id,
|
||||
tenant_id = EXCLUDED.tenant_id,
|
||||
tenant_key = EXCLUDED.tenant_key,
|
||||
default_user_group_id = EXCLUDED.default_user_group_id,
|
||||
roles = EXCLUDED.roles,
|
||||
metadata = EXCLUDED.metadata,
|
||||
status = 'active',
|
||||
deleted_at = NULL,
|
||||
updated_at = now()
|
||||
RETURNING id::text`,
|
||||
userKey,
|
||||
fmt.Sprintf("Local Acceptance %03d", ordinal),
|
||||
tenantID,
|
||||
groupID,
|
||||
roles,
|
||||
ordinal,
|
||||
).Scan(&userID)
|
||||
return userID, err
|
||||
}
|
||||
|
||||
func ensureAcceptanceAPIKey(
|
||||
ctx context.Context,
|
||||
database *store.Store,
|
||||
userID string,
|
||||
groupID string,
|
||||
tenantID string,
|
||||
ordinal int,
|
||||
) (string, string, error) {
|
||||
name := fmt.Sprintf("Local Acceptance %03d", ordinal)
|
||||
var keyID, secret string
|
||||
err := database.Pool().QueryRow(ctx, `
|
||||
SELECT id::text, COALESCE(key_secret, '')
|
||||
FROM gateway_api_keys
|
||||
WHERE gateway_user_id = $1::uuid
|
||||
AND name = $2
|
||||
AND status = 'active'
|
||||
AND deleted_at IS NULL
|
||||
AND COALESCE(key_secret, '') <> ''
|
||||
ORDER BY created_at
|
||||
LIMIT 1`, userID, name).Scan(&keyID, &secret)
|
||||
if err == nil {
|
||||
return keyID, secret, nil
|
||||
}
|
||||
if !store.IsNotFound(err) {
|
||||
return "", "", err
|
||||
}
|
||||
created, err := database.CreateAPIKey(ctx, store.CreateAPIKeyInput{
|
||||
Name: name,
|
||||
Scopes: []string{"image", "video"},
|
||||
}, &auth.User{
|
||||
ID: userID,
|
||||
Username: fmt.Sprintf("local-acceptance-%03d", ordinal),
|
||||
GatewayUserID: userID,
|
||||
GatewayTenantID: tenantID,
|
||||
TenantID: "local-acceptance",
|
||||
TenantKey: "local-acceptance",
|
||||
UserGroupID: groupID,
|
||||
Source: "gateway",
|
||||
})
|
||||
if err != nil {
|
||||
return "", "", err
|
||||
}
|
||||
if _, err := database.Pool().Exec(ctx, `
|
||||
UPDATE gateway_api_keys
|
||||
SET user_group_id = $2::uuid,
|
||||
rate_limit_policy = '{"rules":[]}'::jsonb,
|
||||
quota_policy = '{}'::jsonb,
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid`, created.APIKey.ID, groupID); err != nil {
|
||||
return "", "", err
|
||||
}
|
||||
return created.APIKey.ID, created.Secret, nil
|
||||
}
|
||||
|
||||
func ensureAcceptanceAccessRules(ctx context.Context, database *store.Store, groupID string) error {
|
||||
_, err := database.Pool().Exec(ctx, `
|
||||
WITH selected AS (
|
||||
SELECT 'platform'::text AS resource_type, platform.id AS resource_id
|
||||
FROM integration_platforms platform
|
||||
WHERE COALESCE((platform.config->>'acceptanceSnapshot')::boolean, false)
|
||||
UNION
|
||||
SELECT 'platform_model', model.id
|
||||
FROM platform_models model
|
||||
JOIN integration_platforms platform ON platform.id = model.platform_id
|
||||
WHERE COALESCE((platform.config->>'acceptanceSnapshot')::boolean, false)
|
||||
UNION
|
||||
SELECT 'base_model', model.base_model_id
|
||||
FROM platform_models model
|
||||
JOIN integration_platforms platform ON platform.id = model.platform_id
|
||||
WHERE COALESCE((platform.config->>'acceptanceSnapshot')::boolean, false)
|
||||
AND model.base_model_id IS NOT NULL
|
||||
)
|
||||
INSERT INTO gateway_access_rules (
|
||||
subject_type, subject_id, resource_type, resource_id, effect,
|
||||
priority, min_permission_level, conditions, metadata, status
|
||||
)
|
||||
SELECT 'user_group', $1::uuid, selected.resource_type, selected.resource_id,
|
||||
'allow', 1, 0, '{}'::jsonb,
|
||||
'{"purpose":"local_acceptance"}'::jsonb, 'active'
|
||||
FROM selected
|
||||
ON CONFLICT (subject_type, subject_id, resource_type, resource_id, effect)
|
||||
DO UPDATE SET status = 'active', priority = 1, updated_at = now()`, groupID)
|
||||
return err
|
||||
}
|
||||
|
||||
func selectedModels(ctx context.Context, database *store.Store) (string, string, error) {
|
||||
var gemini, video string
|
||||
err := database.Pool().QueryRow(ctx, `
|
||||
SELECT base_model.invocation_name
|
||||
FROM platform_models model
|
||||
JOIN integration_platforms platform ON platform.id = model.platform_id
|
||||
JOIN base_model_catalog base_model ON base_model.id = model.base_model_id
|
||||
WHERE platform.status = 'enabled'
|
||||
AND COALESCE((platform.config->>'acceptanceSnapshot')::boolean, false)
|
||||
AND model.enabled = true
|
||||
AND model.model_type @> '["image_edit"]'::jsonb
|
||||
ORDER BY platform.priority, model.created_at
|
||||
LIMIT 1`).Scan(&gemini)
|
||||
if err != nil {
|
||||
return "", "", err
|
||||
}
|
||||
err = database.Pool().QueryRow(ctx, `
|
||||
SELECT base_model.invocation_name
|
||||
FROM platform_models model
|
||||
JOIN integration_platforms platform ON platform.id = model.platform_id
|
||||
JOIN base_model_catalog base_model ON base_model.id = model.base_model_id
|
||||
WHERE platform.status = 'enabled'
|
||||
AND COALESCE((platform.config->>'acceptanceSnapshot')::boolean, false)
|
||||
AND model.enabled = true
|
||||
AND model.model_type @> '["omni_video"]'::jsonb
|
||||
ORDER BY platform.priority, model.created_at
|
||||
LIMIT 1`).Scan(&video)
|
||||
return gemini, video, err
|
||||
}
|
||||
|
||||
func randomToken() (string, error) {
|
||||
payload := make([]byte, 32)
|
||||
if _, err := rand.Read(payload); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return hex.EncodeToString(payload), nil
|
||||
}
|
||||
|
||||
func privateOutputPath(path string) (string, error) {
|
||||
absolute, err := filepath.Abs(strings.TrimSpace(path))
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if _, err := os.Lstat(absolute); err == nil {
|
||||
return "", errors.New("runtime output already exists")
|
||||
} else if !os.IsNotExist(err) {
|
||||
return "", err
|
||||
}
|
||||
if err := os.MkdirAll(filepath.Dir(absolute), 0o700); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return absolute, nil
|
||||
}
|
||||
@@ -0,0 +1,120 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
const maxCallbackBodyBytes = 1 << 20
|
||||
|
||||
type collector struct {
|
||||
mu sync.Mutex
|
||||
deliveries int64
|
||||
duplicates int64
|
||||
invalid int64
|
||||
seen map[string]int
|
||||
}
|
||||
|
||||
func main() {
|
||||
address := strings.TrimSpace(os.Getenv("HTTP_ADDR"))
|
||||
if address == "" {
|
||||
address = ":8091"
|
||||
}
|
||||
state := &collector{seen: map[string]int{}}
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("GET /healthz", func(w http.ResponseWriter, _ *http.Request) {
|
||||
writeJSON(w, http.StatusOK, map[string]any{"ok": true})
|
||||
})
|
||||
mux.HandleFunc("GET /report", state.report)
|
||||
mux.HandleFunc("POST /callbacks", state.callback)
|
||||
server := &http.Server{
|
||||
Addr: address,
|
||||
Handler: mux,
|
||||
ReadHeaderTimeout: 10 * time.Second,
|
||||
}
|
||||
slog.Info("acceptance callback collector started", "address", address)
|
||||
if err := server.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
|
||||
slog.Error("acceptance callback collector stopped", "error", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *collector) callback(w http.ResponseWriter, request *http.Request) {
|
||||
body, err := io.ReadAll(io.LimitReader(request.Body, maxCallbackBodyBytes+1))
|
||||
if err != nil || len(body) > maxCallbackBodyBytes {
|
||||
c.recordInvalid()
|
||||
writeJSON(w, http.StatusBadRequest, map[string]any{"error": "invalid callback body"})
|
||||
return
|
||||
}
|
||||
decoder := json.NewDecoder(strings.NewReader(string(body)))
|
||||
decoder.UseNumber()
|
||||
var payload map[string]any
|
||||
if err := decoder.Decode(&payload); err != nil {
|
||||
c.recordInvalid()
|
||||
writeJSON(w, http.StatusBadRequest, map[string]any{"error": "invalid callback JSON"})
|
||||
return
|
||||
}
|
||||
taskID := strings.TrimSpace(stringValue(payload["taskId"]))
|
||||
seq, _ := strconv.ParseInt(stringValue(payload["seq"]), 10, 64)
|
||||
idempotency := strings.TrimSpace(request.Header.Get("Idempotency-Key"))
|
||||
expected := taskID + ":" + strconv.FormatInt(seq, 10)
|
||||
if taskID == "" || seq <= 0 || idempotency != expected {
|
||||
c.recordInvalid()
|
||||
writeJSON(w, http.StatusBadRequest, map[string]any{"error": "callback identity is invalid"})
|
||||
return
|
||||
}
|
||||
c.mu.Lock()
|
||||
c.deliveries++
|
||||
c.seen[expected]++
|
||||
if c.seen[expected] > 1 {
|
||||
c.duplicates++
|
||||
}
|
||||
c.mu.Unlock()
|
||||
writeJSON(w, http.StatusOK, map[string]any{"ok": true})
|
||||
}
|
||||
|
||||
func (c *collector) report(w http.ResponseWriter, _ *http.Request) {
|
||||
c.mu.Lock()
|
||||
report := map[string]any{
|
||||
"schemaVersion": "acceptance-callback-report/v1",
|
||||
"deliveries": c.deliveries,
|
||||
"duplicates": c.duplicates,
|
||||
"invalid": c.invalid,
|
||||
"uniqueEvents": len(c.seen),
|
||||
}
|
||||
c.mu.Unlock()
|
||||
writeJSON(w, http.StatusOK, report)
|
||||
}
|
||||
|
||||
func (c *collector) recordInvalid() {
|
||||
c.mu.Lock()
|
||||
c.invalid++
|
||||
c.mu.Unlock()
|
||||
}
|
||||
|
||||
func stringValue(value any) string {
|
||||
switch typed := value.(type) {
|
||||
case string:
|
||||
return typed
|
||||
case json.Number:
|
||||
return typed.String()
|
||||
case float64:
|
||||
return strconv.FormatFloat(typed, 'f', -1, 64)
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
func writeJSON(w http.ResponseWriter, status int, value any) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(status)
|
||||
_ = json.NewEncoder(w).Encode(value)
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/acceptanceemulator"
|
||||
)
|
||||
|
||||
func main() {
|
||||
address := strings.TrimSpace(os.Getenv("HTTP_ADDR"))
|
||||
if address == "" {
|
||||
address = ":8090"
|
||||
}
|
||||
server := &http.Server{
|
||||
Addr: address,
|
||||
Handler: acceptanceemulator.New(acceptanceemulator.Config{}).Handler(),
|
||||
ReadHeaderTimeout: 10 * time.Second,
|
||||
}
|
||||
slog.Info("acceptance protocol emulator started", "address", address)
|
||||
if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
|
||||
slog.Error("acceptance protocol emulator stopped", "error", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,430 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
)
|
||||
|
||||
type roundTripFunc func(*http.Request) (*http.Response, error)
|
||||
|
||||
func (f roundTripFunc) RoundTrip(request *http.Request) (*http.Response, error) {
|
||||
return f(request)
|
||||
}
|
||||
|
||||
func TestStreamGeminiImageHashDoesNotNeedWholeResponse(t *testing.T) {
|
||||
payload := paddedPNG(4 << 20)
|
||||
response := fmt.Sprintf(`{"candidates":[{"content":{"parts":[{"inlineData":{"mimeType":"image/png","data":"%s"}}]}}]}`,
|
||||
base64.StdEncoding.EncodeToString(payload))
|
||||
size, digest, err := streamGeminiImageHash(bytes.NewBufferString(response))
|
||||
if err != nil {
|
||||
t.Fatalf("stream Gemini output: %v", err)
|
||||
}
|
||||
expected := sha256.Sum256(payload)
|
||||
if size != int64(len(payload)) || digest != hex.EncodeToString(expected[:]) {
|
||||
t.Fatalf("size=%d digest=%s", size, digest)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPaddedPNGVariantsAreExactSizeAndUnique(t *testing.T) {
|
||||
first := paddedPNGVariant(256<<10, 1)
|
||||
second := paddedPNGVariant(256<<10, 2)
|
||||
if len(first) != 256<<10 || len(second) != 256<<10 {
|
||||
t.Fatalf("variant sizes=%d/%d", len(first), len(second))
|
||||
}
|
||||
firstHash := sha256.Sum256(first)
|
||||
secondHash := sha256.Sum256(second)
|
||||
if firstHash == secondHash {
|
||||
t.Fatal("distinct variants have the same SHA-256")
|
||||
}
|
||||
}
|
||||
|
||||
func TestStreamGeminiRequestBodyPreservesMultipleInputs(t *testing.T) {
|
||||
inputs := geminiInputs(3, 2<<20, "multi-image-test", 37)
|
||||
var body struct {
|
||||
Contents []struct {
|
||||
Parts []struct {
|
||||
InlineData *struct {
|
||||
MIMEType string `json:"mimeType"`
|
||||
Data string `json:"data"`
|
||||
} `json:"inlineData"`
|
||||
} `json:"parts"`
|
||||
} `json:"contents"`
|
||||
GenerationConfig struct {
|
||||
ResponseModalities []string `json:"responseModalities"`
|
||||
} `json:"generationConfig"`
|
||||
}
|
||||
stream := streamGeminiRequestBody(inputs)
|
||||
defer stream.Close()
|
||||
if err := json.NewDecoder(stream).Decode(&body); err != nil {
|
||||
t.Fatalf("decode streamed Gemini request: %v", err)
|
||||
}
|
||||
if len(body.Contents) != 1 || len(body.Contents[0].Parts) != 4 {
|
||||
t.Fatalf("unexpected Gemini body structure: %+v", body)
|
||||
}
|
||||
for index, input := range inputs {
|
||||
inlineData := body.Contents[0].Parts[index+1].InlineData
|
||||
if inlineData == nil || inlineData.MIMEType != "image/png" {
|
||||
t.Fatalf("image %d inline data=%+v", index, inlineData)
|
||||
}
|
||||
decoded, err := base64.StdEncoding.DecodeString(inlineData.Data)
|
||||
if err != nil {
|
||||
t.Fatalf("decode streamed input %d: %v", index, err)
|
||||
}
|
||||
if !bytes.Equal(decoded, input) {
|
||||
t.Fatalf("streamed input %d differs from source", index)
|
||||
}
|
||||
}
|
||||
if len(body.GenerationConfig.ResponseModalities) != 1 || body.GenerationConfig.ResponseModalities[0] != "IMAGE" {
|
||||
t.Fatalf("response modalities=%v", body.GenerationConfig.ResponseModalities)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiRequestBodyCanBeReplayedAfterHTTP2Failure(t *testing.T) {
|
||||
inputs := geminiInputs(3, 768<<10, "replay-test", 91)
|
||||
request, err := newGeminiRequest(t.Context(), "https://gateway.example/v1beta/models/test:generateContent", inputs)
|
||||
if err != nil {
|
||||
t.Fatalf("new Gemini request: %v", err)
|
||||
}
|
||||
if request.GetBody == nil {
|
||||
t.Fatal("streaming Gemini request does not provide GetBody")
|
||||
}
|
||||
first, err := io.ReadAll(request.Body)
|
||||
if err != nil {
|
||||
t.Fatalf("read first request body: %v", err)
|
||||
}
|
||||
_ = request.Body.Close()
|
||||
replayedBody, err := request.GetBody()
|
||||
if err != nil {
|
||||
t.Fatalf("replay Gemini request body: %v", err)
|
||||
}
|
||||
replayed, err := io.ReadAll(replayedBody)
|
||||
if err != nil {
|
||||
t.Fatalf("read replayed request body: %v", err)
|
||||
}
|
||||
_ = replayedBody.Close()
|
||||
if !bytes.Equal(first, replayed) {
|
||||
t.Fatal("replayed Gemini request body differs from the original")
|
||||
}
|
||||
if request.ContentLength != int64(len(first)) {
|
||||
t.Fatalf("content length=%d, want %d", request.ContentLength, len(first))
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiMultiImageRequestUsesThreeFileDataReferences(t *testing.T) {
|
||||
request, err := newGeminiFileDataRequest(t.Context(), "https://gateway.example/v1beta/models/test:generateContent", []string{
|
||||
"https://fixtures.example/one.png",
|
||||
"https://fixtures.example/two.png",
|
||||
"https://fixtures.example/three.png",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("new multi-image request: %v", err)
|
||||
}
|
||||
var body struct {
|
||||
Contents []struct {
|
||||
Parts []struct {
|
||||
FileData *struct {
|
||||
FileURI string `json:"fileUri"`
|
||||
} `json:"fileData"`
|
||||
} `json:"parts"`
|
||||
} `json:"contents"`
|
||||
}
|
||||
if err := json.NewDecoder(request.Body).Decode(&body); err != nil {
|
||||
t.Fatalf("decode multi-image request: %v", err)
|
||||
}
|
||||
if len(body.Contents) != 1 || len(body.Contents[0].Parts) != 4 {
|
||||
t.Fatalf("unexpected multi-image body: %+v", body)
|
||||
}
|
||||
for index := 1; index < 4; index++ {
|
||||
if body.Contents[0].Parts[index].FileData == nil || body.Contents[0].Parts[index].FileData.FileURI == "" {
|
||||
t.Fatalf("missing fileData at part %d", index)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestVideoCombinationsProvide128UniqueInputs(t *testing.T) {
|
||||
images := make([]string, 16)
|
||||
for index := range images {
|
||||
images[index] = fmt.Sprintf("https://fixtures.example/image-%02d", index)
|
||||
}
|
||||
combinations := videoCombinations(images, 128)
|
||||
if len(combinations) != 128 {
|
||||
t.Fatalf("combinations=%d", len(combinations))
|
||||
}
|
||||
seen := map[string]struct{}{}
|
||||
for _, combination := range combinations {
|
||||
seen[fmt.Sprint(combination)] = struct{}{}
|
||||
if len(combination) != 3 && len(combination) != 6 && len(combination) != 9 {
|
||||
t.Fatalf("invalid combination size=%d", len(combination))
|
||||
}
|
||||
}
|
||||
if len(seen) != 128 {
|
||||
t.Fatalf("unique combinations=%d", len(seen))
|
||||
}
|
||||
if got := requestedVideoCombinationCount(combinations, "video-capacity", 144, 0, 1); got != 131 {
|
||||
t.Fatalf("requested combinations=%d, want 131", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestVideoProviderQuotaCombinationsUseDistinctNormalSizedImages(t *testing.T) {
|
||||
images := []string{
|
||||
"https://fixtures.example/image-00.png",
|
||||
"https://fixtures.example/image-01.png",
|
||||
"https://fixtures.example/image-02.png",
|
||||
"https://fixtures.example/image-03.png",
|
||||
}
|
||||
combinations := videoProviderQuotaCombinations(images, 24)
|
||||
if got := requestedVideoCombinationCount(combinations, "video-provider-quota", 24, 0, 1); got != 24 {
|
||||
t.Fatalf("quota combinations=%d, want 24", got)
|
||||
}
|
||||
for _, combination := range combinations {
|
||||
if len(combination) != 3 {
|
||||
t.Fatalf("invalid quota combination: %v", combination)
|
||||
}
|
||||
for _, imageURL := range combination {
|
||||
if !strings.Contains(imageURL, "acceptance_quota_variant=") {
|
||||
t.Fatalf("quota image lacks unique variant: %s", imageURL)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiLoadIsSplitAcrossTwoGatewayAPIs(t *testing.T) {
|
||||
output := paddedPNG(256 << 10)
|
||||
encoded := base64.StdEncoding.EncodeToString(output)
|
||||
var firstCalls atomic.Int64
|
||||
var secondCalls atomic.Int64
|
||||
var firstKeyCalls atomic.Int64
|
||||
var secondKeyCalls atomic.Int64
|
||||
newGateway := func(calls *atomic.Int64) *httptest.Server {
|
||||
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
calls.Add(1)
|
||||
switch r.Header.Get("Authorization") {
|
||||
case "Bearer key-1":
|
||||
firstKeyCalls.Add(1)
|
||||
case "Bearer key-2":
|
||||
secondKeyCalls.Add(1)
|
||||
default:
|
||||
t.Errorf("unexpected authorization header")
|
||||
}
|
||||
if r.Header.Get(runHeader) != "run-1" || r.Header.Get(tokenHeader) != "token-1" {
|
||||
t.Errorf("missing acceptance headers")
|
||||
}
|
||||
var body map[string]any
|
||||
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||||
t.Errorf("decode Gemini body: %v", err)
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"candidates": []any{map[string]any{"content": map[string]any{
|
||||
"parts": []any{map[string]any{"inlineData": map[string]any{
|
||||
"mimeType": "image/png", "data": encoded,
|
||||
}}},
|
||||
}}},
|
||||
})
|
||||
}))
|
||||
}
|
||||
first := newGateway(&firstCalls)
|
||||
defer first.Close()
|
||||
second := newGateway(&secondCalls)
|
||||
defer second.Close()
|
||||
|
||||
opts := options{
|
||||
gateways: []string{first.URL, second.URL}, apiKeys: []string{"key-1", "key-2"}, runID: "run-1",
|
||||
runToken: "token-1", geminiModel: "gemini-image-test",
|
||||
}
|
||||
result := runGemini(t.Context(), http.DefaultClient, opts, "dual-api", 8, 1, 256<<10, 256<<10, false, 0, 1)
|
||||
if result.err != nil || result.report.Completed != 8 {
|
||||
t.Fatalf("result=%+v err=%v", result.report, result.err)
|
||||
}
|
||||
if firstCalls.Load() != 4 || secondCalls.Load() != 4 {
|
||||
t.Fatalf("gateway calls=%d/%d", firstCalls.Load(), secondCalls.Load())
|
||||
}
|
||||
if firstKeyCalls.Load() != 4 || secondKeyCalls.Load() != 4 {
|
||||
t.Fatalf("API key calls=%d/%d", firstKeyCalls.Load(), secondKeyCalls.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func TestDistributedShardIndexesCoverWorkloadWithoutOverlap(t *testing.T) {
|
||||
first := options{shardIndex: 0, shardCount: 2}
|
||||
second := options{shardIndex: 1, shardCount: 2}
|
||||
if first.shardRequestCount(1001) != 501 || second.shardRequestCount(1001) != 500 {
|
||||
t.Fatalf("shard counts=%d/%d", first.shardRequestCount(1001), second.shardRequestCount(1001))
|
||||
}
|
||||
seen := map[int]bool{}
|
||||
for local := 0; local < first.shardRequestCount(1001); local++ {
|
||||
seen[first.logicalRequestIndex(local)] = true
|
||||
}
|
||||
for local := 0; local < second.shardRequestCount(1001); local++ {
|
||||
index := second.logicalRequestIndex(local)
|
||||
if seen[index] {
|
||||
t.Fatalf("duplicate logical index %d", index)
|
||||
}
|
||||
seen[index] = true
|
||||
}
|
||||
if len(seen) != 1001 {
|
||||
t.Fatalf("covered indexes=%d", len(seen))
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiInputVariantIsRunScoped(t *testing.T) {
|
||||
if geminiInputVariant("run-a", 7) == geminiInputVariant("run-b", 7) {
|
||||
t.Fatal("different Run IDs produced the same input variant")
|
||||
}
|
||||
if geminiInputVariant("run-a", 7) == geminiInputVariant("run-a", 8) {
|
||||
t.Fatal("different logical indexes produced the same input variant")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAcceptanceIdempotencyKeyIncludesExecutionAndLogicalIndex(t *testing.T) {
|
||||
request := httptest.NewRequest(http.MethodPost, "https://gateway.example/v1", nil)
|
||||
opts := options{
|
||||
apiKeys: []string{"key-1"}, runID: "run-1", runToken: "token-1", executionID: "p24-2-gemini-baseline",
|
||||
}
|
||||
opts.setHeaders(request, 17, false)
|
||||
if got := request.Header.Get("Idempotency-Key"); got != "acceptance-run-1-p24-2-gemini-baseline-17" {
|
||||
t.Fatalf("idempotency key=%q", got)
|
||||
}
|
||||
poll := httptest.NewRequest(http.MethodGet, "https://gateway.example/result", nil)
|
||||
opts.setHeaders(poll, 17, false)
|
||||
if got := poll.Header.Get("Idempotency-Key"); got != "" {
|
||||
t.Fatalf("GET request unexpectedly has idempotency key %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateVideoAssetDownloadsFinalMedia(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
|
||||
if request.Header.Get("Authorization") != "" {
|
||||
t.Fatal("acceptance credentials leaked to external media host")
|
||||
}
|
||||
w.Header().Set("Content-Type", "video/mp4")
|
||||
_, _ = w.Write([]byte{0, 0, 0, 16, 'f', 't', 'y', 'p', 'i', 's', 'o', 'm'})
|
||||
}))
|
||||
defer server.Close()
|
||||
if err := validateVideoAsset(t.Context(), server.Client(), options{}, server.URL+"/result.mp4", 0); err != nil {
|
||||
t.Fatalf("validate video asset: %v", err)
|
||||
}
|
||||
if got := findMediaURL(map[string]any{"content": map[string]any{"video_url": server.URL}}); got != server.URL {
|
||||
t.Fatalf("media URL=%q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateVideoAssetUsesIndependentTransportForExternalMedia(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.Header().Set("Content-Type", "video/mp4")
|
||||
_, _ = w.Write([]byte{0, 0, 0, 16, 'f', 't', 'y', 'p', 'i', 's', 'o', 'm'})
|
||||
}))
|
||||
defer server.Close()
|
||||
poisoned := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
|
||||
return nil, errors.New("gateway-only transport must not be used for external media")
|
||||
})}
|
||||
opts := options{gatewayTLSName: "gateway.easyai.local"}
|
||||
if err := validateVideoAsset(t.Context(), poisoned, opts, server.URL+"/result.mp4", 0); err != nil {
|
||||
t.Fatalf("validate external video with independent transport: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateVideoAssetUsesGatewayForMaterializedPath(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
|
||||
if request.Host != "gateway.easyai.local" || request.Header.Get("Authorization") != "Bearer key-1" {
|
||||
t.Fatalf("host=%q authorization=%q", request.Host, request.Header.Get("Authorization"))
|
||||
}
|
||||
if request.URL.Path != "/static/generated/result.mp4" {
|
||||
t.Fatalf("path=%q", request.URL.Path)
|
||||
}
|
||||
w.Header().Set("Content-Type", "video/mp4")
|
||||
_, _ = w.Write([]byte{0, 0, 0, 16, 'f', 't', 'y', 'p', 'i', 's', 'o', 'm'})
|
||||
}))
|
||||
defer server.Close()
|
||||
opts := options{
|
||||
gateways: []string{server.URL}, apiKeys: []string{"key-1"},
|
||||
runID: "run-1", runToken: "token-1", gatewayTLSName: "gateway.easyai.local",
|
||||
}
|
||||
if err := validateVideoAsset(t.Context(), server.Client(), opts, "/static/generated/result.mp4", 0); err != nil {
|
||||
t.Fatalf("validate materialized video: %v", err)
|
||||
}
|
||||
got := findMediaURL(map[string]any{
|
||||
"raw": map[string]any{"video_url": "http://internal.invalid/video.mp4"},
|
||||
"data": []any{map[string]any{"video_url": "/static/generated/result.mp4"}},
|
||||
})
|
||||
if got != "/static/generated/result.mp4" {
|
||||
t.Fatalf("preferred media URL=%q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAcceptanceReportErrorRedactsSecretsAndURLs(t *testing.T) {
|
||||
got := redactError(
|
||||
`token-1 failed at https://example.invalid/video.mp4?token=signed`,
|
||||
options{runToken: "token-1"},
|
||||
)
|
||||
if got != `[REDACTED] failed at [REDACTED_URL]` {
|
||||
t.Fatalf("redacted error=%q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAcceptanceReportCannotOverwriteExistingRunArtifact(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "load-report.json")
|
||||
if err := writeReportExclusive(path, []byte(`{"runId":"first"}`)); err != nil {
|
||||
t.Fatalf("write first report: %v", err)
|
||||
}
|
||||
if err := writeReportExclusive(path, []byte(`{"runId":"second"}`)); err == nil {
|
||||
t.Fatal("second report unexpectedly overwrote the first report")
|
||||
}
|
||||
payload, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
t.Fatalf("read report: %v", err)
|
||||
}
|
||||
if string(payload) != "{\"runId\":\"first\"}\n" {
|
||||
t.Fatalf("report changed after rejected overwrite: %s", payload)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAcceptanceFailurePreservesOperationAndHTTPStatus(t *testing.T) {
|
||||
err := withOperation("video_poll", &httpStatusError{Status: http.StatusUnauthorized, Body: "unauthorized"})
|
||||
var operationErr *operationError
|
||||
if !errors.As(err, &operationErr) || operationErr.Operation != "video_poll" {
|
||||
t.Fatalf("operation error=%#v", operationErr)
|
||||
}
|
||||
var statusErr *httpStatusError
|
||||
if !errors.As(err, &statusErr) || statusErr.Status != http.StatusUnauthorized {
|
||||
t.Fatalf("HTTP status error=%#v", statusErr)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAcceptanceGatewayTLSNameSetsHostHeader(t *testing.T) {
|
||||
request := httptest.NewRequest(http.MethodGet, "https://127.0.0.1/api/v1/healthz", nil)
|
||||
opts := options{
|
||||
apiKeys: []string{"key-1"}, runID: "run-1", runToken: "token-1",
|
||||
gatewayTLSName: "ai.example.com",
|
||||
}
|
||||
opts.setHeaders(request, 0, false)
|
||||
if request.Host != "ai.example.com" {
|
||||
t.Fatalf("Host=%q", request.Host)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAcceptanceRootCAsRejectsSymlink(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
target := filepath.Join(root, "ca.pem")
|
||||
if err := os.WriteFile(target, []byte("not a certificate"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
link := filepath.Join(root, "ca-link.pem")
|
||||
if err := os.Symlink(target, link); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := acceptanceRootCAs(link); err == nil {
|
||||
t.Fatal("symlink CA file was accepted")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,195 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"flag"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/acceptancesnapshot"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
func main() {
|
||||
if len(os.Args) < 2 {
|
||||
usage()
|
||||
os.Exit(64)
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
|
||||
defer cancel()
|
||||
var err error
|
||||
switch os.Args[1] {
|
||||
case "export":
|
||||
err = runExport(ctx, os.Args[2:])
|
||||
case "validate":
|
||||
err = runValidate(os.Args[2:])
|
||||
case "import":
|
||||
err = runImport(ctx, os.Args[2:])
|
||||
default:
|
||||
usage()
|
||||
os.Exit(64)
|
||||
}
|
||||
if err != nil {
|
||||
fmt.Fprintln(os.Stderr, "acceptance snapshot:", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
}
|
||||
|
||||
func usage() {
|
||||
fmt.Fprintln(os.Stderr, `Usage:
|
||||
easyai-ai-gateway-acceptance-snapshot export --release-sha <SHA> --output <file>
|
||||
easyai-ai-gateway-acceptance-snapshot validate --input <file>
|
||||
easyai-ai-gateway-acceptance-snapshot import --input <file> --local-cluster-id <id>
|
||||
|
||||
export is read-only. import refuses to write unless system_settings contains a
|
||||
matching acceptance_local_cluster_id marker.`)
|
||||
}
|
||||
|
||||
func runExport(ctx context.Context, args []string) error {
|
||||
flags := flag.NewFlagSet("export", flag.ContinueOnError)
|
||||
releaseSHA := flags.String("release-sha", "", "full source release SHA")
|
||||
output := flags.String("output", "", "secret-safe snapshot output path")
|
||||
if err := flags.Parse(args); err != nil {
|
||||
return err
|
||||
}
|
||||
if flags.NArg() != 0 {
|
||||
return errors.New("unexpected export arguments")
|
||||
}
|
||||
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_DATABASE_URL"))
|
||||
if databaseURL == "" {
|
||||
return errors.New("AI_GATEWAY_DATABASE_URL is required")
|
||||
}
|
||||
outputPath, err := safeOutputPath(*output)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
pool, err := pgxpool.New(ctx, databaseURL)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer pool.Close()
|
||||
snapshot, err := acceptancesnapshot.Export(ctx, pool, acceptancesnapshot.ExportOptions{
|
||||
ReleaseSHA: *releaseSHA,
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
payload, err := acceptancesnapshot.Encode(snapshot)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := os.WriteFile(outputPath, payload, 0o600); err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Printf(
|
||||
"acceptance_snapshot_export=PASS schema=%s config_hash=%s snapshot_sha256=%s\n",
|
||||
snapshot.SchemaVersion,
|
||||
snapshot.Source.ConfigHash,
|
||||
snapshot.SnapshotSHA256,
|
||||
)
|
||||
return nil
|
||||
}
|
||||
|
||||
func runValidate(args []string) error {
|
||||
flags := flag.NewFlagSet("validate", flag.ContinueOnError)
|
||||
input := flags.String("input", "", "snapshot input path")
|
||||
if err := flags.Parse(args); err != nil {
|
||||
return err
|
||||
}
|
||||
if flags.NArg() != 0 {
|
||||
return errors.New("unexpected validate arguments")
|
||||
}
|
||||
snapshot, err := readSnapshot(*input)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Printf(
|
||||
"acceptance_snapshot_validate=PASS schema=%s release=%s config_hash=%s snapshot_sha256=%s\n",
|
||||
snapshot.SchemaVersion,
|
||||
snapshot.Source.ReleaseSHA,
|
||||
snapshot.Source.ConfigHash,
|
||||
snapshot.SnapshotSHA256,
|
||||
)
|
||||
return nil
|
||||
}
|
||||
|
||||
func runImport(ctx context.Context, args []string) error {
|
||||
flags := flag.NewFlagSet("import", flag.ContinueOnError)
|
||||
input := flags.String("input", "", "snapshot input path")
|
||||
clusterID := flags.String("local-cluster-id", "", "required local cluster marker")
|
||||
if err := flags.Parse(args); err != nil {
|
||||
return err
|
||||
}
|
||||
if flags.NArg() != 0 {
|
||||
return errors.New("unexpected import arguments")
|
||||
}
|
||||
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_DATABASE_URL"))
|
||||
if databaseURL == "" {
|
||||
return errors.New("AI_GATEWAY_DATABASE_URL is required")
|
||||
}
|
||||
snapshot, err := readSnapshot(*input)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
pool, err := pgxpool.New(ctx, databaseURL)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer pool.Close()
|
||||
if err := acceptancesnapshot.Import(ctx, pool, snapshot, *clusterID); err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Printf(
|
||||
"acceptance_snapshot_import=PASS cluster_id=%s config_hash=%s snapshot_sha256=%s\n",
|
||||
strings.TrimSpace(*clusterID),
|
||||
snapshot.Source.ConfigHash,
|
||||
snapshot.SnapshotSHA256,
|
||||
)
|
||||
return nil
|
||||
}
|
||||
|
||||
func readSnapshot(path string) (acceptancesnapshot.Snapshot, error) {
|
||||
path = strings.TrimSpace(path)
|
||||
if path == "" {
|
||||
return acceptancesnapshot.Snapshot{}, errors.New("snapshot path is required")
|
||||
}
|
||||
info, err := os.Lstat(path)
|
||||
if err != nil {
|
||||
return acceptancesnapshot.Snapshot{}, err
|
||||
}
|
||||
if !info.Mode().IsRegular() || info.Mode()&os.ModeSymlink != 0 {
|
||||
return acceptancesnapshot.Snapshot{}, errors.New("snapshot must be a regular non-symlink file")
|
||||
}
|
||||
payload, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return acceptancesnapshot.Snapshot{}, err
|
||||
}
|
||||
return acceptancesnapshot.Decode(payload)
|
||||
}
|
||||
|
||||
func safeOutputPath(path string) (string, error) {
|
||||
path = strings.TrimSpace(path)
|
||||
if path == "" {
|
||||
return "", errors.New("output path is required")
|
||||
}
|
||||
absolute, err := filepath.Abs(path)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
parent := filepath.Dir(absolute)
|
||||
if err := os.MkdirAll(parent, 0o700); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if info, err := os.Lstat(absolute); err == nil {
|
||||
if !info.Mode().IsRegular() || info.Mode()&os.ModeSymlink != 0 {
|
||||
return "", errors.New("output must be a regular non-symlink file")
|
||||
}
|
||||
} else if !os.IsNotExist(err) {
|
||||
return "", err
|
||||
}
|
||||
return absolute, nil
|
||||
}
|
||||
@@ -0,0 +1,174 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"flag"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/accessruleaudit"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
func main() {
|
||||
if len(os.Args) < 2 {
|
||||
usage()
|
||||
os.Exit(64)
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
|
||||
defer cancel()
|
||||
var err error
|
||||
switch os.Args[1] {
|
||||
case "export":
|
||||
err = runExport(ctx, os.Args[2:])
|
||||
case "verify":
|
||||
err = runVerify(ctx, os.Args[2:])
|
||||
default:
|
||||
usage()
|
||||
os.Exit(64)
|
||||
}
|
||||
if err != nil {
|
||||
fmt.Fprintln(os.Stderr, "access-rule audit:", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
}
|
||||
|
||||
func usage() {
|
||||
fmt.Fprintln(os.Stderr, `Usage:
|
||||
easyai-ai-gateway-access-rule-audit export --output <before.json>
|
||||
easyai-ai-gateway-access-rule-audit verify --before <before.json> --output <after.json>
|
||||
|
||||
Both commands are database read-only and only emit grouped counts and SHA-256
|
||||
digests. Use a SELECT-only AI_GATEWAY_ACCESS_RULE_AUDIT_DATABASE_URL role.`)
|
||||
}
|
||||
|
||||
func runExport(ctx context.Context, args []string) error {
|
||||
flags := flag.NewFlagSet("export", flag.ContinueOnError)
|
||||
output := flags.String("output", "", "secret-safe access-rule audit output")
|
||||
if err := flags.Parse(args); err != nil {
|
||||
return err
|
||||
}
|
||||
if flags.NArg() != 0 {
|
||||
return errors.New("unexpected export arguments")
|
||||
}
|
||||
snapshot, outputPath, err := exportSnapshot(ctx, *output)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := writeSnapshot(outputPath, snapshot); err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Printf("access_rule_audit_export=PASS total=%d allow=%d deny=%d live_sha256=%s\n", snapshot.Live.Total, snapshot.Live.AllowCount, snapshot.Live.DenyCount, snapshot.Live.SHA256)
|
||||
return nil
|
||||
}
|
||||
|
||||
func runVerify(ctx context.Context, args []string) error {
|
||||
flags := flag.NewFlagSet("verify", flag.ContinueOnError)
|
||||
beforePathValue := flags.String("before", "", "pre-migration audit snapshot")
|
||||
output := flags.String("output", "", "post-migration audit output")
|
||||
if err := flags.Parse(args); err != nil {
|
||||
return err
|
||||
}
|
||||
if flags.NArg() != 0 {
|
||||
return errors.New("unexpected verify arguments")
|
||||
}
|
||||
beforePath, err := regularInputPath(*beforePathValue)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
after, outputPath, err := exportSnapshot(ctx, *output)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if beforePath == outputPath {
|
||||
return errors.New("verification output must not overwrite the pre-migration snapshot")
|
||||
}
|
||||
payload, err := os.ReadFile(beforePath)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
before, err := accessruleaudit.Decode(payload)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := accessruleaudit.VerifyMigration(before, after); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := writeSnapshot(outputPath, after); err != nil {
|
||||
return err
|
||||
}
|
||||
fmt.Printf("access_rule_audit_verify=PASS archived_allow=%d live_allow=%d live_deny=%d deny_sha256=%s\n", after.Archive.ArchivedAllowCount, after.Live.AllowCount, after.Live.DenyCount, after.Live.DenySHA256)
|
||||
return nil
|
||||
}
|
||||
|
||||
func exportSnapshot(ctx context.Context, output string) (accessruleaudit.Snapshot, string, error) {
|
||||
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_ACCESS_RULE_AUDIT_DATABASE_URL"))
|
||||
if databaseURL == "" {
|
||||
return accessruleaudit.Snapshot{}, "", errors.New("AI_GATEWAY_ACCESS_RULE_AUDIT_DATABASE_URL is required")
|
||||
}
|
||||
outputPath, err := safeOutputPath(output)
|
||||
if err != nil {
|
||||
return accessruleaudit.Snapshot{}, "", err
|
||||
}
|
||||
pool, err := pgxpool.New(ctx, databaseURL)
|
||||
if err != nil {
|
||||
return accessruleaudit.Snapshot{}, "", err
|
||||
}
|
||||
defer pool.Close()
|
||||
snapshot, err := accessruleaudit.Export(ctx, pool, time.Now())
|
||||
return snapshot, outputPath, err
|
||||
}
|
||||
|
||||
func writeSnapshot(path string, snapshot accessruleaudit.Snapshot) error {
|
||||
payload, err := accessruleaudit.Encode(snapshot)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return os.WriteFile(path, payload, 0o600)
|
||||
}
|
||||
|
||||
func regularInputPath(path string) (string, error) {
|
||||
path = strings.TrimSpace(path)
|
||||
if path == "" {
|
||||
return "", errors.New("input path is required")
|
||||
}
|
||||
absolute, err := filepath.Abs(path)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
info, err := os.Lstat(absolute)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if !info.Mode().IsRegular() || info.Mode()&os.ModeSymlink != 0 {
|
||||
return "", errors.New("input must be a regular non-symlink file")
|
||||
}
|
||||
return absolute, nil
|
||||
}
|
||||
|
||||
func safeOutputPath(path string) (string, error) {
|
||||
path = strings.TrimSpace(path)
|
||||
if path == "" {
|
||||
return "", errors.New("output path is required")
|
||||
}
|
||||
absolute, err := filepath.Abs(path)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
parent := filepath.Dir(absolute)
|
||||
if err := os.MkdirAll(parent, 0o700); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if info, err := os.Lstat(absolute); err == nil {
|
||||
if !info.Mode().IsRegular() || info.Mode()&os.ModeSymlink != 0 {
|
||||
return "", errors.New("output must be a regular non-symlink file")
|
||||
}
|
||||
} else if !os.IsNotExist(err) {
|
||||
return "", err
|
||||
}
|
||||
return absolute, nil
|
||||
}
|
||||
@@ -0,0 +1,133 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"flag"
|
||||
"log/slog"
|
||||
"os"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/runner"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func main() {
|
||||
apply := flag.Bool("apply", false, "persist compacted results; default is dry-run")
|
||||
requireClean := flag.Bool("require-clean", false, "exit non-zero unless a complete dry-run finds no results requiring URL migration")
|
||||
batchSize := flag.Int("batch-size", 100, "rows per batch, maximum 100")
|
||||
maxBatches := flag.Int("max-batches", 10, "maximum batches for one invocation")
|
||||
afterID := flag.String("after-id", "", "resume after this task UUID")
|
||||
flag.Parse()
|
||||
|
||||
logger := slog.New(slog.NewJSONHandler(os.Stdout, nil))
|
||||
cfg := config.Load()
|
||||
if err := cfg.Validate(); err != nil {
|
||||
logger.Error("invalid gateway configuration", "error", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
if *batchSize < 1 || *batchSize > 100 || *maxBatches < 1 {
|
||||
logger.Error("invalid backfill bounds", "batchSize", *batchSize, "maxBatches", *maxBatches)
|
||||
os.Exit(1)
|
||||
}
|
||||
if *apply && *requireClean {
|
||||
logger.Error("--apply and --require-clean cannot be combined")
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
db, err := store.Connect(ctx, cfg.DatabaseURL)
|
||||
if err != nil {
|
||||
logger.Error("connect postgres failed", "error", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
service := runner.New(cfg, db, logger)
|
||||
cursor := *afterID
|
||||
scanned := 0
|
||||
matched := 0
|
||||
blockingMatched := 0
|
||||
updated := 0
|
||||
expired := 0
|
||||
expiredLocalPlaceholders := 0
|
||||
complete := false
|
||||
for batch := 0; batch < *maxBatches; batch++ {
|
||||
items, err := db.ListTaskBinaryResultBackfillBatch(ctx, cursor, *batchSize)
|
||||
if err != nil {
|
||||
logger.Error("list binary result backfill batch failed", "afterId", cursor, "error", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
if len(items) == 0 {
|
||||
complete = true
|
||||
break
|
||||
}
|
||||
for _, item := range items {
|
||||
cursor = item.ID
|
||||
scanned++
|
||||
if !runner.TaskResultNeedsURLMigration(item.Result) {
|
||||
continue
|
||||
}
|
||||
matched++
|
||||
isExpired := item.FinishedAt.Before(time.Now().Add(-time.Duration(localResultTTLHours(cfg)) * time.Hour))
|
||||
hasLocalPlaceholder := runner.TaskResultHasLocalPlaceholder(item.Result)
|
||||
if isExpired && hasLocalPlaceholder {
|
||||
expiredLocalPlaceholders++
|
||||
if *apply {
|
||||
logger.Warn("skip expired local result placeholder without overwriting stored result", "taskId", item.ID)
|
||||
}
|
||||
continue
|
||||
}
|
||||
blockingMatched++
|
||||
if !*apply {
|
||||
continue
|
||||
}
|
||||
persistent, changed, err := service.MigrateTaskResultToURLs(ctx, item.ID, item.Result)
|
||||
if err != nil {
|
||||
logger.Error("materialize historical binary result failed", "taskId", item.ID, "error", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
if !changed {
|
||||
continue
|
||||
}
|
||||
ok, err := db.UpdateTaskBinaryResultBackfill(ctx, item.ID, persistent)
|
||||
if err != nil {
|
||||
logger.Error("update historical binary result failed", "taskId", item.ID, "error", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
updated++
|
||||
if isExpired {
|
||||
expired++
|
||||
}
|
||||
}
|
||||
if len(items) < *batchSize {
|
||||
complete = true
|
||||
break
|
||||
}
|
||||
}
|
||||
logger.Info("binary result backfill completed",
|
||||
"apply", *apply,
|
||||
"scanned", scanned,
|
||||
"matched", matched,
|
||||
"blockingMatched", blockingMatched,
|
||||
"updated", updated,
|
||||
"expired", expired,
|
||||
"expiredLocalPlaceholders", expiredLocalPlaceholders,
|
||||
"complete", complete,
|
||||
"resumeAfterId", cursor,
|
||||
)
|
||||
if *requireClean && (!complete || blockingMatched > 0) {
|
||||
logger.Error("binary result URL migration gate failed", "complete", complete, "blockingMatched", blockingMatched, "expiredLocalPlaceholders", expiredLocalPlaceholders, "resumeAfterId", cursor)
|
||||
os.Exit(1)
|
||||
}
|
||||
}
|
||||
|
||||
func localResultTTLHours(cfg config.Config) int {
|
||||
if cfg.LocalResultTTLHours <= 0 {
|
||||
return 24
|
||||
}
|
||||
return cfg.LocalResultTTLHours
|
||||
}
|
||||
@@ -10,6 +10,7 @@ import (
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/capacitycontroller"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/httpapi"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
@@ -36,30 +37,94 @@ func main() {
|
||||
}
|
||||
|
||||
ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
|
||||
defer stop()
|
||||
|
||||
db, err := store.Connect(ctx, cfg.DatabaseURL)
|
||||
executionMaxConns := cfg.DatabaseMaxConns
|
||||
if cfg.DatabaseCriticalMaxConns > 0 {
|
||||
executionMaxConns -= cfg.DatabaseCriticalMaxConns
|
||||
}
|
||||
if cfg.DatabaseRiverMaxConns > 0 {
|
||||
executionMaxConns -= cfg.DatabaseRiverMaxConns
|
||||
}
|
||||
db, err := store.ConnectWithPoolOptions(ctx, cfg.DatabaseURL, store.PostgresPoolOptions{
|
||||
MaxConns: executionMaxConns,
|
||||
MinIdleConns: cfg.DatabaseMinIdleConns,
|
||||
MaxConnIdleTime: time.Duration(cfg.DatabaseMaxConnIdleSeconds) * time.Second,
|
||||
IdleInTransactionTimeout: time.Duration(cfg.DatabaseIdleInTransactionTimeoutSeconds) * time.Second,
|
||||
LockTimeout: time.Duration(cfg.DatabaseLockTimeoutSeconds) * time.Second,
|
||||
})
|
||||
if err != nil {
|
||||
stop()
|
||||
logger.Error("connect postgres failed", "error", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
defer db.Close()
|
||||
if recovery, err := db.RecoverInterruptedRuntimeState(ctx); err != nil {
|
||||
logger.Error("recover interrupted runtime state failed", "error", err)
|
||||
os.Exit(1)
|
||||
} else if recovery.ReleasedConcurrencyLeases > 0 || recovery.ReleasedRateReservations > 0 || recovery.FailedAttempts > 0 || recovery.FailedTasks > 0 || recovery.RequeuedAsyncTasks > 0 {
|
||||
logger.Warn("interrupted runtime state recovered",
|
||||
"releasedConcurrencyLeases", recovery.ReleasedConcurrencyLeases,
|
||||
"releasedRateReservations", recovery.ReleasedRateReservations,
|
||||
"failedAttempts", recovery.FailedAttempts,
|
||||
"failedTasks", recovery.FailedTasks,
|
||||
"requeuedAsyncTasks", recovery.RequeuedAsyncTasks,
|
||||
)
|
||||
coordinationDB := db
|
||||
if cfg.DatabaseCriticalMaxConns > 0 {
|
||||
coordinationDB, err = store.ConnectWithPoolOptions(ctx, cfg.DatabaseURL, store.PostgresPoolOptions{
|
||||
MaxConns: cfg.DatabaseCriticalMaxConns,
|
||||
MinIdleConns: min(cfg.DatabaseCriticalMaxConns, max(cfg.DatabaseMinIdleConns, 1)),
|
||||
MaxConnIdleTime: time.Duration(cfg.DatabaseMaxConnIdleSeconds) * time.Second,
|
||||
IdleInTransactionTimeout: time.Duration(cfg.DatabaseIdleInTransactionTimeoutSeconds) * time.Second,
|
||||
LockTimeout: time.Duration(cfg.DatabaseLockTimeoutSeconds) * time.Second,
|
||||
})
|
||||
if err != nil {
|
||||
stop()
|
||||
logger.Error("connect critical postgres pool failed", "error", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
defer coordinationDB.Close()
|
||||
}
|
||||
riverDB := db
|
||||
if cfg.DatabaseRiverMaxConns > 0 {
|
||||
riverDB, err = store.ConnectWithPoolOptions(ctx, cfg.DatabaseURL, store.PostgresPoolOptions{
|
||||
MaxConns: cfg.DatabaseRiverMaxConns,
|
||||
MinIdleConns: min(cfg.DatabaseRiverMaxConns, max(cfg.DatabaseMinIdleConns, 1)),
|
||||
MaxConnIdleTime: time.Duration(cfg.DatabaseMaxConnIdleSeconds) * time.Second,
|
||||
IdleInTransactionTimeout: time.Duration(cfg.DatabaseIdleInTransactionTimeoutSeconds) * time.Second,
|
||||
LockTimeout: time.Duration(cfg.DatabaseLockTimeoutSeconds) * time.Second,
|
||||
})
|
||||
if err != nil {
|
||||
stop()
|
||||
logger.Error("connect River postgres pool failed", "error", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
defer riverDB.Close()
|
||||
}
|
||||
// Cancel background LISTEN and worker goroutines before closing the pool.
|
||||
// This also prevents a startup panic from waiting forever in pgxpool.Close.
|
||||
defer stop()
|
||||
if cfg.RunsBackgroundWorkers() {
|
||||
if recovery, recoveryErr := coordinationDB.RecoverInterruptedRuntimeState(ctx); recoveryErr != nil {
|
||||
logger.Error("recover interrupted runtime state failed", "error", recoveryErr)
|
||||
os.Exit(1)
|
||||
} else if recovery.ReleasedConcurrencyLeases > 0 || recovery.ReleasedRateReservations > 0 || recovery.FailedAttempts > 0 || recovery.FailedTasks > 0 || recovery.RequeuedAsyncTasks > 0 || recovery.CleanedTaskAdmissions > 0 {
|
||||
logger.Warn("interrupted runtime state recovered",
|
||||
"releasedConcurrencyLeases", recovery.ReleasedConcurrencyLeases,
|
||||
"releasedRateReservations", recovery.ReleasedRateReservations,
|
||||
"failedAttempts", recovery.FailedAttempts,
|
||||
"failedTasks", recovery.FailedTasks,
|
||||
"requeuedAsyncTasks", recovery.RequeuedAsyncTasks,
|
||||
"cleanedTaskAdmissions", recovery.CleanedTaskAdmissions,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
var handler http.Handler
|
||||
if cfg.RunsCapacityController() {
|
||||
orchestrator, orchestratorErr := capacitycontroller.NewConfiguredAdapter(cfg)
|
||||
if orchestratorErr != nil {
|
||||
logger.Error("initialize capacity orchestrator adapter failed", "error", orchestratorErr)
|
||||
os.Exit(1)
|
||||
}
|
||||
controller := capacitycontroller.New(cfg, coordinationDB, orchestrator, logger)
|
||||
go controller.Run(ctx)
|
||||
handler = controller.Handler()
|
||||
} else {
|
||||
handler = httpapi.NewServerWithStores(ctx, cfg, db, coordinationDB, riverDB, logger)
|
||||
}
|
||||
server := &http.Server{
|
||||
Addr: cfg.HTTPAddr,
|
||||
Handler: httpapi.NewServerWithContext(ctx, cfg, db, logger),
|
||||
Handler: handler,
|
||||
ReadHeaderTimeout: 10 * time.Second,
|
||||
}
|
||||
|
||||
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
"log/slog"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"sort"
|
||||
"strings"
|
||||
|
||||
@@ -13,6 +14,16 @@ import (
|
||||
"github.com/jackc/pgx/v5"
|
||||
)
|
||||
|
||||
const (
|
||||
noTransactionMigrationMarker = "-- easyai:migration:no-transaction"
|
||||
migrationStatementSeparator = "-- easyai:migration:statement"
|
||||
acceptanceImportReplayMarker = "-- easyai:migration:reapply-after-acceptance-import"
|
||||
acceptanceImportReplayEnvironment = "AI_GATEWAY_MIGRATION_REAPPLY_ACCEPTANCE_IMPORT"
|
||||
acceptanceLocalClusterSettingKey = "acceptance_local_cluster_id"
|
||||
)
|
||||
|
||||
var acceptanceLocalClusterMarkerPattern = regexp.MustCompile(`^easyai-local-[0-9a-f]{24}$`)
|
||||
|
||||
func main() {
|
||||
cfg := config.Load()
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
|
||||
@@ -37,6 +48,21 @@ CREATE TABLE IF NOT EXISTS schema_migrations (
|
||||
logger.Error("ensure schema_migrations failed", "error", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
reapplyAcceptanceImport := strings.EqualFold(strings.TrimSpace(os.Getenv(acceptanceImportReplayEnvironment)), "true")
|
||||
if reapplyAcceptanceImport {
|
||||
var localClusterID string
|
||||
if err := conn.QueryRow(ctx, `
|
||||
SELECT COALESCE(value->>'clusterId', '')
|
||||
FROM system_settings
|
||||
WHERE setting_key = $1`, acceptanceLocalClusterSettingKey).Scan(&localClusterID); err != nil {
|
||||
logger.Error("verify local acceptance database failed", "error", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
if !acceptanceLocalClusterMarkerPattern.MatchString(localClusterID) {
|
||||
logger.Error("refusing acceptance migration replay outside a marked local cluster")
|
||||
os.Exit(1)
|
||||
}
|
||||
}
|
||||
|
||||
files, err := filepath.Glob("migrations/*.sql")
|
||||
if err != nil {
|
||||
@@ -52,16 +78,38 @@ CREATE TABLE IF NOT EXISTS schema_migrations (
|
||||
logger.Error("check migration failed", "version", version, "error", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
if exists {
|
||||
logger.Info("migration skipped", "version", version)
|
||||
continue
|
||||
}
|
||||
|
||||
sqlBytes, err := os.ReadFile(file)
|
||||
if err != nil {
|
||||
logger.Error("read migration file failed", "file", file, "error", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
replaying := exists && reapplyAcceptanceImport && hasMigrationMarker(string(sqlBytes), acceptanceImportReplayMarker)
|
||||
if exists && !replaying {
|
||||
logger.Info("migration skipped", "version", version)
|
||||
continue
|
||||
}
|
||||
|
||||
noTransaction, statements := migrationStatements(string(sqlBytes))
|
||||
if noTransaction {
|
||||
for _, statement := range statements {
|
||||
if _, err := conn.Exec(ctx, statement); err != nil {
|
||||
logger.Error("execute non-transaction migration failed", "version", version, "error", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
}
|
||||
if !exists {
|
||||
if _, err := conn.Exec(ctx, "INSERT INTO schema_migrations(version) VALUES($1)", version); err != nil {
|
||||
logger.Error("record non-transaction migration failed", "version", version, "error", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
}
|
||||
if replaying {
|
||||
logger.Info("acceptance migration replayed", "version", version, "transactional", false)
|
||||
} else {
|
||||
logger.Info("migration applied", "version", version, "transactional", false)
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
tx, err := conn.Begin(ctx)
|
||||
if err != nil {
|
||||
@@ -81,17 +129,48 @@ CREATE TABLE IF NOT EXISTS schema_migrations (
|
||||
logger.Error("execute migration failed", "version", version, "error", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
if _, err := tx.Exec(ctx, "INSERT INTO schema_migrations(version) VALUES($1)", version); err != nil {
|
||||
_ = tx.Rollback(ctx)
|
||||
logger.Error("record migration failed", "version", version, "error", err)
|
||||
os.Exit(1)
|
||||
if !exists {
|
||||
if _, err := tx.Exec(ctx, "INSERT INTO schema_migrations(version) VALUES($1)", version); err != nil {
|
||||
_ = tx.Rollback(ctx)
|
||||
logger.Error("record migration failed", "version", version, "error", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
}
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
logger.Error("commit migration failed", "version", version, "error", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
logger.Info("migration applied", "version", version)
|
||||
if replaying {
|
||||
logger.Info("acceptance migration replayed", "version", version)
|
||||
} else {
|
||||
logger.Info("migration applied", "version", version)
|
||||
}
|
||||
}
|
||||
|
||||
fmt.Println("migrations complete")
|
||||
}
|
||||
|
||||
func hasMigrationMarker(sql string, marker string) bool {
|
||||
for _, line := range strings.Split(sql, "\n") {
|
||||
if strings.TrimSpace(line) == marker {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func migrationStatements(sql string) (bool, []string) {
|
||||
trimmed := strings.TrimSpace(sql)
|
||||
if !strings.HasPrefix(trimmed, noTransactionMigrationMarker) {
|
||||
return false, []string{sql}
|
||||
}
|
||||
trimmed = strings.TrimSpace(strings.TrimPrefix(trimmed, noTransactionMigrationMarker))
|
||||
parts := strings.Split(trimmed, migrationStatementSeparator)
|
||||
statements := make([]string, 0, len(parts))
|
||||
for _, part := range parts {
|
||||
if statement := strings.TrimSpace(part); statement != "" {
|
||||
statements = append(statements, statement)
|
||||
}
|
||||
}
|
||||
return true, statements
|
||||
}
|
||||
|
||||
@@ -14,6 +14,102 @@ import (
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
func TestMigrationStatementsSupportsConcurrentIndexes(t *testing.T) {
|
||||
payload := `-- easyai:migration:no-transaction
|
||||
CREATE INDEX CONCURRENTLY IF NOT EXISTS first_index ON first_table(id);
|
||||
-- easyai:migration:statement
|
||||
CREATE INDEX CONCURRENTLY IF NOT EXISTS second_index ON second_table(id);
|
||||
`
|
||||
noTransaction, statements := migrationStatements(payload)
|
||||
if !noTransaction {
|
||||
t.Fatal("expected a non-transaction migration")
|
||||
}
|
||||
if len(statements) != 2 {
|
||||
t.Fatalf("statements=%d, want 2", len(statements))
|
||||
}
|
||||
if !strings.Contains(statements[0], "first_index") || !strings.Contains(statements[1], "second_index") {
|
||||
t.Fatalf("unexpected statements: %#v", statements)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAcceptanceImportReplayMarkerRequiresAnExactCommentLine(t *testing.T) {
|
||||
if !hasMigrationMarker(
|
||||
"-- preface\n"+acceptanceImportReplayMarker+"\nSELECT 1;\n",
|
||||
acceptanceImportReplayMarker,
|
||||
) {
|
||||
t.Fatal("expected exact acceptance import replay marker")
|
||||
}
|
||||
if hasMigrationMarker(
|
||||
"-- mentions "+acceptanceImportReplayMarker+" in prose\nSELECT 1;\n",
|
||||
acceptanceImportReplayMarker,
|
||||
) {
|
||||
t.Fatal("prose mention must not enable acceptance import replay")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSeedanceInputImageConstraintMigrationKeepsCatalogSnapshotsInSync(t *testing.T) {
|
||||
payload, err := os.ReadFile("../../migrations/0078_seedance_input_image_constraints.sql")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
content := string(payload)
|
||||
for _, required := range []string{
|
||||
"easyai:豆包Seedance-2.0",
|
||||
"easyai:豆包Seedance-2.0-fast",
|
||||
"input_image_resolution_range",
|
||||
"input_image_aspect_ratio_range",
|
||||
"'{metadata,rawModel,capabilities}'",
|
||||
"UPDATE platform_models",
|
||||
} {
|
||||
if !strings.Contains(content, required) {
|
||||
t.Fatalf("Seedance input image migration is missing %q", required)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestVolcesSeedanceInputImageConstraintMigrationUsesDocumentedBounds(t *testing.T) {
|
||||
payload, err := os.ReadFile("../../migrations/0098_volces_seedance20_input_image_constraints.sql")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
content := string(payload)
|
||||
for _, required := range []string{
|
||||
"volces:doubao-seedance-2-0-260128",
|
||||
"volces:doubao-seedance-2-0-fast-260128",
|
||||
"volces:doubao-seedance-2-0-mini-260615",
|
||||
`"long_edge":300`,
|
||||
`"long_edge":6000`,
|
||||
`'[0.4,2.5]'::jsonb`,
|
||||
"'{metadata,rawModel,capabilities}'",
|
||||
"UPDATE platform_models",
|
||||
} {
|
||||
if !strings.Contains(content, required) {
|
||||
t.Fatalf("Volces Seedance input image migration is missing %q", required)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestVolcesSeedanceAcceptanceReconciliationIsReplayable(t *testing.T) {
|
||||
payload, err := os.ReadFile("../../migrations/0099_reconcile_volces_seedance20_acceptance_constraints.sql")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
content := string(payload)
|
||||
for _, required := range []string{
|
||||
acceptanceImportReplayMarker,
|
||||
"volces:doubao-seedance-2-0-260128",
|
||||
"volces:doubao-seedance-2-0-fast-260128",
|
||||
"volces:doubao-seedance-2-0-mini-260615",
|
||||
`"long_edge":300`,
|
||||
`"long_edge":6000`,
|
||||
`'[0.4,2.5]'::jsonb`,
|
||||
} {
|
||||
if !strings.Contains(content, required) {
|
||||
t.Fatalf("Volces acceptance reconciliation is missing %q", required)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSecurityEventSchemaMigrationsDefineCurrentLifecycle(t *testing.T) {
|
||||
streamPayload, err := os.ReadFile("../../migrations/0063_oidc_security_events.sql")
|
||||
if err != nil {
|
||||
@@ -279,6 +375,34 @@ SET last_error_category='still-invalid' WHERE id=$1::uuid`, pairingID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOIDCSessionClientIdentityRepairMigrationUpgradesApplied0090Schema(t *testing.T) {
|
||||
pool := newIdentityMigrationPostgresTestSchema(t)
|
||||
ctx := context.Background()
|
||||
if _, err := pool.Exec(ctx, `
|
||||
CREATE TABLE gateway_oidc_sessions (
|
||||
id uuid PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
oidc_user_binding_id uuid
|
||||
)`); err != nil {
|
||||
t.Fatalf("create pre-repair OIDC session schema: %v", err)
|
||||
}
|
||||
|
||||
applyIdentityMigrationTestFile(t, ctx, pool, "../../migrations/0091_oidc_session_client_identity.sql")
|
||||
applyIdentityMigrationTestFile(t, ctx, pool, "../../migrations/0091_oidc_session_client_identity.sql")
|
||||
|
||||
var dataType string
|
||||
if err := pool.QueryRow(ctx, `
|
||||
SELECT data_type
|
||||
FROM information_schema.columns
|
||||
WHERE table_schema=current_schema()
|
||||
AND table_name='gateway_oidc_sessions'
|
||||
AND column_name='oidc_client_id'`).Scan(&dataType); err != nil {
|
||||
t.Fatalf("read repaired OIDC session client identity column: %v", err)
|
||||
}
|
||||
if dataType != "text" {
|
||||
t.Fatalf("oidc_client_id data type=%q, want text", dataType)
|
||||
}
|
||||
}
|
||||
|
||||
func applyIdentityMigrationTestFile(t *testing.T, ctx context.Context, pool *pgxpool.Pool, path string) {
|
||||
t.Helper()
|
||||
payload, err := os.ReadFile(path)
|
||||
|
||||
+3637
-205
File diff suppressed because it is too large
Load Diff
+2476
-175
File diff suppressed because it is too large
Load Diff
@@ -15,6 +15,8 @@ require (
|
||||
github.com/riverqueue/river/riverdriver/riverpgxv5 v0.24.0
|
||||
github.com/riverqueue/river/rivertype v0.24.0
|
||||
golang.org/x/crypto v0.52.0
|
||||
golang.org/x/image v0.43.0
|
||||
golang.org/x/net v0.54.0
|
||||
golang.org/x/oauth2 v0.36.0
|
||||
)
|
||||
|
||||
|
||||
@@ -71,6 +71,10 @@ go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto=
|
||||
go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE=
|
||||
golang.org/x/crypto v0.52.0 h1:RMs7fP2rXdep0CftQlK8Uf+kibLm7qkCcradZWYz988=
|
||||
golang.org/x/crypto v0.52.0/go.mod h1:1QgfPxDqh0T2M/elOJtp9RvuR95kVjir0e6/BvEmGbc=
|
||||
golang.org/x/net v0.54.0 h1:2zJIZAxAHV/OHCDTCOHAYehQzLfSXuf/5SoL/Dv6w/w=
|
||||
golang.org/x/net v0.54.0/go.mod h1:Sj4oj8jK6XmHpBZU/zWHw3BV3abl4Kvi+Ut7cQcY+cQ=
|
||||
golang.org/x/image v0.43.0 h1:FLxcP4ec2350nTfOC8ysKtqYSIFbk/QGjw1ZHNP4tsY=
|
||||
golang.org/x/image v0.43.0/go.mod h1:rrpelvGFt+kLPAjPM4HeWPgrl0FtafueU//e5N0qk/Q=
|
||||
golang.org/x/oauth2 v0.36.0 h1:peZ/1z27fi9hUOFCAZaHyrpWG5lwe0RJEEEeH0ThlIs=
|
||||
golang.org/x/oauth2 v0.36.0/go.mod h1:YDBUJMTkDnJS+A4BP4eZBjCqtokkg1hODuPjwiGPO7Q=
|
||||
golang.org/x/sync v0.21.0 h1:HLII4xRRTtCRkxYp4HNFF0Js/Og6q2i++KXbg0gHCwM=
|
||||
|
||||
@@ -0,0 +1,913 @@
|
||||
package acceptanceemulator
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"encoding/binary"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"hash/crc32"
|
||||
"image"
|
||||
"image/color"
|
||||
"image/jpeg"
|
||||
"image/png"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/acceptanceworkload"
|
||||
_ "golang.org/x/image/webp"
|
||||
)
|
||||
|
||||
const (
|
||||
maxProtocolBodyBytes = 128 << 20
|
||||
maxImageReferenceBytes = 32 << 20
|
||||
)
|
||||
|
||||
type Config struct {
|
||||
Now func() time.Time
|
||||
HTTPClient *http.Client
|
||||
Wait func(context.Context, time.Duration) error
|
||||
}
|
||||
|
||||
type Server struct {
|
||||
now func() time.Time
|
||||
httpClient *http.Client
|
||||
wait func(context.Context, time.Duration) error
|
||||
mu sync.RWMutex
|
||||
tasks map[string]videoTask
|
||||
fixtures map[string]fixture
|
||||
nextID atomic.Uint64
|
||||
report Report
|
||||
callbacks map[string]map[int64]int
|
||||
videoByIdempotency map[string]string
|
||||
geminiIdempotency map[string]struct{}
|
||||
storageObjects map[string]fixture
|
||||
storageAttempts map[string]int
|
||||
pngSmall string
|
||||
pngLarge string
|
||||
pngPeak string
|
||||
}
|
||||
|
||||
type Report struct {
|
||||
GeminiRequests int64 `json:"geminiRequests"`
|
||||
GeminiInvalid int64 `json:"geminiInvalid"`
|
||||
GeminiInputImages int64 `json:"geminiInputImages"`
|
||||
GeminiInputBytes int64 `json:"geminiInputBytes"`
|
||||
GeminiOutputBytes int64 `json:"geminiOutputBytes"`
|
||||
VideoSubmissions int64 `json:"videoSubmissions"`
|
||||
VideoPolls int64 `json:"videoPolls"`
|
||||
VideoInvalid int64 `json:"videoInvalid"`
|
||||
VideoReferenceCounts map[string]int64 `json:"videoReferenceCounts"`
|
||||
VideoRoleCounts map[string]int64 `json:"videoRoleCounts"`
|
||||
UniqueImageHashes int `json:"uniqueImageHashes"`
|
||||
UniqueVideoTasks int `json:"uniqueVideoTasks"`
|
||||
ForcedConversions int64 `json:"forcedConversionRequests"`
|
||||
VerifiedConversions int64 `json:"verifiedConversions"`
|
||||
CallbackEvents int64 `json:"callbackEvents"`
|
||||
DuplicateCallbacks int64 `json:"duplicateCallbacks"`
|
||||
MissingIdempotency int64 `json:"missingIdempotencyKeys"`
|
||||
DuplicateSubmissions int64 `json:"duplicateSubmissionAttempts"`
|
||||
StoragePuts int64 `json:"storagePuts"`
|
||||
StorageGets int64 `json:"storageGets"`
|
||||
StorageHeads int64 `json:"storageHeads"`
|
||||
StorageDeletes int64 `json:"storageDeletes"`
|
||||
StorageFailures int64 `json:"storageFailures"`
|
||||
}
|
||||
|
||||
type videoTask struct {
|
||||
ID string
|
||||
Model string
|
||||
CreatedAt time.Time
|
||||
ReadyAt time.Time
|
||||
UsageTokens int
|
||||
ImageRefs []imageReference
|
||||
}
|
||||
|
||||
type imageReference struct {
|
||||
Role string
|
||||
SHA256 string
|
||||
ContentType string
|
||||
Width int
|
||||
Height int
|
||||
}
|
||||
|
||||
type fixture struct {
|
||||
ContentType string
|
||||
Payload []byte
|
||||
}
|
||||
|
||||
func New(config Config) *Server {
|
||||
now := config.Now
|
||||
if now == nil {
|
||||
now = time.Now
|
||||
}
|
||||
client := config.HTTPClient
|
||||
if client == nil {
|
||||
client = &http.Client{Timeout: 30 * time.Second}
|
||||
}
|
||||
waiter := config.Wait
|
||||
if waiter == nil {
|
||||
waiter = wait
|
||||
}
|
||||
return &Server{
|
||||
now: now, httpClient: client, wait: waiter, tasks: map[string]videoTask{}, fixtures: buildFixtures(),
|
||||
callbacks: map[string]map[int64]int{},
|
||||
videoByIdempotency: map[string]string{},
|
||||
geminiIdempotency: map[string]struct{}{},
|
||||
storageObjects: map[string]fixture{},
|
||||
storageAttempts: map[string]int{},
|
||||
report: Report{
|
||||
VideoReferenceCounts: map[string]int64{},
|
||||
VideoRoleCounts: map[string]int64{},
|
||||
},
|
||||
pngSmall: base64.StdEncoding.EncodeToString(paddedPNG(256 << 10)),
|
||||
pngLarge: base64.StdEncoding.EncodeToString(paddedPNG(4 << 20)),
|
||||
pngPeak: base64.StdEncoding.EncodeToString(paddedPNG(8 << 20)),
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) Handler() http.Handler {
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("GET /healthz", s.health)
|
||||
mux.HandleFunc("GET /report", s.getReport)
|
||||
mux.HandleFunc("POST /v1beta/models/", s.geminiGenerateContent)
|
||||
mux.HandleFunc("POST /v1/models/", s.geminiGenerateContent)
|
||||
mux.HandleFunc("POST /contents/generations/tasks", s.submitVideo)
|
||||
mux.HandleFunc("GET /contents/generations/tasks/{taskID}", s.getVideo)
|
||||
mux.HandleFunc("DELETE /contents/generations/tasks/{taskID}", s.deleteVideo)
|
||||
mux.HandleFunc("GET /media/{asset}", s.getMedia)
|
||||
mux.HandleFunc("GET /fixtures/{asset}", s.getFixture)
|
||||
mux.HandleFunc("POST /callbacks", s.collectCallback)
|
||||
mux.HandleFunc("PUT /storage/{profile}/{object...}", s.objectStorage)
|
||||
mux.HandleFunc("GET /storage/{profile}/{object...}", s.objectStorage)
|
||||
mux.HandleFunc("HEAD /storage/{profile}/{object...}", s.objectStorage)
|
||||
mux.HandleFunc("DELETE /storage/{profile}/{object...}", s.objectStorage)
|
||||
return mux
|
||||
}
|
||||
|
||||
func (s *Server) objectStorage(w http.ResponseWriter, r *http.Request) {
|
||||
profile := strings.ToLower(strings.TrimSpace(r.PathValue("profile")))
|
||||
object := strings.TrimLeft(strings.TrimSpace(r.PathValue("object")), "/")
|
||||
if object == "" {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
baseProfile := strings.TrimSuffix(strings.TrimSuffix(strings.TrimSuffix(profile, "-transient"), "-auth"), "-fail")
|
||||
if baseProfile != "oss" && baseProfile != "s3" {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
key := profile + "/" + object
|
||||
s.mu.Lock()
|
||||
s.storageAttempts[key]++
|
||||
attempt := s.storageAttempts[key]
|
||||
if strings.HasSuffix(profile, "-auth") {
|
||||
s.report.StorageFailures++
|
||||
s.mu.Unlock()
|
||||
http.Error(w, "forbidden", http.StatusForbidden)
|
||||
return
|
||||
}
|
||||
if strings.HasSuffix(profile, "-fail") || (strings.HasSuffix(profile, "-transient") && attempt == 1) {
|
||||
s.report.StorageFailures++
|
||||
s.mu.Unlock()
|
||||
http.Error(w, "temporary storage failure", http.StatusServiceUnavailable)
|
||||
return
|
||||
}
|
||||
switch r.Method {
|
||||
case http.MethodPut:
|
||||
s.mu.Unlock()
|
||||
payload, err := io.ReadAll(io.LimitReader(r.Body, maxProtocolBodyBytes+1))
|
||||
if err != nil || len(payload) > maxProtocolBodyBytes {
|
||||
http.Error(w, "invalid object", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
s.storageObjects[key] = fixture{ContentType: firstNonEmpty(r.Header.Get("Content-Type"), "application/octet-stream"), Payload: payload}
|
||||
s.report.StoragePuts++
|
||||
s.mu.Unlock()
|
||||
w.WriteHeader(http.StatusOK)
|
||||
case http.MethodGet:
|
||||
item, ok := s.storageObjects[key]
|
||||
if ok {
|
||||
s.report.StorageGets++
|
||||
}
|
||||
s.mu.Unlock()
|
||||
if !ok {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", item.ContentType)
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = w.Write(item.Payload)
|
||||
case http.MethodHead:
|
||||
item, ok := s.storageObjects[key]
|
||||
if ok {
|
||||
s.report.StorageHeads++
|
||||
}
|
||||
s.mu.Unlock()
|
||||
if !ok {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", item.ContentType)
|
||||
w.Header().Set("Content-Length", strconv.Itoa(len(item.Payload)))
|
||||
w.WriteHeader(http.StatusOK)
|
||||
case http.MethodDelete:
|
||||
_, ok := s.storageObjects[key]
|
||||
delete(s.storageObjects, key)
|
||||
if ok {
|
||||
s.report.StorageDeletes++
|
||||
}
|
||||
s.mu.Unlock()
|
||||
if !ok {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
w.WriteHeader(http.StatusNoContent)
|
||||
default:
|
||||
s.mu.Unlock()
|
||||
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
|
||||
}
|
||||
}
|
||||
|
||||
func firstNonEmpty(values ...string) string {
|
||||
for _, value := range values {
|
||||
if value = strings.TrimSpace(value); value != "" {
|
||||
return value
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (s *Server) health(w http.ResponseWriter, _ *http.Request) {
|
||||
writeJSON(w, http.StatusOK, map[string]any{"ok": true})
|
||||
}
|
||||
|
||||
func (s *Server) geminiGenerateContent(w http.ResponseWriter, r *http.Request) {
|
||||
if !strings.HasSuffix(r.URL.Path, ":generateContent") {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
idempotencyKey := strings.TrimSpace(r.Header.Get("Idempotency-Key"))
|
||||
if idempotencyKey == "" {
|
||||
s.mu.Lock()
|
||||
s.report.MissingIdempotency++
|
||||
s.mu.Unlock()
|
||||
writeProtocolError(w, http.StatusBadRequest, "Idempotency-Key is required")
|
||||
return
|
||||
}
|
||||
var body map[string]any
|
||||
if err := decodeJSON(r, &body); err != nil {
|
||||
s.recordGeminiInvalid()
|
||||
writeProtocolError(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
inputBytes, inputImages, err := validateGeminiImageRequest(body)
|
||||
if err != nil {
|
||||
s.recordGeminiInvalid()
|
||||
writeProtocolError(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
outputBytes, outputBase64, delay := s.geminiProfile(inputBytes)
|
||||
if err := s.wait(r.Context(), delay); err != nil {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
if _, exists := s.geminiIdempotency[idempotencyKey]; exists {
|
||||
s.report.DuplicateSubmissions++
|
||||
} else {
|
||||
s.geminiIdempotency[idempotencyKey] = struct{}{}
|
||||
}
|
||||
s.report.GeminiRequests++
|
||||
s.report.GeminiInputImages += int64(inputImages)
|
||||
s.report.GeminiInputBytes += int64(inputBytes)
|
||||
s.report.GeminiOutputBytes += int64(outputBytes)
|
||||
s.mu.Unlock()
|
||||
w.Header().Set("X-Request-ID", "acceptance-gemini-"+strconv.FormatUint(s.nextID.Add(1), 10))
|
||||
writeJSON(w, http.StatusOK, map[string]any{
|
||||
"candidates": []any{map[string]any{
|
||||
"index": 0,
|
||||
"content": map[string]any{
|
||||
"role": "model",
|
||||
"parts": []any{map[string]any{"inlineData": map[string]any{
|
||||
"mimeType": "image/png",
|
||||
"data": outputBase64,
|
||||
}}},
|
||||
},
|
||||
"finishReason": "STOP",
|
||||
}},
|
||||
"usageMetadata": map[string]any{
|
||||
"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) geminiProfile(inputBytes int) (int, string, time.Duration) {
|
||||
switch {
|
||||
case inputBytes > acceptanceworkload.GeminiLarge.InputBytes:
|
||||
profile := acceptanceworkload.GeminiPeak
|
||||
return profile.OutputBytes, s.pngPeak, scaledDelay(profile.DelayMin, profile.DelayMax, inputBytes)
|
||||
case inputBytes > acceptanceworkload.GeminiBaseline.InputBytes:
|
||||
profile := acceptanceworkload.GeminiLarge
|
||||
return profile.OutputBytes, s.pngLarge, scaledDelay(profile.DelayMin, profile.DelayMax, inputBytes)
|
||||
default:
|
||||
profile := acceptanceworkload.GeminiBaseline
|
||||
return profile.OutputBytes, s.pngSmall, profile.DelayMin
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Server) submitVideo(w http.ResponseWriter, r *http.Request) {
|
||||
idempotencyKey := strings.TrimSpace(r.Header.Get("Idempotency-Key"))
|
||||
if idempotencyKey == "" {
|
||||
s.mu.Lock()
|
||||
s.report.MissingIdempotency++
|
||||
s.mu.Unlock()
|
||||
writeProtocolError(w, http.StatusBadRequest, "Idempotency-Key is required")
|
||||
return
|
||||
}
|
||||
var body map[string]any
|
||||
if err := decodeJSON(r, &body); err != nil {
|
||||
s.recordVideoInvalid()
|
||||
writeProtocolError(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
refs, longRun, forceConversion, err := s.validateVideoRequest(r.Context(), body)
|
||||
if err != nil {
|
||||
s.recordVideoInvalid()
|
||||
writeProtocolError(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
delay := videoDelay(body, longRun)
|
||||
id := "acceptance-video-" + strconv.FormatUint(s.nextID.Add(1), 10)
|
||||
task := videoTask{
|
||||
ID: id, Model: strings.TrimSpace(stringValue(body["model"])),
|
||||
CreatedAt: s.now(), ReadyAt: s.now().Add(delay), UsageTokens: videoUsageTokens(body), ImageRefs: refs,
|
||||
}
|
||||
s.mu.Lock()
|
||||
if existingID := s.videoByIdempotency[idempotencyKey]; existingID != "" {
|
||||
task = s.tasks[existingID]
|
||||
s.report.DuplicateSubmissions++
|
||||
s.mu.Unlock()
|
||||
w.Header().Set("X-Request-ID", task.ID)
|
||||
writeJSON(w, http.StatusOK, map[string]any{
|
||||
"id": task.ID, "model": task.Model, "status": "queued",
|
||||
"created_at": task.CreatedAt.Unix(),
|
||||
})
|
||||
return
|
||||
}
|
||||
s.tasks[id] = task
|
||||
s.videoByIdempotency[idempotencyKey] = id
|
||||
s.report.VideoSubmissions++
|
||||
s.report.VideoReferenceCounts[strconv.Itoa(len(refs))]++
|
||||
if forceConversion {
|
||||
s.report.ForcedConversions++
|
||||
s.report.VerifiedConversions++
|
||||
}
|
||||
for _, ref := range refs {
|
||||
s.report.VideoRoleCounts[ref.Role]++
|
||||
}
|
||||
s.mu.Unlock()
|
||||
w.Header().Set("X-Request-ID", id)
|
||||
writeJSON(w, http.StatusOK, map[string]any{
|
||||
"id": id, "model": task.Model, "status": "queued",
|
||||
"created_at": task.CreatedAt.Unix(),
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) getVideo(w http.ResponseWriter, r *http.Request) {
|
||||
id := strings.TrimSpace(r.PathValue("taskID"))
|
||||
s.mu.Lock()
|
||||
task, ok := s.tasks[id]
|
||||
if ok {
|
||||
s.report.VideoPolls++
|
||||
}
|
||||
s.mu.Unlock()
|
||||
if !ok {
|
||||
writeProtocolError(w, http.StatusNotFound, "video task not found")
|
||||
return
|
||||
}
|
||||
status := "running"
|
||||
content := map[string]any{}
|
||||
if !s.now().Before(task.ReadyAt) {
|
||||
status = "succeeded"
|
||||
content["video_url"] = requestOrigin(r) + "/media/" + url.PathEscape(id) + ".mp4"
|
||||
}
|
||||
w.Header().Set("X-Request-ID", id)
|
||||
writeJSON(w, http.StatusOK, map[string]any{
|
||||
"id": id, "model": task.Model, "status": status, "content": content,
|
||||
"created_at": task.CreatedAt.Unix(),
|
||||
"usage": map[string]any{"completion_tokens": task.UsageTokens, "total_tokens": task.UsageTokens},
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) deleteVideo(w http.ResponseWriter, r *http.Request) {
|
||||
id := strings.TrimSpace(r.PathValue("taskID"))
|
||||
s.mu.Lock()
|
||||
_, ok := s.tasks[id]
|
||||
delete(s.tasks, id)
|
||||
s.mu.Unlock()
|
||||
if !ok {
|
||||
writeProtocolError(w, http.StatusNotFound, "video task not found")
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]any{"id": id, "status": "cancelled"})
|
||||
}
|
||||
|
||||
func (s *Server) getMedia(w http.ResponseWriter, r *http.Request) {
|
||||
asset := strings.TrimSuffix(strings.TrimSpace(r.PathValue("asset")), ".mp4")
|
||||
s.mu.RLock()
|
||||
_, ok := s.tasks[asset]
|
||||
s.mu.RUnlock()
|
||||
if !ok {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "video/mp4")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = w.Write(minimalMP4)
|
||||
}
|
||||
|
||||
func (s *Server) getFixture(w http.ResponseWriter, r *http.Request) {
|
||||
fixture, ok := s.fixtures[strings.TrimSpace(r.PathValue("asset"))]
|
||||
if !ok {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", fixture.ContentType)
|
||||
w.Header().Set("Cache-Control", "public, max-age=3600")
|
||||
w.Header().Set("X-Content-SHA256", payloadSHA256(fixture.Payload))
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = w.Write(fixture.Payload)
|
||||
}
|
||||
|
||||
func (s *Server) getReport(w http.ResponseWriter, _ *http.Request) {
|
||||
s.mu.RLock()
|
||||
report := s.report
|
||||
tasks := make([]videoTask, 0, len(s.tasks))
|
||||
for _, task := range s.tasks {
|
||||
tasks = append(tasks, task)
|
||||
}
|
||||
s.mu.RUnlock()
|
||||
hashes := map[string]struct{}{}
|
||||
for _, task := range tasks {
|
||||
for _, ref := range task.ImageRefs {
|
||||
hashes[ref.SHA256] = struct{}{}
|
||||
}
|
||||
}
|
||||
report.UniqueImageHashes = len(hashes)
|
||||
report.UniqueVideoTasks = len(tasks)
|
||||
writeJSON(w, http.StatusOK, report)
|
||||
}
|
||||
|
||||
func (s *Server) collectCallback(w http.ResponseWriter, r *http.Request) {
|
||||
var payload map[string]any
|
||||
if err := decodeJSON(r, &payload); err != nil {
|
||||
writeProtocolError(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
taskID := strings.TrimSpace(stringValue(payload["taskId"]))
|
||||
seq := int64(0)
|
||||
switch value := payload["seq"].(type) {
|
||||
case json.Number:
|
||||
seq, _ = value.Int64()
|
||||
case float64:
|
||||
seq = int64(value)
|
||||
}
|
||||
if taskID == "" || seq <= 0 {
|
||||
writeProtocolError(w, http.StatusBadRequest, "callback taskId and seq are required")
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
seen := s.callbacks[taskID]
|
||||
if seen == nil {
|
||||
seen = map[int64]int{}
|
||||
s.callbacks[taskID] = seen
|
||||
}
|
||||
seen[seq]++
|
||||
s.report.CallbackEvents++
|
||||
if seen[seq] > 1 {
|
||||
s.report.DuplicateCallbacks++
|
||||
}
|
||||
s.mu.Unlock()
|
||||
writeJSON(w, http.StatusOK, map[string]any{"ok": true})
|
||||
}
|
||||
|
||||
func (s *Server) validateVideoRequest(ctx context.Context, body map[string]any) ([]imageReference, bool, bool, error) {
|
||||
content, _ := body["content"].([]any)
|
||||
refs := make([]imageReference, 0, 9)
|
||||
firstFrames := 0
|
||||
lastFrames := 0
|
||||
longRun := false
|
||||
forceConversion := false
|
||||
for _, raw := range content {
|
||||
item, _ := raw.(map[string]any)
|
||||
switch strings.TrimSpace(stringValue(item["type"])) {
|
||||
case "text":
|
||||
if strings.Contains(strings.ToLower(stringValue(item["text"])), "acceptance-long-recovery") {
|
||||
longRun = true
|
||||
}
|
||||
if strings.Contains(strings.ToLower(stringValue(item["text"])), "acceptance-force-conversion") {
|
||||
forceConversion = true
|
||||
}
|
||||
case "image_url":
|
||||
role := strings.TrimSpace(stringValue(item["role"]))
|
||||
switch role {
|
||||
case "reference_image":
|
||||
case "first_frame":
|
||||
firstFrames++
|
||||
case "last_frame":
|
||||
lastFrames++
|
||||
default:
|
||||
return nil, false, false, fmt.Errorf("unsupported image role %q", role)
|
||||
}
|
||||
nested, _ := item["image_url"].(map[string]any)
|
||||
rawURL := strings.TrimSpace(stringValue(nested["url"]))
|
||||
ref, err := s.fetchImageReference(ctx, role, rawURL)
|
||||
if err != nil {
|
||||
return nil, false, false, err
|
||||
}
|
||||
refs = append(refs, ref)
|
||||
}
|
||||
}
|
||||
if len(refs) != 3 && len(refs) != 6 && len(refs) != 9 {
|
||||
return nil, false, false, fmt.Errorf("expected 3, 6, or 9 image references, got %d", len(refs))
|
||||
}
|
||||
if firstFrames > 1 || lastFrames > 1 {
|
||||
return nil, false, false, errors.New("first_frame and last_frame must be unique")
|
||||
}
|
||||
if forceConversion {
|
||||
for _, ref := range refs {
|
||||
if !seedanceImageWithinOfficialInputRange(ref.Width, ref.Height) {
|
||||
return nil, false, false, errors.New("forced oversized image did not pass through automatic normalization")
|
||||
}
|
||||
}
|
||||
}
|
||||
return refs, longRun, forceConversion, nil
|
||||
}
|
||||
|
||||
func seedanceImageWithinOfficialInputRange(width int, height int) bool {
|
||||
if width < 300 || width > 6000 || height < 300 || height > 6000 {
|
||||
return false
|
||||
}
|
||||
ratio := float64(width) / float64(height)
|
||||
return ratio >= 0.4 && ratio <= 2.5
|
||||
}
|
||||
|
||||
func (s *Server) fetchImageReference(ctx context.Context, role string, rawURL string) (imageReference, error) {
|
||||
if rawURL == "" {
|
||||
return imageReference{}, errors.New("image reference URL is required")
|
||||
}
|
||||
payload, err := s.readImageReference(ctx, rawURL)
|
||||
if err != nil {
|
||||
return imageReference{}, err
|
||||
}
|
||||
config, format, err := image.DecodeConfig(bytes.NewReader(payload))
|
||||
if err != nil {
|
||||
return imageReference{}, fmt.Errorf("decode image reference: %w", err)
|
||||
}
|
||||
sum := sha256.Sum256(payload)
|
||||
return imageReference{
|
||||
Role: role, SHA256: hex.EncodeToString(sum[:]), ContentType: "image/" + format,
|
||||
Width: config.Width, Height: config.Height,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *Server) readImageReference(ctx context.Context, rawURL string) ([]byte, error) {
|
||||
if strings.HasPrefix(strings.ToLower(rawURL), "data:") {
|
||||
return decodeImageDataURL(rawURL)
|
||||
}
|
||||
request, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
response, err := s.httpClient.Do(request)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("fetch image reference: %w", err)
|
||||
}
|
||||
defer response.Body.Close()
|
||||
if response.StatusCode < 200 || response.StatusCode >= 300 {
|
||||
return nil, fmt.Errorf("fetch image reference: status %d", response.StatusCode)
|
||||
}
|
||||
payload, err := io.ReadAll(io.LimitReader(response.Body, maxImageReferenceBytes+1))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(payload) == 0 {
|
||||
return nil, errors.New("image reference is empty")
|
||||
}
|
||||
if len(payload) > maxImageReferenceBytes {
|
||||
return nil, errors.New("image reference exceeds 32 MiB")
|
||||
}
|
||||
return payload, nil
|
||||
}
|
||||
|
||||
func decodeImageDataURL(raw string) ([]byte, error) {
|
||||
header, encoded, ok := strings.Cut(raw, ",")
|
||||
if !ok || len(header) <= len("data:") || encoded == "" {
|
||||
return nil, errors.New("image data URL is malformed")
|
||||
}
|
||||
metadata := strings.Split(header[len("data:"):], ";")
|
||||
if !strings.HasPrefix(strings.ToLower(strings.TrimSpace(metadata[0])), "image/") {
|
||||
return nil, errors.New("image data URL must declare an image media type")
|
||||
}
|
||||
base64Encoded := false
|
||||
for _, item := range metadata[1:] {
|
||||
if strings.EqualFold(strings.TrimSpace(item), "base64") {
|
||||
base64Encoded = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !base64Encoded {
|
||||
return nil, errors.New("image data URL must use base64 encoding")
|
||||
}
|
||||
if base64.StdEncoding.DecodedLen(len(encoded)) > maxImageReferenceBytes {
|
||||
return nil, errors.New("image reference exceeds 32 MiB")
|
||||
}
|
||||
payload, err := base64.StdEncoding.DecodeString(encoded)
|
||||
if err != nil {
|
||||
return nil, errors.New("image data URL contains invalid base64")
|
||||
}
|
||||
if len(payload) == 0 {
|
||||
return nil, errors.New("image reference is empty")
|
||||
}
|
||||
return payload, nil
|
||||
}
|
||||
|
||||
func validateGeminiImageRequest(body map[string]any) (int, int, error) {
|
||||
generationConfig, _ := body["generationConfig"].(map[string]any)
|
||||
modalities, _ := generationConfig["responseModalities"].([]any)
|
||||
hasImageModality := false
|
||||
for _, modality := range modalities {
|
||||
if strings.EqualFold(strings.TrimSpace(stringValue(modality)), "IMAGE") {
|
||||
hasImageModality = true
|
||||
}
|
||||
}
|
||||
if !hasImageModality {
|
||||
return 0, 0, errors.New("generationConfig.responseModalities must contain IMAGE")
|
||||
}
|
||||
total := 0
|
||||
images := 0
|
||||
contents, _ := body["contents"].([]any)
|
||||
for _, rawContent := range contents {
|
||||
content, _ := rawContent.(map[string]any)
|
||||
parts, _ := content["parts"].([]any)
|
||||
for _, rawPart := range parts {
|
||||
part, _ := rawPart.(map[string]any)
|
||||
inline, _ := part["inlineData"].(map[string]any)
|
||||
if inline == nil {
|
||||
inline, _ = part["inline_data"].(map[string]any)
|
||||
}
|
||||
encoded := strings.TrimSpace(stringValue(inline["data"]))
|
||||
if encoded == "" {
|
||||
fileData, _ := part["fileData"].(map[string]any)
|
||||
if fileData == nil {
|
||||
fileData, _ = part["file_data"].(map[string]any)
|
||||
}
|
||||
fileURI := strings.TrimSpace(stringValue(fileData["fileUri"]))
|
||||
if fileURI == "" {
|
||||
fileURI = strings.TrimSpace(stringValue(fileData["file_uri"]))
|
||||
}
|
||||
if fileURI == "" {
|
||||
fileURI = strings.TrimSpace(stringValue(fileData["uri"]))
|
||||
}
|
||||
if fileURI == "" {
|
||||
continue
|
||||
}
|
||||
if !strings.HasPrefix(fileURI, "http://") && !strings.HasPrefix(fileURI, "https://") {
|
||||
return 0, 0, errors.New("Gemini fileData URI must use HTTP or HTTPS")
|
||||
}
|
||||
images++
|
||||
continue
|
||||
}
|
||||
decoded, err := base64.StdEncoding.DecodeString(encoded)
|
||||
if err != nil {
|
||||
return 0, 0, fmt.Errorf("Gemini inlineData is not valid Base64: %w", err)
|
||||
}
|
||||
if len(decoded) == 0 {
|
||||
return 0, 0, errors.New("Gemini inlineData is empty")
|
||||
}
|
||||
total += len(decoded)
|
||||
images++
|
||||
}
|
||||
}
|
||||
if images == 0 {
|
||||
return 0, 0, errors.New("Gemini request has no inlineData or fileData image")
|
||||
}
|
||||
return total, images, nil
|
||||
}
|
||||
|
||||
func (s *Server) recordGeminiInvalid() {
|
||||
s.mu.Lock()
|
||||
s.report.GeminiInvalid++
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
func (s *Server) recordVideoInvalid() {
|
||||
s.mu.Lock()
|
||||
s.report.VideoInvalid++
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
func decodeJSON(r *http.Request, target any) error {
|
||||
reader := io.LimitReader(r.Body, maxProtocolBodyBytes+1)
|
||||
payload, err := io.ReadAll(reader)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if len(payload) > maxProtocolBodyBytes {
|
||||
return errors.New("protocol request body is too large")
|
||||
}
|
||||
decoder := json.NewDecoder(bytes.NewReader(payload))
|
||||
decoder.UseNumber()
|
||||
return decoder.Decode(target)
|
||||
}
|
||||
|
||||
func writeJSON(w http.ResponseWriter, status int, value any) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(status)
|
||||
_ = json.NewEncoder(w).Encode(value)
|
||||
}
|
||||
|
||||
func writeProtocolError(w http.ResponseWriter, status int, message string) {
|
||||
writeJSON(w, status, map[string]any{"error": map[string]any{
|
||||
"code": "acceptance_protocol_invalid", "message": message,
|
||||
}})
|
||||
}
|
||||
|
||||
func wait(ctx context.Context, delay time.Duration) error {
|
||||
timer := time.NewTimer(delay)
|
||||
defer timer.Stop()
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case <-timer.C:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func scaledDelay(minimum time.Duration, maximum time.Duration, seed int) time.Duration {
|
||||
if maximum <= minimum {
|
||||
return minimum
|
||||
}
|
||||
windowSeconds := int64((maximum - minimum) / time.Second)
|
||||
if windowSeconds <= 0 {
|
||||
return minimum
|
||||
}
|
||||
return minimum + time.Duration(int64(seed)%(windowSeconds+1))*time.Second
|
||||
}
|
||||
|
||||
func videoDelay(body map[string]any, longRun bool) time.Duration {
|
||||
seed := int64(0)
|
||||
switch value := body["seed"].(type) {
|
||||
case json.Number:
|
||||
seed, _ = value.Int64()
|
||||
case float64:
|
||||
seed = int64(value)
|
||||
case int:
|
||||
seed = int64(value)
|
||||
}
|
||||
if seed < 0 {
|
||||
seed = -seed
|
||||
}
|
||||
if longRun {
|
||||
return 2*time.Minute + time.Duration(seed%61)*time.Second
|
||||
}
|
||||
if videoPromptContains(body, "acceptance-provider-quota") {
|
||||
return 5*time.Second + time.Duration(seed%6)*time.Second
|
||||
}
|
||||
if videoPromptContains(body, "acceptance-capacity-ladder") {
|
||||
return 30*time.Second + time.Duration(seed%16)*time.Second
|
||||
}
|
||||
return 5*time.Second + time.Duration(seed%11)*time.Second
|
||||
}
|
||||
|
||||
func videoUsageTokens(body map[string]any) int {
|
||||
if videoPromptContains(body, "acceptance-provider-quota") {
|
||||
return 7
|
||||
}
|
||||
return 1
|
||||
}
|
||||
|
||||
func videoPromptContains(body map[string]any, marker string) bool {
|
||||
content, _ := body["content"].([]any)
|
||||
for _, raw := range content {
|
||||
item, _ := raw.(map[string]any)
|
||||
if strings.TrimSpace(stringValue(item["type"])) == "text" &&
|
||||
strings.Contains(strings.ToLower(stringValue(item["text"])), marker) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func requestOrigin(r *http.Request) string {
|
||||
scheme := strings.TrimSpace(r.Header.Get("X-Forwarded-Proto"))
|
||||
if scheme == "" {
|
||||
scheme = "http"
|
||||
}
|
||||
host := strings.TrimSpace(r.Host)
|
||||
return scheme + "://" + host
|
||||
}
|
||||
|
||||
func stringValue(value any) string {
|
||||
switch typed := value.(type) {
|
||||
case nil:
|
||||
return ""
|
||||
case string:
|
||||
return typed
|
||||
case json.Number:
|
||||
return typed.String()
|
||||
default:
|
||||
return fmt.Sprint(value)
|
||||
}
|
||||
}
|
||||
|
||||
func paddedPNG(size int) []byte {
|
||||
base, _ := base64.StdEncoding.DecodeString("iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII=")
|
||||
if size <= len(base) {
|
||||
return base
|
||||
}
|
||||
iendOffset := len(base) - 12
|
||||
paddingLength := size - len(base) - 12
|
||||
if paddingLength < 0 {
|
||||
paddingLength = 0
|
||||
}
|
||||
chunk := make([]byte, 12+paddingLength)
|
||||
binary.BigEndian.PutUint32(chunk[:4], uint32(paddingLength))
|
||||
copy(chunk[4:8], []byte("teST"))
|
||||
binary.BigEndian.PutUint32(chunk[8+paddingLength:], crc32.ChecksumIEEE(chunk[4:8+paddingLength]))
|
||||
out := make([]byte, 0, len(base)+len(chunk))
|
||||
out = append(out, base[:iendOffset]...)
|
||||
out = append(out, chunk...)
|
||||
out = append(out, base[iendOffset:]...)
|
||||
return out
|
||||
}
|
||||
|
||||
func buildFixtures() map[string]fixture {
|
||||
fixtures := map[string]fixture{}
|
||||
for index := 0; index < 4; index++ {
|
||||
imageValue := patternedImage(768, 512, index+1)
|
||||
var encoded bytes.Buffer
|
||||
_ = png.Encode(&encoded, imageValue)
|
||||
fixtures[fmt.Sprintf("image-%02d.png", index)] = fixture{ContentType: "image/png", Payload: encoded.Bytes()}
|
||||
}
|
||||
for index := 0; index < 4; index++ {
|
||||
imageValue := patternedImage(1024, 768, index+11)
|
||||
var encoded bytes.Buffer
|
||||
_ = jpeg.Encode(&encoded, imageValue, &jpeg.Options{Quality: 88})
|
||||
fixtures[fmt.Sprintf("image-%02d.jpg", index+4)] = fixture{ContentType: "image/jpeg", Payload: encoded.Bytes()}
|
||||
}
|
||||
webpFixtures := []string{
|
||||
"524946465a00000057454250565038204e000000f003009d012a400040003e9148a04c25a42322220800b012096900d3ca8000103b93c116da67710000feeea63fff80dd7c5b4cbfff7381ff7381ff7381fc6d4326178e9af219737c3e2990000000",
|
||||
"5249464650000000574542505650382044000000d003009d012a400040003e9148a04c25a42322220800b012096900760000206ea6a00af10b720000feeed3dfffc5b9fb75d21ffff8b43ad0eb43447bcc63798400000000",
|
||||
"52494646500000005745425056503820440000009003009d012a400040003e9148a04c25a42322220800b0120969000010685716c42e300a2000fef558bffffb9c0fffd9c0fffd9c0fe3afff5ea96b93a3f9696c1c000000",
|
||||
"524946464e0000005745425056503820420000009003009d012a400040003e9148a04c25a42322220800b0120969000010375350057885b90000fef4351ffffba7afffda7afffda7afe60fff239de951800000000000",
|
||||
}
|
||||
for index, encoded := range webpFixtures {
|
||||
payload, _ := hex.DecodeString(encoded)
|
||||
fixtures[fmt.Sprintf("image-%02d.webp", index+8)] = fixture{ContentType: "image/webp", Payload: payload}
|
||||
}
|
||||
for index := 0; index < 4; index++ {
|
||||
imageValue := patternedImage(6144, 2160, index+21)
|
||||
var encoded bytes.Buffer
|
||||
_ = jpeg.Encode(&encoded, imageValue, &jpeg.Options{Quality: 90})
|
||||
fixtures[fmt.Sprintf("image-%02d-oversized.jpg", index+12)] = fixture{ContentType: "image/jpeg", Payload: encoded.Bytes()}
|
||||
}
|
||||
return fixtures
|
||||
}
|
||||
|
||||
func patternedImage(width int, height int, seed int) image.Image {
|
||||
out := image.NewNRGBA(image.Rect(0, 0, width, height))
|
||||
for y := 0; y < height; y++ {
|
||||
for x := 0; x < width; x++ {
|
||||
out.SetNRGBA(x, y, color.NRGBA{
|
||||
R: uint8((x + seed*31) % 256),
|
||||
G: uint8((y + seed*47) % 256),
|
||||
B: uint8((x/8 + y/8 + seed*59) % 256),
|
||||
A: 255,
|
||||
})
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func payloadSHA256(payload []byte) string {
|
||||
sum := sha256.Sum256(payload)
|
||||
return hex.EncodeToString(sum[:])
|
||||
}
|
||||
|
||||
var minimalMP4 = []byte{
|
||||
0, 0, 0, 24, 'f', 't', 'y', 'p', 'i', 's', 'o', 'm',
|
||||
0, 0, 0, 0, 'i', 's', 'o', 'm', 'm', 'p', '4', '2',
|
||||
0, 0, 0, 8, 'm', 'd', 'a', 't',
|
||||
}
|
||||
@@ -0,0 +1,406 @@
|
||||
package acceptanceemulator
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"image"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestGeminiAndVolcesProtocolEmulation(t *testing.T) {
|
||||
var mu sync.Mutex
|
||||
now := time.Unix(1_800_000_000, 0)
|
||||
server := New(Config{
|
||||
Now: func() time.Time {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
return now
|
||||
},
|
||||
Wait: func(context.Context, time.Duration) error { return nil },
|
||||
})
|
||||
httpServer := httptest.NewServer(server.Handler())
|
||||
defer httpServer.Close()
|
||||
|
||||
input := paddedPNG(256 << 10)
|
||||
geminiBody, _ := json.Marshal(map[string]any{
|
||||
"contents": []any{map[string]any{"parts": []any{
|
||||
map[string]any{"text": "edit image"},
|
||||
map[string]any{"inlineData": map[string]any{
|
||||
"mimeType": "image/png", "data": base64.StdEncoding.EncodeToString(input),
|
||||
}},
|
||||
}}},
|
||||
"generationConfig": map[string]any{"responseModalities": []any{"IMAGE"}},
|
||||
})
|
||||
response, err := postIdempotent(
|
||||
httpServer.URL+"/v1beta/models/gemini-test:generateContent",
|
||||
geminiBody,
|
||||
"gemini-task",
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("Gemini request: %v", err)
|
||||
}
|
||||
if response.StatusCode != http.StatusOK {
|
||||
payload, _ := io.ReadAll(response.Body)
|
||||
t.Fatalf("Gemini status=%d body=%s", response.StatusCode, payload)
|
||||
}
|
||||
var geminiResult map[string]any
|
||||
if err := json.NewDecoder(response.Body).Decode(&geminiResult); err != nil {
|
||||
t.Fatalf("decode Gemini response: %v", err)
|
||||
}
|
||||
_ = response.Body.Close()
|
||||
|
||||
fixtureResponse, err := http.Get(httpServer.URL + "/fixtures/image-12-oversized.jpg")
|
||||
if err != nil {
|
||||
t.Fatalf("get oversized fixture: %v", err)
|
||||
}
|
||||
fixturePayload, readErr := io.ReadAll(fixtureResponse.Body)
|
||||
_ = fixtureResponse.Body.Close()
|
||||
if readErr != nil {
|
||||
t.Fatalf("read oversized fixture: %v", readErr)
|
||||
}
|
||||
config, format, err := image.DecodeConfig(bytes.NewReader(fixturePayload))
|
||||
if err != nil || format != "jpeg" || config.Width != 6144 || config.Height != 2160 {
|
||||
t.Fatalf("oversized fixture format=%s config=%+v err=%v", format, config, err)
|
||||
}
|
||||
if _, format, err = image.Decode(bytes.NewReader(fixturePayload)); err != nil || format != "jpeg" {
|
||||
t.Fatalf("fully decode oversized fixture format=%s err=%v", format, err)
|
||||
}
|
||||
|
||||
content := []any{map[string]any{"type": "text", "text": "video"}}
|
||||
for index, name := range []string{"image-00.png", "image-04.jpg", "image-08.webp"} {
|
||||
role := "reference_image"
|
||||
if index == 0 {
|
||||
role = "first_frame"
|
||||
}
|
||||
content = append(content, map[string]any{
|
||||
"type": "image_url", "role": role,
|
||||
"image_url": map[string]any{"url": httpServer.URL + "/fixtures/" + name},
|
||||
})
|
||||
}
|
||||
videoBody, _ := json.Marshal(map[string]any{"model": "seedance-test", "content": content, "seed": 1})
|
||||
response, err = postIdempotent(
|
||||
httpServer.URL+"/contents/generations/tasks",
|
||||
videoBody,
|
||||
"video-task",
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("submit video: %v", err)
|
||||
}
|
||||
if response.StatusCode != http.StatusOK {
|
||||
t.Fatalf("submit video status=%d", response.StatusCode)
|
||||
}
|
||||
var submitted map[string]any
|
||||
_ = json.NewDecoder(response.Body).Decode(&submitted)
|
||||
_ = response.Body.Close()
|
||||
taskID := strings.TrimSpace(stringValue(submitted["id"]))
|
||||
if taskID == "" {
|
||||
t.Fatal("video emulator returned no task ID")
|
||||
}
|
||||
mu.Lock()
|
||||
now = now.Add(20 * time.Second)
|
||||
mu.Unlock()
|
||||
response, err = http.Get(httpServer.URL + "/contents/generations/tasks/" + taskID)
|
||||
if err != nil {
|
||||
t.Fatalf("poll video: %v", err)
|
||||
}
|
||||
var polled map[string]any
|
||||
_ = json.NewDecoder(response.Body).Decode(&polled)
|
||||
_ = response.Body.Close()
|
||||
if polled["status"] != "succeeded" {
|
||||
t.Fatalf("video status=%v", polled["status"])
|
||||
}
|
||||
|
||||
response, err = http.Get(httpServer.URL + "/report")
|
||||
if err != nil {
|
||||
t.Fatalf("get report: %v", err)
|
||||
}
|
||||
var report Report
|
||||
_ = json.NewDecoder(response.Body).Decode(&report)
|
||||
_ = response.Body.Close()
|
||||
if report.GeminiRequests != 1 || report.VideoSubmissions != 1 || report.VideoInvalid != 0 ||
|
||||
report.VideoReferenceCounts["3"] != 1 || report.UniqueImageHashes != 3 {
|
||||
t.Fatalf("unexpected emulator report: %+v", report)
|
||||
}
|
||||
}
|
||||
|
||||
func TestForcedConversionRejectsImagesOutsideOfficialSeedanceRange(t *testing.T) {
|
||||
server := New(Config{Wait: func(context.Context, time.Duration) error { return nil }})
|
||||
httpServer := httptest.NewServer(server.Handler())
|
||||
defer httpServer.Close()
|
||||
|
||||
content := []any{map[string]any{"type": "text", "text": "acceptance-force-conversion"}}
|
||||
for _, name := range []string{"image-12-oversized.jpg", "image-00.png", "image-04.jpg"} {
|
||||
content = append(content, map[string]any{
|
||||
"type": "image_url", "role": "reference_image",
|
||||
"image_url": map[string]any{"url": httpServer.URL + "/fixtures/" + name},
|
||||
})
|
||||
}
|
||||
body, _ := json.Marshal(map[string]any{"model": "seedance-test", "content": content})
|
||||
response, err := postIdempotent(httpServer.URL+"/contents/generations/tasks", body, "oversized")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer response.Body.Close()
|
||||
if response.StatusCode != http.StatusBadRequest {
|
||||
t.Fatalf("status=%d, want %d", response.StatusCode, http.StatusBadRequest)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiProtocolAcceptsMultipleFileDataImages(t *testing.T) {
|
||||
server := New(Config{Wait: func(context.Context, time.Duration) error { return nil }})
|
||||
httpServer := httptest.NewServer(server.Handler())
|
||||
defer httpServer.Close()
|
||||
parts := []any{map[string]any{"text": "combine references"}}
|
||||
for index := 0; index < 3; index++ {
|
||||
parts = append(parts, map[string]any{"fileData": map[string]any{
|
||||
"mimeType": "image/png",
|
||||
"fileUri": fmt.Sprintf("https://fixtures.example/%d.png", index),
|
||||
}})
|
||||
}
|
||||
body, _ := json.Marshal(map[string]any{
|
||||
"contents": []any{map[string]any{"parts": parts}},
|
||||
"generationConfig": map[string]any{"responseModalities": []any{"IMAGE"}},
|
||||
})
|
||||
response, err := postIdempotent(httpServer.URL+"/v1beta/models/gemini-test:generateContent", body, "gemini-multi")
|
||||
if err != nil {
|
||||
t.Fatalf("Gemini multi-image request: %v", err)
|
||||
}
|
||||
defer response.Body.Close()
|
||||
if response.StatusCode != http.StatusOK {
|
||||
payload, _ := io.ReadAll(response.Body)
|
||||
t.Fatalf("Gemini multi-image status=%d body=%s", response.StatusCode, payload)
|
||||
}
|
||||
reportResponse, err := http.Get(httpServer.URL + "/report")
|
||||
if err != nil {
|
||||
t.Fatalf("get report: %v", err)
|
||||
}
|
||||
defer reportResponse.Body.Close()
|
||||
var report Report
|
||||
if err := json.NewDecoder(reportResponse.Body).Decode(&report); err != nil {
|
||||
t.Fatalf("decode report: %v", err)
|
||||
}
|
||||
if report.GeminiRequests != 1 || report.GeminiInputImages != 3 || report.GeminiInvalid != 0 {
|
||||
t.Fatalf("unexpected multi-image report: %+v", report)
|
||||
}
|
||||
}
|
||||
|
||||
func TestImageFixturesFullyDecode(t *testing.T) {
|
||||
for name, fixture := range buildFixtures() {
|
||||
if !strings.HasPrefix(fixture.ContentType, "image/") {
|
||||
continue
|
||||
}
|
||||
decoded, format, err := image.Decode(bytes.NewReader(fixture.Payload))
|
||||
if err != nil {
|
||||
t.Fatalf("fixture %s does not fully decode: %v", name, err)
|
||||
}
|
||||
bounds := decoded.Bounds()
|
||||
if format == "" || bounds.Dx() <= 0 || bounds.Dy() <= 0 {
|
||||
t.Fatalf("fixture %s format=%q bounds=%v", name, format, bounds)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPaddedPNGIsValidAndExact(t *testing.T) {
|
||||
for _, size := range []int{256 << 10, 4 << 20, 8 << 20} {
|
||||
payload := paddedPNG(size)
|
||||
if len(payload) != size {
|
||||
t.Fatalf("PNG bytes=%d, want=%d", len(payload), size)
|
||||
}
|
||||
if _, _, err := image.DecodeConfig(bytes.NewReader(payload)); err != nil {
|
||||
t.Fatalf("decode padded PNG: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestScaledDelayStaysOnWholeSecondsWithinProfile(t *testing.T) {
|
||||
delay := scaledDelay(8*time.Second, 15*time.Second, 2<<20)
|
||||
if delay < 8*time.Second || delay > 15*time.Second || delay%time.Second != 0 {
|
||||
t.Fatalf("scaled delay=%s", delay)
|
||||
}
|
||||
}
|
||||
|
||||
func TestVideoCapacityDelay(t *testing.T) {
|
||||
body := map[string]any{
|
||||
"seed": json.Number("7"),
|
||||
"content": []any{map[string]any{
|
||||
"type": "text",
|
||||
"text": "capacity acceptance-capacity-ladder",
|
||||
}},
|
||||
}
|
||||
delay := videoDelay(body, false)
|
||||
if delay < 30*time.Second || delay > 45*time.Second {
|
||||
t.Fatalf("video capacity delay=%s, want 30s..45s", delay)
|
||||
}
|
||||
}
|
||||
|
||||
func TestVideoProviderQuotaDelayAndUsage(t *testing.T) {
|
||||
body := map[string]any{
|
||||
"seed": json.Number("5"),
|
||||
"content": []any{map[string]any{
|
||||
"type": "text",
|
||||
"text": "acceptance-provider-quota",
|
||||
}},
|
||||
}
|
||||
delay := videoDelay(body, false)
|
||||
if delay < 5*time.Second || delay > 10*time.Second {
|
||||
t.Fatalf("video provider quota delay=%s, want 5s..10s", delay)
|
||||
}
|
||||
if tokens := videoUsageTokens(body); tokens != 7 {
|
||||
t.Fatalf("video provider quota usage tokens=%d, want 7", tokens)
|
||||
}
|
||||
}
|
||||
|
||||
func TestVolcesProtocolAcceptsThreeSixAndNineReferenceImages(t *testing.T) {
|
||||
server := New(Config{Wait: func(context.Context, time.Duration) error { return nil }})
|
||||
httpServer := httptest.NewServer(server.Handler())
|
||||
defer httpServer.Close()
|
||||
dataImage := "data:image/png;base64," + base64.StdEncoding.EncodeToString(paddedPNG(0))
|
||||
|
||||
for _, imageCount := range []int{3, 6, 9} {
|
||||
content := []any{map[string]any{"type": "text", "text": "multi-reference video"}}
|
||||
for index := 0; index < imageCount; index++ {
|
||||
role := "reference_image"
|
||||
if index == 0 {
|
||||
role = "first_frame"
|
||||
}
|
||||
if imageCount > 3 && index == imageCount-1 {
|
||||
role = "last_frame"
|
||||
}
|
||||
imageURL := fmt.Sprintf("%s/fixtures/image-%02d.png", httpServer.URL, index%4)
|
||||
if index%2 == 0 {
|
||||
imageURL = dataImage
|
||||
}
|
||||
content = append(content, map[string]any{
|
||||
"type": "image_url", "role": role,
|
||||
"image_url": map[string]any{"url": imageURL},
|
||||
})
|
||||
}
|
||||
body, _ := json.Marshal(map[string]any{"model": "seedance-test", "content": content, "seed": imageCount})
|
||||
response, err := postIdempotent(
|
||||
httpServer.URL+"/contents/generations/tasks",
|
||||
body,
|
||||
fmt.Sprintf("video-%d", imageCount),
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("%d-image submit: %v", imageCount, err)
|
||||
}
|
||||
_ = response.Body.Close()
|
||||
if response.StatusCode != http.StatusOK {
|
||||
t.Fatalf("%d-image status=%d", imageCount, response.StatusCode)
|
||||
}
|
||||
}
|
||||
|
||||
response, err := http.Get(httpServer.URL + "/report")
|
||||
if err != nil {
|
||||
t.Fatalf("get report: %v", err)
|
||||
}
|
||||
defer response.Body.Close()
|
||||
var report Report
|
||||
if err := json.NewDecoder(response.Body).Decode(&report); err != nil {
|
||||
t.Fatalf("decode report: %v", err)
|
||||
}
|
||||
for _, imageCount := range []string{"3", "6", "9"} {
|
||||
if report.VideoReferenceCounts[imageCount] != 1 {
|
||||
t.Fatalf("reference counts=%v", report.VideoReferenceCounts)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestObjectStorageEmulationSupportsRetryLifecycleAndFaults(t *testing.T) {
|
||||
server := httptest.NewServer(New(Config{}).Handler())
|
||||
defer server.Close()
|
||||
objectURL := server.URL + "/storage/s3-transient/bucket/media/result.png"
|
||||
|
||||
request, _ := http.NewRequest(http.MethodPut, objectURL, strings.NewReader("image-bytes"))
|
||||
request.Header.Set("Content-Type", "image/png")
|
||||
response, err := http.DefaultClient.Do(request)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_ = response.Body.Close()
|
||||
if response.StatusCode != http.StatusServiceUnavailable {
|
||||
t.Fatalf("first transient PUT status=%d", response.StatusCode)
|
||||
}
|
||||
request, _ = http.NewRequest(http.MethodPut, objectURL, strings.NewReader("image-bytes"))
|
||||
request.Header.Set("Content-Type", "image/png")
|
||||
response, err = http.DefaultClient.Do(request)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_ = response.Body.Close()
|
||||
if response.StatusCode != http.StatusOK {
|
||||
t.Fatalf("retried PUT status=%d", response.StatusCode)
|
||||
}
|
||||
|
||||
response, err = http.Get(objectURL)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
payload, _ := io.ReadAll(response.Body)
|
||||
_ = response.Body.Close()
|
||||
if response.StatusCode != http.StatusOK || string(payload) != "image-bytes" {
|
||||
t.Fatalf("GET status=%d payload=%q", response.StatusCode, payload)
|
||||
}
|
||||
request, _ = http.NewRequest(http.MethodHead, objectURL, nil)
|
||||
response, err = http.DefaultClient.Do(request)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_ = response.Body.Close()
|
||||
if response.StatusCode != http.StatusOK {
|
||||
t.Fatalf("HEAD status=%d", response.StatusCode)
|
||||
}
|
||||
request, _ = http.NewRequest(http.MethodDelete, objectURL, nil)
|
||||
response, err = http.DefaultClient.Do(request)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_ = response.Body.Close()
|
||||
if response.StatusCode != http.StatusNoContent {
|
||||
t.Fatalf("DELETE status=%d", response.StatusCode)
|
||||
}
|
||||
|
||||
for profile, wantStatus := range map[string]int{"oss-auth": http.StatusForbidden, "s3-fail": http.StatusServiceUnavailable} {
|
||||
request, _ = http.NewRequest(http.MethodPut, server.URL+"/storage/"+profile+"/probe.bin", strings.NewReader("probe"))
|
||||
response, err = http.DefaultClient.Do(request)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_ = response.Body.Close()
|
||||
if response.StatusCode != wantStatus {
|
||||
t.Fatalf("%s status=%d, want %d", profile, response.StatusCode, wantStatus)
|
||||
}
|
||||
}
|
||||
|
||||
response, err = http.Get(server.URL + "/report")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer response.Body.Close()
|
||||
var report Report
|
||||
if err := json.NewDecoder(response.Body).Decode(&report); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if report.StoragePuts != 1 || report.StorageGets != 1 || report.StorageHeads != 1 || report.StorageDeletes != 1 || report.StorageFailures != 3 {
|
||||
t.Fatalf("unexpected storage report: %+v", report)
|
||||
}
|
||||
}
|
||||
|
||||
func postIdempotent(url string, body []byte, key string) (*http.Response, error) {
|
||||
request, err := http.NewRequest(http.MethodPost, url, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
request.Header.Set("Idempotency-Key", key)
|
||||
return http.DefaultClient.Do(request)
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,180 @@
|
||||
package acceptancesnapshot
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestSnapshotValidationRejectsSecretsAndDetectsTampering(t *testing.T) {
|
||||
snapshot := validSnapshot(t)
|
||||
if err := Validate(snapshot); err != nil {
|
||||
t.Fatalf("validate fixture: %v", err)
|
||||
}
|
||||
|
||||
tampered := snapshot
|
||||
tampered.Candidates = append([]Candidate(nil), snapshot.Candidates...)
|
||||
tampered.Candidates[0].Platform.Name = "tampered"
|
||||
if err := Validate(tampered); err == nil {
|
||||
t.Fatal("tampered snapshot passed validation")
|
||||
}
|
||||
|
||||
unsafe := snapshot
|
||||
unsafe.Candidates = append([]Candidate(nil), snapshot.Candidates...)
|
||||
unsafe.Candidates[0].Platform.Config = map[string]any{"apiToken": "must-not-leak"}
|
||||
refreshHashes(t, &unsafe)
|
||||
if err := Validate(unsafe); err == nil {
|
||||
t.Fatal("secret-like snapshot field passed validation")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSanitizeMapRemovesCredentialAndProxyFields(t *testing.T) {
|
||||
input := map[string]any{
|
||||
"specType": "gemini",
|
||||
"credentialEnv": "PRODUCTION_TOKEN",
|
||||
"networkProxy": map[string]any{
|
||||
"proxyUrl": "http://user:password@example.invalid",
|
||||
},
|
||||
"requestAssetImageURLFormat": "base64",
|
||||
}
|
||||
got := sanitizeMap(input)
|
||||
if got["specType"] != "gemini" || got["requestAssetImageURLFormat"] != "base64" {
|
||||
t.Fatalf("safe protocol fields were removed: %#v", got)
|
||||
}
|
||||
if _, ok := got["credentialEnv"]; ok {
|
||||
t.Fatal("credentialEnv was not removed")
|
||||
}
|
||||
nested, _ := got["networkProxy"].(map[string]any)
|
||||
if _, ok := nested["proxyUrl"]; ok {
|
||||
t.Fatal("nested proxy URL was not removed")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDecodeRejectsUnknownSnapshotFields(t *testing.T) {
|
||||
snapshot := validSnapshot(t)
|
||||
payload, err := json.Marshal(snapshot)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var object map[string]any
|
||||
if err := json.Unmarshal(payload, &object); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
object["unexpected"] = true
|
||||
payload, _ = json.Marshal(object)
|
||||
if _, err := Decode(payload); err == nil {
|
||||
t.Fatal("snapshot with unknown field decoded successfully")
|
||||
}
|
||||
}
|
||||
|
||||
func validSnapshot(t *testing.T) Snapshot {
|
||||
t.Helper()
|
||||
candidates := []Candidate{
|
||||
{
|
||||
Workload: "gemini_image_edit",
|
||||
Provider: Provider{
|
||||
ProviderKey: "gemini", ProviderCode: "google-gemini",
|
||||
DisplayName: "Gemini", ProviderType: "gemini",
|
||||
CapabilitySchema: map[string]any{}, DefaultRateLimitPolicy: map[string]any{},
|
||||
},
|
||||
BaseModel: BaseModel{
|
||||
CanonicalModelKey: "gemini:image", InvocationName: "gemini-image",
|
||||
ProviderModelName: "gemini-image", ModelType: []string{"image_edit"},
|
||||
DisplayName: "Gemini Image", Capabilities: map[string]any{},
|
||||
BaseBillingConfig: map[string]any{}, DefaultRateLimitPolicy: map[string]any{},
|
||||
RuntimePolicyOverride: map[string]any{}, PricingVersion: 1,
|
||||
},
|
||||
Platform: Platform{
|
||||
Provider: "gemini", PlatformKey: "gemini-production", Name: "Gemini Production",
|
||||
BaseURLPath: "/", AuthType: "api_key", Config: map[string]any{"specType": "gemini"},
|
||||
DefaultPricingMode: "inherit_discount", DefaultDiscountFactor: "1",
|
||||
RetryPolicy: map[string]any{}, RateLimitPolicy: map[string]any{}, Priority: 10,
|
||||
},
|
||||
PlatformModel: PlatformModel{
|
||||
ModelName: "gemini-image", ProviderModelName: "gemini-image",
|
||||
ModelType: []string{"image_edit"}, DisplayName: "Gemini Image",
|
||||
CapabilityOverride: map[string]any{}, Capabilities: map[string]any{},
|
||||
PricingMode: "inherit_discount", BillingConfigOverride: map[string]any{},
|
||||
BillingConfig: map[string]any{}, PermissionConfig: map[string]any{},
|
||||
RetryPolicy: map[string]any{}, RateLimitPolicy: map[string]any{},
|
||||
RuntimePolicyOverride: map[string]any{},
|
||||
},
|
||||
RuntimePolicy: RuntimePolicy{
|
||||
RateLimitPolicy: map[string]any{}, RetryPolicy: map[string]any{},
|
||||
AutoDisablePolicy: map[string]any{}, DegradePolicy: map[string]any{},
|
||||
},
|
||||
PricingRules: []PricingRule{},
|
||||
Metadata: CandidateSource{
|
||||
PlatformID: "00000000-0000-0000-0000-000000000001",
|
||||
PlatformModelID: "00000000-0000-0000-0000-000000000002",
|
||||
BaseModelID: "00000000-0000-0000-0000-000000000003",
|
||||
},
|
||||
},
|
||||
{
|
||||
Workload: "multi_reference_video",
|
||||
Provider: Provider{
|
||||
ProviderKey: "volces", ProviderCode: "volces",
|
||||
DisplayName: "Volces", ProviderType: "volces",
|
||||
CapabilitySchema: map[string]any{}, DefaultRateLimitPolicy: map[string]any{},
|
||||
},
|
||||
BaseModel: BaseModel{
|
||||
CanonicalModelKey: "volces:seedance", InvocationName: "seedance-2-fast",
|
||||
ProviderModelName: "seedance-2-fast", ModelType: []string{"omni_video"},
|
||||
DisplayName: "Seedance", Capabilities: map[string]any{"omni_video": map[string]any{"max_images": 9}},
|
||||
BaseBillingConfig: map[string]any{}, DefaultRateLimitPolicy: map[string]any{},
|
||||
RuntimePolicyOverride: map[string]any{}, PricingVersion: 1,
|
||||
},
|
||||
Platform: Platform{
|
||||
Provider: "volces", PlatformKey: "volces-production", Name: "Volces Production",
|
||||
BaseURLPath: "/", AuthType: "bearer", Config: map[string]any{"specType": "volces"},
|
||||
DefaultPricingMode: "inherit_discount", DefaultDiscountFactor: "1",
|
||||
RetryPolicy: map[string]any{}, RateLimitPolicy: map[string]any{}, Priority: 10,
|
||||
},
|
||||
PlatformModel: PlatformModel{
|
||||
ModelName: "seedance-2-fast", ProviderModelName: "seedance-2-fast",
|
||||
ModelType: []string{"omni_video"}, DisplayName: "Seedance",
|
||||
CapabilityOverride: map[string]any{"omni_video": map[string]any{"max_images": 9}},
|
||||
Capabilities: map[string]any{}, PricingMode: "inherit_discount",
|
||||
BillingConfigOverride: map[string]any{}, BillingConfig: map[string]any{},
|
||||
PermissionConfig: map[string]any{}, RetryPolicy: map[string]any{},
|
||||
RateLimitPolicy: map[string]any{}, RuntimePolicyOverride: map[string]any{},
|
||||
},
|
||||
RuntimePolicy: RuntimePolicy{
|
||||
RateLimitPolicy: map[string]any{}, RetryPolicy: map[string]any{},
|
||||
AutoDisablePolicy: map[string]any{}, DegradePolicy: map[string]any{},
|
||||
},
|
||||
PricingRules: []PricingRule{},
|
||||
Metadata: CandidateSource{
|
||||
PlatformID: "00000000-0000-0000-0000-000000000004",
|
||||
PlatformModelID: "00000000-0000-0000-0000-000000000005",
|
||||
BaseModelID: "00000000-0000-0000-0000-000000000006",
|
||||
},
|
||||
},
|
||||
}
|
||||
snapshot := Snapshot{
|
||||
SchemaVersion: SchemaVersion,
|
||||
Source: Source{
|
||||
ReleaseSHA: "0123456789abcdef0123456789abcdef01234567",
|
||||
CreatedAt: time.Date(2026, 7, 31, 0, 0, 0, 0, time.UTC),
|
||||
},
|
||||
Candidates: candidates,
|
||||
SecretSafe: true,
|
||||
}
|
||||
refreshHashes(t, &snapshot)
|
||||
return snapshot
|
||||
}
|
||||
|
||||
func refreshHashes(t *testing.T, snapshot *Snapshot) {
|
||||
t.Helper()
|
||||
configHash, err := candidateConfigHash(snapshot.Candidates)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
snapshot.Source.ConfigHash = configHash
|
||||
snapshot.SnapshotSHA256 = ""
|
||||
hash, err := snapshotHash(*snapshot)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
snapshot.SnapshotSHA256 = hash
|
||||
}
|
||||
@@ -0,0 +1,59 @@
|
||||
package acceptanceworkload
|
||||
|
||||
import "time"
|
||||
|
||||
type GeminiProfile struct {
|
||||
Name string
|
||||
Requests int
|
||||
InputImages int
|
||||
InputBytes int
|
||||
OutputBytes int
|
||||
DelayMin time.Duration
|
||||
DelayMax time.Duration
|
||||
}
|
||||
|
||||
var (
|
||||
GeminiBaseline = GeminiProfile{
|
||||
Name: "gemini-baseline", Requests: 1000, InputImages: 1, InputBytes: 256 << 10, OutputBytes: 256 << 10,
|
||||
DelayMin: 4 * time.Second, DelayMax: 4 * time.Second,
|
||||
}
|
||||
GeminiMultiImage = GeminiProfile{
|
||||
Name: "gemini-multi-image", Requests: 192, InputImages: 3, InputBytes: 768 << 10, OutputBytes: 0,
|
||||
DelayMin: 8 * time.Second, DelayMax: 15 * time.Second,
|
||||
}
|
||||
GeminiLarge = GeminiProfile{
|
||||
Name: "gemini-large", Requests: 128, InputImages: 1, InputBytes: 2 << 20, OutputBytes: 4 << 20,
|
||||
DelayMin: 8 * time.Second, DelayMax: 15 * time.Second,
|
||||
}
|
||||
GeminiPeak = GeminiProfile{
|
||||
Name: "gemini-peak", Requests: 32, InputImages: 1, InputBytes: 8 << 20, OutputBytes: 8 << 20,
|
||||
DelayMin: 15 * time.Second, DelayMax: 30 * time.Second,
|
||||
}
|
||||
)
|
||||
|
||||
func GeminiProfileByName(name string) (GeminiProfile, bool) {
|
||||
switch name {
|
||||
case GeminiBaseline.Name:
|
||||
return GeminiBaseline, true
|
||||
case GeminiMultiImage.Name:
|
||||
return GeminiMultiImage, true
|
||||
case GeminiLarge.Name:
|
||||
return GeminiLarge, true
|
||||
case GeminiPeak.Name:
|
||||
return GeminiPeak, true
|
||||
default:
|
||||
return GeminiProfile{}, false
|
||||
}
|
||||
}
|
||||
|
||||
// VideoImageCount provides a deterministic 60/30/10 split for 3/6/9-image tasks.
|
||||
func VideoImageCount(index int) int {
|
||||
switch index % 10 {
|
||||
case 0, 1, 2, 3, 4, 5:
|
||||
return 3
|
||||
case 6, 7, 8:
|
||||
return 6
|
||||
default:
|
||||
return 9
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,26 @@
|
||||
package acceptanceworkload
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestGeminiProfilesAndVideoDistribution(t *testing.T) {
|
||||
for _, profile := range []GeminiProfile{GeminiBaseline, GeminiMultiImage, GeminiLarge, GeminiPeak} {
|
||||
resolved, ok := GeminiProfileByName(profile.Name)
|
||||
if !ok || resolved != profile {
|
||||
t.Fatalf("profile %q resolved to %+v, ok=%v", profile.Name, resolved, ok)
|
||||
}
|
||||
if profile.Requests <= 0 || profile.InputImages <= 0 || profile.InputBytes < profile.InputImages || profile.OutputBytes < 0 ||
|
||||
profile.DelayMin <= 0 || profile.DelayMax < profile.DelayMin {
|
||||
t.Fatalf("invalid Gemini profile: %+v", profile)
|
||||
}
|
||||
if profile.Name != GeminiMultiImage.Name && profile.OutputBytes == 0 {
|
||||
t.Fatalf("fixed-output Gemini profile has no output size: %+v", profile)
|
||||
}
|
||||
}
|
||||
counts := map[int]int{}
|
||||
for index := 0; index < 100; index++ {
|
||||
counts[VideoImageCount(index)]++
|
||||
}
|
||||
if counts[3] != 60 || counts[6] != 30 || counts[9] != 10 {
|
||||
t.Fatalf("video image distribution=%v", counts)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,323 @@
|
||||
package accessruleaudit
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
const (
|
||||
SchemaVersion = "access-rule-audit/v1"
|
||||
MigrationBatch = "0102_access_rule_allow_whitelist_semantics"
|
||||
)
|
||||
|
||||
type Snapshot struct {
|
||||
SchemaVersion string `json:"schemaVersion"`
|
||||
GeneratedAt time.Time `json:"generatedAt"`
|
||||
SecretSafe bool `json:"secretSafe"`
|
||||
Live RuleSet `json:"live"`
|
||||
Archive *ArchiveManifest `json:"archive,omitempty"`
|
||||
}
|
||||
|
||||
type RuleSet struct {
|
||||
Total int64 `json:"total"`
|
||||
SHA256 string `json:"sha256"`
|
||||
AllowCount int64 `json:"allowCount"`
|
||||
AllowSHA256 string `json:"allowSha256"`
|
||||
DenyCount int64 `json:"denyCount"`
|
||||
DenySHA256 string `json:"denySha256"`
|
||||
Counts []RuleCount `json:"counts"`
|
||||
}
|
||||
|
||||
type RuleCount struct {
|
||||
SubjectType string `json:"subjectType"`
|
||||
Effect string `json:"effect"`
|
||||
ResourceType string `json:"resourceType"`
|
||||
Status string `json:"status"`
|
||||
Count int64 `json:"count"`
|
||||
}
|
||||
|
||||
type ArchiveManifest struct {
|
||||
MigrationBatch string `json:"migrationBatch"`
|
||||
ManifestAllowCount int64 `json:"manifestAllowCount"`
|
||||
ManifestAllowSHA string `json:"manifestAllowSha256"`
|
||||
ArchivedAllowCount int64 `json:"archivedAllowCount"`
|
||||
ArchivedAllowSHA string `json:"archivedAllowSha256"`
|
||||
ManifestDenyCount int64 `json:"manifestDenyCount"`
|
||||
ManifestDenySHA string `json:"manifestDenySha256"`
|
||||
Consistent bool `json:"consistent"`
|
||||
}
|
||||
|
||||
type ruleDigestRow struct {
|
||||
id string
|
||||
subjectType string
|
||||
subjectID string
|
||||
resourceType string
|
||||
resourceID string
|
||||
effect string
|
||||
priority string
|
||||
minPermissionLevel string
|
||||
conditions string
|
||||
metadata string
|
||||
status string
|
||||
createdEpoch string
|
||||
updatedEpoch string
|
||||
}
|
||||
|
||||
func Export(ctx context.Context, pool *pgxpool.Pool, now time.Time) (Snapshot, error) {
|
||||
live, err := loadLiveRules(ctx, pool)
|
||||
if err != nil {
|
||||
return Snapshot{}, err
|
||||
}
|
||||
archive, err := loadArchiveManifest(ctx, pool)
|
||||
if err != nil {
|
||||
return Snapshot{}, err
|
||||
}
|
||||
if now.IsZero() {
|
||||
now = time.Now()
|
||||
}
|
||||
snapshot := Snapshot{
|
||||
SchemaVersion: SchemaVersion,
|
||||
GeneratedAt: now.UTC(),
|
||||
SecretSafe: true,
|
||||
Live: live,
|
||||
Archive: archive,
|
||||
}
|
||||
if err := Validate(snapshot); err != nil {
|
||||
return Snapshot{}, err
|
||||
}
|
||||
return snapshot, nil
|
||||
}
|
||||
|
||||
func Validate(snapshot Snapshot) error {
|
||||
if snapshot.SchemaVersion != SchemaVersion {
|
||||
return fmt.Errorf("unsupported access-rule audit schema %q", snapshot.SchemaVersion)
|
||||
}
|
||||
if !snapshot.SecretSafe {
|
||||
return errors.New("access-rule audit snapshot is not marked secret-safe")
|
||||
}
|
||||
for name, value := range map[string]string{
|
||||
"live": snapshot.Live.SHA256, "allow": snapshot.Live.AllowSHA256, "deny": snapshot.Live.DenySHA256,
|
||||
} {
|
||||
if !validSHA256(value) {
|
||||
return fmt.Errorf("%s SHA-256 is invalid", name)
|
||||
}
|
||||
}
|
||||
if snapshot.Live.Total != snapshot.Live.AllowCount+snapshot.Live.DenyCount {
|
||||
return errors.New("live access-rule counts are inconsistent")
|
||||
}
|
||||
var groupedTotal int64
|
||||
for _, count := range snapshot.Live.Counts {
|
||||
if count.SubjectType == "" || count.Effect == "" || count.ResourceType == "" || count.Status == "" || count.Count < 1 {
|
||||
return errors.New("access-rule grouped count is invalid")
|
||||
}
|
||||
groupedTotal += count.Count
|
||||
}
|
||||
if groupedTotal != snapshot.Live.Total {
|
||||
return errors.New("access-rule grouped counts do not match total")
|
||||
}
|
||||
if snapshot.Archive != nil {
|
||||
archive := snapshot.Archive
|
||||
if archive.MigrationBatch != MigrationBatch || !validSHA256(archive.ManifestAllowSHA) || !validSHA256(archive.ArchivedAllowSHA) || !validSHA256(archive.ManifestDenySHA) {
|
||||
return errors.New("access-rule archive manifest is invalid")
|
||||
}
|
||||
consistent := archive.ManifestAllowCount == archive.ArchivedAllowCount && archive.ManifestAllowSHA == archive.ArchivedAllowSHA
|
||||
if archive.Consistent != consistent {
|
||||
return errors.New("access-rule archive consistency marker is incorrect")
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func VerifyMigration(before Snapshot, after Snapshot) error {
|
||||
if err := Validate(before); err != nil {
|
||||
return fmt.Errorf("before snapshot: %w", err)
|
||||
}
|
||||
if err := Validate(after); err != nil {
|
||||
return fmt.Errorf("after snapshot: %w", err)
|
||||
}
|
||||
if after.Archive == nil || !after.Archive.Consistent {
|
||||
return errors.New("verified allow archive is unavailable")
|
||||
}
|
||||
if after.Live.AllowCount != 0 {
|
||||
return fmt.Errorf("live legacy allow count=%d, want 0", after.Live.AllowCount)
|
||||
}
|
||||
if before.Live.AllowCount != after.Archive.ManifestAllowCount || before.Live.AllowSHA256 != after.Archive.ManifestAllowSHA {
|
||||
return errors.New("pre-migration allow snapshot does not match archived manifest")
|
||||
}
|
||||
if before.Live.DenyCount != after.Live.DenyCount || before.Live.DenySHA256 != after.Live.DenySHA256 {
|
||||
return errors.New("deny rules changed during allow migration")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func Encode(snapshot Snapshot) ([]byte, error) {
|
||||
if err := Validate(snapshot); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
payload, err := json.MarshalIndent(snapshot, "", " ")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return append(payload, '\n'), nil
|
||||
}
|
||||
|
||||
func Decode(payload []byte) (Snapshot, error) {
|
||||
decoder := json.NewDecoder(strings.NewReader(string(payload)))
|
||||
decoder.DisallowUnknownFields()
|
||||
var snapshot Snapshot
|
||||
if err := decoder.Decode(&snapshot); err != nil {
|
||||
return Snapshot{}, err
|
||||
}
|
||||
if err := Validate(snapshot); err != nil {
|
||||
return Snapshot{}, err
|
||||
}
|
||||
return snapshot, nil
|
||||
}
|
||||
|
||||
func loadLiveRules(ctx context.Context, pool *pgxpool.Pool) (RuleSet, error) {
|
||||
rows, err := pool.Query(ctx, `
|
||||
SELECT id::text, subject_type, subject_id::text, resource_type,
|
||||
resource_id::text, effect, priority::text,
|
||||
min_permission_level::text, conditions::text, metadata::text,
|
||||
status, extract(epoch FROM created_at)::text,
|
||||
extract(epoch FROM updated_at)::text
|
||||
FROM gateway_access_rules
|
||||
ORDER BY id`)
|
||||
if err != nil {
|
||||
return RuleSet{}, err
|
||||
}
|
||||
defer rows.Close()
|
||||
allDigests := make([]string, 0)
|
||||
allowDigests := make([]string, 0)
|
||||
denyDigests := make([]string, 0)
|
||||
counts := map[string]int64{}
|
||||
for rows.Next() {
|
||||
var row ruleDigestRow
|
||||
if err := rows.Scan(
|
||||
&row.id, &row.subjectType, &row.subjectID, &row.resourceType,
|
||||
&row.resourceID, &row.effect, &row.priority,
|
||||
&row.minPermissionLevel, &row.conditions, &row.metadata,
|
||||
&row.status, &row.createdEpoch, &row.updatedEpoch,
|
||||
); err != nil {
|
||||
return RuleSet{}, err
|
||||
}
|
||||
digest := rowSHA256(row)
|
||||
allDigests = append(allDigests, digest)
|
||||
if row.effect == "allow" {
|
||||
allowDigests = append(allowDigests, digest)
|
||||
}
|
||||
if row.effect == "deny" {
|
||||
denyDigests = append(denyDigests, digest)
|
||||
}
|
||||
counts[strings.Join([]string{row.subjectType, row.effect, row.resourceType, row.status}, "\x00")]++
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return RuleSet{}, err
|
||||
}
|
||||
grouped := make([]RuleCount, 0, len(counts))
|
||||
for key, count := range counts {
|
||||
parts := strings.Split(key, "\x00")
|
||||
grouped = append(grouped, RuleCount{SubjectType: parts[0], Effect: parts[1], ResourceType: parts[2], Status: parts[3], Count: count})
|
||||
}
|
||||
sort.Slice(grouped, func(i, j int) bool {
|
||||
left := grouped[i]
|
||||
right := grouped[j]
|
||||
return strings.Join([]string{left.SubjectType, left.Effect, left.ResourceType, left.Status}, "\x00") <
|
||||
strings.Join([]string{right.SubjectType, right.Effect, right.ResourceType, right.Status}, "\x00")
|
||||
})
|
||||
return RuleSet{
|
||||
Total: int64(len(allDigests)),
|
||||
SHA256: aggregateSHA256(allDigests),
|
||||
AllowCount: int64(len(allowDigests)),
|
||||
AllowSHA256: aggregateSHA256(allowDigests),
|
||||
DenyCount: int64(len(denyDigests)),
|
||||
DenySHA256: aggregateSHA256(denyDigests),
|
||||
Counts: grouped,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func loadArchiveManifest(ctx context.Context, pool *pgxpool.Pool) (*ArchiveManifest, error) {
|
||||
var available bool
|
||||
if err := pool.QueryRow(ctx, `
|
||||
SELECT to_regclass('gateway_access_rule_migration_batches') IS NOT NULL
|
||||
AND to_regclass('gateway_access_rule_allow_archive') IS NOT NULL`).Scan(&available); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !available {
|
||||
return nil, nil
|
||||
}
|
||||
manifest := &ArchiveManifest{MigrationBatch: MigrationBatch}
|
||||
if err := pool.QueryRow(ctx, `
|
||||
SELECT allow_count, allow_sha256, deny_count, deny_sha256
|
||||
FROM gateway_access_rule_migration_batches
|
||||
WHERE migration_batch = $1`, MigrationBatch).Scan(
|
||||
&manifest.ManifestAllowCount,
|
||||
&manifest.ManifestAllowSHA,
|
||||
&manifest.ManifestDenyCount,
|
||||
&manifest.ManifestDenySHA,
|
||||
); err != nil {
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
rows, err := pool.Query(ctx, `
|
||||
SELECT row_sha256
|
||||
FROM gateway_access_rule_allow_archive
|
||||
WHERE migration_batch = $1
|
||||
ORDER BY rule_id`, MigrationBatch)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
digests := make([]string, 0)
|
||||
for rows.Next() {
|
||||
var digest string
|
||||
if err := rows.Scan(&digest); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
digests = append(digests, digest)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
manifest.ArchivedAllowCount = int64(len(digests))
|
||||
manifest.ArchivedAllowSHA = aggregateSHA256(digests)
|
||||
manifest.Consistent = manifest.ManifestAllowCount == manifest.ArchivedAllowCount && manifest.ManifestAllowSHA == manifest.ArchivedAllowSHA
|
||||
return manifest, nil
|
||||
}
|
||||
|
||||
func rowSHA256(row ruleDigestRow) string {
|
||||
return sha256Hex(strings.Join([]string{
|
||||
row.id, row.subjectType, row.subjectID, row.resourceType,
|
||||
row.resourceID, row.effect, row.priority, row.minPermissionLevel,
|
||||
row.conditions, row.metadata, row.status, row.createdEpoch, row.updatedEpoch,
|
||||
}, "\x1f"))
|
||||
}
|
||||
|
||||
func aggregateSHA256(rowDigests []string) string {
|
||||
return sha256Hex(strings.Join(rowDigests, "\n"))
|
||||
}
|
||||
|
||||
func sha256Hex(value string) string {
|
||||
digest := sha256.Sum256([]byte(value))
|
||||
return hex.EncodeToString(digest[:])
|
||||
}
|
||||
|
||||
func validSHA256(value string) bool {
|
||||
if len(value) != 64 {
|
||||
return false
|
||||
}
|
||||
_, err := hex.DecodeString(value)
|
||||
return err == nil && value == strings.ToLower(value)
|
||||
}
|
||||
@@ -0,0 +1,69 @@
|
||||
package accessruleaudit
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestVerifyMigrationAcceptsMatchingArchiveAndUnchangedDeny(t *testing.T) {
|
||||
before := auditSnapshot(3, "allow-before", 2, "deny-before", nil)
|
||||
after := auditSnapshot(0, "", 2, "deny-before", &ArchiveManifest{
|
||||
MigrationBatch: MigrationBatch,
|
||||
ManifestAllowCount: 3,
|
||||
ManifestAllowSHA: sha256Hex("allow-before"),
|
||||
ArchivedAllowCount: 3,
|
||||
ArchivedAllowSHA: sha256Hex("allow-before"),
|
||||
ManifestDenyCount: 2,
|
||||
ManifestDenySHA: sha256Hex("deny-before"),
|
||||
Consistent: true,
|
||||
})
|
||||
if err := VerifyMigration(before, after); err != nil {
|
||||
t.Fatalf("verify matching migration: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestVerifyMigrationRejectsChangedDenyOrMissingArchive(t *testing.T) {
|
||||
before := auditSnapshot(1, "allow", 1, "deny", nil)
|
||||
withoutArchive := auditSnapshot(0, "", 1, "deny", nil)
|
||||
if err := VerifyMigration(before, withoutArchive); err == nil {
|
||||
t.Fatal("missing archive must fail verification")
|
||||
}
|
||||
changedDeny := auditSnapshot(0, "", 1, "changed-deny", &ArchiveManifest{
|
||||
MigrationBatch: MigrationBatch,
|
||||
ManifestAllowCount: 1,
|
||||
ManifestAllowSHA: sha256Hex("allow"),
|
||||
ArchivedAllowCount: 1,
|
||||
ArchivedAllowSHA: sha256Hex("allow"),
|
||||
ManifestDenyCount: 1,
|
||||
ManifestDenySHA: sha256Hex("deny"),
|
||||
Consistent: true,
|
||||
})
|
||||
if err := VerifyMigration(before, changedDeny); err == nil {
|
||||
t.Fatal("changed deny hash must fail verification")
|
||||
}
|
||||
}
|
||||
|
||||
func auditSnapshot(allowCount int64, allowSeed string, denyCount int64, denySeed string, archive *ArchiveManifest) Snapshot {
|
||||
counts := make([]RuleCount, 0, 2)
|
||||
if allowCount > 0 {
|
||||
counts = append(counts, RuleCount{SubjectType: "api_key", Effect: "allow", ResourceType: "platform_model", Status: "active", Count: allowCount})
|
||||
}
|
||||
if denyCount > 0 {
|
||||
counts = append(counts, RuleCount{SubjectType: "api_key", Effect: "deny", ResourceType: "platform_model", Status: "active", Count: denyCount})
|
||||
}
|
||||
return Snapshot{
|
||||
SchemaVersion: SchemaVersion,
|
||||
GeneratedAt: time.Date(2026, 8, 3, 0, 0, 0, 0, time.UTC),
|
||||
SecretSafe: true,
|
||||
Live: RuleSet{
|
||||
Total: allowCount + denyCount,
|
||||
SHA256: sha256Hex("all-" + allowSeed + "-" + denySeed),
|
||||
AllowCount: allowCount,
|
||||
AllowSHA256: sha256Hex(allowSeed),
|
||||
DenyCount: denyCount,
|
||||
DenySHA256: sha256Hex(denySeed),
|
||||
Counts: counts,
|
||||
},
|
||||
Archive: archive,
|
||||
}
|
||||
}
|
||||
@@ -34,27 +34,36 @@ const (
|
||||
)
|
||||
|
||||
type User struct {
|
||||
ID string `json:"sub"`
|
||||
Username string `json:"username"`
|
||||
Roles []string `json:"role,omitempty"`
|
||||
TenantID string `json:"tenantId,omitempty"`
|
||||
GatewayTenantID string `json:"gatewayTenantId,omitempty"`
|
||||
TenantKey string `json:"tenantKey,omitempty"`
|
||||
SSOID string `json:"sso_id,omitempty"`
|
||||
Source string `json:"source,omitempty"`
|
||||
GatewayUserID string `json:"gatewayUserId,omitempty"`
|
||||
UserGroupID string `json:"userGroupId,omitempty"`
|
||||
UserGroupKey string `json:"userGroupKey,omitempty"`
|
||||
UserGroupKeys []string `json:"userGroupKeys,omitempty"`
|
||||
APIKeyID string `json:"apiKeyId,omitempty"`
|
||||
APIKeySecret string `json:"apiKeySecret,omitempty"`
|
||||
APIKeyName string `json:"apiKeyName,omitempty"`
|
||||
APIKeyPrefix string `json:"apiKeyPrefix,omitempty"`
|
||||
APIKeyScopes []string `json:"apiKeyScopes,omitempty"`
|
||||
TokenExpiresAt time.Time `json:"-"`
|
||||
TokenIssuedAt time.Time `json:"-"`
|
||||
Issuer string `json:"-"`
|
||||
TokenPurpose string `json:"-"`
|
||||
ID string `json:"sub"`
|
||||
Username string `json:"username"`
|
||||
DisplayName string `json:"-"`
|
||||
Email string `json:"-"`
|
||||
Phone string `json:"-"`
|
||||
AvatarURL string `json:"-"`
|
||||
Roles []string `json:"role,omitempty"`
|
||||
ContextType string `json:"contextType,omitempty"`
|
||||
TenantID string `json:"tenantId,omitempty"`
|
||||
TenantName string `json:"tenantName,omitempty"`
|
||||
GatewayTenantID string `json:"gatewayTenantId,omitempty"`
|
||||
TenantKey string `json:"tenantKey,omitempty"`
|
||||
SSOID string `json:"sso_id,omitempty"`
|
||||
Source string `json:"source,omitempty"`
|
||||
GatewayUserID string `json:"gatewayUserId,omitempty"`
|
||||
UserGroupID string `json:"userGroupId,omitempty"`
|
||||
UserGroupKey string `json:"userGroupKey,omitempty"`
|
||||
UserGroupKeys []string `json:"userGroupKeys,omitempty"`
|
||||
APIKeyID string `json:"apiKeyId,omitempty"`
|
||||
APIKeySecret string `json:"apiKeySecret,omitempty"`
|
||||
APIKeyName string `json:"apiKeyName,omitempty"`
|
||||
APIKeyPrefix string `json:"apiKeyPrefix,omitempty"`
|
||||
APIKeyScopes []string `json:"apiKeyScopes,omitempty"`
|
||||
TokenExpiresAt time.Time `json:"-"`
|
||||
TokenIssuedAt time.Time `json:"-"`
|
||||
Issuer string `json:"-"`
|
||||
ApplicationID string `json:"-"`
|
||||
OIDCClientID string `json:"-"`
|
||||
OIDCUserBindingID string `json:"-"`
|
||||
TokenPurpose string `json:"-"`
|
||||
}
|
||||
|
||||
type contextKey string
|
||||
@@ -64,9 +73,10 @@ const userContextKey contextKey = "easyai-auth-user"
|
||||
var ErrUnauthorized = errors.New("unauthorized")
|
||||
|
||||
type RequestAuthError struct {
|
||||
Status int
|
||||
Code string
|
||||
Message string
|
||||
Status int
|
||||
Code string
|
||||
Message string
|
||||
RetryAfterSeconds int
|
||||
}
|
||||
|
||||
func (e *RequestAuthError) Error() string { return e.Code }
|
||||
@@ -75,6 +85,12 @@ func NewRequestAuthError(status int, code, message string) error {
|
||||
return &RequestAuthError{Status: status, Code: code, Message: message}
|
||||
}
|
||||
|
||||
func NewRetryableRequestAuthError(status int, code, message string, retryAfterSeconds int) error {
|
||||
return &RequestAuthError{
|
||||
Status: status, Code: code, Message: message, RetryAfterSeconds: retryAfterSeconds,
|
||||
}
|
||||
}
|
||||
|
||||
type Authenticator struct {
|
||||
JWTSecret string
|
||||
ServerMainBaseURL string
|
||||
@@ -122,6 +138,9 @@ func (a *Authenticator) Require(permission Permission, next http.Handler) http.H
|
||||
}
|
||||
var requestError *RequestAuthError
|
||||
if errors.As(err, &requestError) {
|
||||
if requestError.RetryAfterSeconds > 0 {
|
||||
w.Header().Set("Retry-After", strconv.Itoa(requestError.RetryAfterSeconds))
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json; charset=utf-8")
|
||||
w.Header().Set("Cache-Control", "no-store")
|
||||
w.WriteHeader(requestError.Status)
|
||||
@@ -313,7 +332,12 @@ func (a *Authenticator) verifyAPIKey(ctx context.Context, apiKey string) (*User,
|
||||
return user, nil
|
||||
}
|
||||
if !errors.Is(err, ErrUnauthorized) {
|
||||
return nil, err
|
||||
return nil, NewRetryableRequestAuthError(
|
||||
http.StatusServiceUnavailable,
|
||||
"AUTH_STORE_UNAVAILABLE",
|
||||
"authentication service temporarily unavailable",
|
||||
2,
|
||||
)
|
||||
}
|
||||
if strings.HasPrefix(apiKey, "sk-gw-") {
|
||||
return nil, ErrUnauthorized
|
||||
|
||||
@@ -0,0 +1,31 @@
|
||||
package auth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestAPIKeyStoreFailureReturnsRetryableServiceUnavailable(t *testing.T) {
|
||||
authenticator := New("test-secret", "", "")
|
||||
authenticator.LocalAPIKeyVerifier = func(context.Context, string) (*User, error) {
|
||||
return nil, errors.New("database unavailable")
|
||||
}
|
||||
handler := authenticator.Require(PermissionBasic, http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
|
||||
t.Fatal("unavailable authentication store must not call the protected handler")
|
||||
}))
|
||||
request := httptest.NewRequest(http.MethodGet, "/protected", nil)
|
||||
request.Header.Set("Authorization", "Bearer sk-gw-local")
|
||||
recorder := httptest.NewRecorder()
|
||||
|
||||
handler.ServeHTTP(recorder, request)
|
||||
|
||||
if recorder.Code != http.StatusServiceUnavailable {
|
||||
t.Fatalf("status=%d, want 503; body=%s", recorder.Code, recorder.Body.String())
|
||||
}
|
||||
if recorder.Header().Get("Retry-After") != "2" {
|
||||
t.Fatalf("Retry-After=%q, want 2", recorder.Header().Get("Retry-After"))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
package auth
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestGatewaySourceDoesNotDependOnAuthenticationCoreVocabulary(t *testing.T) {
|
||||
_, currentFile, _, ok := runtime.Caller(0)
|
||||
if !ok {
|
||||
t.Fatal("resolve current test file")
|
||||
}
|
||||
repositoryRoot := filepath.Clean(filepath.Join(filepath.Dir(currentFile), "..", "..", "..", ".."))
|
||||
appsRoot := filepath.Join(repositoryRoot, "apps")
|
||||
forbiddenMarkers := []string{
|
||||
strings.ToLower("key" + "cloak"),
|
||||
"/rea" + "lms/",
|
||||
}
|
||||
sourceExtensions := []string{".go", ".js", ".json", ".ts", ".tsx", ".yaml", ".yml"}
|
||||
skippedDirectories := []string{"coverage", "dist", "node_modules", "test-results"}
|
||||
var violations []string
|
||||
|
||||
err := filepath.WalkDir(appsRoot, func(filePath string, entry os.DirEntry, walkErr error) error {
|
||||
if walkErr != nil {
|
||||
return walkErr
|
||||
}
|
||||
if entry.IsDir() {
|
||||
if filePath != appsRoot && slices.Contains(skippedDirectories, entry.Name()) {
|
||||
return filepath.SkipDir
|
||||
}
|
||||
return nil
|
||||
}
|
||||
if !slices.Contains(sourceExtensions, strings.ToLower(filepath.Ext(entry.Name()))) {
|
||||
return nil
|
||||
}
|
||||
content, readErr := os.ReadFile(filePath)
|
||||
if readErr != nil {
|
||||
return readErr
|
||||
}
|
||||
lowerContent := strings.ToLower(string(content))
|
||||
for _, marker := range forbiddenMarkers {
|
||||
if strings.Contains(lowerContent, marker) {
|
||||
relativePath, relativeErr := filepath.Rel(repositoryRoot, filePath)
|
||||
if relativeErr != nil {
|
||||
return relativeErr
|
||||
}
|
||||
violations = append(violations, filepath.ToSlash(relativePath))
|
||||
break
|
||||
}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("scan Gateway sources: %v", err)
|
||||
}
|
||||
if len(violations) > 0 {
|
||||
t.Fatalf("Gateway source contains authentication-core-specific vocabulary: %s", strings.Join(violations, ", "))
|
||||
}
|
||||
}
|
||||
+164
-29
@@ -20,6 +20,7 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/golang-jwt/jwt/v5"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
const maxOIDCResponseBytes = 1 << 20
|
||||
@@ -30,6 +31,9 @@ type OIDCConfig struct {
|
||||
Issuer string
|
||||
Audience string
|
||||
TenantID string
|
||||
TenantMode string
|
||||
ApplicationID string
|
||||
ClientID string
|
||||
RolePrefix string
|
||||
RequiredScopes []string
|
||||
JWKSCacheTTL time.Duration
|
||||
@@ -44,10 +48,12 @@ type OIDCConfig struct {
|
||||
}
|
||||
|
||||
type OIDCSecurityEventIdentity struct {
|
||||
Issuer string
|
||||
TenantID string
|
||||
Subject string
|
||||
IssuedAt time.Time
|
||||
Issuer string
|
||||
ApplicationID string
|
||||
ContextType string
|
||||
TenantID string
|
||||
Subject string
|
||||
IssuedAt time.Time
|
||||
}
|
||||
|
||||
type OIDCSecurityEventEvaluation struct {
|
||||
@@ -91,8 +97,16 @@ func NewOIDCVerifier(config OIDCConfig) (*OIDCVerifier, error) {
|
||||
config.Issuer = strings.TrimRight(strings.TrimSpace(config.Issuer), "/")
|
||||
config.Audience = strings.TrimSpace(config.Audience)
|
||||
config.TenantID = strings.TrimSpace(config.TenantID)
|
||||
config.TenantMode = strings.TrimSpace(config.TenantMode)
|
||||
config.ApplicationID = strings.TrimSpace(config.ApplicationID)
|
||||
config.ClientID = strings.TrimSpace(config.ClientID)
|
||||
config.RolePrefix = strings.TrimSpace(config.RolePrefix)
|
||||
if err := validatePublicURL(config.Issuer, config.AppEnv); err != nil || config.Audience == "" || config.TenantID == "" || config.RolePrefix == "" {
|
||||
if config.TenantMode == "" {
|
||||
config.TenantMode = "single_tenant"
|
||||
}
|
||||
validTenantMode := config.TenantMode == "single_tenant" && config.TenantID != "" ||
|
||||
config.TenantMode == "multi_tenant" && config.TenantID == "" && config.ApplicationID != ""
|
||||
if err := validatePublicURL(config.Issuer, config.AppEnv); err != nil || config.Audience == "" || !validTenantMode || config.RolePrefix == "" {
|
||||
return nil, errors.New("issuer, audience, tenant and role prefix are required")
|
||||
}
|
||||
if config.IntrospectionEnabled && config.IntrospectionCredentialProvider == nil &&
|
||||
@@ -132,15 +146,28 @@ func (v *OIDCVerifier) Verify(ctx context.Context, raw string) (*User, error) {
|
||||
parser := jwt.NewParser(jwt.WithValidMethods([]string{"RS256", "ES256"}))
|
||||
unverified, _, err := parser.ParseUnverified(raw, jwt.MapClaims{})
|
||||
if err != nil || unverified == nil {
|
||||
return nil, oidcUnauthorized("token envelope is invalid", err)
|
||||
return nil, oidcUnauthorized("TOKEN_ENVELOPE_INVALID", "token envelope is invalid", err)
|
||||
}
|
||||
unverifiedClaims, ok := unverified.Claims.(jwt.MapClaims)
|
||||
if !ok {
|
||||
return nil, oidcUnauthorized("TOKEN_ENVELOPE_INVALID", "token envelope is invalid", nil)
|
||||
}
|
||||
if stringClaim(unverifiedClaims, "iss") == "" {
|
||||
return nil, oidcUnauthorized("ISSUER_MISSING", "issuer is missing", nil)
|
||||
}
|
||||
if len(stringSliceClaim(unverifiedClaims, "aud")) == 0 {
|
||||
return nil, oidcUnauthorized("AUDIENCE_MISSING", "audience is missing", nil)
|
||||
}
|
||||
if _, exists := unverifiedClaims["exp"]; !exists {
|
||||
return nil, oidcUnauthorized("EXP_MISSING", "exp is missing", nil)
|
||||
}
|
||||
kid, _ := unverified.Header["kid"].(string)
|
||||
if kid == "" {
|
||||
return nil, oidcUnauthorized("kid is missing", nil)
|
||||
return nil, oidcUnauthorized("KID_MISSING", "kid is missing", nil)
|
||||
}
|
||||
key, err := v.key(ctx, kid)
|
||||
if err != nil {
|
||||
return nil, oidcUnauthorized("signing key lookup failed", err)
|
||||
return nil, oidcUnauthorized("SIGNING_KEY_LOOKUP_FAILED", "signing key lookup failed", err)
|
||||
}
|
||||
token, err := jwt.Parse(raw, func(token *jwt.Token) (any, error) {
|
||||
if token.Header["kid"] != kid {
|
||||
@@ -150,43 +177,60 @@ func (v *OIDCVerifier) Verify(ctx context.Context, raw string) (*User, error) {
|
||||
}, jwt.WithValidMethods([]string{"RS256", "ES256"}), jwt.WithIssuer(v.config.Issuer),
|
||||
jwt.WithAudience(v.config.Audience), jwt.WithExpirationRequired(), jwt.WithLeeway(30*time.Second))
|
||||
if err != nil || !token.Valid {
|
||||
return nil, oidcUnauthorized("signature or registered claims are invalid", err)
|
||||
return nil, oidcUnauthorized(registeredClaimsValidationCategory(err), "signature or registered claims are invalid", err)
|
||||
}
|
||||
claims, ok := token.Claims.(jwt.MapClaims)
|
||||
if !ok || stringClaim(claims, "sub") == "" || stringClaim(claims, "tid") != v.config.TenantID {
|
||||
return nil, oidcUnauthorized("stable identity claims are invalid", nil)
|
||||
contextType := stringClaim(claims, "context_type")
|
||||
tenantID := stringClaim(claims, "tid")
|
||||
validContext := contextType == "tenant" && tenantID == v.config.TenantID
|
||||
if v.config.TenantMode == "multi_tenant" {
|
||||
switch contextType {
|
||||
case "platform":
|
||||
validContext = tenantID == ""
|
||||
case "tenant":
|
||||
validContext = uuid.Validate(tenantID) == nil
|
||||
default:
|
||||
validContext = false
|
||||
}
|
||||
}
|
||||
clientID := stringClaim(claims, "client_id")
|
||||
if !ok || stringClaim(claims, "sub") == "" || !validContext ||
|
||||
v.config.ClientID != "" && clientID != v.config.ClientID {
|
||||
return nil, oidcUnauthorized("STABLE_IDENTITY_CLAIMS_INVALID", "stable identity claims are invalid", nil)
|
||||
}
|
||||
expiresAt, ok := numericDateClaim(claims["exp"])
|
||||
if !ok {
|
||||
return nil, oidcUnauthorized("exp is invalid", nil)
|
||||
return nil, oidcUnauthorized("EXP_INVALID", "exp is invalid", nil)
|
||||
}
|
||||
if _, ok := numericDateClaim(claims["nbf"]); !ok {
|
||||
return nil, oidcUnauthorized("nbf is missing", nil)
|
||||
return nil, oidcUnauthorized("NBF_MISSING", "nbf is missing", nil)
|
||||
}
|
||||
issuedAt, hasIssuedAt := numericDateClaim(claims["iat"])
|
||||
scopes := scopeClaims(claims)
|
||||
if !containsAll(scopes, v.config.RequiredScopes) {
|
||||
return nil, oidcUnauthorized("required scope is missing", nil)
|
||||
return nil, oidcUnauthorized("REQUIRED_SCOPE_MISSING", "required scope is missing", nil)
|
||||
}
|
||||
roles := gatewayRoles(stringSliceClaim(claims, "roles"), v.config.RolePrefix)
|
||||
if len(roles) == 0 {
|
||||
roles = gatewayRoles(stringSliceClaim(claims, "role"), v.config.RolePrefix)
|
||||
}
|
||||
if len(roles) == 0 {
|
||||
return nil, oidcUnauthorized("mapped role is missing", nil)
|
||||
return nil, oidcUnauthorized("MAPPED_ROLE_MISSING", "mapped role is missing", nil)
|
||||
}
|
||||
if v.config.SecurityEventEvaluator != nil {
|
||||
evaluation, evaluateErr := v.config.SecurityEventEvaluator(ctx, OIDCSecurityEventIdentity{
|
||||
Issuer: v.config.Issuer, TenantID: v.config.TenantID, Subject: stringClaim(claims, "sub"), IssuedAt: issuedAt,
|
||||
Issuer: v.config.Issuer, ApplicationID: v.config.ApplicationID,
|
||||
ContextType: contextType, TenantID: tenantID,
|
||||
Subject: stringClaim(claims, "sub"), IssuedAt: issuedAt,
|
||||
})
|
||||
if evaluateErr != nil {
|
||||
return nil, NewRequestAuthError(http.StatusServiceUnavailable, "OIDC_SECURITY_EVENT_STATE_UNAVAILABLE", "认证撤销状态暂时不可用")
|
||||
}
|
||||
if evaluation.Enabled && !hasIssuedAt {
|
||||
return nil, oidcUnauthorized("iat is required when security events are enabled", nil)
|
||||
return nil, oidcUnauthorized("SECURITY_EVENT_IAT_MISSING", "iat is required when security events are enabled", nil)
|
||||
}
|
||||
if evaluation.Revoked {
|
||||
return nil, oidcUnauthorized("token was issued before the revocation watermark", nil)
|
||||
return nil, oidcUnauthorized("TOKEN_REVOKED", "token was issued before the revocation watermark", nil)
|
||||
}
|
||||
if evaluation.RequireIntrospection {
|
||||
active, introspectionErr := v.introspect(ctx, raw)
|
||||
@@ -194,35 +238,126 @@ func (v *OIDCVerifier) Verify(ctx context.Context, raw string) (*User, error) {
|
||||
return nil, NewRequestAuthError(http.StatusServiceUnavailable, "OIDC_INTROSPECTION_UNAVAILABLE", "认证中心内省暂时不可用")
|
||||
}
|
||||
if !active {
|
||||
return nil, oidcUnauthorized("token is inactive", nil)
|
||||
return nil, oidcUnauthorized("TOKEN_INACTIVE", "token is inactive", nil)
|
||||
}
|
||||
} else if !evaluation.Enabled && v.config.IntrospectionEnabled {
|
||||
active, introspectionErr := v.introspect(ctx, raw)
|
||||
if introspectionErr != nil || !active {
|
||||
return nil, oidcUnauthorized("token is inactive", introspectionErr)
|
||||
return nil, oidcUnauthorized("TOKEN_INACTIVE", "token is inactive", introspectionErr)
|
||||
}
|
||||
}
|
||||
} else if v.config.IntrospectionEnabled {
|
||||
active, err := v.introspect(ctx, raw)
|
||||
if err != nil || !active {
|
||||
return nil, oidcUnauthorized("token is inactive", err)
|
||||
return nil, oidcUnauthorized("TOKEN_INACTIVE", "token is inactive", err)
|
||||
}
|
||||
}
|
||||
username := stringClaim(claims, "preferred_username")
|
||||
username := oidcProfileText(claims, "preferred_username", 320)
|
||||
if username == "" {
|
||||
username = stringClaim(claims, "username")
|
||||
username = oidcProfileText(claims, "username", 320)
|
||||
}
|
||||
return &User{
|
||||
ID: stringClaim(claims, "sub"), Username: username, Roles: roles,
|
||||
TenantID: v.config.TenantID, Source: "oidc", TokenExpiresAt: expiresAt, TokenIssuedAt: issuedAt, Issuer: v.config.Issuer,
|
||||
ID: stringClaim(claims, "sub"), Username: username,
|
||||
DisplayName: oidcProfileText(claims, "name", 200),
|
||||
Email: oidcVerifiedProfileText(claims, "email", "email_verified", 320),
|
||||
Phone: oidcVerifiedProfileText(claims, "phone_number", "phone_number_verified", 64),
|
||||
AvatarURL: safeOIDCProfileURL(oidcProfileText(claims, "picture", 2048)),
|
||||
Roles: roles,
|
||||
ContextType: contextType, TenantID: tenantID, Source: "oidc",
|
||||
TokenExpiresAt: expiresAt, TokenIssuedAt: issuedAt, Issuer: v.config.Issuer,
|
||||
ApplicationID: v.config.ApplicationID, OIDCClientID: clientID,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func oidcUnauthorized(reason string, cause error) error {
|
||||
if cause != nil {
|
||||
return fmt.Errorf("%w: %s: %v", ErrUnauthorized, reason, cause)
|
||||
func oidcProfileText(claims jwt.MapClaims, key string, limit int) string {
|
||||
value := strings.TrimSpace(stringClaim(claims, key))
|
||||
if value == "" {
|
||||
return ""
|
||||
}
|
||||
runes := []rune(value)
|
||||
if len(runes) > limit {
|
||||
return string(runes[:limit])
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
func oidcVerifiedProfileText(claims jwt.MapClaims, key, verifiedKey string, limit int) string {
|
||||
verified, ok := claims[verifiedKey].(bool)
|
||||
if !ok || !verified {
|
||||
return ""
|
||||
}
|
||||
return oidcProfileText(claims, key, limit)
|
||||
}
|
||||
|
||||
func safeOIDCProfileURL(value string) string {
|
||||
parsed, err := url.Parse(value)
|
||||
if err != nil || parsed.Scheme != "https" || parsed.Host == "" || parsed.User != nil {
|
||||
return ""
|
||||
}
|
||||
return parsed.String()
|
||||
}
|
||||
|
||||
type oidcValidationError struct {
|
||||
category string
|
||||
reason string
|
||||
cause error
|
||||
}
|
||||
|
||||
func (err *oidcValidationError) Error() string {
|
||||
if err.cause != nil {
|
||||
return fmt.Sprintf("%v: %s: %v", ErrUnauthorized, err.reason, err.cause)
|
||||
}
|
||||
return fmt.Sprintf("%v: %s", ErrUnauthorized, err.reason)
|
||||
}
|
||||
|
||||
func (err *oidcValidationError) Unwrap() []error {
|
||||
if err.cause == nil {
|
||||
return []error{ErrUnauthorized}
|
||||
}
|
||||
return []error{ErrUnauthorized, err.cause}
|
||||
}
|
||||
|
||||
func OIDCValidationCategory(err error) string {
|
||||
var validationError *oidcValidationError
|
||||
if errors.As(err, &validationError) {
|
||||
return validationError.category
|
||||
}
|
||||
return "UNCLASSIFIED"
|
||||
}
|
||||
|
||||
func registeredClaimsValidationCategory(err error) string {
|
||||
switch {
|
||||
case errors.Is(err, jwt.ErrTokenRequiredClaimMissing):
|
||||
return "REQUIRED_CLAIM_MISSING"
|
||||
case errors.Is(err, jwt.ErrTokenInvalidAudience):
|
||||
return "AUDIENCE_INVALID"
|
||||
case errors.Is(err, jwt.ErrTokenInvalidIssuer):
|
||||
return "ISSUER_INVALID"
|
||||
case errors.Is(err, jwt.ErrTokenExpired):
|
||||
return "TOKEN_EXPIRED"
|
||||
case errors.Is(err, jwt.ErrTokenNotValidYet):
|
||||
return "TOKEN_NOT_VALID_YET"
|
||||
case errors.Is(err, jwt.ErrTokenUsedBeforeIssued):
|
||||
return "TOKEN_USED_BEFORE_ISSUED"
|
||||
case errors.Is(err, jwt.ErrTokenSignatureInvalid):
|
||||
return "SIGNATURE_INVALID"
|
||||
case errors.Is(err, jwt.ErrTokenMalformed):
|
||||
return "TOKEN_MALFORMED"
|
||||
case errors.Is(err, jwt.ErrTokenUnverifiable):
|
||||
return "TOKEN_UNVERIFIABLE"
|
||||
case errors.Is(err, jwt.ErrTokenInvalidClaims):
|
||||
return "CLAIMS_INVALID"
|
||||
default:
|
||||
return "REGISTERED_CLAIMS_INVALID"
|
||||
}
|
||||
}
|
||||
|
||||
func oidcUnauthorized(category, reason string, cause error) error {
|
||||
return &oidcValidationError{
|
||||
category: category,
|
||||
reason: reason,
|
||||
cause: cause,
|
||||
}
|
||||
return fmt.Errorf("%w: %s", ErrUnauthorized, reason)
|
||||
}
|
||||
|
||||
func (v *OIDCVerifier) key(ctx context.Context, kid string) (any, error) {
|
||||
|
||||
@@ -11,6 +11,7 @@ import (
|
||||
"sync"
|
||||
|
||||
"github.com/coreos/go-oidc/v3/oidc"
|
||||
"github.com/google/uuid"
|
||||
"golang.org/x/oauth2"
|
||||
)
|
||||
|
||||
@@ -77,15 +78,38 @@ func (c *OIDCPublicClient) ValidateConfiguration(ctx context.Context) error {
|
||||
return err
|
||||
}
|
||||
|
||||
func (c *OIDCPublicClient) AuthorizationURL(ctx context.Context, state, nonce, pkceVerifier string) (string, error) {
|
||||
func (c *OIDCPublicClient) AuthorizationURL(
|
||||
ctx context.Context,
|
||||
state, nonce, pkceVerifier, contextType, tenantHint string,
|
||||
) (string, error) {
|
||||
if strings.TrimSpace(state) == "" || strings.TrimSpace(nonce) == "" || !validPKCEVerifier(pkceVerifier) {
|
||||
return "", errors.New("state, nonce and PKCE verifier are required")
|
||||
}
|
||||
contextType = strings.TrimSpace(contextType)
|
||||
if contextType != "" && contextType != "platform" && contextType != "tenant" {
|
||||
return "", errors.New("context type must be empty, platform or tenant")
|
||||
}
|
||||
if strings.TrimSpace(tenantHint) != "" && uuid.Validate(strings.TrimSpace(tenantHint)) != nil {
|
||||
return "", errors.New("tenant hint must be a UUID")
|
||||
}
|
||||
if contextType != "tenant" && strings.TrimSpace(tenantHint) != "" {
|
||||
return "", errors.New("only tenant context can bind a tenant hint")
|
||||
}
|
||||
config, _, err := c.configuration(ctx)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return config.AuthCodeURL(state, oidc.Nonce(nonce), oauth2.S256ChallengeOption(pkceVerifier)), nil
|
||||
options := []oauth2.AuthCodeOption{
|
||||
oidc.Nonce(nonce),
|
||||
oauth2.S256ChallengeOption(pkceVerifier),
|
||||
}
|
||||
if contextType != "" {
|
||||
options = append(options, oauth2.SetAuthURLParam("context_type", contextType))
|
||||
}
|
||||
if strings.TrimSpace(tenantHint) != "" {
|
||||
options = append(options, oauth2.SetAuthURLParam("tenant_hint", strings.TrimSpace(tenantHint)))
|
||||
}
|
||||
return config.AuthCodeURL(state, options...), nil
|
||||
}
|
||||
|
||||
func (c *OIDCPublicClient) ExchangeCode(ctx context.Context, code, verifier string) (OIDCTokenResponse, error) {
|
||||
@@ -121,30 +145,30 @@ func (c *OIDCPublicClient) Refresh(ctx context.Context, refreshToken string) (OI
|
||||
func (c *OIDCPublicClient) VerifyIDToken(ctx context.Context, raw, expectedNonce string) (string, error) {
|
||||
expectedNonce = strings.TrimSpace(expectedNonce)
|
||||
if strings.TrimSpace(raw) == "" || expectedNonce == "" {
|
||||
return "", oidcUnauthorized("ID token validation context is invalid", nil)
|
||||
return "", oidcUnauthorized("ID_TOKEN_CONTEXT_INVALID", "ID token validation context is invalid", nil)
|
||||
}
|
||||
if _, _, err := c.configuration(ctx); err != nil {
|
||||
return "", oidcUnauthorized("ID token provider discovery failed", err)
|
||||
return "", oidcUnauthorized("ID_TOKEN_DISCOVERY_FAILED", "ID token provider discovery failed", err)
|
||||
}
|
||||
c.mu.Lock()
|
||||
verifier := c.idTokenVerifier
|
||||
c.mu.Unlock()
|
||||
if verifier == nil {
|
||||
return "", oidcUnauthorized("ID token verifier is unavailable", nil)
|
||||
return "", oidcUnauthorized("ID_TOKEN_VERIFIER_UNAVAILABLE", "ID token verifier is unavailable", nil)
|
||||
}
|
||||
token, err := verifier.Verify(c.requestContext(ctx), raw)
|
||||
if err != nil {
|
||||
return "", oidcUnauthorized("ID token signature or registered claims are invalid", nil)
|
||||
return "", oidcUnauthorized("ID_TOKEN_REGISTERED_CLAIMS_INVALID", "ID token signature or registered claims are invalid", nil)
|
||||
}
|
||||
if token.Subject == "" || token.Nonce != expectedNonce {
|
||||
return "", oidcUnauthorized("ID token subject or nonce is invalid", nil)
|
||||
return "", oidcUnauthorized("ID_TOKEN_SUBJECT_OR_NONCE_INVALID", "ID token subject or nonce is invalid", nil)
|
||||
}
|
||||
var claims map[string]any
|
||||
if err := token.Claims(&claims); err != nil {
|
||||
return "", oidcUnauthorized("ID token claims are invalid", nil)
|
||||
return "", oidcUnauthorized("ID_TOKEN_CLAIMS_INVALID", "ID token claims are invalid", nil)
|
||||
}
|
||||
if _, ok := numericDateClaim(claims["nbf"]); !ok {
|
||||
return "", oidcUnauthorized("ID token nbf is missing", nil)
|
||||
return "", oidcUnauthorized("ID_TOKEN_NBF_MISSING", "ID token nbf is missing", nil)
|
||||
}
|
||||
return token.Subject, nil
|
||||
}
|
||||
|
||||
@@ -67,7 +67,10 @@ func TestOIDCPublicClientUsesAuthorizationCodePKCES256WithoutSecret(t *testing.T
|
||||
if err := client.ValidateConfiguration(context.Background()); err != nil {
|
||||
t.Fatalf("ValidateConfiguration() error = %v", err)
|
||||
}
|
||||
authorizationURL, err := client.AuthorizationURL(context.Background(), "state", "nonce", pkceVerifier)
|
||||
tenantHint := "aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa"
|
||||
authorizationURL, err := client.AuthorizationURL(
|
||||
context.Background(), "state", "nonce", pkceVerifier, "tenant", tenantHint,
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
@@ -78,11 +81,54 @@ func TestOIDCPublicClientUsesAuthorizationCodePKCES256WithoutSecret(t *testing.T
|
||||
if query.Get("response_type") != "code" || query.Get("code_challenge_method") != "S256" || query.Get("code_challenge") != expectedChallenge {
|
||||
t.Fatalf("authorization request is not PKCE S256: %v", query)
|
||||
}
|
||||
if query.Get("tenant_hint") != tenantHint {
|
||||
t.Fatalf("tenant_hint = %q, want %q", query.Get("tenant_hint"), tenantHint)
|
||||
}
|
||||
if query.Get("context_type") != "tenant" {
|
||||
t.Fatalf("context_type = %q, want tenant", query.Get("context_type"))
|
||||
}
|
||||
if _, err := client.ExchangeCode(context.Background(), "authorization-code", pkceVerifier); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOIDCPublicClientOmitsOptionalContextForUnifiedLogin(t *testing.T) {
|
||||
var issuer string
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/.well-known/openid-configuration" {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(map[string]string{
|
||||
"issuer": issuer, "authorization_endpoint": issuer + "/authorize",
|
||||
"token_endpoint": issuer + "/token", "jwks_uri": issuer + "/jwks",
|
||||
})
|
||||
}))
|
||||
defer server.Close()
|
||||
issuer = server.URL
|
||||
|
||||
client, err := NewOIDCPublicClient(OIDCPublicClientConfig{
|
||||
AppEnv: "test", Issuer: issuer, ClientID: "gateway-public",
|
||||
RedirectURI: "https://gateway.example.com/callback",
|
||||
PostLogoutRedirectURI: "https://gateway.example.com/",
|
||||
HTTPClient: server.Client(),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
authorizationURL, err := client.AuthorizationURL(
|
||||
context.Background(), "state", "nonce",
|
||||
"test-pkce-verifier-with-at-least-43-characters-1234", "", "",
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
parsed, _ := url.Parse(authorizationURL)
|
||||
if parsed.Query().Has("context_type") || parsed.Query().Has("tenant_hint") {
|
||||
t.Fatalf("unified login leaked optional context binding: %v", parsed.Query())
|
||||
}
|
||||
}
|
||||
|
||||
func TestOIDCPublicClientRejectsOfflineAccess(t *testing.T) {
|
||||
_, err := NewOIDCPublicClient(OIDCPublicClientConfig{
|
||||
AppEnv: "production",
|
||||
|
||||
@@ -66,11 +66,27 @@ func TestOIDCVerifierAcceptsRS256AndES256StableClaims(t *testing.T) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if user.ID != "platform-subject" || user.TenantID != "tenant-1" || user.Source != "oidc" ||
|
||||
user.Username != "acceptance" || user.DisplayName != "王小明" ||
|
||||
user.Email != "real.user@example.test" || user.Phone != "+8613800000000" ||
|
||||
user.AvatarURL != "https://static.example.test/avatar.png" ||
|
||||
len(user.Roles) != 1 || user.Roles[0] != "admin" {
|
||||
t.Fatalf("unexpected OIDC user: %#v", user)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
unverified := signedOIDCToken(t, issuer, "rsa-key", jwt.SigningMethodRS256, rsaKey, func(claims jwt.MapClaims) {
|
||||
claims["email_verified"] = false
|
||||
claims["phone_number_verified"] = false
|
||||
claims["picture"] = "javascript:alert(1)"
|
||||
})
|
||||
user, err := verifier.Verify(context.Background(), unverified)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if user.Email != "" || user.Phone != "" || user.AvatarURL != "" {
|
||||
t.Fatalf("untrusted profile claims were accepted: %#v", user)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOIDCVerifierRejectsMissingOrMismatchedSecurityClaims(t *testing.T) {
|
||||
@@ -91,20 +107,116 @@ func TestOIDCVerifierRejectsMissingOrMismatchedSecurityClaims(t *testing.T) {
|
||||
RequiredScopes: []string{"gateway.access"}, HTTPClient: server.Client(),
|
||||
})
|
||||
tests := []struct {
|
||||
name string
|
||||
mutate func(jwt.MapClaims)
|
||||
name, wantCategory string
|
||||
mutate func(jwt.MapClaims)
|
||||
}{
|
||||
{"missing nbf", func(claims jwt.MapClaims) { delete(claims, "nbf") }},
|
||||
{"wrong audience", func(claims jwt.MapClaims) { claims["aud"] = "other-api" }},
|
||||
{"wrong tenant", func(claims jwt.MapClaims) { claims["tid"] = "tenant-2" }},
|
||||
{"missing scope", func(claims jwt.MapClaims) { claims["scope"] = "openid" }},
|
||||
{"unmapped role", func(claims jwt.MapClaims) { claims["roles"] = []string{"other.admin"} }},
|
||||
{"missing nbf", "NBF_MISSING", func(claims jwt.MapClaims) { delete(claims, "nbf") }},
|
||||
{"missing exp", "EXP_MISSING", func(claims jwt.MapClaims) { delete(claims, "exp") }},
|
||||
{"missing issuer", "ISSUER_MISSING", func(claims jwt.MapClaims) { delete(claims, "iss") }},
|
||||
{"missing audience", "AUDIENCE_MISSING", func(claims jwt.MapClaims) { delete(claims, "aud") }},
|
||||
{"wrong audience", "AUDIENCE_INVALID", func(claims jwt.MapClaims) { claims["aud"] = "other-api" }},
|
||||
{"missing context type", "STABLE_IDENTITY_CLAIMS_INVALID", func(claims jwt.MapClaims) { delete(claims, "context_type") }},
|
||||
{"unknown context type", "STABLE_IDENTITY_CLAIMS_INVALID", func(claims jwt.MapClaims) { claims["context_type"] = "account" }},
|
||||
{"wrong tenant", "STABLE_IDENTITY_CLAIMS_INVALID", func(claims jwt.MapClaims) { claims["tid"] = "tenant-2" }},
|
||||
{"missing scope", "REQUIRED_SCOPE_MISSING", func(claims jwt.MapClaims) { claims["scope"] = "openid" }},
|
||||
{"unmapped role", "MAPPED_ROLE_MISSING", func(claims jwt.MapClaims) { claims["roles"] = []string{"other.admin"} }},
|
||||
}
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
raw := signedOIDCToken(t, issuer, "ec-key", jwt.SigningMethodES256, key, test.mutate)
|
||||
if _, err := verifier.Verify(context.Background(), raw); err == nil {
|
||||
t.Fatal("expected token to be rejected")
|
||||
} else {
|
||||
if !errors.Is(err, ErrUnauthorized) {
|
||||
t.Fatalf("validation error no longer wraps ErrUnauthorized: %v", err)
|
||||
}
|
||||
if category := OIDCValidationCategory(err); category != test.wantCategory {
|
||||
t.Fatalf("validation category=%q, want %q", category, test.wantCategory)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOIDCVerifierAcceptsExplicitPlatformAndTenantContextsForMultiTenantApplication(t *testing.T) {
|
||||
key, _ := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
|
||||
var issuer string
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
|
||||
if request.URL.Path == "/.well-known/openid-configuration" {
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"issuer": issuer,
|
||||
"jwks_uri": issuer + "/jwks",
|
||||
})
|
||||
return
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"keys": []any{ecJWK("ec-key", &key.PublicKey)},
|
||||
})
|
||||
}))
|
||||
defer server.Close()
|
||||
issuer = server.URL
|
||||
verifier, err := NewOIDCVerifier(OIDCConfig{
|
||||
AppEnv: "test",
|
||||
Issuer: issuer, Audience: "gateway-api",
|
||||
TenantMode: "multi_tenant",
|
||||
ApplicationID: "11111111-1111-4111-8111-111111111111",
|
||||
RolePrefix: "gateway.", HTTPClient: server.Client(),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
platformToken := signedOIDCToken(
|
||||
t,
|
||||
issuer,
|
||||
"ec-key",
|
||||
jwt.SigningMethodES256,
|
||||
key,
|
||||
func(claims jwt.MapClaims) {
|
||||
claims["context_type"] = "platform"
|
||||
delete(claims, "tid")
|
||||
},
|
||||
)
|
||||
platformUser, err := verifier.Verify(context.Background(), platformToken)
|
||||
if err != nil {
|
||||
t.Fatalf("platform context rejected: %v", err)
|
||||
}
|
||||
if platformUser.ContextType != "platform" || platformUser.TenantID != "" {
|
||||
t.Fatalf("platform user=%#v", platformUser)
|
||||
}
|
||||
tenantID := "22222222-2222-4222-8222-222222222222"
|
||||
tenantToken := signedOIDCToken(
|
||||
t,
|
||||
issuer,
|
||||
"ec-key",
|
||||
jwt.SigningMethodES256,
|
||||
key,
|
||||
func(claims jwt.MapClaims) {
|
||||
claims["context_type"] = "tenant"
|
||||
claims["tid"] = tenantID
|
||||
},
|
||||
)
|
||||
tenantUser, err := verifier.Verify(context.Background(), tenantToken)
|
||||
if err != nil {
|
||||
t.Fatalf("tenant context rejected: %v", err)
|
||||
}
|
||||
if tenantUser.ContextType != "tenant" || tenantUser.TenantID != tenantID {
|
||||
t.Fatalf("tenant user=%#v", tenantUser)
|
||||
}
|
||||
for name, mutate := range map[string]func(jwt.MapClaims){
|
||||
"platform with tenant": func(claims jwt.MapClaims) {
|
||||
claims["context_type"] = "platform"
|
||||
},
|
||||
"tenant without tenant": func(claims jwt.MapClaims) {
|
||||
claims["context_type"] = "tenant"
|
||||
delete(claims, "tid")
|
||||
},
|
||||
} {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
raw := signedOIDCToken(
|
||||
t, issuer, "ec-key", jwt.SigningMethodES256, key, mutate,
|
||||
)
|
||||
if _, err := verifier.Verify(context.Background(), raw); err == nil {
|
||||
t.Fatal("context/tenant mismatch was accepted")
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -302,8 +414,12 @@ func signedOIDCToken(t *testing.T, issuer, kid string, method jwt.SigningMethod,
|
||||
now := time.Now()
|
||||
claims := jwt.MapClaims{
|
||||
"iss": issuer, "aud": "gateway-api", "sub": "platform-subject", "tid": "tenant-1",
|
||||
"context_type": "tenant",
|
||||
"preferred_username": "acceptance", "roles": []string{"gateway.admin"},
|
||||
"scope": "openid gateway.access", "iat": now.Unix(), "nbf": now.Add(-time.Second).Unix(),
|
||||
"name": "王小明", "email": "real.user@example.test", "email_verified": true,
|
||||
"phone_number": "+8613800000000", "phone_number_verified": true,
|
||||
"picture": "https://static.example.test/avatar.png",
|
||||
"scope": "openid gateway.access", "iat": now.Unix(), "nbf": now.Add(-time.Second).Unix(),
|
||||
"exp": now.Add(time.Hour).Unix(),
|
||||
}
|
||||
if mutate != nil {
|
||||
|
||||
@@ -0,0 +1,40 @@
|
||||
package capacitycontroller
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
|
||||
)
|
||||
|
||||
func NewConfiguredAdapter(cfg config.Config) (OrchestratorAdapter, error) {
|
||||
switch strings.ToLower(strings.TrimSpace(cfg.CapacityOrchestratorAdapter)) {
|
||||
case "static":
|
||||
pools, err := cfg.CapacityPools()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
initial := make(map[string]int, len(pools))
|
||||
for _, pool := range pools {
|
||||
initial[pool.ID] = pool.BootstrapReplicas
|
||||
}
|
||||
return NewStaticAdapter(initial), nil
|
||||
case "", "kubernetes":
|
||||
return NewKubernetesClient(KubernetesConfig{
|
||||
Namespace: adapterEnv("AI_GATEWAY_CAPACITY_CONTROLLER_NAMESPACE", adapterEnv("POD_NAMESPACE", "easyai")),
|
||||
APIServer: adapterEnv("AI_GATEWAY_CAPACITY_CONTROLLER_API_SERVER", "https://kubernetes.default.svc"),
|
||||
TokenFile: adapterEnv("AI_GATEWAY_CAPACITY_CONTROLLER_TOKEN_FILE", "/var/run/secrets/kubernetes.io/serviceaccount/token"),
|
||||
CAFile: adapterEnv("AI_GATEWAY_CAPACITY_CONTROLLER_CA_FILE", "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt"),
|
||||
})
|
||||
default:
|
||||
return nil, errors.New("unsupported capacity orchestrator adapter")
|
||||
}
|
||||
}
|
||||
|
||||
func adapterEnv(name, fallback string) string {
|
||||
if value := strings.TrimSpace(os.Getenv(name)); value != "" {
|
||||
return value
|
||||
}
|
||||
return fallback
|
||||
}
|
||||
@@ -0,0 +1,454 @@
|
||||
package capacitycontroller
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"os"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/executionpool"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
"github.com/jackc/pgx/v5"
|
||||
)
|
||||
|
||||
const (
|
||||
controllerReconcileInterval = 10 * time.Second
|
||||
controllerLeadershipRetry = 5 * time.Second
|
||||
controllerDeletionCost = -1000
|
||||
)
|
||||
|
||||
type capacityStore interface {
|
||||
TryAcquireCapacityControllerLeadership(context.Context) (store.Leadership, bool, error)
|
||||
WorkerQueueRuntime(context.Context) (store.WorkerQueueRuntime, error)
|
||||
ListPoolQueueRuntime(context.Context) ([]store.PoolQueueRuntime, error)
|
||||
ListWorkerInstanceRuntime(context.Context) ([]store.WorkerInstanceRuntime, error)
|
||||
CapacityDatabaseHealth(context.Context) (store.CapacityDatabaseHealth, error)
|
||||
MarkWorkerDraining(context.Context, string) error
|
||||
ReactivateWorkerInstance(context.Context, string) error
|
||||
PublishDesiredCapacity(context.Context, executionpool.DesiredCapacity) error
|
||||
}
|
||||
|
||||
type OrchestratorAdapter interface {
|
||||
PoolState(context.Context, string, string) (PoolInfrastructureState, error)
|
||||
ScalePool(context.Context, string, string, int) error
|
||||
SetInstanceTerminationPriority(context.Context, string, int) error
|
||||
}
|
||||
|
||||
type Status struct {
|
||||
Leader bool `json:"leader"`
|
||||
LastRunAt time.Time `json:"lastRunAt,omitempty"`
|
||||
LastError string `json:"lastError,omitempty"`
|
||||
Queue store.WorkerQueueRuntime `json:"queue"`
|
||||
Plan Plan `json:"plan"`
|
||||
ScaleActions uint64 `json:"scaleActions"`
|
||||
}
|
||||
|
||||
type Controller struct {
|
||||
cfg config.Config
|
||||
store capacityStore
|
||||
orchestrator OrchestratorAdapter
|
||||
pools []config.ExecutionPoolCapacityConfig
|
||||
logger *slog.Logger
|
||||
expectedRevision string
|
||||
now func() time.Time
|
||||
highSince time.Time
|
||||
lowSince time.Time
|
||||
statusMu sync.RWMutex
|
||||
status Status
|
||||
}
|
||||
|
||||
func New(
|
||||
cfg config.Config,
|
||||
db capacityStore,
|
||||
orchestrator OrchestratorAdapter,
|
||||
logger *slog.Logger,
|
||||
) *Controller {
|
||||
return &Controller{
|
||||
cfg: cfg, store: db, orchestrator: orchestrator, logger: logger,
|
||||
expectedRevision: strings.TrimSpace(os.Getenv("AI_GATEWAY_REVISION")),
|
||||
now: time.Now,
|
||||
pools: mustCapacityPools(cfg),
|
||||
}
|
||||
}
|
||||
|
||||
func mustCapacityPools(cfg config.Config) []config.ExecutionPoolCapacityConfig {
|
||||
pools, _ := cfg.CapacityPools()
|
||||
return pools
|
||||
}
|
||||
|
||||
func (controller *Controller) Run(ctx context.Context) {
|
||||
for ctx.Err() == nil {
|
||||
leadership, acquired, err := controller.store.TryAcquireCapacityControllerLeadership(ctx)
|
||||
if err != nil {
|
||||
controller.setError(err)
|
||||
if !waitContext(ctx, controllerLeadershipRetry) {
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
if !acquired {
|
||||
controller.setLeader(false)
|
||||
controller.clearError()
|
||||
if !waitContext(ctx, controllerLeadershipRetry) {
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
controller.setLeader(true)
|
||||
controller.clearError()
|
||||
controller.runLeader(ctx, leadership)
|
||||
leadership.Release()
|
||||
controller.setLeader(false)
|
||||
}
|
||||
}
|
||||
|
||||
func (controller *Controller) runLeader(ctx context.Context, leadership store.Leadership) {
|
||||
ticker := time.NewTicker(controllerReconcileInterval)
|
||||
defer ticker.Stop()
|
||||
if err := controller.reconcile(ctx); err != nil {
|
||||
controller.logError("initial capacity reconciliation failed", err)
|
||||
}
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
keepAliveCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
err := leadership.KeepAlive(keepAliveCtx)
|
||||
cancel()
|
||||
if err != nil {
|
||||
controller.logError("capacity controller leadership lost", err)
|
||||
return
|
||||
}
|
||||
if err := controller.reconcile(ctx); err != nil {
|
||||
controller.logError("capacity reconciliation failed", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (controller *Controller) reconcile(ctx context.Context) error {
|
||||
queue, err := controller.store.WorkerQueueRuntime(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
poolQueues, err := controller.store.ListPoolQueueRuntime(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
demandByPool := make(map[string]int, len(poolQueues))
|
||||
for _, poolQueue := range poolQueues {
|
||||
demandByPool[poolQueue.PoolID] = poolQueue.Queued + poolQueue.Running
|
||||
}
|
||||
instances, err := controller.store.ListWorkerInstanceRuntime(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
database, err := controller.store.CapacityDatabaseHealth(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
now := controller.now()
|
||||
poolResources := make([]PoolResources, 0, len(controller.pools))
|
||||
infrastructureStates := make(map[string]PoolInfrastructureState, len(controller.pools))
|
||||
for _, pool := range controller.pools {
|
||||
if pool.MaxReplicas == 0 {
|
||||
continue
|
||||
}
|
||||
state, stateErr := controller.orchestrator.PoolState(ctx, pool.ID, pool.AdapterRef)
|
||||
if stateErr != nil {
|
||||
return stateErr
|
||||
}
|
||||
infrastructureStates[pool.ID] = state
|
||||
poolResources = append(poolResources, PoolResources{
|
||||
PoolID: pool.ID, CurrentReplicas: state.CurrentReplicas,
|
||||
MinReplicas: pool.MinReplicas, MaxReplicas: pool.MaxReplicas,
|
||||
Demand: demandByPool[pool.ID],
|
||||
AllocatableMemoryBytes: state.AllocatableMemoryBytes,
|
||||
UsedMemoryBytes: state.UsedMemoryBytes,
|
||||
WorkerRequestMemoryBytes: state.WorkerRequestMemoryBytes,
|
||||
AllocatableMilliCPU: state.AllocatableMilliCPU,
|
||||
UsedMilliCPU: state.UsedMilliCPU,
|
||||
WorkerRequestMilliCPU: state.WorkerRequestMilliCPU,
|
||||
MemoryPressure: state.MemoryPressure,
|
||||
Nodes: nodeResources(state.Nodes),
|
||||
})
|
||||
}
|
||||
currentTotal := 0
|
||||
for _, pool := range poolResources {
|
||||
currentTotal += pool.CurrentReplicas
|
||||
}
|
||||
target := controller.cfg.WorkerTargetOutstandingPerReplica
|
||||
if target < 1 {
|
||||
target = 2 * controller.cfg.AsyncWorkerInstanceHardLimit
|
||||
}
|
||||
rawDesired := (max(queue.Queued+queue.Running, 0) + target - 1) / target
|
||||
if rawDesired > currentTotal {
|
||||
if controller.highSince.IsZero() {
|
||||
controller.highSince = now
|
||||
}
|
||||
controller.lowSince = time.Time{}
|
||||
} else if rawDesired < currentTotal && queue.Queued == 0 {
|
||||
if controller.lowSince.IsZero() {
|
||||
controller.lowSince = now
|
||||
}
|
||||
controller.highSince = time.Time{}
|
||||
} else {
|
||||
controller.highSince = time.Time{}
|
||||
controller.lowSince = time.Time{}
|
||||
}
|
||||
revisionHealthy := controller.revisionsMatch(instances, infrastructureStates)
|
||||
scaleUpEligible := !controller.highSince.IsZero() &&
|
||||
now.Sub(controller.highSince) >= time.Duration(controller.cfg.WorkerScaleUpWindowSeconds)*time.Second &&
|
||||
revisionHealthy
|
||||
scaleDownEligible := !controller.lowSince.IsZero() &&
|
||||
now.Sub(controller.lowSince) >= time.Duration(controller.cfg.WorkerScaleDownStabilizationSeconds)*time.Second &&
|
||||
revisionHealthy
|
||||
plan := CalculatePlan(PlanInput{
|
||||
Queued: queue.Queued, Running: queue.Running,
|
||||
InstanceSlots: controller.cfg.AsyncWorkerInstanceHardLimit,
|
||||
TargetOutstandingPerReplica: target,
|
||||
MemoryTargetPercent: controller.cfg.NodeMemoryTargetPercent,
|
||||
MemoryHardPercent: controller.cfg.NodeMemoryHardPercent,
|
||||
CPUTargetPercent: controller.cfg.NodeCPUTargetPercent,
|
||||
DatabaseConnections: database.Connections,
|
||||
DatabaseConnectionBudget: controller.cfg.PostgresConnectionBudget,
|
||||
NonWorkerConnectionBudget: controller.cfg.PostgresNonWorkerConnectionBudget,
|
||||
WorkerDatabasePoolMax: controller.cfg.WorkerDatabaseMaxConns,
|
||||
SynchronousDatabasePeers: database.SynchronousPeers,
|
||||
ScaleUpEligible: scaleUpEligible,
|
||||
ScaleDownEligible: scaleDownEligible,
|
||||
Pools: poolResources,
|
||||
})
|
||||
if !revisionHealthy && plan.FrozenReason == "" {
|
||||
plan.FrozenReason = "release_revision_mismatch"
|
||||
}
|
||||
for _, poolPlan := range plan.Pools {
|
||||
reason := "capacity_plan"
|
||||
if plan.FrozenReason != "" {
|
||||
reason = "frozen:" + plan.FrozenReason
|
||||
}
|
||||
if err := controller.store.PublishDesiredCapacity(ctx, executionpool.DesiredCapacity{
|
||||
PoolID: poolPlan.PoolID, Desired: poolPlan.DesiredReplicas,
|
||||
Reason: reason, ValidUntil: now.Add(2 * controllerReconcileInterval),
|
||||
}); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if controller.cfg.WorkerAutoscalingEnabled {
|
||||
for _, poolPlan := range plan.Pools {
|
||||
switch {
|
||||
case poolPlan.DesiredReplicas > poolPlan.CurrentReplicas:
|
||||
if err := controller.orchestrator.ScalePool(ctx, poolPlan.PoolID, controller.adapterRef(poolPlan.PoolID), poolPlan.DesiredReplicas); err != nil {
|
||||
return err
|
||||
}
|
||||
controller.observeScale(poolPlan.PoolID, poolPlan.CurrentReplicas, poolPlan.DesiredReplicas, "scale_up")
|
||||
case poolPlan.DesiredReplicas < poolPlan.CurrentReplicas:
|
||||
if err := controller.reconcileScaleDown(ctx, now, poolPlan, instances); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
controller.statusMu.Lock()
|
||||
controller.status.LastRunAt = now
|
||||
controller.status.LastError = ""
|
||||
controller.status.Queue = queue
|
||||
controller.status.Plan = plan
|
||||
controller.statusMu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
func nodeResources(states []KubernetesNodeState) []NodeResources {
|
||||
nodes := make([]NodeResources, 0, len(states))
|
||||
for _, state := range states {
|
||||
nodes = append(nodes, NodeResources{
|
||||
NodeName: state.NodeName,
|
||||
CurrentReplicas: state.CurrentReplicas,
|
||||
AllocatableMemoryBytes: state.AllocatableMemoryBytes,
|
||||
UsedMemoryBytes: state.UsedMemoryBytes,
|
||||
WorkerUsedMemoryBytes: state.WorkerUsedMemoryBytes,
|
||||
AllocatableMilliCPU: state.AllocatableMilliCPU,
|
||||
UsedMilliCPU: state.UsedMilliCPU,
|
||||
WorkerUsedMilliCPU: state.WorkerUsedMilliCPU,
|
||||
MemoryPressure: state.MemoryPressure,
|
||||
})
|
||||
}
|
||||
return nodes
|
||||
}
|
||||
|
||||
func (controller *Controller) reconcileScaleDown(
|
||||
ctx context.Context,
|
||||
now time.Time,
|
||||
poolPlan PoolPlan,
|
||||
instances []store.WorkerInstanceRuntime,
|
||||
) error {
|
||||
for _, instance := range instances {
|
||||
if workerPoolID(instance) != poolPlan.PoolID || instance.Status != "draining" {
|
||||
continue
|
||||
}
|
||||
if instance.RunningTasks == 0 && instance.ActiveLeases == 0 {
|
||||
if strings.TrimSpace(instance.OrchestratorInstanceRef) == "" {
|
||||
return errors.New("worker orchestrator instance reference is required for scale down")
|
||||
}
|
||||
if err := controller.orchestrator.SetInstanceTerminationPriority(ctx, instance.OrchestratorInstanceRef, controllerDeletionCost); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := controller.orchestrator.ScalePool(ctx, poolPlan.PoolID, controller.adapterRef(poolPlan.PoolID), poolPlan.CurrentReplicas-1); err != nil {
|
||||
return err
|
||||
}
|
||||
controller.observeScale(poolPlan.PoolID, poolPlan.CurrentReplicas, poolPlan.CurrentReplicas-1, "drained_scale_down")
|
||||
return nil
|
||||
}
|
||||
if instance.DrainingAt != nil &&
|
||||
now.Sub(*instance.DrainingAt) >= time.Duration(controller.cfg.WorkerDrainTimeoutSeconds)*time.Second {
|
||||
if err := controller.store.ReactivateWorkerInstance(ctx, instance.InstanceID); err != nil && !errors.Is(err, pgx.ErrNoRows) {
|
||||
return err
|
||||
}
|
||||
controller.observeScale(poolPlan.PoolID, poolPlan.CurrentReplicas, poolPlan.CurrentReplicas, "drain_timeout")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
var candidate *store.WorkerInstanceRuntime
|
||||
for index := range instances {
|
||||
instance := &instances[index]
|
||||
if workerPoolID(*instance) != poolPlan.PoolID || instance.Status != "active" {
|
||||
continue
|
||||
}
|
||||
if candidate == nil ||
|
||||
instance.RunningTasks+instance.ActiveLeases < candidate.RunningTasks+candidate.ActiveLeases {
|
||||
candidate = instance
|
||||
}
|
||||
}
|
||||
if candidate == nil {
|
||||
return nil
|
||||
}
|
||||
if err := controller.store.MarkWorkerDraining(ctx, candidate.InstanceID); err != nil && !errors.Is(err, pgx.ErrNoRows) {
|
||||
return err
|
||||
}
|
||||
controller.observeScale(poolPlan.PoolID, poolPlan.CurrentReplicas, poolPlan.CurrentReplicas, "drain_started")
|
||||
return nil
|
||||
}
|
||||
|
||||
func workerPoolID(instance store.WorkerInstanceRuntime) string {
|
||||
if strings.TrimSpace(instance.PoolID) != "" {
|
||||
return instance.PoolID
|
||||
}
|
||||
return instance.Site
|
||||
}
|
||||
|
||||
func (controller *Controller) adapterRef(poolID string) string {
|
||||
for _, pool := range controller.pools {
|
||||
if pool.ID == poolID {
|
||||
return pool.AdapterRef
|
||||
}
|
||||
}
|
||||
return poolID
|
||||
}
|
||||
|
||||
func (controller *Controller) revisionsMatch(
|
||||
instances []store.WorkerInstanceRuntime,
|
||||
states map[string]PoolInfrastructureState,
|
||||
) bool {
|
||||
if controller.expectedRevision == "" {
|
||||
return true
|
||||
}
|
||||
for _, state := range states {
|
||||
if state.Revision != "" && state.Revision != controller.expectedRevision {
|
||||
return false
|
||||
}
|
||||
}
|
||||
for _, instance := range instances {
|
||||
if instance.Revision != "" && instance.Revision != controller.expectedRevision {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (controller *Controller) Status() Status {
|
||||
controller.statusMu.RLock()
|
||||
defer controller.statusMu.RUnlock()
|
||||
return controller.status
|
||||
}
|
||||
|
||||
func (controller *Controller) Handler() http.Handler {
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("GET /healthz", func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"ok":true,"service":"easyai-capacity-controller"}`))
|
||||
})
|
||||
mux.HandleFunc("GET /readyz", func(w http.ResponseWriter, _ *http.Request) {
|
||||
status := controller.Status()
|
||||
// Followers are intentionally idle while the database advisory lock is
|
||||
// held by the elected leader. They are still ready to take over and must
|
||||
// not make a two-replica Deployment permanently fail its rollout.
|
||||
if status.LastError != "" {
|
||||
http.Error(w, `{"ok":false}`, http.StatusServiceUnavailable)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"ok":true}`))
|
||||
})
|
||||
mux.HandleFunc("GET /status", func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(controller.Status())
|
||||
})
|
||||
return mux
|
||||
}
|
||||
|
||||
func (controller *Controller) setLeader(leader bool) {
|
||||
controller.statusMu.Lock()
|
||||
controller.status.Leader = leader
|
||||
controller.statusMu.Unlock()
|
||||
}
|
||||
|
||||
func (controller *Controller) setError(err error) {
|
||||
controller.statusMu.Lock()
|
||||
controller.status.LastError = err.Error()
|
||||
controller.statusMu.Unlock()
|
||||
}
|
||||
|
||||
func (controller *Controller) clearError() {
|
||||
controller.statusMu.Lock()
|
||||
controller.status.LastError = ""
|
||||
controller.statusMu.Unlock()
|
||||
}
|
||||
|
||||
func (controller *Controller) logError(message string, err error) {
|
||||
controller.setError(err)
|
||||
if controller.logger != nil {
|
||||
controller.logger.Error(message, "error", err)
|
||||
}
|
||||
}
|
||||
|
||||
func (controller *Controller) observeScale(site string, from int, to int, reason string) {
|
||||
controller.statusMu.Lock()
|
||||
controller.status.ScaleActions++
|
||||
controller.statusMu.Unlock()
|
||||
if controller.logger != nil {
|
||||
controller.logger.Info("worker capacity action",
|
||||
"site", site,
|
||||
"fromReplicas", from,
|
||||
"toReplicas", to,
|
||||
"reason", reason,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
func waitContext(ctx context.Context, duration time.Duration) bool {
|
||||
timer := time.NewTimer(duration)
|
||||
defer timer.Stop()
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return false
|
||||
case <-timer.C:
|
||||
return true
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,482 @@
|
||||
package capacitycontroller
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"crypto/x509"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
type KubernetesConfig struct {
|
||||
Namespace string
|
||||
APIServer string
|
||||
TokenFile string
|
||||
CAFile string
|
||||
HTTPClient *http.Client
|
||||
}
|
||||
|
||||
type KubernetesClient struct {
|
||||
namespace string
|
||||
baseURL string
|
||||
tokenFile string
|
||||
client *http.Client
|
||||
}
|
||||
|
||||
type PoolInfrastructureState struct {
|
||||
PoolID string
|
||||
NodeName string
|
||||
CurrentReplicas int
|
||||
Revision string
|
||||
Image string
|
||||
AllocatableMemoryBytes int64
|
||||
UsedMemoryBytes int64
|
||||
WorkerRequestMemoryBytes int64
|
||||
AllocatableMilliCPU int64
|
||||
UsedMilliCPU int64
|
||||
WorkerRequestMilliCPU int64
|
||||
MemoryPressure bool
|
||||
Nodes []PoolNodeState
|
||||
}
|
||||
|
||||
type PoolNodeState struct {
|
||||
NodeName string
|
||||
CurrentReplicas int
|
||||
AllocatableMemoryBytes int64
|
||||
UsedMemoryBytes int64
|
||||
WorkerUsedMemoryBytes int64
|
||||
AllocatableMilliCPU int64
|
||||
UsedMilliCPU int64
|
||||
WorkerUsedMilliCPU int64
|
||||
MemoryPressure bool
|
||||
}
|
||||
|
||||
type KubernetesSiteState = PoolInfrastructureState
|
||||
type KubernetesNodeState = PoolNodeState
|
||||
|
||||
var _ OrchestratorAdapter = (*KubernetesClient)(nil)
|
||||
|
||||
func (client *KubernetesClient) PoolState(ctx context.Context, poolID, adapterRef string) (PoolInfrastructureState, error) {
|
||||
state, err := client.SiteState(ctx, adapterRef)
|
||||
state.PoolID = poolID
|
||||
return state, err
|
||||
}
|
||||
|
||||
func (client *KubernetesClient) ScalePool(ctx context.Context, _, adapterRef string, replicas int) error {
|
||||
return client.ScaleWorkerDeployment(ctx, adapterRef, replicas)
|
||||
}
|
||||
|
||||
func (client *KubernetesClient) SetInstanceTerminationPriority(ctx context.Context, instanceRef string, priority int) error {
|
||||
return client.SetPodDeletionCost(ctx, instanceRef, priority)
|
||||
}
|
||||
|
||||
func NewKubernetesClient(config KubernetesConfig) (*KubernetesClient, error) {
|
||||
if strings.TrimSpace(config.Namespace) == "" {
|
||||
return nil, errors.New("capacity controller Kubernetes namespace is required")
|
||||
}
|
||||
if config.APIServer == "" {
|
||||
config.APIServer = "https://kubernetes.default.svc"
|
||||
}
|
||||
if config.TokenFile == "" {
|
||||
config.TokenFile = "/var/run/secrets/kubernetes.io/serviceaccount/token"
|
||||
}
|
||||
parsed, err := url.Parse(strings.TrimRight(config.APIServer, "/"))
|
||||
if err != nil || parsed.Host == "" || parsed.User != nil || parsed.RawQuery != "" || parsed.Fragment != "" ||
|
||||
(parsed.Scheme != "https" && config.HTTPClient == nil) {
|
||||
return nil, errors.New("capacity controller Kubernetes API server URL is invalid")
|
||||
}
|
||||
client := config.HTTPClient
|
||||
if client == nil {
|
||||
if config.CAFile == "" {
|
||||
config.CAFile = "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt"
|
||||
}
|
||||
certificate, readErr := os.ReadFile(config.CAFile)
|
||||
if readErr != nil {
|
||||
return nil, fmt.Errorf("read Kubernetes service account CA: %w", readErr)
|
||||
}
|
||||
roots := x509.NewCertPool()
|
||||
if !roots.AppendCertsFromPEM(certificate) {
|
||||
return nil, errors.New("Kubernetes service account CA is invalid")
|
||||
}
|
||||
transport := http.DefaultTransport.(*http.Transport).Clone()
|
||||
transport.Proxy = nil
|
||||
transport.DisableCompression = true
|
||||
transport.TLSClientConfig = &tls.Config{MinVersion: tls.VersionTLS12, RootCAs: roots}
|
||||
client = &http.Client{
|
||||
Timeout: 10 * time.Second,
|
||||
Transport: transport,
|
||||
CheckRedirect: func(*http.Request, []*http.Request) error {
|
||||
return http.ErrUseLastResponse
|
||||
},
|
||||
}
|
||||
}
|
||||
return &KubernetesClient{
|
||||
namespace: config.Namespace,
|
||||
baseURL: strings.TrimRight(config.APIServer, "/"),
|
||||
tokenFile: config.TokenFile,
|
||||
client: client,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (client *KubernetesClient) SiteState(ctx context.Context, site string) (KubernetesSiteState, error) {
|
||||
site = strings.TrimSpace(site)
|
||||
if site == "" {
|
||||
return KubernetesSiteState{}, errors.New("site is required")
|
||||
}
|
||||
var deployment struct {
|
||||
Metadata struct {
|
||||
Name string `json:"name"`
|
||||
} `json:"metadata"`
|
||||
Spec struct {
|
||||
Replicas int `json:"replicas"`
|
||||
Template struct {
|
||||
Spec struct {
|
||||
Containers []struct {
|
||||
Name string `json:"name"`
|
||||
Image string `json:"image"`
|
||||
Env []struct {
|
||||
Name string `json:"name"`
|
||||
Value string `json:"value"`
|
||||
} `json:"env"`
|
||||
Resources struct {
|
||||
Requests map[string]string `json:"requests"`
|
||||
} `json:"resources"`
|
||||
} `json:"containers"`
|
||||
} `json:"spec"`
|
||||
} `json:"template"`
|
||||
} `json:"spec"`
|
||||
}
|
||||
deploymentPath := fmt.Sprintf(
|
||||
"/apis/apps/v1/namespaces/%s/deployments/easyai-worker-%s",
|
||||
url.PathEscape(client.namespace),
|
||||
url.PathEscape(site),
|
||||
)
|
||||
if err := client.getJSON(ctx, deploymentPath, &deployment); err != nil {
|
||||
return KubernetesSiteState{}, err
|
||||
}
|
||||
state := KubernetesSiteState{PoolID: site, CurrentReplicas: deployment.Spec.Replicas}
|
||||
for _, container := range deployment.Spec.Template.Spec.Containers {
|
||||
if container.Name != "worker" {
|
||||
continue
|
||||
}
|
||||
state.Image = container.Image
|
||||
state.WorkerRequestMemoryBytes, _ = parseBinaryQuantity(container.Resources.Requests["memory"])
|
||||
state.WorkerRequestMilliCPU, _ = parseCPUQuantity(container.Resources.Requests["cpu"])
|
||||
for _, environment := range container.Env {
|
||||
if environment.Name == "AI_GATEWAY_REVISION" {
|
||||
state.Revision = environment.Value
|
||||
}
|
||||
}
|
||||
}
|
||||
var nodes struct {
|
||||
Items []struct {
|
||||
Metadata struct {
|
||||
Name string `json:"name"`
|
||||
} `json:"metadata"`
|
||||
Status struct {
|
||||
Allocatable map[string]string `json:"allocatable"`
|
||||
Conditions []struct {
|
||||
Type string `json:"type"`
|
||||
Status string `json:"status"`
|
||||
} `json:"conditions"`
|
||||
} `json:"status"`
|
||||
} `json:"items"`
|
||||
}
|
||||
nodePath := "/api/v1/nodes?labelSelector=" + url.QueryEscape(
|
||||
"easyai.io/site="+site+",easyai.io/worker=true",
|
||||
)
|
||||
if err := client.getJSON(ctx, nodePath, &nodes); err != nil {
|
||||
return KubernetesSiteState{}, err
|
||||
}
|
||||
readyNodes := nodes.Items[:0]
|
||||
for _, node := range nodes.Items {
|
||||
ready := false
|
||||
for _, condition := range node.Status.Conditions {
|
||||
if condition.Type == "Ready" && condition.Status == "True" {
|
||||
ready = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if ready {
|
||||
readyNodes = append(readyNodes, node)
|
||||
}
|
||||
}
|
||||
nodes.Items = readyNodes
|
||||
if len(nodes.Items) == 0 {
|
||||
// A zero-replica Deployment with no eligible nodes represents an
|
||||
// explicitly disabled site. It must remain visible to the planner, but it
|
||||
// has no node or metrics budget to collect. A non-zero Deployment still
|
||||
// fails closed so replicas cannot disappear from capacity accounting.
|
||||
if state.CurrentReplicas == 0 {
|
||||
return state, nil
|
||||
}
|
||||
return KubernetesSiteState{}, fmt.Errorf("site %s has no ready matching nodes", site)
|
||||
}
|
||||
var pods struct {
|
||||
Items []struct {
|
||||
Metadata struct {
|
||||
Name string `json:"name"`
|
||||
DeletionTimestamp *time.Time `json:"deletionTimestamp"`
|
||||
} `json:"metadata"`
|
||||
Spec struct {
|
||||
NodeName string `json:"nodeName"`
|
||||
} `json:"spec"`
|
||||
} `json:"items"`
|
||||
}
|
||||
podPath := fmt.Sprintf(
|
||||
"/api/v1/namespaces/%s/pods?labelSelector=%s",
|
||||
url.PathEscape(client.namespace),
|
||||
url.QueryEscape("app.kubernetes.io/name=easyai-worker,easyai.io/site="+site),
|
||||
)
|
||||
if err := client.getJSON(ctx, podPath, &pods); err != nil {
|
||||
return KubernetesSiteState{}, err
|
||||
}
|
||||
replicasByNode := make(map[string]int, len(nodes.Items))
|
||||
nodeByPod := make(map[string]string, len(pods.Items))
|
||||
for _, pod := range pods.Items {
|
||||
if pod.Metadata.DeletionTimestamp == nil && pod.Spec.NodeName != "" {
|
||||
replicasByNode[pod.Spec.NodeName]++
|
||||
nodeByPod[pod.Metadata.Name] = pod.Spec.NodeName
|
||||
}
|
||||
}
|
||||
var podMetrics struct {
|
||||
Items []struct {
|
||||
Metadata struct {
|
||||
Name string `json:"name"`
|
||||
} `json:"metadata"`
|
||||
Containers []struct {
|
||||
Usage map[string]string `json:"usage"`
|
||||
} `json:"containers"`
|
||||
} `json:"items"`
|
||||
}
|
||||
podMetricsPath := fmt.Sprintf(
|
||||
"/apis/metrics.k8s.io/v1beta1/namespaces/%s/pods?labelSelector=%s",
|
||||
url.PathEscape(client.namespace),
|
||||
url.QueryEscape("app.kubernetes.io/name=easyai-worker,easyai.io/site="+site),
|
||||
)
|
||||
if err := client.getJSON(ctx, podMetricsPath, &podMetrics); err != nil {
|
||||
return KubernetesSiteState{}, err
|
||||
}
|
||||
workerMemoryByNode := make(map[string]int64, len(nodes.Items))
|
||||
workerCPUByNode := make(map[string]int64, len(nodes.Items))
|
||||
measuredPods := make(map[string]bool, len(podMetrics.Items))
|
||||
for _, pod := range podMetrics.Items {
|
||||
nodeName := nodeByPod[pod.Metadata.Name]
|
||||
if nodeName == "" {
|
||||
continue
|
||||
}
|
||||
measuredPods[pod.Metadata.Name] = true
|
||||
for _, container := range pod.Containers {
|
||||
memory, memoryErr := parseBinaryQuantity(container.Usage["memory"])
|
||||
cpu, cpuErr := parseCPUQuantity(container.Usage["cpu"])
|
||||
if memoryErr != nil || cpuErr != nil {
|
||||
return KubernetesSiteState{}, fmt.Errorf(
|
||||
"site %s pod %s returned incomplete resource metrics",
|
||||
site,
|
||||
pod.Metadata.Name,
|
||||
)
|
||||
}
|
||||
workerMemoryByNode[nodeName] += memory
|
||||
workerCPUByNode[nodeName] += cpu
|
||||
}
|
||||
}
|
||||
for podName := range nodeByPod {
|
||||
if !measuredPods[podName] {
|
||||
return KubernetesSiteState{}, fmt.Errorf(
|
||||
"site %s pod %s has no resource metrics yet",
|
||||
site,
|
||||
podName,
|
||||
)
|
||||
}
|
||||
}
|
||||
for _, node := range nodes.Items {
|
||||
nodeState := KubernetesNodeState{
|
||||
NodeName: node.Metadata.Name,
|
||||
CurrentReplicas: replicasByNode[node.Metadata.Name],
|
||||
WorkerUsedMemoryBytes: workerMemoryByNode[node.Metadata.Name],
|
||||
WorkerUsedMilliCPU: workerCPUByNode[node.Metadata.Name],
|
||||
}
|
||||
nodeState.AllocatableMemoryBytes, _ = parseBinaryQuantity(node.Status.Allocatable["memory"])
|
||||
nodeState.AllocatableMilliCPU, _ = parseCPUQuantity(node.Status.Allocatable["cpu"])
|
||||
for _, condition := range node.Status.Conditions {
|
||||
if condition.Type == "MemoryPressure" && condition.Status == "True" {
|
||||
nodeState.MemoryPressure = true
|
||||
state.MemoryPressure = true
|
||||
}
|
||||
}
|
||||
var metrics struct {
|
||||
Usage map[string]string `json:"usage"`
|
||||
}
|
||||
metricsPath := "/apis/metrics.k8s.io/v1beta1/nodes/" + url.PathEscape(nodeState.NodeName)
|
||||
if err := client.getJSON(ctx, metricsPath, &metrics); err != nil {
|
||||
return KubernetesSiteState{}, err
|
||||
}
|
||||
nodeState.UsedMemoryBytes, _ = parseBinaryQuantity(metrics.Usage["memory"])
|
||||
nodeState.UsedMilliCPU, _ = parseCPUQuantity(metrics.Usage["cpu"])
|
||||
if nodeState.AllocatableMemoryBytes <= 0 || nodeState.AllocatableMilliCPU <= 0 {
|
||||
return KubernetesSiteState{}, fmt.Errorf(
|
||||
"site %s node %s returned incomplete resource quantities",
|
||||
site,
|
||||
nodeState.NodeName,
|
||||
)
|
||||
}
|
||||
state.Nodes = append(state.Nodes, nodeState)
|
||||
state.AllocatableMemoryBytes += nodeState.AllocatableMemoryBytes
|
||||
state.AllocatableMilliCPU += nodeState.AllocatableMilliCPU
|
||||
state.UsedMemoryBytes += nodeState.UsedMemoryBytes
|
||||
state.UsedMilliCPU += nodeState.UsedMilliCPU
|
||||
}
|
||||
if len(state.Nodes) == 1 {
|
||||
state.NodeName = state.Nodes[0].NodeName
|
||||
}
|
||||
if state.WorkerRequestMemoryBytes <= 0 || state.WorkerRequestMilliCPU <= 0 ||
|
||||
state.AllocatableMemoryBytes <= 0 || state.AllocatableMilliCPU <= 0 {
|
||||
return KubernetesSiteState{}, fmt.Errorf("site %s returned incomplete resource quantities", site)
|
||||
}
|
||||
return state, nil
|
||||
}
|
||||
|
||||
func (client *KubernetesClient) ScaleWorkerDeployment(ctx context.Context, site string, replicas int) error {
|
||||
if replicas < 0 || replicas > 64 {
|
||||
return errors.New("worker replica target must be between 0 and 64")
|
||||
}
|
||||
path := fmt.Sprintf(
|
||||
"/apis/apps/v1/namespaces/%s/deployments/easyai-worker-%s/scale",
|
||||
url.PathEscape(client.namespace),
|
||||
url.PathEscape(strings.TrimSpace(site)),
|
||||
)
|
||||
return client.patchJSON(ctx, path, map[string]any{"spec": map[string]any{"replicas": replicas}})
|
||||
}
|
||||
|
||||
func (client *KubernetesClient) SetPodDeletionCost(ctx context.Context, podName string, cost int) error {
|
||||
podName = strings.TrimSpace(podName)
|
||||
if podName == "" {
|
||||
return errors.New("worker pod name is required")
|
||||
}
|
||||
path := fmt.Sprintf(
|
||||
"/api/v1/namespaces/%s/pods/%s",
|
||||
url.PathEscape(client.namespace),
|
||||
url.PathEscape(podName),
|
||||
)
|
||||
return client.patchJSON(ctx, path, map[string]any{
|
||||
"metadata": map[string]any{
|
||||
"annotations": map[string]any{
|
||||
"controller.kubernetes.io/pod-deletion-cost": strconv.Itoa(cost),
|
||||
},
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func (client *KubernetesClient) getJSON(ctx context.Context, path string, target any) error {
|
||||
request, err := client.request(ctx, http.MethodGet, path, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
response, err := client.client.Do(request)
|
||||
if err != nil {
|
||||
return fmt.Errorf("request Kubernetes API: %w", err)
|
||||
}
|
||||
defer response.Body.Close()
|
||||
if response.StatusCode != http.StatusOK {
|
||||
_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 64*1024))
|
||||
return fmt.Errorf("Kubernetes API %s returned HTTP %d", path, response.StatusCode)
|
||||
}
|
||||
decoder := json.NewDecoder(io.LimitReader(response.Body, 4*1024*1024))
|
||||
if err := decoder.Decode(target); err != nil {
|
||||
return fmt.Errorf("decode Kubernetes API %s: %w", path, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (client *KubernetesClient) patchJSON(ctx context.Context, path string, payload any) error {
|
||||
encoded, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
request, err := client.request(ctx, http.MethodPatch, path, bytes.NewReader(encoded))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
request.Header.Set("Content-Type", "application/merge-patch+json")
|
||||
response, err := client.client.Do(request)
|
||||
if err != nil {
|
||||
return fmt.Errorf("patch Kubernetes API: %w", err)
|
||||
}
|
||||
defer response.Body.Close()
|
||||
_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 64*1024))
|
||||
if response.StatusCode < 200 || response.StatusCode >= 300 {
|
||||
return fmt.Errorf("Kubernetes API %s returned HTTP %d", path, response.StatusCode)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (client *KubernetesClient) request(ctx context.Context, method string, path string, body io.Reader) (*http.Request, error) {
|
||||
token, err := os.ReadFile(client.tokenFile)
|
||||
if err != nil || strings.TrimSpace(string(token)) == "" {
|
||||
clear(token)
|
||||
return nil, errors.New("Kubernetes service account token is unavailable")
|
||||
}
|
||||
request, err := http.NewRequestWithContext(ctx, method, client.baseURL+path, body)
|
||||
if err != nil {
|
||||
clear(token)
|
||||
return nil, err
|
||||
}
|
||||
request.Header.Set("Authorization", "Bearer "+strings.TrimSpace(string(token)))
|
||||
request.Header.Set("Accept", "application/json")
|
||||
clear(token)
|
||||
return request, nil
|
||||
}
|
||||
|
||||
func parseBinaryQuantity(value string) (int64, error) {
|
||||
value = strings.TrimSpace(value)
|
||||
if value == "" {
|
||||
return 0, errors.New("quantity is empty")
|
||||
}
|
||||
units := []struct {
|
||||
suffix string
|
||||
multiplier int64
|
||||
}{
|
||||
{"Ei", 1 << 60},
|
||||
{"Pi", 1 << 50},
|
||||
{"Ti", 1 << 40},
|
||||
{"Gi", 1 << 30},
|
||||
{"Mi", 1 << 20},
|
||||
{"Ki", 1 << 10},
|
||||
{"G", 1_000_000_000},
|
||||
{"M", 1_000_000},
|
||||
{"K", 1_000},
|
||||
}
|
||||
for _, unit := range units {
|
||||
if strings.HasSuffix(value, unit.suffix) {
|
||||
number, err := strconv.ParseFloat(strings.TrimSuffix(value, unit.suffix), 64)
|
||||
return int64(number * float64(unit.multiplier)), err
|
||||
}
|
||||
}
|
||||
return strconv.ParseInt(value, 10, 64)
|
||||
}
|
||||
|
||||
func parseCPUQuantity(value string) (int64, error) {
|
||||
value = strings.TrimSpace(value)
|
||||
if strings.HasSuffix(value, "n") {
|
||||
nanocores, err := strconv.ParseInt(strings.TrimSuffix(value, "n"), 10, 64)
|
||||
return nanocores / 1_000_000, err
|
||||
}
|
||||
if strings.HasSuffix(value, "u") {
|
||||
microcores, err := strconv.ParseInt(strings.TrimSuffix(value, "u"), 10, 64)
|
||||
return microcores / 1_000, err
|
||||
}
|
||||
if strings.HasSuffix(value, "m") {
|
||||
return strconv.ParseInt(strings.TrimSuffix(value, "m"), 10, 64)
|
||||
}
|
||||
cores, err := strconv.ParseFloat(value, 64)
|
||||
return int64(cores * 1000), err
|
||||
}
|
||||
@@ -0,0 +1,172 @@
|
||||
package capacitycontroller
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestKubernetesClientReadsSiteAndMutatesOnlyWorkerScale(t *testing.T) {
|
||||
var patches []string
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
|
||||
if request.Header.Get("Authorization") != "Bearer test-token" {
|
||||
http.Error(w, "unauthorized", http.StatusUnauthorized)
|
||||
return
|
||||
}
|
||||
switch {
|
||||
case request.Method == http.MethodGet && strings.Contains(request.URL.Path, "/deployments/easyai-worker-hongkong"):
|
||||
_, _ = w.Write([]byte(`{
|
||||
"spec":{"replicas":1,"template":{"spec":{"containers":[{
|
||||
"name":"worker","image":"registry.invalid/gateway@sha256:abc",
|
||||
"env":[{"name":"AI_GATEWAY_REVISION","value":"release-sha"}],
|
||||
"resources":{"requests":{"memory":"512Mi","cpu":"250m"}}
|
||||
}]}}}
|
||||
}`))
|
||||
case request.Method == http.MethodGet && request.URL.Path == "/api/v1/nodes":
|
||||
if request.URL.Query().Get("labelSelector") != "easyai.io/site=hongkong,easyai.io/worker=true" {
|
||||
http.Error(w, "unexpected Worker node selector", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
_, _ = w.Write([]byte(`{"items":[
|
||||
{"metadata":{"name":"easyai-hongkong"},"status":{
|
||||
"allocatable":{"memory":"8Gi","cpu":"4"},
|
||||
"conditions":[{"type":"Ready","status":"False"}]
|
||||
}},
|
||||
{"metadata":{"name":"easyai-hongkong-worker-2"},"status":{
|
||||
"allocatable":{"memory":"8Gi","cpu":"4"},
|
||||
"conditions":[{"type":"Ready","status":"True"},{"type":"MemoryPressure","status":"False"}]
|
||||
}}]}`))
|
||||
case request.Method == http.MethodGet && strings.Contains(request.URL.Path, "/api/v1/namespaces/easyai/pods"):
|
||||
_, _ = w.Write([]byte(`{"items":[{"metadata":{"name":"worker-hongkong-1"},"spec":{"nodeName":"easyai-hongkong-worker-2"}}]}`))
|
||||
case request.Method == http.MethodGet && strings.Contains(request.URL.Path, "/metrics.k8s.io/") &&
|
||||
strings.HasSuffix(request.URL.Path, "/nodes/easyai-hongkong-worker-2"):
|
||||
_, _ = w.Write([]byte(`{"usage":{"memory":"3Gi","cpu":"500m"}}`))
|
||||
case request.Method == http.MethodGet && strings.Contains(request.URL.Path, "/metrics.k8s.io/") &&
|
||||
strings.HasSuffix(request.URL.Path, "/nodes/easyai-hongkong"):
|
||||
t.Error("NotReady node metrics must not be queried")
|
||||
http.Error(w, "NotReady node", http.StatusInternalServerError)
|
||||
case request.Method == http.MethodGet && strings.Contains(request.URL.Path, "/metrics.k8s.io/") &&
|
||||
strings.HasSuffix(request.URL.Path, "/pods"):
|
||||
_, _ = w.Write([]byte(`{"items":[{
|
||||
"metadata":{"name":"worker-hongkong-1"},
|
||||
"containers":[{"usage":{"memory":"384Mi","cpu":"125m"}}]
|
||||
}]}`))
|
||||
case request.Method == http.MethodPatch:
|
||||
patches = append(patches, request.URL.Path)
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{}`))
|
||||
default:
|
||||
http.NotFound(w, request)
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
tokenFile := filepath.Join(t.TempDir(), "token")
|
||||
if err := os.WriteFile(tokenFile, []byte("test-token"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
client, err := NewKubernetesClient(KubernetesConfig{
|
||||
Namespace: "easyai", APIServer: server.URL, TokenFile: tokenFile, HTTPClient: server.Client(),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
state, err := client.SiteState(context.Background(), "hongkong")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if state.CurrentReplicas != 1 || state.NodeName != "easyai-hongkong-worker-2" ||
|
||||
state.WorkerRequestMemoryBytes != 512<<20 || state.WorkerRequestMilliCPU != 250 ||
|
||||
state.AllocatableMemoryBytes != 8<<30 || state.UsedMemoryBytes != 3<<30 ||
|
||||
state.Nodes[0].WorkerUsedMemoryBytes != 384<<20 || state.Nodes[0].WorkerUsedMilliCPU != 125 {
|
||||
t.Fatalf("site state=%+v", state)
|
||||
}
|
||||
if err := client.SetPodDeletionCost(context.Background(), "worker-pod", -1000); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := client.ScaleWorkerDeployment(context.Background(), "hongkong", 2); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(patches) != 2 ||
|
||||
!strings.Contains(patches[0], "/pods/worker-pod") ||
|
||||
!strings.Contains(patches[1], "/deployments/easyai-worker-hongkong/scale") {
|
||||
t.Fatalf("patches=%v", patches)
|
||||
}
|
||||
}
|
||||
|
||||
func TestKubernetesClientAllowsDisabledSiteWithoutWorkerNodes(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
|
||||
switch {
|
||||
case request.Method == http.MethodGet && strings.Contains(request.URL.Path, "/deployments/easyai-worker-ningbo"):
|
||||
_, _ = w.Write([]byte(`{
|
||||
"spec":{"replicas":0,"template":{"spec":{"containers":[{
|
||||
"name":"worker","image":"registry.invalid/gateway@sha256:abc",
|
||||
"env":[{"name":"AI_GATEWAY_REVISION","value":"release-sha"}],
|
||||
"resources":{"requests":{"memory":"1536Mi","cpu":"500m"}}
|
||||
}]}}}
|
||||
}`))
|
||||
case request.Method == http.MethodGet && request.URL.Path == "/api/v1/nodes":
|
||||
_, _ = w.Write([]byte(`{"items":[]}`))
|
||||
default:
|
||||
http.NotFound(w, request)
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
tokenFile := filepath.Join(t.TempDir(), "token")
|
||||
if err := os.WriteFile(tokenFile, []byte("test-token"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
client, err := NewKubernetesClient(KubernetesConfig{
|
||||
Namespace: "easyai", APIServer: server.URL, TokenFile: tokenFile, HTTPClient: server.Client(),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
state, err := client.SiteState(context.Background(), "ningbo")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if state.PoolID != "ningbo" || state.CurrentReplicas != 0 || len(state.Nodes) != 0 {
|
||||
t.Fatalf("site state=%+v", state)
|
||||
}
|
||||
}
|
||||
|
||||
func TestKubernetesClientRejectsActiveSiteWithoutWorkerNodes(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
|
||||
switch {
|
||||
case request.Method == http.MethodGet && strings.Contains(request.URL.Path, "/deployments/easyai-worker-ningbo"):
|
||||
_, _ = w.Write([]byte(`{"spec":{"replicas":1,"template":{"spec":{"containers":[]}}}}`))
|
||||
case request.Method == http.MethodGet && request.URL.Path == "/api/v1/nodes":
|
||||
_, _ = w.Write([]byte(`{"items":[]}`))
|
||||
default:
|
||||
http.NotFound(w, request)
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
tokenFile := filepath.Join(t.TempDir(), "token")
|
||||
if err := os.WriteFile(tokenFile, []byte("test-token"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
client, err := NewKubernetesClient(KubernetesConfig{
|
||||
Namespace: "easyai", APIServer: server.URL, TokenFile: tokenFile, HTTPClient: server.Client(),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := client.SiteState(context.Background(), "ningbo"); err == nil ||
|
||||
!strings.Contains(err.Error(), "site ningbo has no ready matching nodes") {
|
||||
t.Fatalf("expected missing Worker node error, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestQuantityParsing(t *testing.T) {
|
||||
if value, err := parseBinaryQuantity("1536Mi"); err != nil || value != 1536<<20 {
|
||||
t.Fatalf("memory=%d err=%v", value, err)
|
||||
}
|
||||
if value, err := parseCPUQuantity("250000000n"); err != nil || value != 250 {
|
||||
t.Fatalf("cpu=%d err=%v", value, err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,289 @@
|
||||
package capacitycontroller
|
||||
|
||||
import (
|
||||
"math"
|
||||
"sort"
|
||||
)
|
||||
|
||||
type PoolResources struct {
|
||||
PoolID string
|
||||
Demand int
|
||||
CurrentReplicas int
|
||||
MinReplicas int
|
||||
MaxReplicas int
|
||||
AllocatableMemoryBytes int64
|
||||
UsedMemoryBytes int64
|
||||
WorkerRequestMemoryBytes int64
|
||||
AllocatableMilliCPU int64
|
||||
UsedMilliCPU int64
|
||||
WorkerRequestMilliCPU int64
|
||||
MemoryPressure bool
|
||||
Nodes []NodeResources
|
||||
}
|
||||
|
||||
type NodeResources struct {
|
||||
NodeName string
|
||||
CurrentReplicas int
|
||||
AllocatableMemoryBytes int64
|
||||
UsedMemoryBytes int64
|
||||
WorkerUsedMemoryBytes int64
|
||||
AllocatableMilliCPU int64
|
||||
UsedMilliCPU int64
|
||||
WorkerUsedMilliCPU int64
|
||||
MemoryPressure bool
|
||||
}
|
||||
|
||||
type PlanInput struct {
|
||||
Queued int
|
||||
Running int
|
||||
InstanceSlots int
|
||||
TargetOutstandingPerReplica int
|
||||
MemoryTargetPercent int
|
||||
MemoryHardPercent int
|
||||
CPUTargetPercent int
|
||||
DatabaseConnections int
|
||||
DatabaseConnectionBudget int
|
||||
NonWorkerConnectionBudget int
|
||||
WorkerDatabasePoolMax int
|
||||
SynchronousDatabasePeers int
|
||||
ScaleUpEligible bool
|
||||
ScaleDownEligible bool
|
||||
Pools []PoolResources
|
||||
}
|
||||
|
||||
type PoolPlan struct {
|
||||
PoolID string `json:"poolId"`
|
||||
Demand int `json:"demand"`
|
||||
CurrentReplicas int `json:"currentReplicas"`
|
||||
DesiredReplicas int `json:"desiredReplicas"`
|
||||
ResourceMax int `json:"resourceMax"`
|
||||
MemoryPercent float64 `json:"memoryPercent"`
|
||||
CPUPercent float64 `json:"cpuPercent"`
|
||||
}
|
||||
|
||||
type Plan struct {
|
||||
RawDesired int `json:"rawDesired"`
|
||||
DesiredTotal int `json:"desiredTotal"`
|
||||
CurrentTotal int `json:"currentTotal"`
|
||||
FrozenReason string `json:"frozenReason,omitempty"`
|
||||
Pools []PoolPlan `json:"pools"`
|
||||
}
|
||||
|
||||
func CalculatePlan(input PlanInput) Plan {
|
||||
if input.InstanceSlots < 1 {
|
||||
input.InstanceSlots = 1
|
||||
}
|
||||
target := input.TargetOutstandingPerReplica
|
||||
if target < 1 {
|
||||
target = 2 * input.InstanceSlots
|
||||
}
|
||||
if input.MemoryTargetPercent < 1 {
|
||||
input.MemoryTargetPercent = 75
|
||||
}
|
||||
if input.MemoryHardPercent < 1 {
|
||||
input.MemoryHardPercent = 85
|
||||
}
|
||||
if input.CPUTargetPercent < 1 {
|
||||
input.CPUTargetPercent = 70
|
||||
}
|
||||
rawDesired := int(math.Ceil(float64(max(input.Queued+input.Running, 0)) / float64(target)))
|
||||
plan := Plan{RawDesired: rawDesired}
|
||||
minTotal := 0
|
||||
resourceMaxTotal := 0
|
||||
for _, pool := range input.Pools {
|
||||
for _, node := range pool.Nodes {
|
||||
if node.MemoryPressure {
|
||||
pool.MemoryPressure = true
|
||||
}
|
||||
}
|
||||
current := max(pool.CurrentReplicas, 0)
|
||||
minReplicas := max(pool.MinReplicas, 0)
|
||||
configMax := max(pool.MaxReplicas, minReplicas)
|
||||
resourceMax, memoryPercent, cpuPercent := poolResourceMaximum(
|
||||
pool,
|
||||
input.MemoryTargetPercent,
|
||||
input.CPUTargetPercent,
|
||||
)
|
||||
resourceMax = min(resourceMax, configMax)
|
||||
if resourceMax < minReplicas && plan.FrozenReason == "" {
|
||||
plan.FrozenReason = "pool_resource_budget"
|
||||
}
|
||||
resourceMax = max(resourceMax, minReplicas)
|
||||
plan.Pools = append(plan.Pools, PoolPlan{
|
||||
PoolID: pool.PoolID, CurrentReplicas: current, DesiredReplicas: minReplicas,
|
||||
Demand: pool.Demand,
|
||||
ResourceMax: resourceMax, MemoryPercent: memoryPercent, CPUPercent: cpuPercent,
|
||||
})
|
||||
plan.CurrentTotal += current
|
||||
minTotal += minReplicas
|
||||
resourceMaxTotal += resourceMax
|
||||
if pool.MemoryPressure || memoryPercent >= float64(input.MemoryHardPercent) {
|
||||
plan.FrozenReason = "node_memory_pressure"
|
||||
}
|
||||
}
|
||||
if input.DatabaseConnectionBudget > 0 && input.WorkerDatabasePoolMax > 0 {
|
||||
workerReplicaBudget := max(
|
||||
(input.DatabaseConnectionBudget-input.NonWorkerConnectionBudget)/input.WorkerDatabasePoolMax,
|
||||
0,
|
||||
)
|
||||
capPoolResourceMaxima(plan.Pools, workerReplicaBudget)
|
||||
resourceMaxTotal = 0
|
||||
for _, pool := range plan.Pools {
|
||||
resourceMaxTotal += pool.ResourceMax
|
||||
}
|
||||
if workerReplicaBudget < minTotal && plan.FrozenReason == "" {
|
||||
plan.FrozenReason = "database_connection_budget"
|
||||
}
|
||||
}
|
||||
desired := max(rawDesired, minTotal)
|
||||
desired = min(desired, resourceMaxTotal)
|
||||
|
||||
if input.SynchronousDatabasePeers < 1 && plan.FrozenReason == "" {
|
||||
plan.FrozenReason = "database_not_synchronous"
|
||||
}
|
||||
if desired > plan.CurrentTotal {
|
||||
if input.DatabaseConnectionBudget > 0 && input.WorkerDatabasePoolMax > 0 {
|
||||
additional := max(
|
||||
(input.DatabaseConnectionBudget-input.DatabaseConnections)/input.WorkerDatabasePoolMax,
|
||||
0,
|
||||
)
|
||||
desired = min(desired, plan.CurrentTotal+additional)
|
||||
}
|
||||
if !input.ScaleUpEligible || plan.FrozenReason != "" ||
|
||||
(input.DatabaseConnectionBudget > 0 && input.DatabaseConnections >= input.DatabaseConnectionBudget) {
|
||||
desired = plan.CurrentTotal
|
||||
if plan.FrozenReason == "" {
|
||||
if !input.ScaleUpEligible {
|
||||
plan.FrozenReason = "scale_up_window"
|
||||
} else {
|
||||
plan.FrozenReason = "database_connection_budget"
|
||||
}
|
||||
}
|
||||
} else {
|
||||
desired = min(desired, max(plan.CurrentTotal*2, plan.CurrentTotal+2))
|
||||
}
|
||||
} else if desired < plan.CurrentTotal && !input.ScaleDownEligible {
|
||||
desired = plan.CurrentTotal
|
||||
if plan.FrozenReason == "" {
|
||||
plan.FrozenReason = "scale_down_stabilization"
|
||||
}
|
||||
}
|
||||
desired = max(desired, minTotal)
|
||||
plan.DesiredTotal = desired
|
||||
distributeDesiredReplicas(plan.Pools, desired)
|
||||
return plan
|
||||
}
|
||||
|
||||
func capPoolResourceMaxima(pools []PoolPlan, total int) {
|
||||
minimum := 0
|
||||
original := make(map[string]int, len(pools))
|
||||
for index := range pools {
|
||||
original[pools[index].PoolID] = pools[index].ResourceMax
|
||||
pools[index].ResourceMax = pools[index].DesiredReplicas
|
||||
minimum += pools[index].ResourceMax
|
||||
}
|
||||
target := max(total, minimum)
|
||||
assigned := minimum
|
||||
for assigned < target {
|
||||
sort.SliceStable(pools, func(left, right int) bool {
|
||||
leftRoom := original[pools[left].PoolID] - pools[left].ResourceMax
|
||||
rightRoom := original[pools[right].PoolID] - pools[right].ResourceMax
|
||||
if leftRoom != rightRoom {
|
||||
return leftRoom > rightRoom
|
||||
}
|
||||
return pools[left].PoolID < pools[right].PoolID
|
||||
})
|
||||
if original[pools[0].PoolID] <= pools[0].ResourceMax {
|
||||
break
|
||||
}
|
||||
pools[0].ResourceMax++
|
||||
assigned++
|
||||
}
|
||||
}
|
||||
|
||||
func poolResourceMaximum(pool PoolResources, memoryTargetPercent int, cpuTargetPercent int) (int, float64, float64) {
|
||||
if len(pool.Nodes) > 0 {
|
||||
if pool.WorkerRequestMemoryBytes <= 0 || pool.WorkerRequestMilliCPU <= 0 {
|
||||
return 0, 0, 0
|
||||
}
|
||||
memoryMax := 0
|
||||
cpuMax := 0
|
||||
memoryPercent := float64(0)
|
||||
cpuPercent := float64(0)
|
||||
for _, node := range pool.Nodes {
|
||||
nodeMemoryPercent := usagePercent(node.UsedMemoryBytes, node.AllocatableMemoryBytes)
|
||||
nodeCPUPercent := usagePercent(node.UsedMilliCPU, node.AllocatableMilliCPU)
|
||||
memoryPercent = max(memoryPercent, nodeMemoryPercent)
|
||||
cpuPercent = max(cpuPercent, nodeCPUPercent)
|
||||
nonWorkerMemory := max(
|
||||
node.UsedMemoryBytes-node.WorkerUsedMemoryBytes,
|
||||
0,
|
||||
)
|
||||
memoryBudget := node.AllocatableMemoryBytes*int64(memoryTargetPercent)/100 - nonWorkerMemory
|
||||
memoryMax += int(max(memoryBudget, 0) / pool.WorkerRequestMemoryBytes)
|
||||
nonWorkerCPU := max(
|
||||
node.UsedMilliCPU-node.WorkerUsedMilliCPU,
|
||||
0,
|
||||
)
|
||||
cpuBudget := node.AllocatableMilliCPU*int64(cpuTargetPercent)/100 - nonWorkerCPU
|
||||
cpuMax += int(max(cpuBudget, 0) / pool.WorkerRequestMilliCPU)
|
||||
}
|
||||
return min(pool.MaxReplicas, min(memoryMax, cpuMax)), memoryPercent, cpuPercent
|
||||
}
|
||||
memoryPercent := usagePercent(pool.UsedMemoryBytes, pool.AllocatableMemoryBytes)
|
||||
cpuPercent := usagePercent(pool.UsedMilliCPU, pool.AllocatableMilliCPU)
|
||||
memoryMax := pool.MaxReplicas
|
||||
if pool.AllocatableMemoryBytes > 0 && pool.WorkerRequestMemoryBytes > 0 {
|
||||
nonWorker := max(pool.UsedMemoryBytes-int64(pool.CurrentReplicas)*pool.WorkerRequestMemoryBytes, 0)
|
||||
budget := pool.AllocatableMemoryBytes*int64(memoryTargetPercent)/100 - nonWorker
|
||||
memoryMax = int(max(budget, 0) / pool.WorkerRequestMemoryBytes)
|
||||
}
|
||||
cpuMax := pool.MaxReplicas
|
||||
if pool.AllocatableMilliCPU > 0 && pool.WorkerRequestMilliCPU > 0 {
|
||||
nonWorker := max(pool.UsedMilliCPU-int64(pool.CurrentReplicas)*pool.WorkerRequestMilliCPU, 0)
|
||||
budget := pool.AllocatableMilliCPU*int64(cpuTargetPercent)/100 - nonWorker
|
||||
cpuMax = int(max(budget, 0) / pool.WorkerRequestMilliCPU)
|
||||
}
|
||||
return min(pool.MaxReplicas, min(memoryMax, cpuMax)), memoryPercent, cpuPercent
|
||||
}
|
||||
|
||||
func distributeDesiredReplicas(pools []PoolPlan, desired int) {
|
||||
if len(pools) == 0 {
|
||||
return
|
||||
}
|
||||
assigned := 0
|
||||
for index := range pools {
|
||||
assigned += pools[index].DesiredReplicas
|
||||
}
|
||||
for assigned < desired {
|
||||
sort.SliceStable(pools, func(left, right int) bool {
|
||||
leftDemandPerReplica := float64(max(pools[left].Demand, 0)) / float64(max(pools[left].DesiredReplicas, 1))
|
||||
rightDemandPerReplica := float64(max(pools[right].Demand, 0)) / float64(max(pools[right].DesiredReplicas, 1))
|
||||
if leftDemandPerReplica != rightDemandPerReplica {
|
||||
return leftDemandPerReplica > rightDemandPerReplica
|
||||
}
|
||||
leftRoom := pools[left].ResourceMax - pools[left].DesiredReplicas
|
||||
rightRoom := pools[right].ResourceMax - pools[right].DesiredReplicas
|
||||
if leftRoom != rightRoom {
|
||||
return leftRoom > rightRoom
|
||||
}
|
||||
if pools[left].DesiredReplicas != pools[right].DesiredReplicas {
|
||||
return pools[left].DesiredReplicas < pools[right].DesiredReplicas
|
||||
}
|
||||
return pools[left].PoolID < pools[right].PoolID
|
||||
})
|
||||
if pools[0].DesiredReplicas >= pools[0].ResourceMax {
|
||||
break
|
||||
}
|
||||
pools[0].DesiredReplicas++
|
||||
assigned++
|
||||
}
|
||||
sort.Slice(pools, func(left, right int) bool { return pools[left].PoolID < pools[right].PoolID })
|
||||
}
|
||||
|
||||
func usagePercent(used int64, allocatable int64) float64 {
|
||||
if allocatable <= 0 {
|
||||
return 0
|
||||
}
|
||||
return float64(used) * 100 / float64(allocatable)
|
||||
}
|
||||
@@ -0,0 +1,164 @@
|
||||
package capacitycontroller
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestCalculatePlanRespectsResourceDatabaseAndStepLimits(t *testing.T) {
|
||||
input := PlanInput{
|
||||
Queued: 1000, Running: 48, InstanceSlots: 24,
|
||||
MemoryTargetPercent: 75, MemoryHardPercent: 85, CPUTargetPercent: 70,
|
||||
DatabaseConnections: 80, DatabaseConnectionBudget: 150, WorkerDatabasePoolMax: 24,
|
||||
SynchronousDatabasePeers: 1, ScaleUpEligible: true,
|
||||
Pools: []PoolResources{
|
||||
{
|
||||
PoolID: "ningbo", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4,
|
||||
AllocatableMemoryBytes: 8 << 30, UsedMemoryBytes: 6 << 30,
|
||||
WorkerRequestMemoryBytes: 1 << 30,
|
||||
AllocatableMilliCPU: 4000, UsedMilliCPU: 2000, WorkerRequestMilliCPU: 500,
|
||||
},
|
||||
{
|
||||
PoolID: "hongkong", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4,
|
||||
AllocatableMemoryBytes: 8 << 30, UsedMemoryBytes: 3 << 30,
|
||||
WorkerRequestMemoryBytes: 1 << 30,
|
||||
AllocatableMilliCPU: 4000, UsedMilliCPU: 500, WorkerRequestMilliCPU: 500,
|
||||
},
|
||||
},
|
||||
}
|
||||
plan := CalculatePlan(input)
|
||||
if plan.RawDesired != 22 {
|
||||
t.Fatalf("raw desired=%d, want 22", plan.RawDesired)
|
||||
}
|
||||
if plan.DesiredTotal != 4 {
|
||||
t.Fatalf("desired total=%d, want step-limited 4: %+v", plan.DesiredTotal, plan)
|
||||
}
|
||||
if plan.Pools[0].PoolID != "hongkong" || plan.Pools[0].DesiredReplicas != 3 {
|
||||
t.Fatalf("pool allocation=%+v, want hongkong=3 ningbo=1", plan.Pools)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCalculatePlanFreezesScaleUpOnHardMemoryPressure(t *testing.T) {
|
||||
plan := CalculatePlan(PlanInput{
|
||||
Queued: 200, InstanceSlots: 24,
|
||||
DatabaseConnections: 10, DatabaseConnectionBudget: 150, WorkerDatabasePoolMax: 24,
|
||||
SynchronousDatabasePeers: 1, ScaleUpEligible: true,
|
||||
Pools: []PoolResources{{
|
||||
PoolID: "ningbo", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4,
|
||||
AllocatableMemoryBytes: 8 << 30, UsedMemoryBytes: 7 << 30,
|
||||
WorkerRequestMemoryBytes: 1 << 30,
|
||||
}},
|
||||
})
|
||||
if plan.DesiredTotal != 1 || plan.FrozenReason != "node_memory_pressure" {
|
||||
t.Fatalf("plan=%+v, want frozen at one replica", plan)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCalculatePlanReportsConfiguredMinimumOutsideResourceBudget(t *testing.T) {
|
||||
plan := CalculatePlan(PlanInput{
|
||||
Queued: 200, InstanceSlots: 24,
|
||||
DatabaseConnections: 10, DatabaseConnectionBudget: 150, WorkerDatabasePoolMax: 24,
|
||||
SynchronousDatabasePeers: 1, ScaleUpEligible: true,
|
||||
Pools: []PoolResources{{
|
||||
PoolID: "ningbo", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4,
|
||||
AllocatableMemoryBytes: 8 << 30, UsedMemoryBytes: 13 << 29,
|
||||
WorkerRequestMemoryBytes: 2 << 30,
|
||||
AllocatableMilliCPU: 4000, UsedMilliCPU: 2800, WorkerRequestMilliCPU: 500,
|
||||
}},
|
||||
})
|
||||
if plan.DesiredTotal != 1 || plan.FrozenReason != "pool_resource_budget" {
|
||||
t.Fatalf("plan=%+v, want the existing minimum preserved and expansion frozen", plan)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCalculatePlanWaitsForSafeScaleDownWindow(t *testing.T) {
|
||||
input := PlanInput{
|
||||
InstanceSlots: 24, DatabaseConnections: 20, DatabaseConnectionBudget: 150,
|
||||
WorkerDatabasePoolMax: 24, SynchronousDatabasePeers: 1,
|
||||
ScaleDownEligible: false,
|
||||
Pools: []PoolResources{
|
||||
{PoolID: "ningbo", CurrentReplicas: 2, MinReplicas: 1, MaxReplicas: 4},
|
||||
{PoolID: "hongkong", CurrentReplicas: 2, MinReplicas: 1, MaxReplicas: 4},
|
||||
},
|
||||
}
|
||||
plan := CalculatePlan(input)
|
||||
if plan.DesiredTotal != 4 || plan.FrozenReason != "scale_down_stabilization" {
|
||||
t.Fatalf("plan=%+v, want current replicas during stabilization", plan)
|
||||
}
|
||||
input.ScaleDownEligible = true
|
||||
plan = CalculatePlan(input)
|
||||
if plan.DesiredTotal != 2 {
|
||||
t.Fatalf("desired total=%d, want min total 2", plan.DesiredTotal)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCalculatePlanAddsCapacityAcrossLabeledSiteNodes(t *testing.T) {
|
||||
plan := CalculatePlan(PlanInput{
|
||||
Queued: 500, InstanceSlots: 24,
|
||||
MemoryTargetPercent: 75, MemoryHardPercent: 85, CPUTargetPercent: 70,
|
||||
DatabaseConnections: 20, DatabaseConnectionBudget: 150, WorkerDatabasePoolMax: 20,
|
||||
SynchronousDatabasePeers: 1, ScaleUpEligible: true,
|
||||
Pools: []PoolResources{{
|
||||
PoolID: "hongkong", CurrentReplicas: 2, MinReplicas: 1, MaxReplicas: 8,
|
||||
WorkerRequestMemoryBytes: 1 << 30, WorkerRequestMilliCPU: 500,
|
||||
Nodes: []NodeResources{
|
||||
{
|
||||
NodeName: "hk-a", CurrentReplicas: 1,
|
||||
AllocatableMemoryBytes: 8 << 30, UsedMemoryBytes: 3 << 30,
|
||||
AllocatableMilliCPU: 4000, UsedMilliCPU: 1000,
|
||||
},
|
||||
{
|
||||
NodeName: "hk-b", CurrentReplicas: 1,
|
||||
AllocatableMemoryBytes: 8 << 30, UsedMemoryBytes: 2 << 30,
|
||||
AllocatableMilliCPU: 4000, UsedMilliCPU: 500,
|
||||
},
|
||||
},
|
||||
}},
|
||||
})
|
||||
if plan.Pools[0].ResourceMax != 7 {
|
||||
t.Fatalf("resource max=%d, want database-budgeted multi-node maximum: %+v", plan.Pools[0].ResourceMax, plan)
|
||||
}
|
||||
if plan.DesiredTotal != 4 {
|
||||
t.Fatalf("desired=%d, want two-wave step from 2 to 4", plan.DesiredTotal)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCalculatePlanCapsReplicaMaximumByDeclaredDatabasePools(t *testing.T) {
|
||||
plan := CalculatePlan(PlanInput{
|
||||
Queued: 500, InstanceSlots: 24,
|
||||
DatabaseConnections: 20, DatabaseConnectionBudget: 150,
|
||||
NonWorkerConnectionBudget: 72, WorkerDatabasePoolMax: 32,
|
||||
SynchronousDatabasePeers: 1, ScaleUpEligible: true,
|
||||
Pools: []PoolResources{
|
||||
{PoolID: "ningbo", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4},
|
||||
{PoolID: "hongkong", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4},
|
||||
},
|
||||
})
|
||||
if plan.DesiredTotal != 2 {
|
||||
t.Fatalf("desired=%d, want the current 1+1 database-safe topology: %+v", plan.DesiredTotal, plan)
|
||||
}
|
||||
resourceMax := 0
|
||||
for _, pool := range plan.Pools {
|
||||
resourceMax += pool.ResourceMax
|
||||
}
|
||||
if resourceMax != 2 {
|
||||
t.Fatalf("resource maximum=%d, want floor((150-72)/32)=2: %+v", resourceMax, plan)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPlanJSONUsesAcceptanceContractFieldNames(t *testing.T) {
|
||||
payload, err := json.Marshal(Plan{
|
||||
RawDesired: 3,
|
||||
Pools: []PoolPlan{{PoolID: "hongkong", ResourceMax: 2}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
text := string(payload)
|
||||
for _, field := range []string{`"rawDesired":3`, `"pools"`, `"poolId":"hongkong"`, `"resourceMax":2`} {
|
||||
if !strings.Contains(text, field) {
|
||||
t.Fatalf("plan JSON %s does not contain %s", text, field)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,42 @@
|
||||
package capacitycontroller
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// StaticAdapter proves that capacity coordination does not require Kubernetes.
|
||||
// It reports configured fixed replicas and intentionally ignores scale writes.
|
||||
type StaticAdapter struct {
|
||||
mu sync.RWMutex
|
||||
replicas map[string]int
|
||||
}
|
||||
|
||||
var _ OrchestratorAdapter = (*StaticAdapter)(nil)
|
||||
|
||||
func NewStaticAdapter(initial map[string]int) *StaticAdapter {
|
||||
copyOfInitial := make(map[string]int, len(initial))
|
||||
for poolID, replicas := range initial {
|
||||
copyOfInitial[poolID] = replicas
|
||||
}
|
||||
return &StaticAdapter{replicas: copyOfInitial}
|
||||
}
|
||||
|
||||
func (adapter *StaticAdapter) PoolState(_ context.Context, poolID, _ string) (PoolInfrastructureState, error) {
|
||||
adapter.mu.RLock()
|
||||
replicas := adapter.replicas[poolID]
|
||||
adapter.mu.RUnlock()
|
||||
return PoolInfrastructureState{
|
||||
PoolID: poolID, CurrentReplicas: replicas,
|
||||
AllocatableMemoryBytes: 1 << 60, WorkerRequestMemoryBytes: 1,
|
||||
AllocatableMilliCPU: 1 << 50, WorkerRequestMilliCPU: 1,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (adapter *StaticAdapter) ScalePool(context.Context, string, string, int) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (adapter *StaticAdapter) SetInstanceTerminationPriority(context.Context, string, int) error {
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
package capacitycontroller
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestStaticAdapterDoesNotRequireKubernetes(t *testing.T) {
|
||||
adapter := NewStaticAdapter(map[string]int{"pool-a": 2})
|
||||
state, err := adapter.PoolState(context.Background(), "pool-a", "ignored")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if state.PoolID != "pool-a" || state.CurrentReplicas != 2 {
|
||||
t.Fatalf("state=%+v", state)
|
||||
}
|
||||
if err := adapter.ScalePool(context.Background(), "pool-a", "ignored", 4); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
state, _ = adapter.PoolState(context.Background(), "pool-a", "ignored")
|
||||
if state.CurrentReplicas != 2 {
|
||||
t.Fatalf("static adapter scaled to %d", state.CurrentReplicas)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,282 @@
|
||||
package clients
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func TestVectorizerClientMultipartAndPrivateResumeState(t *testing.T) {
|
||||
var receivedPath string
|
||||
var receivedFields map[string]string
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
receivedPath = r.URL.Path
|
||||
if username, password, ok := r.BasicAuth(); !ok || username != "id" || password != "secret" {
|
||||
t.Fatalf("unexpected vectorizer auth")
|
||||
}
|
||||
if err := r.ParseMultipartForm(1 << 20); err != nil {
|
||||
t.Fatalf("parse multipart: %v", err)
|
||||
}
|
||||
receivedFields = map[string]string{}
|
||||
for key, values := range r.MultipartForm.Value {
|
||||
receivedFields[key] = values[0]
|
||||
}
|
||||
w.Header().Set("Content-Type", "image/svg+xml")
|
||||
w.Header().Set("X-Image-Token", "private-image-token")
|
||||
w.Header().Set("X-Receipt", "private-receipt")
|
||||
_, _ = io.WriteString(w, `<svg xmlns="http://www.w3.org/2000/svg"/>`)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
privateState := map[string]any{}
|
||||
response, err := (VectorizerClient{LookupIP: func(context.Context, string) ([]net.IPAddr, error) {
|
||||
return []net.IPAddr{{IP: net.ParseIP("93.184.216.34")}}, nil
|
||||
}}).Run(context.Background(), Request{
|
||||
Kind: "images.vectorize", Model: "easy-image-vectorizer-1",
|
||||
Body: map[string]any{"source": map[string]any{"url": "https://example.com/input.png"}, "format": "svg", "maxColors": 16, "cleanupLevel": "strong"},
|
||||
Candidate: storeCandidate(server.URL, "vectorizer", map[string]any{"accessKey": "id", "secretKey": "secret"}),
|
||||
OnRemoteTaskSubmitted: func(_ string, payload map[string]any) error { privateState = payload; return nil },
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("run vectorizer: %v", err)
|
||||
}
|
||||
if receivedPath != "/vectorize" || receivedFields["image.url"] == "" || receivedFields["output.file_format"] != "svg" || receivedFields["processing.max_colors"] != "16" {
|
||||
t.Fatalf("unexpected multipart request: path=%s fields=%+v", receivedPath, receivedFields)
|
||||
}
|
||||
if privateState["imageToken"] != "private-image-token" || privateState["receipt"] != "private-receipt" {
|
||||
t.Fatalf("private resume state missing: %+v", privateState)
|
||||
}
|
||||
serialized := strings.ToLower(toJSONForTest(response.Result))
|
||||
if strings.Contains(serialized, "private-image-token") || strings.Contains(serialized, "private-receipt") {
|
||||
t.Fatalf("private vectorizer state leaked into response: %s", serialized)
|
||||
}
|
||||
}
|
||||
|
||||
func TestVectorizerClientRejectsPrivateSourceBeforeProviderCall(t *testing.T) {
|
||||
providerCalls := 0
|
||||
server := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { providerCalls++ }))
|
||||
defer server.Close()
|
||||
_, err := (VectorizerClient{LookupIP: func(context.Context, string) ([]net.IPAddr, error) {
|
||||
return []net.IPAddr{{IP: net.ParseIP("127.0.0.1")}}, nil
|
||||
}}).Run(context.Background(), Request{
|
||||
Body: map[string]any{"source": map[string]any{"url": "https://assets.example.test/input.png"}},
|
||||
Candidate: storeCandidate(server.URL, "vectorizer", map[string]any{"apiKey": "key"}),
|
||||
})
|
||||
if err == nil || !strings.Contains(err.Error(), "blocked network") {
|
||||
t.Fatalf("expected blocked source error, got %v", err)
|
||||
}
|
||||
if providerCalls != 0 {
|
||||
t.Fatalf("provider received %d calls for blocked source", providerCalls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestVectorizerClientReusesImageTokenForExtraFormat(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/download" {
|
||||
t.Fatalf("unexpected path %s", r.URL.Path)
|
||||
}
|
||||
if err := r.ParseMultipartForm(1 << 20); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if r.FormValue("image.token") != "token" || r.FormValue("receipt") != "receipt" || r.FormValue("output.file_format") != "pdf" {
|
||||
t.Fatalf("unexpected reuse form: %+v", r.MultipartForm.Value)
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/pdf")
|
||||
_, _ = w.Write([]byte("%PDF-1.7\n"))
|
||||
}))
|
||||
defer server.Close()
|
||||
response, err := (VectorizerClient{}).Run(context.Background(), Request{
|
||||
Body: map[string]any{"_vectorizer_image_token": "token", "_vectorizer_receipt": "receipt", "format": "pdf"},
|
||||
Candidate: storeCandidate(server.URL, "vectorizer", map[string]any{"apiKey": "key"}),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
item := response.Result["data"].([]any)[0].(map[string]any)
|
||||
if item["type"] != "file" || item["mime_type"] != "application/pdf" {
|
||||
t.Fatalf("unexpected pdf result: %+v", item)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTopazClientUploadsPartsPollsAndValidatesOutput(t *testing.T) {
|
||||
var mu sync.Mutex
|
||||
called := map[string]int{}
|
||||
var server *httptest.Server
|
||||
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
mu.Lock()
|
||||
called[r.Method+" "+r.URL.Path]++
|
||||
mu.Unlock()
|
||||
switch {
|
||||
case r.Method == http.MethodGet && r.URL.Path == "/source.mp4":
|
||||
_, _ = w.Write([]byte("small-video-source"))
|
||||
case r.Method == http.MethodPost && r.URL.Path == "/video/":
|
||||
if r.Header.Get("X-API-Key") != "topaz-key" {
|
||||
t.Fatalf("missing Topaz API key")
|
||||
}
|
||||
_, _ = io.WriteString(w, `{"requestId":"job-1"}`)
|
||||
case r.Method == http.MethodPatch && r.URL.Path == "/video/job-1/accept":
|
||||
_, _ = io.WriteString(w, `{"urls":["`+server.URL+`/upload/1","`+server.URL+`/upload/2"]}`)
|
||||
case r.Method == http.MethodPut && strings.HasPrefix(r.URL.Path, "/upload/"):
|
||||
w.Header().Set("ETag", `"etag-`+strings.TrimPrefix(r.URL.Path, "/upload/")+`"`)
|
||||
w.WriteHeader(http.StatusOK)
|
||||
case r.Method == http.MethodPatch && r.URL.Path == "/video/job-1/complete-upload/":
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = io.WriteString(w, `{}`)
|
||||
case r.Method == http.MethodGet && r.URL.Path == "/video/job-1/status":
|
||||
_, _ = io.WriteString(w, `{"status":"completed","download":{"url":"`+server.URL+`/output.mp4"}}`)
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
phases := []string{}
|
||||
client := TopazClient{Probe: func(_ context.Context, target string) (TopazVideoMetadata, error) {
|
||||
if strings.HasPrefix(target, "http") {
|
||||
return TopazVideoMetadata{Width: 640, Height: 360, Duration: 3, FrameRate: 24, HasAudio: true}, nil
|
||||
}
|
||||
return TopazVideoMetadata{Width: 320, Height: 180, Duration: 3, FrameRate: 24, FrameCount: 72, HasAudio: true}, nil
|
||||
}}
|
||||
response, err := client.Run(context.Background(), Request{
|
||||
Model: "easy-proteus-standard-4",
|
||||
Body: map[string]any{"video_url": server.URL + "/source.mp4", "output_width": 640, "output_height": 360, "preserve_audio": true},
|
||||
Candidate: storeCandidateWithConfig(server.URL, "topaz", map[string]any{"apiKey": "topaz-key"}, map[string]any{"allowPrivateSourceDownloads": true, "pollIntervalMs": 1, "pollTimeoutMs": 100}),
|
||||
OnRemoteTaskSubmitted: func(_ string, payload map[string]any) error {
|
||||
phases = append(phases, payload["phase"].(string))
|
||||
return nil
|
||||
},
|
||||
OnRemoteTaskPolled: func(_ string, payload map[string]any) error {
|
||||
phases = append(phases, payload["phase"].(string))
|
||||
return nil
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("run Topaz: %v", err)
|
||||
}
|
||||
if response.RequestID != "job-1" || !containsStringForTest(phases, "created") || !containsStringForTest(phases, "uploaded") {
|
||||
t.Fatalf("unexpected Topaz response/phases: response=%+v phases=%+v", response, phases)
|
||||
}
|
||||
if called["PUT /upload/1"] != 1 || called["PUT /upload/2"] != 1 || called["PATCH /video/job-1/complete-upload/"] != 1 {
|
||||
t.Fatalf("upload lifecycle incomplete: %+v", called)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTopazClientResumesUploadedTaskWithoutDownloadingSource(t *testing.T) {
|
||||
var sourceCalls int
|
||||
var server *httptest.Server
|
||||
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path == "/source.mp4" {
|
||||
sourceCalls++
|
||||
}
|
||||
if r.URL.Path == "/video/job-resume/status" {
|
||||
_, _ = io.WriteString(w, `{"status":"completed","download":{"url":"`+server.URL+`/result.mp4"}}`)
|
||||
return
|
||||
}
|
||||
http.NotFound(w, r)
|
||||
}))
|
||||
defer server.Close()
|
||||
_, err := (TopazClient{Probe: func(context.Context, string) (TopazVideoMetadata, error) {
|
||||
return TopazVideoMetadata{Width: 640, Height: 360, Duration: 3, FrameRate: 24}, nil
|
||||
}}).Run(context.Background(), Request{
|
||||
RemoteTaskID: "job-resume", RemoteTaskPayload: map[string]any{"phase": "uploaded", "targetResolution": map[string]any{"width": 640, "height": 360}},
|
||||
Body: map[string]any{"video_url": server.URL + "/source.mp4"},
|
||||
Candidate: storeCandidateWithConfig(server.URL, "topaz", map[string]any{"apiKey": "topaz-key"}, map[string]any{"pollIntervalMs": 1, "pollTimeoutMs": 100}),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if sourceCalls != 0 {
|
||||
t.Fatalf("resume downloaded source %d times", sourceCalls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTopazClientReportsProviderFailureAndPollingTimeout(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
statusBody string
|
||||
timeoutMS int
|
||||
want string
|
||||
}{
|
||||
{name: "provider failure", statusBody: `{"status":"failed","message":"upstream rejected"}`, timeoutMS: 100, want: "upstream rejected"},
|
||||
{name: "poll timeout", statusBody: `{"status":"processing"}`, timeoutMS: 3, want: "polling timed out"},
|
||||
}
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/video/job/status" {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
_, _ = io.WriteString(w, test.statusBody)
|
||||
}))
|
||||
defer server.Close()
|
||||
_, err := (TopazClient{}).Run(context.Background(), Request{
|
||||
RemoteTaskID: "job", RemoteTaskPayload: map[string]any{"phase": "uploaded", "targetResolution": map[string]any{"width": 640, "height": 360}},
|
||||
Body: map[string]any{"video_url": "https://assets.example.test/source.mp4"},
|
||||
Candidate: storeCandidateWithConfig(server.URL, "topaz", map[string]any{"apiKey": "topaz-key"}, map[string]any{
|
||||
"pollIntervalMs": 1, "pollTimeoutMs": test.timeoutMS,
|
||||
}),
|
||||
})
|
||||
if err == nil || !strings.Contains(err.Error(), test.want) {
|
||||
t.Fatalf("expected %q error, got %v", test.want, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestTopazClientRejectsOversizedInputBeforeProbe(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path == "/source.mp4" {
|
||||
_, _ = io.WriteString(w, "too-large")
|
||||
return
|
||||
}
|
||||
http.NotFound(w, r)
|
||||
}))
|
||||
defer server.Close()
|
||||
probeCalls := 0
|
||||
_, err := (TopazClient{Probe: func(context.Context, string) (TopazVideoMetadata, error) {
|
||||
probeCalls++
|
||||
return TopazVideoMetadata{}, nil
|
||||
}}).Run(context.Background(), Request{
|
||||
Body: map[string]any{"video_url": server.URL + "/source.mp4"},
|
||||
Candidate: storeCandidateWithConfig(server.URL, "topaz", map[string]any{"apiKey": "topaz-key"}, map[string]any{
|
||||
"allowPrivateSourceDownloads": true, "maxInputBytes": 3,
|
||||
}),
|
||||
})
|
||||
if err == nil || !strings.Contains(err.Error(), "size limit") {
|
||||
t.Fatalf("expected input size limit error, got %v", err)
|
||||
}
|
||||
if probeCalls != 0 {
|
||||
t.Fatalf("oversized input was probed %d times", probeCalls)
|
||||
}
|
||||
}
|
||||
|
||||
func storeCandidate(baseURL, specType string, credentials map[string]any) store.RuntimeModelCandidate {
|
||||
return storeCandidateWithConfig(baseURL, specType, credentials, nil)
|
||||
}
|
||||
|
||||
func storeCandidateWithConfig(baseURL, specType string, credentials, config map[string]any) store.RuntimeModelCandidate {
|
||||
return store.RuntimeModelCandidate{BaseURL: baseURL, SpecType: specType, Provider: specType, Credentials: credentials, PlatformConfig: config}
|
||||
}
|
||||
|
||||
func toJSONForTest(value any) string {
|
||||
raw, _ := json.Marshal(value)
|
||||
return string(raw)
|
||||
}
|
||||
|
||||
func containsStringForTest(values []string, expected string) bool {
|
||||
for _, value := range values {
|
||||
if value == expected {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -12,7 +12,8 @@ import (
|
||||
const OpenAIReasoningEffortValidationMessage = "reasoning_effort must be one of: none, minimal, low, medium, high, xhigh, max"
|
||||
|
||||
var (
|
||||
openAIReasoningEfforts = map[string]struct{}{
|
||||
chatReasoningEffortOrder = []string{"none", "minimal", "low", "medium", "high", "xhigh", "max"}
|
||||
openAIReasoningEfforts = map[string]struct{}{
|
||||
"none": {},
|
||||
"minimal": {},
|
||||
"low": {},
|
||||
@@ -49,9 +50,48 @@ func ValidateOpenAIReasoningEffort(value any) error {
|
||||
}
|
||||
|
||||
func applyOpenAIChatReasoningParams(body map[string]any, candidate store.RuntimeModelCandidate) {
|
||||
_ = applyOpenAIChatReasoningParamsWithSource(body, candidate, nil)
|
||||
}
|
||||
|
||||
func applyOpenAIChatReasoningParamsWithSource(body map[string]any, candidate store.RuntimeModelCandidate, original map[string]any) error {
|
||||
effort := normalizedReasoningString(body["reasoning_effort"])
|
||||
_, explicitEffort := original["reasoning_effort"]
|
||||
_, explicitTemperature := original["temperature"]
|
||||
model := chatReasoningModelName(body, candidate)
|
||||
if isAliyunBailianOpenAI(candidate) && isAliyunQwen38MaxPreview(model) {
|
||||
if explicitEffort && qwen38MaxPreviewReasoningEffort(effort) != effort {
|
||||
return explicitParameterAdaptationError("reasoning_effort", "the selected upstream only supports low, medium, or xhigh without changing the requested reasoning semantics")
|
||||
}
|
||||
if temperature, ok := finiteFloatFromAny(body["temperature"]); explicitTemperature && ok && temperature < 0.6 {
|
||||
return explicitParameterAdaptationError("temperature", "the selected upstream requires temperature >= 0.6")
|
||||
}
|
||||
applyAliyunQwen38Reasoning(body, effort)
|
||||
if temperature, ok := finiteFloatFromAny(body["temperature"]); ok && temperature < 0.6 {
|
||||
body["temperature"] = 0.6
|
||||
}
|
||||
return nil
|
||||
}
|
||||
if effort == "" || !isOpenAIReasoningEffort(effort) {
|
||||
return
|
||||
return nil
|
||||
}
|
||||
resolved, state := resolveCandidateReasoningEffort(effort, candidate)
|
||||
if state == reasoningCapabilityUnsupported {
|
||||
if explicitEffort {
|
||||
return explicitParameterAdaptationError("reasoning_effort", "the selected upstream does not support reasoning_effort")
|
||||
}
|
||||
delete(body, "reasoning_effort")
|
||||
return nil
|
||||
}
|
||||
if resolved != "" {
|
||||
if explicitEffort && resolved != effort {
|
||||
return explicitParameterAdaptationError("reasoning_effort", "the selected upstream does not support the requested reasoning effort exactly")
|
||||
}
|
||||
effort = resolved
|
||||
}
|
||||
if explicitEffort {
|
||||
if err := validateExplicitProviderReasoningEffort(effort, candidate, model); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
body["reasoning_effort"] = effort
|
||||
|
||||
@@ -65,10 +105,239 @@ func applyOpenAIChatReasoningParams(body map[string]any, candidate store.Runtime
|
||||
case isVolcesOpenAI(candidate):
|
||||
applyVolcesReasoning(body, candidate, effort)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func validateExplicitProviderReasoningEffort(effort string, candidate store.RuntimeModelCandidate, model string) error {
|
||||
if effort == "none" {
|
||||
return nil
|
||||
}
|
||||
unsupported := func() error {
|
||||
return explicitParameterAdaptationError("reasoning_effort", "the selected upstream cannot preserve the requested reasoning effort exactly")
|
||||
}
|
||||
switch {
|
||||
case isAliyunBailianOpenAI(candidate):
|
||||
if !isAliyunHighMaxReasoningModel(model) || highMaxReasoningEffort(effort) != effort {
|
||||
return unsupported()
|
||||
}
|
||||
case isDeepSeekOpenAI(candidate):
|
||||
if highMaxReasoningEffort(effort) != effort {
|
||||
return unsupported()
|
||||
}
|
||||
case isZhipuOpenAI(candidate):
|
||||
if !isZhipuReasoningEffortModel(model) || zhipuReasoningEffort(effort) != effort {
|
||||
return unsupported()
|
||||
}
|
||||
case isVolcesOpenAI(candidate):
|
||||
if !isVolcesReasoningEffortModel(model) || volcesChatReasoningEffort(effort) != effort {
|
||||
return unsupported()
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func applyOpenAIResponsesReasoningParams(body map[string]any, candidate store.RuntimeModelCandidate) {
|
||||
_ = applyOpenAIResponsesReasoningParamsWithSource(body, candidate, nil)
|
||||
}
|
||||
|
||||
func applyOpenAIResponsesReasoningParamsWithSource(body map[string]any, candidate store.RuntimeModelCandidate, original map[string]any) error {
|
||||
reasoning, _ := body["reasoning"].(map[string]any)
|
||||
if reasoning != nil {
|
||||
reasoning = cloneBody(reasoning)
|
||||
}
|
||||
effort := ""
|
||||
if reasoning != nil {
|
||||
effort = normalizedReasoningString(reasoning["effort"])
|
||||
}
|
||||
originalReasoning, _ := original["reasoning"].(map[string]any)
|
||||
_, explicitEffort := originalReasoning["effort"]
|
||||
_, explicitTemperature := original["temperature"]
|
||||
model := chatReasoningModelName(body, candidate)
|
||||
qwen38 := isAliyunBailianOpenAI(candidate) && isAliyunQwen38MaxPreview(model)
|
||||
if qwen38 {
|
||||
if effort == "" {
|
||||
if enabled, ok := body["enable_thinking"].(bool); ok && !enabled {
|
||||
effort = "low"
|
||||
}
|
||||
} else {
|
||||
mapped := qwen38MaxPreviewReasoningEffort(effort)
|
||||
if explicitEffort && mapped != effort {
|
||||
return explicitParameterAdaptationError("reasoning.effort", "the selected upstream only supports low, medium, or xhigh without changing the requested reasoning semantics")
|
||||
}
|
||||
effort = mapped
|
||||
}
|
||||
} else if effort != "" && isOpenAIReasoningEffort(effort) {
|
||||
resolved, state := resolveCandidateReasoningEffort(effort, candidate)
|
||||
if state == reasoningCapabilityUnsupported {
|
||||
if explicitEffort {
|
||||
return explicitParameterAdaptationError("reasoning.effort", "the selected upstream does not support reasoning.effort")
|
||||
}
|
||||
effort = ""
|
||||
} else if resolved != "" {
|
||||
if explicitEffort && resolved != effort {
|
||||
return explicitParameterAdaptationError("reasoning.effort", "the selected upstream does not support the requested reasoning effort exactly")
|
||||
}
|
||||
effort = resolved
|
||||
}
|
||||
}
|
||||
|
||||
if effort != "" {
|
||||
if reasoning == nil {
|
||||
reasoning = map[string]any{}
|
||||
}
|
||||
reasoning["effort"] = effort
|
||||
} else if reasoning != nil {
|
||||
delete(reasoning, "effort")
|
||||
}
|
||||
if len(reasoning) > 0 {
|
||||
body["reasoning"] = reasoning
|
||||
} else {
|
||||
delete(body, "reasoning")
|
||||
}
|
||||
delete(body, "reasoning_effort")
|
||||
delete(body, "enable_thinking")
|
||||
if qwen38 {
|
||||
if temperature, ok := finiteFloatFromAny(body["temperature"]); ok && temperature < 0.6 {
|
||||
if explicitTemperature {
|
||||
return explicitParameterAdaptationError("temperature", "the selected upstream requires temperature >= 0.6")
|
||||
}
|
||||
body["temperature"] = 0.6
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func explicitParameterAdaptationError(param string, message string) error {
|
||||
return &ClientError{Code: "invalid_parameter", Message: message, Param: param, StatusCode: 400, Retryable: false}
|
||||
}
|
||||
|
||||
func applyAliyunQwen38Reasoning(body map[string]any, effort string) {
|
||||
defer delete(body, "thinking_budget_tokens")
|
||||
requestedDisable := false
|
||||
if enabled, ok := body["enable_thinking"].(bool); ok && !enabled {
|
||||
requestedDisable = true
|
||||
}
|
||||
body["enable_thinking"] = true
|
||||
budget, hasBudget := nonNegativeIntFromAny(body["thinking_budget"])
|
||||
if !hasBudget {
|
||||
budget, hasBudget = nonNegativeIntFromAny(body["thinking_budget_tokens"])
|
||||
}
|
||||
if hasBudget {
|
||||
if budget > 262144 {
|
||||
budget = 262144
|
||||
}
|
||||
body["thinking_budget"] = budget
|
||||
delete(body, "reasoning_effort")
|
||||
return
|
||||
}
|
||||
delete(body, "thinking_budget")
|
||||
if effort == "" {
|
||||
if requestedDisable {
|
||||
body["reasoning_effort"] = "low"
|
||||
}
|
||||
return
|
||||
}
|
||||
body["reasoning_effort"] = qwen38MaxPreviewReasoningEffort(effort)
|
||||
}
|
||||
|
||||
func qwen38MaxPreviewReasoningEffort(effort string) string {
|
||||
switch effort {
|
||||
case "none", "minimal", "low":
|
||||
return "low"
|
||||
case "medium":
|
||||
return "medium"
|
||||
default:
|
||||
return "xhigh"
|
||||
}
|
||||
}
|
||||
|
||||
type reasoningCapabilityState int
|
||||
|
||||
const (
|
||||
reasoningCapabilityUnknown reasoningCapabilityState = iota
|
||||
reasoningCapabilityUnsupported
|
||||
reasoningCapabilitySupported
|
||||
)
|
||||
|
||||
func resolveCandidateReasoningEffort(effort string, candidate store.RuntimeModelCandidate) (string, reasoningCapabilityState) {
|
||||
supported := map[string]struct{}{}
|
||||
declaredLevels := false
|
||||
explicitlyUnsupported := false
|
||||
explicitlySupported := false
|
||||
for _, key := range []string{"text_generate", "tools_call", "image_analysis", "video_understanding", "audio_understanding", "omni"} {
|
||||
capability, ok := candidate.Capabilities[key].(map[string]any)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if value, ok := capability["supportThinking"].(bool); ok && !value {
|
||||
explicitlyUnsupported = true
|
||||
continue
|
||||
}
|
||||
if value, ok := capability["supportThinking"].(bool); ok && value {
|
||||
explicitlySupported = true
|
||||
}
|
||||
rawLevels, ok := capability["thinkingEffortLevels"].([]any)
|
||||
if !ok {
|
||||
if stringsValue, stringsOK := capability["thinkingEffortLevels"].([]string); stringsOK {
|
||||
rawLevels = make([]any, len(stringsValue))
|
||||
for index, value := range stringsValue {
|
||||
rawLevels[index] = value
|
||||
}
|
||||
ok = true
|
||||
}
|
||||
}
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
declaredLevels = true
|
||||
for _, raw := range rawLevels {
|
||||
level := normalizedReasoningString(raw)
|
||||
if isOpenAIReasoningEffort(level) {
|
||||
supported[level] = struct{}{}
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(supported) == 0 {
|
||||
if declaredLevels || (explicitlyUnsupported && !explicitlySupported) {
|
||||
return "", reasoningCapabilityUnsupported
|
||||
}
|
||||
return effort, reasoningCapabilityUnknown
|
||||
}
|
||||
requestedIndex := reasoningEffortIndex(effort)
|
||||
best := ""
|
||||
bestDistance := len(chatReasoningEffortOrder) + 1
|
||||
for index, candidateEffort := range chatReasoningEffortOrder {
|
||||
if _, ok := supported[candidateEffort]; !ok {
|
||||
continue
|
||||
}
|
||||
distance := index - requestedIndex
|
||||
if distance < 0 {
|
||||
distance = -distance
|
||||
}
|
||||
if distance < bestDistance {
|
||||
best = candidateEffort
|
||||
bestDistance = distance
|
||||
}
|
||||
}
|
||||
return best, reasoningCapabilitySupported
|
||||
}
|
||||
|
||||
func reasoningEffortIndex(effort string) int {
|
||||
for index, value := range chatReasoningEffortOrder {
|
||||
if value == effort {
|
||||
return index
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func applyAliyunReasoning(body map[string]any, candidate store.RuntimeModelCandidate, effort string) {
|
||||
defer delete(body, "thinking_budget_tokens")
|
||||
budget, hasBudget := positiveIntFromAny(body["thinking_budget"])
|
||||
if !hasBudget {
|
||||
budget, hasBudget = positiveIntFromAny(body["thinking_budget_tokens"])
|
||||
}
|
||||
delete(body, "thinking_budget")
|
||||
|
||||
if effort == "none" {
|
||||
body["enable_thinking"] = false
|
||||
@@ -78,6 +347,15 @@ func applyAliyunReasoning(body map[string]any, candidate store.RuntimeModelCandi
|
||||
|
||||
body["enable_thinking"] = true
|
||||
model := chatReasoningModelName(body, candidate)
|
||||
if isAliyunQwen38MaxPreview(model) {
|
||||
if budget, ok := positiveIntFromAny(body["thinking_budget_tokens"]); ok {
|
||||
body["thinking_budget"] = budget
|
||||
delete(body, "reasoning_effort")
|
||||
return
|
||||
}
|
||||
body["reasoning_effort"] = qwen38MaxPreviewReasoningEffort(effort)
|
||||
return
|
||||
}
|
||||
if isAliyunHighMaxReasoningModel(model) {
|
||||
body["reasoning_effort"] = highMaxReasoningEffort(effort)
|
||||
return
|
||||
@@ -85,7 +363,7 @@ func applyAliyunReasoning(body map[string]any, candidate store.RuntimeModelCandi
|
||||
|
||||
delete(body, "reasoning_effort")
|
||||
if isAliyunThinkingBudgetModel(model) {
|
||||
if budget, ok := positiveIntFromAny(body["thinking_budget_tokens"]); ok {
|
||||
if hasBudget {
|
||||
body["thinking_budget"] = budget
|
||||
}
|
||||
}
|
||||
@@ -200,6 +478,10 @@ func isAliyunHighMaxReasoningModel(model string) bool {
|
||||
return strings.Contains(model, "deepseek-v4") || strings.HasPrefix(model, "glm-")
|
||||
}
|
||||
|
||||
func isAliyunQwen38MaxPreview(model string) bool {
|
||||
return strings.Contains(model, "qwen3.8-max-preview")
|
||||
}
|
||||
|
||||
func isAliyunThinkingBudgetModel(model string) bool {
|
||||
return strings.Contains(model, "qwen") || strings.Contains(model, "qwq") || strings.Contains(model, "qvq") || strings.Contains(model, "kimi")
|
||||
}
|
||||
@@ -282,6 +564,49 @@ func positiveIntFromAny(value any) (int, bool) {
|
||||
return int(math.Floor(number)), true
|
||||
}
|
||||
|
||||
func nonNegativeIntFromAny(value any) (int, bool) {
|
||||
number, ok := finiteFloatFromAny(value)
|
||||
if !ok || number < 0 {
|
||||
return 0, false
|
||||
}
|
||||
return int(math.Floor(number)), true
|
||||
}
|
||||
|
||||
func finiteFloatFromAny(value any) (float64, bool) {
|
||||
if value == nil {
|
||||
return 0, false
|
||||
}
|
||||
var number float64
|
||||
switch typed := value.(type) {
|
||||
case int:
|
||||
number = float64(typed)
|
||||
case int64:
|
||||
number = float64(typed)
|
||||
case float64:
|
||||
number = typed
|
||||
case float32:
|
||||
number = float64(typed)
|
||||
case jsonNumber:
|
||||
parsed, err := typed.Float64()
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
number = parsed
|
||||
case string:
|
||||
parsed, err := strconv.ParseFloat(strings.TrimSpace(typed), 64)
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
number = parsed
|
||||
default:
|
||||
return 0, false
|
||||
}
|
||||
if math.IsNaN(number) || math.IsInf(number, 0) {
|
||||
return 0, false
|
||||
}
|
||||
return number, true
|
||||
}
|
||||
|
||||
type jsonNumber interface {
|
||||
Float64() (float64, error)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,184 @@
|
||||
package clients
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"os"
|
||||
"testing"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func TestQwen38MatchesCrossServiceBehaviorVectors(t *testing.T) {
|
||||
raw, err := os.ReadFile("testdata/qwen38-reasoning-vectors.json")
|
||||
if err != nil {
|
||||
t.Fatalf("read behavior vectors: %v", err)
|
||||
}
|
||||
var vectors []struct {
|
||||
Name string `json:"name"`
|
||||
Input map[string]any `json:"input"`
|
||||
Expected map[string]any `json:"expected"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &vectors); err != nil {
|
||||
t.Fatalf("decode behavior vectors: %v", err)
|
||||
}
|
||||
candidate := store.RuntimeModelCandidate{Provider: "aliyun-bailian-openai", ProviderModelName: "qwen3.8-max-preview"}
|
||||
for _, vector := range vectors {
|
||||
t.Run(vector.Name, func(t *testing.T) {
|
||||
body := cloneBody(vector.Input)
|
||||
body["model"] = "qwen3.8-max-preview"
|
||||
applyOpenAIChatReasoningParams(body, candidate)
|
||||
for _, key := range []string{"enable_thinking", "reasoning_effort", "thinking_budget", "temperature"} {
|
||||
expected, expectedPresent := vector.Expected[key]
|
||||
actual, actualPresent := body[key]
|
||||
if expectedPresent != actualPresent || (expectedPresent && normalizedJSONScalar(actual) != normalizedJSONScalar(expected)) {
|
||||
t.Fatalf("%s mismatch: expected=%#v actual=%#v body=%+v", key, expected, actual, body)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func normalizedJSONScalar(value any) any {
|
||||
if number, ok := value.(int); ok {
|
||||
return float64(number)
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
func TestQwen38MaxPreviewReasoningNormalization(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
body map[string]any
|
||||
effort any
|
||||
budget any
|
||||
temp any
|
||||
enabled any
|
||||
noEffort bool
|
||||
}{
|
||||
{name: "none", body: map[string]any{"reasoning_effort": "none"}, effort: "low"},
|
||||
{name: "minimal", body: map[string]any{"reasoning_effort": "minimal"}, effort: "low"},
|
||||
{name: "medium", body: map[string]any{"reasoning_effort": "medium"}, effort: "medium"},
|
||||
{name: "high", body: map[string]any{"reasoning_effort": "high"}, effort: "xhigh"},
|
||||
{name: "max", body: map[string]any{"reasoning_effort": "max"}, effort: "xhigh"},
|
||||
{name: "unspecified", body: map[string]any{}, noEffort: true},
|
||||
{name: "forced on", body: map[string]any{"enable_thinking": false}, effort: "low"},
|
||||
{name: "budget wins", body: map[string]any{"reasoning_effort": "high", "thinking_budget": 999999}, budget: 262144, noEffort: true},
|
||||
{name: "zero budget", body: map[string]any{"reasoning_effort": "medium", "thinking_budget_tokens": 0}, budget: 0, noEffort: true},
|
||||
{name: "temperature floor", body: map[string]any{"temperature": 0.1}, temp: 0.6, noEffort: true},
|
||||
}
|
||||
candidate := store.RuntimeModelCandidate{
|
||||
Provider: "aliyun-bailian-openai",
|
||||
ProviderModelName: "qwen3.8-max-preview",
|
||||
}
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
body := cloneBody(test.body)
|
||||
body["model"] = "qwen3.8-max-preview"
|
||||
applyOpenAIChatReasoningParams(body, candidate)
|
||||
if body["enable_thinking"] != true {
|
||||
t.Fatalf("Qwen3.8 must always enable thinking: %+v", body)
|
||||
}
|
||||
if test.noEffort {
|
||||
if _, ok := body["reasoning_effort"]; ok {
|
||||
t.Fatalf("reasoning_effort must be omitted: %+v", body)
|
||||
}
|
||||
} else if body["reasoning_effort"] != test.effort {
|
||||
t.Fatalf("expected effort %#v, got %+v", test.effort, body)
|
||||
}
|
||||
if test.budget != nil && body["thinking_budget"] != test.budget {
|
||||
t.Fatalf("expected budget %#v, got %+v", test.budget, body)
|
||||
}
|
||||
if test.temp != nil && body["temperature"] != test.temp {
|
||||
t.Fatalf("expected temperature %#v, got %+v", test.temp, body)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGenericReasoningCapabilityUsesNearestTierAndIgnoresModeSwitch(t *testing.T) {
|
||||
body := map[string]any{"reasoning_effort": "medium"}
|
||||
applyOpenAIChatReasoningParams(body, store.RuntimeModelCandidate{
|
||||
Provider: "openai",
|
||||
Capabilities: map[string]any{
|
||||
"text_generate": map[string]any{
|
||||
"supportThinking": true,
|
||||
"supportThinkingModeSwitch": false,
|
||||
"thinkingEffortLevels": []any{"low", "high"},
|
||||
},
|
||||
},
|
||||
})
|
||||
if body["reasoning_effort"] != "low" {
|
||||
t.Fatalf("equidistant effort must choose lower tier, got %+v", body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGenericReasoningKeepsUnknownAndRemovesExplicitUnsupported(t *testing.T) {
|
||||
unknown := map[string]any{"reasoning_effort": "high"}
|
||||
applyOpenAIChatReasoningParams(unknown, store.RuntimeModelCandidate{Provider: "openai"})
|
||||
if unknown["reasoning_effort"] != "high" {
|
||||
t.Fatalf("unknown capability should preserve effort: %+v", unknown)
|
||||
}
|
||||
unsupported := map[string]any{"reasoning_effort": "high"}
|
||||
applyOpenAIChatReasoningParams(unsupported, store.RuntimeModelCandidate{
|
||||
Provider: "openai",
|
||||
Capabilities: map[string]any{
|
||||
"text_generate": map[string]any{"supportThinking": false},
|
||||
},
|
||||
})
|
||||
if _, ok := unsupported["reasoning_effort"]; ok {
|
||||
t.Fatalf("explicit unsupported capability should remove effort: %+v", unsupported)
|
||||
}
|
||||
}
|
||||
|
||||
func TestQwen38ResponsesReasoningNormalization(t *testing.T) {
|
||||
candidate := store.RuntimeModelCandidate{
|
||||
Provider: "aliyun-bailian-openai",
|
||||
ProviderModelName: "qwen3.8-max-preview",
|
||||
}
|
||||
tests := []struct {
|
||||
name string
|
||||
requested string
|
||||
expected string
|
||||
}{
|
||||
{name: "none", requested: "none", expected: "low"},
|
||||
{name: "minimal", requested: "minimal", expected: "low"},
|
||||
{name: "medium", requested: "medium", expected: "medium"},
|
||||
{name: "high", requested: "high", expected: "xhigh"},
|
||||
{name: "max", requested: "max", expected: "xhigh"},
|
||||
}
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
body := map[string]any{
|
||||
"model": "qwen3.8-max-preview",
|
||||
"reasoning": map[string]any{"effort": test.requested, "summary": "auto"},
|
||||
"temperature": 0.1,
|
||||
}
|
||||
applyOpenAIResponsesReasoningParams(body, candidate)
|
||||
reasoning := body["reasoning"].(map[string]any)
|
||||
if reasoning["effort"] != test.expected || reasoning["summary"] != "auto" {
|
||||
t.Fatalf("unexpected Responses reasoning: %+v", body)
|
||||
}
|
||||
if body["temperature"] != 0.6 {
|
||||
t.Fatalf("Qwen3.8 Responses temperature was not floored: %+v", body)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestQwen38ResponsesForcedThinkingAndDefault(t *testing.T) {
|
||||
candidate := store.RuntimeModelCandidate{Provider: "aliyun-bailian-openai", ProviderModelName: "qwen3.8-max-preview"}
|
||||
forced := map[string]any{"enable_thinking": false}
|
||||
applyOpenAIResponsesReasoningParams(forced, candidate)
|
||||
if forced["reasoning"].(map[string]any)["effort"] != "low" {
|
||||
t.Fatalf("disabled Qwen3.8 Responses request must become low: %+v", forced)
|
||||
}
|
||||
if _, ok := forced["enable_thinking"]; ok {
|
||||
t.Fatalf("Chat-only alias leaked into Responses body: %+v", forced)
|
||||
}
|
||||
defaulted := map[string]any{"reasoning": map[string]any{"summary": "auto"}}
|
||||
applyOpenAIResponsesReasoningParams(defaulted, candidate)
|
||||
reasoning := defaulted["reasoning"].(map[string]any)
|
||||
if _, ok := reasoning["effort"]; ok || reasoning["summary"] != "auto" {
|
||||
t.Fatalf("unspecified effort must preserve upstream xhigh default: %+v", defaulted)
|
||||
}
|
||||
}
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
@@ -344,6 +345,453 @@ func TestOpenAIClientChatContract(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIClientDelayedStream429ReturnsBeforeAnyDelta(t *testing.T) {
|
||||
requestStarted := make(chan struct{})
|
||||
releaseResponse := make(chan struct{})
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
close(requestStarted)
|
||||
<-releaseResponse
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusTooManyRequests)
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"error": map[string]any{"message": "controlled delayed 429"},
|
||||
})
|
||||
}))
|
||||
defer server.Close()
|
||||
release := func() {
|
||||
select {
|
||||
case <-releaseResponse:
|
||||
default:
|
||||
close(releaseResponse)
|
||||
}
|
||||
}
|
||||
defer release()
|
||||
|
||||
type runResult struct {
|
||||
response Response
|
||||
err error
|
||||
}
|
||||
deltaCount := 0
|
||||
done := make(chan runResult, 1)
|
||||
go func() {
|
||||
response, err := (OpenAIClient{HTTPClient: server.Client()}).Run(context.Background(), Request{
|
||||
Kind: "chat.completions",
|
||||
Model: "delayed-429-model",
|
||||
Stream: true,
|
||||
Body: map[string]any{
|
||||
"model": "delayed-429-model",
|
||||
"stream": true,
|
||||
"messages": []any{map[string]any{"role": "user", "content": "ping"}},
|
||||
},
|
||||
Candidate: store.RuntimeModelCandidate{
|
||||
BaseURL: server.URL,
|
||||
ProviderModelName: "delayed-429-model",
|
||||
Credentials: map[string]any{"apiKey": "test-key"},
|
||||
},
|
||||
StreamDelta: func(StreamDeltaEvent) error {
|
||||
deltaCount++
|
||||
return nil
|
||||
},
|
||||
})
|
||||
done <- runResult{response: response, err: err}
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-requestStarted:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("upstream request did not start")
|
||||
}
|
||||
select {
|
||||
case result := <-done:
|
||||
t.Fatalf("stream returned before upstream status arrived: response=%+v err=%v", result.response, result.err)
|
||||
case <-time.After(50 * time.Millisecond):
|
||||
}
|
||||
release()
|
||||
|
||||
var result runResult
|
||||
select {
|
||||
case result = <-done:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("stream did not return after delayed 429")
|
||||
}
|
||||
clientErr, ok := result.err.(*ClientError)
|
||||
if !ok || clientErr.Code != "rate_limit" || clientErr.StatusCode != http.StatusTooManyRequests || !clientErr.Retryable {
|
||||
t.Fatalf("unexpected delayed 429 error: %T %+v", result.err, result.err)
|
||||
}
|
||||
if deltaCount != 0 {
|
||||
t.Fatalf("delayed 429 emitted %d stream deltas, want 0", deltaCount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIClientAliyunChatSendsStandardInputAudioAndNormalizesLegacyAudioURL(t *testing.T) {
|
||||
var captured map[string]any
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if err := json.NewDecoder(r.Body).Decode(&captured); err != nil {
|
||||
t.Fatalf("decode request: %v", err)
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"id": "chatcmpl-audio",
|
||||
"object": "chat.completion",
|
||||
"model": captured["model"],
|
||||
"choices": []any{map[string]any{
|
||||
"message": map[string]any{"role": "assistant", "content": "ok"},
|
||||
}},
|
||||
})
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
_, err := (OpenAIClient{HTTPClient: server.Client()}).Run(context.Background(), Request{
|
||||
Kind: "chat.completions",
|
||||
Model: "qwen-omni",
|
||||
Body: map[string]any{
|
||||
"model": "qwen-omni",
|
||||
"messages": []any{
|
||||
map[string]any{
|
||||
"role": "user",
|
||||
"content": []any{
|
||||
map[string]any{
|
||||
"type": "input_audio",
|
||||
"input_audio": map[string]any{
|
||||
"data": "https://cdn.example.com/audio/standard.wav",
|
||||
"format": "wav",
|
||||
},
|
||||
},
|
||||
map[string]any{
|
||||
"type": "audio_url",
|
||||
"audio_url": map[string]any{"url": "https://cdn.example.com/audio/legacy.m4a"},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
Candidate: store.RuntimeModelCandidate{
|
||||
Provider: "aliyun-bailian-openai",
|
||||
BaseURL: server.URL,
|
||||
ProviderModelName: "qwen-omni",
|
||||
Credentials: map[string]any{"apiKey": "test-key"},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("run aliyun openai-compatible client: %v", err)
|
||||
}
|
||||
|
||||
messages, _ := captured["messages"].([]any)
|
||||
message, _ := messages[0].(map[string]any)
|
||||
content, _ := message["content"].([]any)
|
||||
standard, _ := content[0].(map[string]any)
|
||||
if standard["type"] != "input_audio" {
|
||||
t.Fatalf("standard input_audio was not preserved: %+v", standard)
|
||||
}
|
||||
standardAudio, _ := standard["input_audio"].(map[string]any)
|
||||
if standardAudio["data"] != "https://cdn.example.com/audio/standard.wav" || standardAudio["format"] != "wav" {
|
||||
t.Fatalf("unexpected standard input_audio: %+v", standardAudio)
|
||||
}
|
||||
legacy, _ := content[1].(map[string]any)
|
||||
if legacy["type"] != "input_audio" {
|
||||
t.Fatalf("legacy audio_url was not normalized: %+v", legacy)
|
||||
}
|
||||
legacyAudio, _ := legacy["input_audio"].(map[string]any)
|
||||
if legacyAudio["data"] != "https://cdn.example.com/audio/legacy.m4a" || legacyAudio["format"] != "m4a" {
|
||||
t.Fatalf("unexpected normalized legacy audio: %+v", legacyAudio)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIClientImageEditUsesMultipartWithMultipleImages(t *testing.T) {
|
||||
var contentType string
|
||||
var receivedFields map[string][]string
|
||||
var receivedImageCount int
|
||||
var receivedImageFields []string
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
contentType = r.Header.Get("Content-Type")
|
||||
if err := r.ParseMultipartForm(1 << 20); err != nil {
|
||||
t.Fatalf("parse multipart image edit: %v", err)
|
||||
}
|
||||
receivedFields = r.MultipartForm.Value
|
||||
for field, files := range r.MultipartForm.File {
|
||||
receivedImageFields = append(receivedImageFields, field)
|
||||
receivedImageCount += len(files)
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"data": []any{map[string]any{"b64_json": "aW1hZ2U="}},
|
||||
})
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
_, err := (OpenAIClient{HTTPClient: server.Client()}).Run(context.Background(), Request{
|
||||
Kind: "images.edits",
|
||||
ModelType: "image_edit",
|
||||
Model: "gpt-image-2",
|
||||
Body: map[string]any{
|
||||
"model": "gpt-image-2",
|
||||
"prompt": "combine the references",
|
||||
"size": "1024x1536",
|
||||
"quality": "medium",
|
||||
"aspect_ratio": "2:3",
|
||||
"resolution": "2K",
|
||||
"width": 1024,
|
||||
"height": 1536,
|
||||
"images": []any{
|
||||
"data:image/png;base64," + base64.StdEncoding.EncodeToString([]byte("image one")),
|
||||
"data:image/jpeg;base64," + base64.StdEncoding.EncodeToString([]byte("image two")),
|
||||
},
|
||||
"_metadata": map[string]any{"private": true},
|
||||
},
|
||||
Candidate: store.RuntimeModelCandidate{
|
||||
BaseURL: server.URL,
|
||||
Provider: "openai",
|
||||
ProviderModelName: "gpt-image-2-vip",
|
||||
ModelType: "image_edit",
|
||||
Credentials: map[string]any{"apiKey": "openai-key"},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("run OpenAI image edit: %v", err)
|
||||
}
|
||||
if !strings.HasPrefix(contentType, "multipart/form-data; boundary=") {
|
||||
t.Fatalf("image edit should use multipart/form-data, got %q", contentType)
|
||||
}
|
||||
if receivedImageCount != 2 || len(receivedImageFields) != 1 || receivedImageFields[0] != "image[]" {
|
||||
t.Fatalf("unexpected multipart image files: fields=%v count=%d", receivedImageFields, receivedImageCount)
|
||||
}
|
||||
if receivedFields["model"][0] != "gpt-image-2-vip" ||
|
||||
receivedFields["prompt"][0] != "combine the references" ||
|
||||
receivedFields["size"][0] != "1024x1536" ||
|
||||
receivedFields["quality"][0] != "medium" {
|
||||
t.Fatalf("unexpected multipart fields: %+v", receivedFields)
|
||||
}
|
||||
if _, ok := receivedFields["_metadata"]; ok {
|
||||
t.Fatalf("internal metadata must not be forwarded: %+v", receivedFields)
|
||||
}
|
||||
for _, field := range []string{"aspect_ratio", "resolution", "width", "height"} {
|
||||
if _, ok := receivedFields[field]; ok {
|
||||
t.Fatalf("generic image geometry field %q must not be forwarded: %+v", field, receivedFields)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIClientImageEditDownloadsURLWhenBuildingMultipart(t *testing.T) {
|
||||
imagePayload := []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n', 1, 2, 3, 4}
|
||||
fetchCount := 0
|
||||
imageServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
fetchCount++
|
||||
w.Header().Set("Content-Type", "image/png")
|
||||
_, _ = w.Write(imagePayload)
|
||||
}))
|
||||
defer imageServer.Close()
|
||||
|
||||
var uploaded []byte
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if err := r.ParseMultipartForm(1 << 20); err != nil {
|
||||
t.Fatalf("parse multipart image edit: %v", err)
|
||||
}
|
||||
files := r.MultipartForm.File["image"]
|
||||
if len(files) != 1 {
|
||||
t.Fatalf("expected one multipart image, got %+v", r.MultipartForm.File)
|
||||
}
|
||||
file, err := files[0].Open()
|
||||
if err != nil {
|
||||
t.Fatalf("open multipart image: %v", err)
|
||||
}
|
||||
defer file.Close()
|
||||
uploaded, err = io.ReadAll(file)
|
||||
if err != nil {
|
||||
t.Fatalf("read multipart image: %v", err)
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"data": []any{map[string]any{"b64_json": "aW1hZ2U="}},
|
||||
})
|
||||
}))
|
||||
defer upstream.Close()
|
||||
|
||||
_, err := (OpenAIClient{HTTPClient: upstream.Client()}).Run(context.Background(), Request{
|
||||
Kind: "images.edits",
|
||||
ModelType: "image_edit",
|
||||
Body: map[string]any{
|
||||
"image": imageServer.URL + "/source.png",
|
||||
"prompt": "edit it",
|
||||
},
|
||||
Candidate: store.RuntimeModelCandidate{
|
||||
BaseURL: upstream.URL,
|
||||
Provider: "openai",
|
||||
ProviderModelName: "gpt-image-compatible",
|
||||
ModelType: "image_edit",
|
||||
Credentials: map[string]any{"apiKey": "openai-key"},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("run OpenAI URL multipart image edit: %v", err)
|
||||
}
|
||||
if fetchCount != 1 || string(uploaded) != string(imagePayload) {
|
||||
t.Fatalf("image URL should be fetched once during multipart construction: fetches=%d uploaded=%v", fetchCount, uploaded)
|
||||
}
|
||||
}
|
||||
|
||||
func TestKelingAndMinimaxMaterializeRemoteMediaInsideClient(t *testing.T) {
|
||||
payload := []byte("remote media payload")
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "audio/mpeg")
|
||||
_, _ = w.Write(payload)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
encoded, err := kelingImageToBase64(context.Background(), Request{HTTPClient: server.Client()}, server.URL+"/image.png")
|
||||
if err != nil {
|
||||
t.Fatalf("materialize Keling image URL: %v", err)
|
||||
}
|
||||
if encoded != base64.StdEncoding.EncodeToString(payload) {
|
||||
t.Fatalf("unexpected Keling image base64: %q", encoded)
|
||||
}
|
||||
|
||||
audio, filename, contentType, err := minimaxVoiceCloneFilePayload(context.Background(), server.Client(), server.URL+"/audio.mp3", "voice_clone")
|
||||
if err != nil {
|
||||
t.Fatalf("materialize Minimax audio URL: %v", err)
|
||||
}
|
||||
if string(audio) != string(payload) || filename != "voice_clone.mp3" || contentType != "audio/mpeg" {
|
||||
t.Fatalf("unexpected Minimax audio materialization: filename=%q contentType=%q payload=%q", filename, contentType, audio)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIClientImageEditUsesJSONURLsOnlyWhenExplicitlyConfigured(t *testing.T) {
|
||||
var contentType string
|
||||
var received map[string]any
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
contentType = r.Header.Get("Content-Type")
|
||||
if err := json.NewDecoder(r.Body).Decode(&received); err != nil {
|
||||
t.Fatalf("decode image edit JSON: %v", err)
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"data": []any{map[string]any{"url": "https://cdn.example.com/output.png"}},
|
||||
})
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
_, err := (OpenAIClient{HTTPClient: server.Client()}).Run(context.Background(), Request{
|
||||
Kind: "images.edits",
|
||||
ModelType: "image_edit",
|
||||
Model: "compatible-image-edit",
|
||||
Body: map[string]any{
|
||||
"images": []any{
|
||||
"https://cdn.example.com/first.png",
|
||||
map[string]any{"url": "https://cdn.example.com/second.png"},
|
||||
},
|
||||
"mask_image": map[string]any{"image_url": "https://cdn.example.com/mask.png"},
|
||||
"prompt": "combine the references",
|
||||
"_metadata": map[string]any{"private": true},
|
||||
},
|
||||
Candidate: store.RuntimeModelCandidate{
|
||||
BaseURL: server.URL,
|
||||
Provider: "compatible-openai",
|
||||
SpecType: "openai",
|
||||
ProviderModelName: "compatible-image-edit-v2",
|
||||
ModelType: "image_edit",
|
||||
Credentials: map[string]any{"apiKey": "openai-key"},
|
||||
PlatformConfig: map[string]any{"imageEditRequestFormat": "json_url"},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("run OpenAI-compatible JSON image edit: %v", err)
|
||||
}
|
||||
if !strings.HasPrefix(contentType, "application/json") {
|
||||
t.Fatalf("explicit JSON URL mode should use application/json, got %q", contentType)
|
||||
}
|
||||
images, _ := received["image"].([]any)
|
||||
if len(images) != 2 || images[0] != "https://cdn.example.com/first.png" || images[1] != "https://cdn.example.com/second.png" {
|
||||
t.Fatalf("unexpected JSON image URLs: %+v", received)
|
||||
}
|
||||
if received["mask"] != "https://cdn.example.com/mask.png" || received["model"] != "compatible-image-edit-v2" {
|
||||
t.Fatalf("unexpected JSON image edit body: %+v", received)
|
||||
}
|
||||
if _, ok := received["images"]; ok {
|
||||
t.Fatalf("images alias must be normalized to image: %+v", received)
|
||||
}
|
||||
if _, ok := received["_metadata"]; ok {
|
||||
t.Fatalf("internal metadata must not be forwarded: %+v", received)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIClientImageEditJSONModeRejectsInlineImageData(t *testing.T) {
|
||||
_, _, err := openAIRequestPayload(context.Background(), "images.edits", map[string]any{
|
||||
"image": "data:image/png;base64,aW1hZ2U=",
|
||||
}, store.RuntimeModelCandidate{PlatformConfig: map[string]any{"imageEditRequestFormat": "json_url"}})
|
||||
clientErr, ok := err.(*ClientError)
|
||||
if !ok || clientErr.Code != "invalid_parameter" || clientErr.Param != "image" || clientErr.StatusCode != http.StatusBadRequest {
|
||||
t.Fatalf("expected non-retryable image URL validation error, got %#v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIImageEditFieldNameUsesArraySyntaxForMultipleImages(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
candidate store.RuntimeModelCandidate
|
||||
imageCount int
|
||||
want string
|
||||
}{
|
||||
{name: "compatible single image", candidate: store.RuntimeModelCandidate{BaseURL: "https://compatible.example/v1"}, imageCount: 1, want: "image"},
|
||||
{name: "compatible multiple images", candidate: store.RuntimeModelCandidate{BaseURL: "https://compatible.example/v1"}, imageCount: 2, want: "image[]"},
|
||||
{name: "official single image", candidate: store.RuntimeModelCandidate{BaseURL: "https://api.openai.com/v1"}, imageCount: 1, want: "image[]"},
|
||||
{name: "explicit scalar override", candidate: store.RuntimeModelCandidate{BaseURL: "https://compatible.example/v1", PlatformConfig: map[string]any{"imageEditMultipartFieldName": "image"}}, imageCount: 2, want: "image"},
|
||||
{name: "explicit array override", candidate: store.RuntimeModelCandidate{BaseURL: "https://compatible.example/v1", PlatformConfig: map[string]any{"imageEditMultipartFieldName": "image[]"}}, imageCount: 1, want: "image[]"},
|
||||
}
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
if got := openAIImageEditFieldName(test.candidate, test.imageCount); got != test.want {
|
||||
t.Fatalf("field name = %q, want %q", got, test.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIClientImageGenerationUsesSizeWithoutGenericGeometryFields(t *testing.T) {
|
||||
var received map[string]any
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if err := json.NewDecoder(r.Body).Decode(&received); err != nil {
|
||||
t.Fatalf("decode image generation body: %v", err)
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"data": []any{map[string]any{"b64_json": "aW1hZ2U="}},
|
||||
})
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
_, err := (OpenAIClient{HTTPClient: server.Client()}).Run(context.Background(), Request{
|
||||
Kind: "images.generations",
|
||||
ModelType: "image_generate",
|
||||
Model: "gpt-image-2",
|
||||
Body: map[string]any{
|
||||
"model": "gpt-image-2",
|
||||
"prompt": "draw",
|
||||
"size": "2048x1152",
|
||||
"quality": "high",
|
||||
"aspect_ratio": "16:9",
|
||||
"aspectRatio": "16:9",
|
||||
"resolution": "2K",
|
||||
"width": 2048,
|
||||
"height": 1152,
|
||||
},
|
||||
Candidate: store.RuntimeModelCandidate{
|
||||
BaseURL: server.URL,
|
||||
Provider: "openai",
|
||||
ProviderModelName: "gpt-image-2-vip",
|
||||
ModelType: "image_generate",
|
||||
Credentials: map[string]any{"apiKey": "openai-key"},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("run OpenAI image generation: %v", err)
|
||||
}
|
||||
if received["model"] != "gpt-image-2-vip" ||
|
||||
received["prompt"] != "draw" ||
|
||||
received["size"] != "2048x1152" ||
|
||||
received["quality"] != "high" {
|
||||
t.Fatalf("unexpected image generation fields: %+v", received)
|
||||
}
|
||||
for _, field := range []string{"aspect_ratio", "aspectRatio", "resolution", "width", "height"} {
|
||||
if _, ok := received[field]; ok {
|
||||
t.Fatalf("generic image geometry field %q must not be forwarded: %+v", field, received)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIClientChatReasoningParamsByProvider(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
@@ -405,6 +853,36 @@ func TestOpenAIClientChatReasoningParamsByProvider(t *testing.T) {
|
||||
}
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "aliyun qwen 3.8 preview preserves supported effort",
|
||||
provider: "aliyun-bailian-openai",
|
||||
model: "qwen3.8-max-preview",
|
||||
body: map[string]any{
|
||||
"reasoning_effort": "high",
|
||||
},
|
||||
assertion: func(t *testing.T, body map[string]any) {
|
||||
if body["enable_thinking"] != true || body["reasoning_effort"] != "xhigh" {
|
||||
t.Fatalf("aliyun qwen 3.8 should use xhigh reasoning effort, got %+v", body)
|
||||
}
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "aliyun qwen 3.8 preview uses explicit budget exclusively",
|
||||
provider: "aliyun-bailian-openai",
|
||||
model: "qwen3.8-max-preview",
|
||||
body: map[string]any{
|
||||
"reasoning_effort": "medium",
|
||||
"thinking_budget_tokens": 16384,
|
||||
},
|
||||
assertion: func(t *testing.T, body map[string]any) {
|
||||
if body["enable_thinking"] != true || body["thinking_budget"] != float64(16384) {
|
||||
t.Fatalf("aliyun qwen 3.8 should use the explicit thinking budget, got %+v", body)
|
||||
}
|
||||
if _, ok := body["reasoning_effort"]; ok {
|
||||
t.Fatalf("aliyun qwen 3.8 must not combine reasoning effort with thinking budget: %+v", body)
|
||||
}
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "same deepseek model on aliyun uses aliyun protocol",
|
||||
provider: "aliyun-bailian-openai",
|
||||
@@ -707,13 +1185,11 @@ func TestOpenAIClientChatRequestNormalizesToolContext(t *testing.T) {
|
||||
}
|
||||
assistant, _ := messages[0].(map[string]any)
|
||||
if _, ok := assistant["functionCall"]; ok {
|
||||
t.Fatalf("functionCall should be converted away: %+v", assistant)
|
||||
t.Fatalf("functionCall alias should be converted away: %+v", assistant)
|
||||
}
|
||||
toolCalls, _ := assistant["tool_calls"].([]any)
|
||||
toolCall, _ := toolCalls[0].(map[string]any)
|
||||
function, _ := toolCall["function"].(map[string]any)
|
||||
function, _ := assistant["function_call"].(map[string]any)
|
||||
if function["name"] != "lookup" || function["arguments"] != `{"q":"weather"}` {
|
||||
t.Fatalf("unexpected normalized tool call: %+v", assistant)
|
||||
t.Fatalf("unexpected normalized legacy function call: %+v", assistant)
|
||||
}
|
||||
toolMessage, _ := messages[1].(map[string]any)
|
||||
if toolMessage["tool_call_id"] != "call_0" || toolMessage["toolCallId"] != nil {
|
||||
@@ -811,14 +1287,18 @@ func TestOpenAIClientChatResponseNormalizesToolCallFormats(t *testing.T) {
|
||||
if message["content"] != "calling tools" {
|
||||
t.Fatalf("tool_use block should be removed from content: %+v", message)
|
||||
}
|
||||
for _, key := range []string{"toolCalls", "function_call"} {
|
||||
for _, key := range []string{"toolCalls"} {
|
||||
if _, ok := message[key]; ok {
|
||||
t.Fatalf("%s should be converted away: %+v", key, message)
|
||||
}
|
||||
}
|
||||
legacyFunction, _ := message["function_call"].(map[string]any)
|
||||
if legacyFunction["name"] != "legacy_lookup" || legacyFunction["arguments"] != "{\"city\":\"NYC\"}" {
|
||||
t.Fatalf("canonical legacy function_call should be preserved: %+v", message)
|
||||
}
|
||||
toolCalls, _ := message["tool_calls"].([]any)
|
||||
if len(toolCalls) != 3 {
|
||||
t.Fatalf("expected 3 normalized tool calls, got %+v", message)
|
||||
if len(toolCalls) != 2 {
|
||||
t.Fatalf("expected 2 normalized tool calls plus legacy function_call, got %+v", message)
|
||||
}
|
||||
assertToolCall := func(index int, id string, name string, arguments string) {
|
||||
t.Helper()
|
||||
@@ -829,8 +1309,7 @@ func TestOpenAIClientChatResponseNormalizesToolCallFormats(t *testing.T) {
|
||||
}
|
||||
}
|
||||
assertToolCall(0, "call_camel", "camel_lookup", "{\"city\":\"SF\"}")
|
||||
assertToolCall(1, "call_1", "legacy_lookup", "{\"city\":\"NYC\"}")
|
||||
assertToolCall(2, "toolu_1", "anthropic_lookup", "{\"city\":\"Boston\"}")
|
||||
assertToolCall(1, "toolu_1", "anthropic_lookup", "{\"city\":\"Boston\"}")
|
||||
}
|
||||
|
||||
func TestOpenAIClientChatStreamContract(t *testing.T) {
|
||||
@@ -1088,19 +1567,21 @@ func TestOpenAIClientChatStreamNormalizesToolCallFormats(t *testing.T) {
|
||||
if len(captured) != 3 {
|
||||
t.Fatalf("unexpected captured events: %+v", captured)
|
||||
}
|
||||
for _, event := range captured {
|
||||
for index, event := range captured {
|
||||
choices, _ := event.Event["choices"].([]any)
|
||||
choice, _ := choices[0].(map[string]any)
|
||||
delta, _ := choice["delta"].(map[string]any)
|
||||
if _, ok := delta["function_call"]; ok {
|
||||
t.Fatalf("function_call should be converted away: %+v", event.Event)
|
||||
}
|
||||
if _, ok := delta["functionCall"]; ok {
|
||||
t.Fatalf("functionCall should be converted away: %+v", event.Event)
|
||||
}
|
||||
if _, ok := delta["toolCall"]; ok {
|
||||
t.Fatalf("toolCall should be converted away: %+v", event.Event)
|
||||
}
|
||||
if index < 2 {
|
||||
if _, ok := delta["function_call"]; !ok {
|
||||
t.Fatalf("canonical legacy function_call should be preserved: %+v", event.Event)
|
||||
}
|
||||
}
|
||||
}
|
||||
choices, _ := response.Result["choices"].([]any)
|
||||
choice, _ := choices[0].(map[string]any)
|
||||
@@ -1304,6 +1785,104 @@ func TestGeminiClientImageGenerateBuildsNativeImageBody(t *testing.T) {
|
||||
if data[0].(map[string]any)["b64_json"] != "aW1hZ2U=" {
|
||||
t.Fatalf("unexpected image response: %+v", response.Result)
|
||||
}
|
||||
if response.Wire == nil || len(response.Wire.RawJSON) != 0 {
|
||||
t.Fatalf("Gemini image response retained duplicate raw JSON: %+v", response.Wire)
|
||||
}
|
||||
candidates, _ := response.Wire.Body["candidates"].([]any)
|
||||
content, _ := candidates[0].(map[string]any)["content"].(map[string]any)
|
||||
wireParts, _ := content["parts"].([]any)
|
||||
inline, _ := wireParts[0].(map[string]any)["inlineData"].(map[string]any)
|
||||
if inline["data"] != "aW1hZ2U=" {
|
||||
t.Fatalf("decoded Gemini wire body lost Base64 output: %+v", response.Wire.Body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiClientImageGenerateRejectsPromptBlockedWithoutImage(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.Header().Set("X-Request-Id", "gemini-blocked-request")
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"promptFeedback": map[string]any{
|
||||
"blockReason": "SAFETY",
|
||||
"blockReasonMessage": "The request was blocked by the image safety policy.",
|
||||
},
|
||||
"usageMetadata": map[string]any{"promptTokenCount": 12, "totalTokenCount": 12},
|
||||
})
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
_, err := (GeminiClient{HTTPClient: server.Client()}).Run(context.Background(), Request{
|
||||
Kind: "images.generations",
|
||||
ModelType: "image_generate",
|
||||
Model: "gemini-image",
|
||||
Body: map[string]any{"model": "gemini-image", "prompt": "blocked prompt"},
|
||||
Candidate: store.RuntimeModelCandidate{
|
||||
BaseURL: server.URL,
|
||||
ProviderModelName: "gemini-image",
|
||||
ModelType: "image_generate",
|
||||
Credentials: map[string]any{"apiKey": "gemini-key"},
|
||||
},
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("missing Gemini image must fail instead of returning a placeholder")
|
||||
}
|
||||
if ErrorCode(err) != "gemini_image_output_missing" || IsRetryable(err) {
|
||||
t.Fatalf("unexpected missing image error: code=%q retryable=%v err=%v", ErrorCode(err), IsRetryable(err), err)
|
||||
}
|
||||
meta := ErrorResponseMetadata(err)
|
||||
if meta.RequestID != "gemini-blocked-request" || meta.ResponseDurationMS < 0 {
|
||||
t.Fatalf("missing image error lost response metadata: %+v", meta)
|
||||
}
|
||||
details := ErrorDetails(err)
|
||||
if details["reason"] != "prompt_blocked" || details["promptBlockReason"] != "SAFETY" || details["candidateCount"] != 0 {
|
||||
t.Fatalf("unexpected prompt block details: %+v", details)
|
||||
}
|
||||
if !strings.Contains(err.Error(), "blockReason=SAFETY") || strings.Contains(err.Error(), "gemini-image-placeholder") {
|
||||
t.Fatalf("missing image error should explain the upstream block without a placeholder: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiMissingImageDetailsExplainEmptyImagePayload(t *testing.T) {
|
||||
details := geminiMissingImageDetails(map[string]any{
|
||||
"candidates": []any{map[string]any{
|
||||
"finishReason": "STOP",
|
||||
"content": map[string]any{"parts": []any{
|
||||
map[string]any{"inlineData": map[string]any{"mimeType": "image/png", "data": ""}},
|
||||
map[string]any{"text": "The model completed without producing an image."},
|
||||
}},
|
||||
}},
|
||||
})
|
||||
message := geminiMissingImageMessage(details)
|
||||
|
||||
if details["reason"] != "empty_image_payload" || details["candidateCount"] != 1 || details["partCount"] != 2 {
|
||||
t.Fatalf("unexpected empty image payload details: %+v", details)
|
||||
}
|
||||
if !strings.Contains(message, "missing inline data or file URI") ||
|
||||
!strings.Contains(message, "finishReason=STOP") ||
|
||||
!strings.Contains(message, "The model completed without producing an image.") {
|
||||
t.Fatalf("missing image message should preserve the concrete extraction reason: %s", message)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiMissingImageDetailsPreserveUpstreamErrorEnvelope(t *testing.T) {
|
||||
details := geminiMissingImageDetails(map[string]any{
|
||||
"error": map[string]any{
|
||||
"code": 500,
|
||||
"status": "INTERNAL",
|
||||
"message": "image backend returned no generated asset",
|
||||
},
|
||||
})
|
||||
message := geminiMissingImageMessage(details)
|
||||
|
||||
if details["reason"] != "upstream_error" ||
|
||||
details["upstreamErrorCode"] != "500" ||
|
||||
details["upstreamErrorStatus"] != "INTERNAL" {
|
||||
t.Fatalf("unexpected upstream error details: %+v", details)
|
||||
}
|
||||
if !strings.Contains(message, "upstreamErrorCode=500") ||
|
||||
!strings.Contains(message, "upstreamErrorStatus=INTERNAL") ||
|
||||
!strings.Contains(message, "image backend returned no generated asset") {
|
||||
t.Fatalf("upstream error message should remain diagnostic: %s", message)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiGenerationConfigInitializesMissingImageConfig(t *testing.T) {
|
||||
@@ -1977,7 +2556,7 @@ func TestVolcesClientVideoSubmitsAndPollsTask(t *testing.T) {
|
||||
}
|
||||
data, _ := response.Result["data"].([]any)
|
||||
item, _ := data[0].(map[string]any)
|
||||
if item["url"] != "https://example.com/out.mp4" || response.Usage.TotalTokens != 9 {
|
||||
if item["url"] != "https://example.com/out.mp4" || response.Result["content"] != nil || response.Result["usage"] != nil || response.Usage.TotalTokens != 9 {
|
||||
t.Fatalf("unexpected response: %+v usage=%+v", response.Result, response.Usage)
|
||||
}
|
||||
}
|
||||
@@ -2021,8 +2600,9 @@ func TestVolcesClientVideoRetriesTransientPollAndKeepsOfficialResult(t *testing.
|
||||
if polls != 2 || len(persisted) != 1 || persisted[0] != "cgt-retry:succeeded" {
|
||||
t.Fatalf("unexpected poll state polls=%d persisted=%+v", polls, persisted)
|
||||
}
|
||||
if response.Result["updated_at"] != float64(124) || response.Result["seed"] != float64(7) || response.Result["raw"] == nil {
|
||||
t.Fatalf("official result fields lost: %+v", response.Result)
|
||||
if response.Result["seed"] != float64(7) || response.Result["raw"] != nil || response.Result["updated_at"] != nil ||
|
||||
response.Result["content"] != nil || response.Result["usage"] != nil || response.Result["upstream_task_id"] != nil {
|
||||
t.Fatalf("provider response should be reduced to the canonical result: %+v", response.Result)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2109,6 +2689,8 @@ func TestVolcesVideoBodyAllowsOnlyTaskPayloadFields(t *testing.T) {
|
||||
"tools": []any{map[string]any{"type": "web_search"}},
|
||||
"task_id": "local-task-id",
|
||||
"runMode": "simulation",
|
||||
"modelType": "omni_video",
|
||||
"model_type": "omni_video",
|
||||
"fps": 24,
|
||||
"content": []any{
|
||||
map[string]any{"type": "text", "text": "Use <<<element_1>>> in a product reveal"},
|
||||
@@ -2300,7 +2882,7 @@ func TestVolcesClientVideoResumePollsExistingTaskID(t *testing.T) {
|
||||
}
|
||||
data, _ := response.Result["data"].([]any)
|
||||
item, _ := data[0].(map[string]any)
|
||||
if response.Result["upstream_task_id"] != "cgt-existing" || item["url"] != "https://example.com/resumed.mp4" {
|
||||
if response.Result["upstream_task_id"] != nil || item["url"] != "https://example.com/resumed.mp4" {
|
||||
t.Fatalf("unexpected resumed response: %+v", response.Result)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,8 +1,10 @@
|
||||
package clients
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"mime"
|
||||
@@ -11,6 +13,8 @@ import (
|
||||
"path"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
type GeminiClient struct {
|
||||
@@ -22,25 +26,73 @@ func (c GeminiClient) Run(ctx context.Context, request Request) (Response, error
|
||||
if apiKey == "" {
|
||||
return Response{}, &ClientError{Code: "missing_credentials", Message: "gemini api key is required", Retryable: false}
|
||||
}
|
||||
if geminiVeoRequest(request) {
|
||||
return c.runVeo(ctx, request, apiKey)
|
||||
}
|
||||
body := geminiBody(request)
|
||||
_, nativeContents := request.Body["contents"]
|
||||
if !nativeContents && GeminiUsesOfficialAPI(request.Candidate) {
|
||||
var err error
|
||||
body, err = c.prepareOfficialGeminiImageFiles(ctx, request, apiKey, body)
|
||||
if err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
} else if !nativeContents {
|
||||
var err error
|
||||
body, err = prepareCompatibleGeminiInlineImages(ctx, body)
|
||||
if err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
}
|
||||
raw, _ := json.Marshal(body)
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, geminiURL(request.Candidate.BaseURL, upstreamModelName(request.Candidate), apiKey), bytes.NewReader(raw))
|
||||
endpoint := geminiURL(request.Candidate.BaseURL, upstreamModelName(request.Candidate), apiKey)
|
||||
if request.Stream {
|
||||
endpoint = geminiStreamURL(request.Candidate.BaseURL, upstreamModelName(request.Candidate), apiKey)
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(raw))
|
||||
if err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
applyUpstreamIdempotency(req, request)
|
||||
responseStartedAt := time.Now()
|
||||
if err := notifySubmissionStarted(request); err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
resp, err := httpClient(request.HTTPClient, c.HTTPClient).Do(req)
|
||||
if err != nil {
|
||||
return Response{}, &ClientError{Code: "network", Message: err.Error(), Retryable: true}
|
||||
return Response{}, transportClientError(err)
|
||||
}
|
||||
if err := notifyResponseReceived(request); err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
requestID := requestIDFromHTTPResponse(resp)
|
||||
result, err := decodeHTTPResponse(resp)
|
||||
var result map[string]any
|
||||
var wire *WireResponse
|
||||
if request.Stream && resp.StatusCode >= 200 && resp.StatusCode < 300 {
|
||||
wire = &WireResponse{Protocol: ProtocolGeminiGenerateContent, StatusCode: resp.StatusCode, Headers: compatibleResponseHeaders(resp.Header)}
|
||||
result, err = decodeGeminiStreamResponse(resp, request.StreamDelta, wire)
|
||||
} else {
|
||||
result, wire, err = decodeHTTPResponseForProtocol(resp, ProtocolGeminiGenerateContent)
|
||||
}
|
||||
responseFinishedAt := time.Now()
|
||||
if err != nil {
|
||||
return Response{}, annotateResponseError(err, requestID, responseStartedAt, responseFinishedAt)
|
||||
}
|
||||
output := geminiResult(request, result)
|
||||
output, err := geminiResult(request, result)
|
||||
if err != nil {
|
||||
if requestID == "" {
|
||||
requestID = firstNonEmptyString(result["responseId"], result["response_id"])
|
||||
}
|
||||
return Response{}, annotateResponseError(err, requestID, responseStartedAt, responseFinishedAt)
|
||||
}
|
||||
if wire != nil && geminiImageModelType(request.ModelType) {
|
||||
// Gemini image responses can contain multi-megabyte Base64 fields. The
|
||||
// decoded wire body is still used for native synchronous passthrough, so
|
||||
// retaining the original JSON bytes would keep a second full copy alive
|
||||
// until the client response finishes.
|
||||
wire.RawJSON = nil
|
||||
}
|
||||
if requestID == "" {
|
||||
requestID = requestIDFromResult(output)
|
||||
}
|
||||
@@ -52,10 +104,244 @@ func (c GeminiClient) Run(ctx context.Context, request Request) (Response, error
|
||||
ResponseStartedAt: responseStartedAt,
|
||||
ResponseFinishedAt: responseFinishedAt,
|
||||
ResponseDurationMS: responseDurationMS(responseStartedAt, responseFinishedAt),
|
||||
UpstreamProtocol: ProtocolGeminiGenerateContent,
|
||||
Wire: wire,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// GeminiUsesOfficialAPI distinguishes the native Google Gemini endpoint from
|
||||
// Gemini-shaped compatibility services and Google's OpenAI compatibility path.
|
||||
func GeminiUsesOfficialAPI(candidate store.RuntimeModelCandidate) bool {
|
||||
base := strings.TrimSpace(candidate.BaseURL)
|
||||
if base == "" {
|
||||
return true
|
||||
}
|
||||
parsed, err := url.Parse(base)
|
||||
if err != nil || !strings.EqualFold(parsed.Hostname(), "generativelanguage.googleapis.com") {
|
||||
return false
|
||||
}
|
||||
pathValue := strings.ToLower(strings.TrimRight(parsed.Path, "/"))
|
||||
return pathValue != "/openai" && !strings.HasSuffix(pathValue, "/openai")
|
||||
}
|
||||
|
||||
func (c GeminiClient) prepareOfficialGeminiImageFiles(ctx context.Context, request Request, apiKey string, body map[string]any) (map[string]any, error) {
|
||||
client := httpClient(request.HTTPClient, c.HTTPClient)
|
||||
cache := map[string]map[string]any{}
|
||||
prepared, err := prepareOfficialGeminiImageFileValue(ctx, client, request.Candidate, apiKey, body, cache)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out, _ := prepared.(map[string]any)
|
||||
if out == nil {
|
||||
return body, nil
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func prepareOfficialGeminiImageFileValue(ctx context.Context, client *http.Client, candidate store.RuntimeModelCandidate, apiKey string, value any, cache map[string]map[string]any) (any, error) {
|
||||
switch typed := value.(type) {
|
||||
case map[string]any:
|
||||
next := make(map[string]any, len(typed))
|
||||
for key, item := range typed {
|
||||
prepared, err := prepareOfficialGeminiImageFileValue(ctx, client, candidate, apiKey, item, cache)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
next[key] = prepared
|
||||
}
|
||||
for _, key := range []string{"fileData", "file_data"} {
|
||||
fileData, ok := next[key].(map[string]any)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
fileURI := firstNonEmptyString(fileData["fileUri"], fileData["file_uri"], fileData["uri"])
|
||||
mimeType := firstNonEmptyString(fileData["mimeType"], fileData["mime_type"], mimeFromURI(fileURI))
|
||||
if !requestAssetStringIsHTTPURL(fileURI) || !strings.HasPrefix(strings.ToLower(mimeType), "image/") || geminiOfficialFileURI(fileURI) {
|
||||
continue
|
||||
}
|
||||
cacheKey := fileURI + "\x00" + mimeType
|
||||
uploaded := cache[cacheKey]
|
||||
if uploaded == nil {
|
||||
var err error
|
||||
uploaded, err = uploadOfficialGeminiImageFile(ctx, client, candidate, apiKey, fileURI, mimeType)
|
||||
if err != nil {
|
||||
fallback, fallbackErr := compatibleGeminiInlineImage(ctx, fileURI, mimeType)
|
||||
if fallbackErr != nil {
|
||||
return nil, err
|
||||
}
|
||||
delete(next, key)
|
||||
next["inlineData"] = fallback
|
||||
continue
|
||||
}
|
||||
cache[cacheKey] = uploaded
|
||||
}
|
||||
next[key] = cloneBody(uploaded)
|
||||
}
|
||||
return next, nil
|
||||
case []any:
|
||||
next := make([]any, 0, len(typed))
|
||||
for _, item := range typed {
|
||||
prepared, err := prepareOfficialGeminiImageFileValue(ctx, client, candidate, apiKey, item, cache)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
next = append(next, prepared)
|
||||
}
|
||||
return next, nil
|
||||
default:
|
||||
return value, nil
|
||||
}
|
||||
}
|
||||
|
||||
func prepareCompatibleGeminiInlineImages(ctx context.Context, body map[string]any) (map[string]any, error) {
|
||||
prepared, err := prepareCompatibleGeminiInlineImageValue(ctx, body, map[string]map[string]any{})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out, _ := prepared.(map[string]any)
|
||||
if out == nil {
|
||||
return body, nil
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func prepareCompatibleGeminiInlineImageValue(ctx context.Context, value any, cache map[string]map[string]any) (any, error) {
|
||||
switch typed := value.(type) {
|
||||
case map[string]any:
|
||||
next := make(map[string]any, len(typed))
|
||||
for key, item := range typed {
|
||||
prepared, err := prepareCompatibleGeminiInlineImageValue(ctx, item, cache)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
next[key] = prepared
|
||||
}
|
||||
for _, key := range []string{"fileData", "file_data"} {
|
||||
fileData, ok := next[key].(map[string]any)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
fileURI := firstNonEmptyString(fileData["fileUri"], fileData["file_uri"], fileData["uri"])
|
||||
mimeType := firstNonEmptyString(fileData["mimeType"], fileData["mime_type"], mimeFromURI(fileURI))
|
||||
if !requestAssetStringIsHTTPURL(fileURI) || !strings.HasPrefix(strings.ToLower(mimeType), "image/") {
|
||||
continue
|
||||
}
|
||||
cacheKey := fileURI + "\x00" + mimeType
|
||||
inline := cache[cacheKey]
|
||||
if inline == nil {
|
||||
var err error
|
||||
inline, err = compatibleGeminiInlineImage(ctx, fileURI, mimeType)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
cache[cacheKey] = inline
|
||||
}
|
||||
delete(next, key)
|
||||
next["inlineData"] = cloneBody(inline)
|
||||
}
|
||||
return next, nil
|
||||
case []any:
|
||||
next := make([]any, 0, len(typed))
|
||||
for _, item := range typed {
|
||||
prepared, err := prepareCompatibleGeminiInlineImageValue(ctx, item, cache)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
next = append(next, prepared)
|
||||
}
|
||||
return next, nil
|
||||
default:
|
||||
return value, nil
|
||||
}
|
||||
}
|
||||
|
||||
func compatibleGeminiInlineImage(ctx context.Context, sourceURL string, declaredMimeType string) (map[string]any, error) {
|
||||
fetchedMimeType, payload, err := fetchRemoteMediaInputPayload(ctx, sourceURL, 256<<20)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return map[string]any{
|
||||
"mimeType": geminiMediaMime(firstNonEmptyString(declaredMimeType, fetchedMimeType), "image"),
|
||||
"data": base64.StdEncoding.EncodeToString(payload),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func uploadOfficialGeminiImageFile(ctx context.Context, client *http.Client, candidate store.RuntimeModelCandidate, apiKey string, sourceURL string, declaredMimeType string) (map[string]any, error) {
|
||||
fetchedMimeType, payload, err := fetchRemoteMediaInputPayload(ctx, sourceURL, 256<<20)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
mimeType := geminiMediaMime(firstNonEmptyString(declaredMimeType, fetchedMimeType), "image")
|
||||
base := strings.TrimRight(strings.TrimSpace(candidate.BaseURL), "/")
|
||||
if base == "" {
|
||||
base = "https://generativelanguage.googleapis.com"
|
||||
}
|
||||
parsed, err := url.Parse(base)
|
||||
if err != nil || parsed.Scheme == "" || parsed.Host == "" {
|
||||
return nil, &ClientError{Code: "invalid_parameter", Message: "invalid official Gemini base URL", Retryable: false}
|
||||
}
|
||||
uploadStartURL := parsed.Scheme + "://" + parsed.Host + "/upload/v1beta/files"
|
||||
metadata, _ := json.Marshal(map[string]any{"file": map[string]any{"display_name": fmt.Sprintf("easyai-image-%d", time.Now().UnixMilli())}})
|
||||
startRequest, err := http.NewRequestWithContext(ctx, http.MethodPost, uploadStartURL, bytes.NewReader(metadata))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
startRequest.Header.Set("x-goog-api-key", apiKey)
|
||||
startRequest.Header.Set("X-Goog-Upload-Protocol", "resumable")
|
||||
startRequest.Header.Set("X-Goog-Upload-Command", "start")
|
||||
startRequest.Header.Set("X-Goog-Upload-Header-Content-Length", fmt.Sprintf("%d", len(payload)))
|
||||
startRequest.Header.Set("X-Goog-Upload-Header-Content-Type", mimeType)
|
||||
startRequest.Header.Set("Content-Type", "application/json")
|
||||
startResponse, err := client.Do(startRequest)
|
||||
if err != nil {
|
||||
return nil, transportClientError(err)
|
||||
}
|
||||
if startResponse.StatusCode < http.StatusOK || startResponse.StatusCode >= http.StatusMultipleChoices {
|
||||
_, responseErr := decodeHTTPResponse(startResponse)
|
||||
return nil, responseErr
|
||||
}
|
||||
uploadURL := strings.TrimSpace(startResponse.Header.Get("X-Goog-Upload-URL"))
|
||||
_ = startResponse.Body.Close()
|
||||
if uploadURL == "" {
|
||||
return nil, &ClientError{Code: "invalid_response", Message: "Gemini Files API did not return an upload URL", Retryable: false}
|
||||
}
|
||||
uploadRequest, err := http.NewRequestWithContext(ctx, http.MethodPost, uploadURL, bytes.NewReader(payload))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
uploadRequest.Header.Set("X-Goog-Upload-Offset", "0")
|
||||
uploadRequest.Header.Set("X-Goog-Upload-Command", "upload, finalize")
|
||||
uploadRequest.Header.Set("Content-Type", mimeType)
|
||||
uploadResponse, err := client.Do(uploadRequest)
|
||||
if err != nil {
|
||||
return nil, transportClientError(err)
|
||||
}
|
||||
result, err := decodeHTTPResponse(uploadResponse)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
file := mapFromAny(result["file"])
|
||||
fileURI := firstNonEmptyString(file["uri"], result["uri"])
|
||||
if fileURI == "" {
|
||||
return nil, &ClientError{Code: "invalid_response", Message: "Gemini Files API response did not include file.uri", Retryable: false}
|
||||
}
|
||||
return map[string]any{"fileUri": fileURI, "mimeType": firstNonEmptyString(file["mimeType"], file["mime_type"], mimeType)}, nil
|
||||
}
|
||||
|
||||
func geminiOfficialFileURI(value string) bool {
|
||||
parsed, err := url.Parse(strings.TrimSpace(value))
|
||||
return err == nil && strings.EqualFold(parsed.Hostname(), "generativelanguage.googleapis.com") && strings.Contains(strings.ToLower(parsed.Path), "/files/")
|
||||
}
|
||||
|
||||
func geminiURL(baseURL string, model string, apiKey string) string {
|
||||
return geminiActionURL(baseURL, model, "generateContent", apiKey, false)
|
||||
}
|
||||
|
||||
func geminiStreamURL(baseURL string, model string, apiKey string) string {
|
||||
return geminiActionURL(baseURL, model, "streamGenerateContent", apiKey, true)
|
||||
}
|
||||
|
||||
func geminiActionURL(baseURL string, model string, action string, apiKey string, stream bool) string {
|
||||
base := strings.TrimRight(strings.TrimSpace(baseURL), "/")
|
||||
if base == "" {
|
||||
base = "https://generativelanguage.googleapis.com"
|
||||
@@ -67,7 +353,53 @@ func geminiURL(baseURL string, model string, apiKey string) string {
|
||||
base += "/v1beta"
|
||||
}
|
||||
escapedModel := url.PathEscape(model)
|
||||
return fmt.Sprintf("%s/models/%s:generateContent?key=%s", base, escapedModel, url.QueryEscape(apiKey))
|
||||
endpoint := fmt.Sprintf("%s/models/%s:%s?key=%s", base, escapedModel, action, url.QueryEscape(apiKey))
|
||||
if stream {
|
||||
endpoint += "&alt=sse"
|
||||
}
|
||||
return endpoint
|
||||
}
|
||||
|
||||
func decodeGeminiStreamResponse(resp *http.Response, onDelta StreamDelta, wire *WireResponse) (map[string]any, error) {
|
||||
defer resp.Body.Close()
|
||||
scanner := bufio.NewScanner(resp.Body)
|
||||
scanner.Buffer(make([]byte, 64*1024), 16*1024*1024)
|
||||
result := map[string]any{}
|
||||
for scanner.Scan() {
|
||||
line := strings.TrimSpace(scanner.Text())
|
||||
if !strings.HasPrefix(line, "data:") {
|
||||
continue
|
||||
}
|
||||
data := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
|
||||
if data == "" || data == "[DONE]" {
|
||||
continue
|
||||
}
|
||||
var event map[string]any
|
||||
if err := json.Unmarshal([]byte(data), &event); err != nil {
|
||||
return nil, &ClientError{Code: "invalid_response", Message: err.Error(), StatusCode: resp.StatusCode, Retryable: false, Wire: wire}
|
||||
}
|
||||
result = event
|
||||
if onDelta != nil {
|
||||
if err := onDelta(StreamDeltaEvent{
|
||||
Text: geminiText(event),
|
||||
Event: event,
|
||||
WireProtocol: ProtocolGeminiGenerateContent,
|
||||
WireStatusCode: resp.StatusCode,
|
||||
WireHeaders: compatibleResponseHeaders(resp.Header),
|
||||
}); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
}
|
||||
if err := scanner.Err(); err != nil {
|
||||
clientErr := transportClientError(err)
|
||||
clientErr.Wire = wire
|
||||
return nil, clientErr
|
||||
}
|
||||
if len(result) == 0 {
|
||||
return nil, &ClientError{Code: "invalid_response", Message: "gemini stream returned no events", StatusCode: resp.StatusCode, Retryable: false, Wire: wire}
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func geminiBody(request Request) map[string]any {
|
||||
@@ -200,24 +532,39 @@ func geminiGenerationConfig(body map[string]any, imageResponse bool) map[string]
|
||||
if out == nil {
|
||||
out = map[string]any{}
|
||||
}
|
||||
if aspectRatio := firstNonEmptyString(body["aspect_ratio"], body["aspectRatio"]); aspectRatio != "" {
|
||||
imageConfig := cloneMapAny(mapFromAny(firstPresent(out["imageConfig"], out["image_config"])))
|
||||
imageConfig := cloneMapAny(mapFromAny(firstPresent(out["imageConfig"], out["image_config"])))
|
||||
if imageConfig != nil {
|
||||
aspectRatio := normalizedProviderAspectRatio(firstNonEmptyString(imageConfig["aspectRatio"], imageConfig["aspect_ratio"]))
|
||||
imageSize := normalizedProviderImageResolution(firstNonEmptyString(imageConfig["imageSize"], imageConfig["image_size"]))
|
||||
delete(imageConfig, "aspectRatio")
|
||||
delete(imageConfig, "aspect_ratio")
|
||||
delete(imageConfig, "imageSize")
|
||||
delete(imageConfig, "image_size")
|
||||
if aspectRatio != "" {
|
||||
imageConfig["aspectRatio"] = aspectRatio
|
||||
}
|
||||
if imageSize != "" {
|
||||
imageConfig["imageSize"] = imageSize
|
||||
}
|
||||
}
|
||||
if aspectRatio := normalizedProviderAspectRatio(firstNonEmptyString(body["aspect_ratio"], body["aspectRatio"], body["ratio"])); aspectRatio != "" {
|
||||
if imageConfig == nil {
|
||||
imageConfig = map[string]any{}
|
||||
}
|
||||
imageConfig["aspectRatio"] = aspectRatio
|
||||
out["imageConfig"] = imageConfig
|
||||
delete(out, "image_config")
|
||||
}
|
||||
if imageSize := firstNonEmptyString(body["resolution"], body["imageSize"], body["image_size"], body["size"]); imageSize != "" {
|
||||
imageConfig := cloneMapAny(mapFromAny(firstPresent(out["imageConfig"], out["image_config"])))
|
||||
if imageSize := normalizedProviderImageResolution(firstNonEmptyString(body["resolution"], body["imageSize"], body["image_size"], body["size"])); imageSize != "" {
|
||||
if imageConfig == nil {
|
||||
imageConfig = map[string]any{}
|
||||
}
|
||||
imageConfig["imageSize"] = imageSize
|
||||
out["imageConfig"] = imageConfig
|
||||
delete(out, "image_config")
|
||||
}
|
||||
if len(imageConfig) > 0 {
|
||||
out["imageConfig"] = imageConfig
|
||||
} else {
|
||||
delete(out, "imageConfig")
|
||||
}
|
||||
delete(out, "image_config")
|
||||
if count := intFromAny(firstPresent(body["n"], body["candidateCount"], body["candidate_count"])); count > 0 {
|
||||
out["candidateCount"] = count
|
||||
}
|
||||
@@ -544,19 +891,24 @@ func geminiToolsFromOpenAITools(value any) []any {
|
||||
return []any{map[string]any{"functionDeclarations": declarations}}
|
||||
}
|
||||
|
||||
func geminiResult(request Request, raw map[string]any) map[string]any {
|
||||
func geminiResult(request Request, raw map[string]any) (map[string]any, error) {
|
||||
if geminiImageModelType(request.ModelType) {
|
||||
data := geminiImageData(raw)
|
||||
if len(data) == 0 {
|
||||
data = []any{map[string]any{"url": "/static/provider/gemini-image-placeholder.png"}}
|
||||
details := geminiMissingImageDetails(raw)
|
||||
return nil, &ClientError{
|
||||
Code: "gemini_image_output_missing",
|
||||
Message: geminiMissingImageMessage(details),
|
||||
Details: details,
|
||||
Retryable: false,
|
||||
}
|
||||
}
|
||||
return map[string]any{
|
||||
"id": "gemini-image",
|
||||
"created": nowUnix(),
|
||||
"model": request.Model,
|
||||
"data": data,
|
||||
"raw": raw,
|
||||
}
|
||||
}, nil
|
||||
}
|
||||
message, finishReason := geminiChatMessage(raw)
|
||||
return map[string]any{
|
||||
@@ -570,8 +922,7 @@ func geminiResult(request Request, raw map[string]any) map[string]any {
|
||||
"message": message,
|
||||
}},
|
||||
"usage": geminiUsageMap(raw),
|
||||
"raw": raw,
|
||||
}
|
||||
}, nil
|
||||
}
|
||||
|
||||
func textFromMessages(body map[string]any) string {
|
||||
@@ -686,6 +1037,250 @@ func geminiImageData(raw map[string]any) []any {
|
||||
return out
|
||||
}
|
||||
|
||||
func geminiMissingImageDetails(raw map[string]any) map[string]any {
|
||||
details := map[string]any{
|
||||
"provider": "gemini",
|
||||
"resourceType": "image",
|
||||
}
|
||||
upstreamError, _ := raw["error"].(map[string]any)
|
||||
upstreamErrorCode := stringFromPathValue(upstreamError["code"])
|
||||
upstreamErrorStatus := strings.TrimSpace(stringFromAny(upstreamError["status"]))
|
||||
upstreamErrorMessage := compactGeminiDiagnosticText(stringFromAny(upstreamError["message"]), 320)
|
||||
if upstreamErrorCode != "" {
|
||||
details["upstreamErrorCode"] = upstreamErrorCode
|
||||
}
|
||||
if upstreamErrorStatus != "" {
|
||||
details["upstreamErrorStatus"] = upstreamErrorStatus
|
||||
}
|
||||
if upstreamErrorMessage != "" {
|
||||
details["upstreamErrorMessage"] = upstreamErrorMessage
|
||||
}
|
||||
promptFeedback, _ := raw["promptFeedback"].(map[string]any)
|
||||
if promptFeedback == nil {
|
||||
promptFeedback, _ = raw["prompt_feedback"].(map[string]any)
|
||||
}
|
||||
promptBlockReason := firstNonEmptyString(promptFeedback["blockReason"], promptFeedback["block_reason"])
|
||||
promptBlockMessage := compactGeminiDiagnosticText(firstNonEmptyString(
|
||||
promptFeedback["blockReasonMessage"],
|
||||
promptFeedback["block_reason_message"],
|
||||
promptFeedback["message"],
|
||||
), 320)
|
||||
if promptBlockReason != "" {
|
||||
details["promptBlockReason"] = promptBlockReason
|
||||
}
|
||||
if promptBlockMessage != "" {
|
||||
details["promptBlockMessage"] = promptBlockMessage
|
||||
}
|
||||
|
||||
candidates, _ := raw["candidates"].([]any)
|
||||
details["candidateCount"] = len(candidates)
|
||||
finishReasons := make([]string, 0)
|
||||
safetyCategories := make([]string, 0)
|
||||
partKinds := make([]string, 0)
|
||||
providerMessages := make([]string, 0)
|
||||
partCount := 0
|
||||
emptyImagePayload := false
|
||||
for _, rawCandidate := range candidates {
|
||||
candidate, _ := rawCandidate.(map[string]any)
|
||||
if finishReason := strings.TrimSpace(stringFromAny(firstPresent(candidate["finishReason"], candidate["finish_reason"]))); finishReason != "" {
|
||||
finishReasons = appendGeminiDiagnosticValue(finishReasons, finishReason)
|
||||
}
|
||||
safetyRatings, _ := candidate["safetyRatings"].([]any)
|
||||
if safetyRatings == nil {
|
||||
safetyRatings, _ = candidate["safety_ratings"].([]any)
|
||||
}
|
||||
for _, rawRating := range safetyRatings {
|
||||
rating, _ := rawRating.(map[string]any)
|
||||
blocked, _ := firstPresent(rating["blocked"], rating["isBlocked"], rating["is_blocked"]).(bool)
|
||||
if !blocked {
|
||||
continue
|
||||
}
|
||||
if category := strings.TrimSpace(stringFromAny(rating["category"])); category != "" {
|
||||
safetyCategories = appendGeminiDiagnosticValue(safetyCategories, category)
|
||||
}
|
||||
}
|
||||
content, _ := candidate["content"].(map[string]any)
|
||||
parts, _ := content["parts"].([]any)
|
||||
partCount += len(parts)
|
||||
for _, rawPart := range parts {
|
||||
part, _ := rawPart.(map[string]any)
|
||||
if text := compactGeminiDiagnosticText(stringFromAny(part["text"]), 320); text != "" {
|
||||
partKinds = appendGeminiDiagnosticValue(partKinds, "text")
|
||||
providerMessages = appendGeminiDiagnosticValue(providerMessages, text)
|
||||
}
|
||||
inline, inlinePresent := part["inlineData"].(map[string]any)
|
||||
if !inlinePresent {
|
||||
inline, inlinePresent = part["inline_data"].(map[string]any)
|
||||
}
|
||||
if inlinePresent {
|
||||
partKinds = appendGeminiDiagnosticValue(partKinds, "inlineData")
|
||||
if strings.TrimSpace(stringFromAny(inline["data"])) == "" {
|
||||
emptyImagePayload = true
|
||||
}
|
||||
}
|
||||
fileData, filePresent := part["fileData"].(map[string]any)
|
||||
if !filePresent {
|
||||
fileData, filePresent = part["file_data"].(map[string]any)
|
||||
}
|
||||
if filePresent {
|
||||
partKinds = appendGeminiDiagnosticValue(partKinds, "fileData")
|
||||
if firstNonEmptyString(fileData["fileUri"], fileData["file_uri"], fileData["uri"]) == "" {
|
||||
emptyImagePayload = true
|
||||
}
|
||||
}
|
||||
for _, item := range []struct {
|
||||
key string
|
||||
kind string
|
||||
}{
|
||||
{key: "thought", kind: "thought"},
|
||||
{key: "functionCall", kind: "functionCall"},
|
||||
{key: "function_call", kind: "functionCall"},
|
||||
{key: "executableCode", kind: "executableCode"},
|
||||
{key: "executable_code", kind: "executableCode"},
|
||||
{key: "codeExecutionResult", kind: "codeExecutionResult"},
|
||||
{key: "code_execution_result", kind: "codeExecutionResult"},
|
||||
} {
|
||||
if _, ok := part[item.key]; ok {
|
||||
partKinds = appendGeminiDiagnosticValue(partKinds, item.kind)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
details["partCount"] = partCount
|
||||
if len(finishReasons) > 0 {
|
||||
details["candidateFinishReasons"] = finishReasons
|
||||
}
|
||||
if len(safetyCategories) > 0 {
|
||||
details["blockedSafetyCategories"] = safetyCategories
|
||||
}
|
||||
if len(partKinds) > 0 {
|
||||
details["observedPartKinds"] = partKinds
|
||||
}
|
||||
if len(providerMessages) > 0 {
|
||||
details["providerMessages"] = providerMessages
|
||||
}
|
||||
|
||||
reason := "no_supported_image_parts"
|
||||
switch {
|
||||
case upstreamErrorCode != "" || upstreamErrorStatus != "" || upstreamErrorMessage != "":
|
||||
reason = "upstream_error"
|
||||
case promptBlockReason != "":
|
||||
reason = "prompt_blocked"
|
||||
case len(safetyCategories) > 0 || geminiFinishReasonsIndicateFiltering(finishReasons):
|
||||
reason = "candidate_filtered"
|
||||
case len(candidates) == 0:
|
||||
reason = "no_candidates"
|
||||
case emptyImagePayload:
|
||||
reason = "empty_image_payload"
|
||||
case partCount == 0:
|
||||
reason = "no_content_parts"
|
||||
}
|
||||
details["reason"] = reason
|
||||
return details
|
||||
}
|
||||
|
||||
func geminiMissingImageMessage(details map[string]any) string {
|
||||
message := "Gemini image response contains no extractable image"
|
||||
switch stringFromAny(details["reason"]) {
|
||||
case "upstream_error":
|
||||
message += ": upstream returned an error envelope"
|
||||
case "prompt_blocked":
|
||||
message += ": prompt blocked"
|
||||
case "candidate_filtered":
|
||||
message += ": candidate filtered"
|
||||
case "no_candidates":
|
||||
message += ": upstream response contains no candidates"
|
||||
case "empty_image_payload":
|
||||
message += ": image part is missing inline data or file URI"
|
||||
case "no_content_parts":
|
||||
message += ": candidate content contains no parts"
|
||||
default:
|
||||
message += ": candidate parts contain no supported image resource"
|
||||
}
|
||||
if reason := stringFromAny(details["promptBlockReason"]); reason != "" {
|
||||
message += " (blockReason=" + reason + ")"
|
||||
}
|
||||
if code := stringFromAny(details["upstreamErrorCode"]); code != "" {
|
||||
message += "; upstreamErrorCode=" + code
|
||||
}
|
||||
if status := stringFromAny(details["upstreamErrorStatus"]); status != "" {
|
||||
message += "; upstreamErrorStatus=" + status
|
||||
}
|
||||
if values := geminiDiagnosticStrings(details["candidateFinishReasons"]); len(values) > 0 {
|
||||
message += "; finishReason=" + strings.Join(values, ",")
|
||||
}
|
||||
if values := geminiDiagnosticStrings(details["blockedSafetyCategories"]); len(values) > 0 {
|
||||
message += "; blockedSafetyCategories=" + strings.Join(values, ",")
|
||||
}
|
||||
if values := geminiDiagnosticStrings(details["providerMessages"]); len(values) > 0 {
|
||||
message += "; providerMessage=" + values[0]
|
||||
} else if upstreamMessage := stringFromAny(details["upstreamErrorMessage"]); upstreamMessage != "" {
|
||||
message += "; providerMessage=" + upstreamMessage
|
||||
} else if blockMessage := stringFromAny(details["promptBlockMessage"]); blockMessage != "" {
|
||||
message += "; providerMessage=" + blockMessage
|
||||
}
|
||||
if values := geminiDiagnosticStrings(details["observedPartKinds"]); len(values) > 0 {
|
||||
message += "; observedPartKinds=" + strings.Join(values, ",")
|
||||
}
|
||||
return compactGeminiDiagnosticText(message, 1800)
|
||||
}
|
||||
|
||||
func geminiFinishReasonsIndicateFiltering(values []string) bool {
|
||||
for _, value := range values {
|
||||
normalized := strings.ToUpper(strings.TrimSpace(value))
|
||||
switch normalized {
|
||||
case "SAFETY", "IMAGE_SAFETY", "RECITATION", "BLOCKLIST", "PROHIBITED_CONTENT", "SPII":
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func geminiDiagnosticStrings(value any) []string {
|
||||
switch typed := value.(type) {
|
||||
case []string:
|
||||
return typed
|
||||
case []any:
|
||||
values := make([]string, 0, len(typed))
|
||||
for _, item := range typed {
|
||||
if text := strings.TrimSpace(stringFromAny(item)); text != "" {
|
||||
values = append(values, text)
|
||||
}
|
||||
}
|
||||
return values
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func appendGeminiDiagnosticValue(values []string, value string) []string {
|
||||
value = strings.TrimSpace(value)
|
||||
if value == "" {
|
||||
return values
|
||||
}
|
||||
for _, existing := range values {
|
||||
if existing == value {
|
||||
return values
|
||||
}
|
||||
}
|
||||
if len(values) >= 8 {
|
||||
return values
|
||||
}
|
||||
return append(values, value)
|
||||
}
|
||||
|
||||
func compactGeminiDiagnosticText(value string, maxRunes int) string {
|
||||
value = strings.Join(strings.Fields(value), " ")
|
||||
if maxRunes <= 0 {
|
||||
return ""
|
||||
}
|
||||
runes := []rune(value)
|
||||
if len(runes) <= maxRunes {
|
||||
return value
|
||||
}
|
||||
return string(runes[:maxRunes]) + "…"
|
||||
}
|
||||
|
||||
func geminiUsage(raw map[string]any) Usage {
|
||||
usageMap := geminiUsageMap(raw)
|
||||
input := intFromAny(usageMap["prompt_tokens"])
|
||||
|
||||
@@ -0,0 +1,667 @@
|
||||
package clients
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
const geminiVeoMaxVideoBytes int64 = 256 << 20
|
||||
|
||||
type geminiVeoImage struct {
|
||||
URI string
|
||||
MimeType string
|
||||
}
|
||||
|
||||
func geminiVeoRequest(request Request) bool {
|
||||
return request.Kind == "videos.generations" &&
|
||||
strings.Contains(strings.ToLower(strings.TrimSpace(upstreamModelName(request.Candidate))), "veo")
|
||||
}
|
||||
|
||||
func (c GeminiClient) runVeo(ctx context.Context, request Request, apiKey string) (Response, error) {
|
||||
startedAt := time.Now()
|
||||
operationName := strings.TrimSpace(request.RemoteTaskID)
|
||||
lastRequestID := operationName
|
||||
var operation map[string]any
|
||||
var wire *WireResponse
|
||||
|
||||
if operationName == "" {
|
||||
body, err := geminiVeoBody(ctx, request)
|
||||
if err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
operation, lastRequestID, wire, err = c.geminiVeoPost(ctx, request, apiKey, body)
|
||||
if err != nil {
|
||||
return Response{}, annotateResponseError(err, lastRequestID, startedAt, time.Now())
|
||||
}
|
||||
operationName = strings.TrimSpace(stringFromAny(operation["name"]))
|
||||
if err := validateGeminiVeoOperationName(operationName); err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
if request.OnRemoteTaskSubmitted != nil {
|
||||
if err := request.OnRemoteTaskSubmitted(operationName, geminiVeoOperationCheckpoint(operation)); err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
}
|
||||
} else if err := validateGeminiVeoOperationName(operationName); err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
|
||||
interval := durationFromConfig(request.Candidate.PlatformConfig, 10*time.Second, "pollIntervalMs", "poll_interval_ms")
|
||||
timeout := durationFromConfig(request.Candidate.PlatformConfig, ProviderRequestTimeout(request.Kind), "pollTimeoutMs", "poll_timeout_ms", "timeoutMs")
|
||||
deadline := time.NewTimer(timeout)
|
||||
defer deadline.Stop()
|
||||
nextPoll := time.NewTimer(0)
|
||||
defer nextPoll.Stop()
|
||||
transientFailures := 0
|
||||
|
||||
for {
|
||||
if operation != nil && boolFromAny(operation["done"]) {
|
||||
return c.geminiVeoCompletedResponse(ctx, request, apiKey, operationName, operation, wire, lastRequestID, startedAt)
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return Response{}, &ClientError{Code: "cancelled", Message: ctx.Err().Error(), RequestID: lastRequestID, Retryable: true}
|
||||
case <-deadline.C:
|
||||
return Response{}, &ClientError{
|
||||
Code: "timeout",
|
||||
Message: fmt.Sprintf("gemini Veo operation %s did not finish before timeout", operationName),
|
||||
RequestID: lastRequestID,
|
||||
Retryable: true,
|
||||
}
|
||||
case <-nextPoll.C:
|
||||
pollStartedAt := time.Now()
|
||||
result, requestID, pollWire, err := c.geminiVeoGetOperation(ctx, request, apiKey, operationName)
|
||||
if requestID != "" {
|
||||
lastRequestID = requestID
|
||||
}
|
||||
if err != nil {
|
||||
err = annotateResponseError(err, lastRequestID, pollStartedAt, time.Now())
|
||||
if !IsRetryable(err) {
|
||||
return Response{}, err
|
||||
}
|
||||
transientFailures++
|
||||
resetGeminiVeoPollTimer(nextPoll, geminiVeoRetryInterval(interval, transientFailures))
|
||||
continue
|
||||
}
|
||||
transientFailures = 0
|
||||
operation = result
|
||||
wire = pollWire
|
||||
if request.OnRemoteTaskPolled != nil {
|
||||
if err := request.OnRemoteTaskPolled(operationName, geminiVeoOperationCheckpoint(operation)); err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
}
|
||||
if boolFromAny(operation["done"]) {
|
||||
return c.geminiVeoCompletedResponse(ctx, request, apiKey, operationName, operation, wire, lastRequestID, startedAt)
|
||||
}
|
||||
resetGeminiVeoPollTimer(nextPoll, interval)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func geminiVeoBody(ctx context.Context, request Request) (map[string]any, error) {
|
||||
prompt := firstNonEmptyPrompt(request.Body, "")
|
||||
if prompt == "" {
|
||||
return nil, &ClientError{Code: "invalid_parameter", Message: "gemini Veo prompt is required", Param: "prompt", StatusCode: http.StatusBadRequest, Retryable: false}
|
||||
}
|
||||
|
||||
firstFrame, lastFrame, references := geminiVeoImageInputs(request.Body)
|
||||
instance := map[string]any{"prompt": prompt}
|
||||
if firstFrame.URI != "" {
|
||||
image, err := geminiVeoInlineImage(ctx, firstFrame)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
instance["image"] = image
|
||||
}
|
||||
if lastFrame.URI != "" {
|
||||
image, err := geminiVeoInlineImage(ctx, lastFrame)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
instance["lastFrame"] = image
|
||||
}
|
||||
if len(references) > 3 {
|
||||
return nil, &ClientError{Code: "invalid_parameter", Message: "gemini Veo supports at most 3 reference images", Param: "reference_images", StatusCode: http.StatusBadRequest, Retryable: false}
|
||||
}
|
||||
if len(references) > 0 {
|
||||
items := make([]any, 0, len(references))
|
||||
for _, reference := range references {
|
||||
image, err := geminiVeoInlineImage(ctx, reference)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, map[string]any{"image": image, "referenceType": "asset"})
|
||||
}
|
||||
instance["referenceImages"] = items
|
||||
}
|
||||
|
||||
parameters, err := geminiVeoParameters(request.Body, lastFrame.URI != "", len(references) > 0)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
body := map[string]any{"instances": []any{instance}}
|
||||
if len(parameters) > 0 {
|
||||
body["parameters"] = parameters
|
||||
}
|
||||
return body, nil
|
||||
}
|
||||
|
||||
func geminiVeoParameters(body map[string]any, hasLastFrame bool, hasReferences bool) (map[string]any, error) {
|
||||
parameters := map[string]any{}
|
||||
count := intFromAny(firstPresent(body["n"], body["numberOfVideos"], body["number_of_videos"]))
|
||||
if count == 0 {
|
||||
count = 1
|
||||
}
|
||||
if count != 1 {
|
||||
return nil, &ClientError{Code: "invalid_parameter", Message: "gemini Veo supports exactly 1 output video", Param: "n", StatusCode: http.StatusBadRequest, Retryable: false}
|
||||
}
|
||||
parameters["sampleCount"] = count
|
||||
|
||||
duration := intFromAny(firstPresent(body["duration"], body["duration_seconds"], body["durationSeconds"]))
|
||||
if duration != 0 {
|
||||
if duration != 4 && duration != 6 && duration != 8 {
|
||||
return nil, &ClientError{Code: "invalid_parameter", Message: "gemini Veo duration must be 4, 6, or 8 seconds", Param: "duration", StatusCode: http.StatusBadRequest, Retryable: false}
|
||||
}
|
||||
parameters["durationSeconds"] = duration
|
||||
}
|
||||
|
||||
aspectRatio := strings.TrimSpace(firstNonEmptyString(body["aspect_ratio"], body["aspectRatio"], body["ratio"]))
|
||||
if aspectRatio != "" {
|
||||
if aspectRatio != "16:9" && aspectRatio != "9:16" {
|
||||
return nil, &ClientError{Code: "invalid_parameter", Message: "gemini Veo aspect ratio must be 16:9 or 9:16", Param: "aspect_ratio", StatusCode: http.StatusBadRequest, Retryable: false}
|
||||
}
|
||||
parameters["aspectRatio"] = aspectRatio
|
||||
}
|
||||
|
||||
resolution := geminiVeoResolution(firstNonEmptyString(body["resolution"], body["size"]))
|
||||
if resolution == "invalid" {
|
||||
return nil, &ClientError{Code: "invalid_parameter", Message: "gemini Veo resolution must be 720p, 1080p, or 4k", Param: "resolution", StatusCode: http.StatusBadRequest, Retryable: false}
|
||||
}
|
||||
if resolution != "" {
|
||||
parameters["resolution"] = resolution
|
||||
}
|
||||
if (resolution == "1080p" || resolution == "4k" || hasLastFrame || hasReferences) && duration != 0 && duration != 8 {
|
||||
return nil, &ClientError{Code: "invalid_parameter", Message: "gemini Veo requires an 8-second duration for 1080p, 4k, last-frame, or reference-image generation", Param: "duration", StatusCode: http.StatusBadRequest, Retryable: false}
|
||||
}
|
||||
|
||||
for _, item := range []struct {
|
||||
to string
|
||||
from []string
|
||||
}{
|
||||
{to: "personGeneration", from: []string{"personGeneration", "person_generation"}},
|
||||
{to: "negativePrompt", from: []string{"negativePrompt", "negative_prompt"}},
|
||||
} {
|
||||
if value := strings.TrimSpace(firstNonEmptyStringValue(body, item.from...)); value != "" {
|
||||
parameters[item.to] = value
|
||||
}
|
||||
}
|
||||
for _, key := range []string{"enhancePrompt", "enhance_prompt"} {
|
||||
if value, ok := body[key].(bool); ok {
|
||||
parameters["enhancePrompt"] = value
|
||||
break
|
||||
}
|
||||
}
|
||||
return parameters, nil
|
||||
}
|
||||
|
||||
func geminiVeoResolution(value string) string {
|
||||
normalized := strings.ToLower(strings.TrimSpace(value))
|
||||
switch normalized {
|
||||
case "":
|
||||
return ""
|
||||
case "720p", "1080p", "4k":
|
||||
return normalized
|
||||
case "2160p", "3840x2160", "2160x3840":
|
||||
return "4k"
|
||||
default:
|
||||
return "invalid"
|
||||
}
|
||||
}
|
||||
|
||||
func geminiVeoImageInputs(body map[string]any) (geminiVeoImage, geminiVeoImage, []geminiVeoImage) {
|
||||
first := firstGeminiVeoImage(body["first_frame"], body["firstFrame"], body["first_frame_image"], body["firstFrameImage"])
|
||||
last := firstGeminiVeoImage(body["last_frame"], body["lastFrame"], body["last_frame_image"], body["lastFrameImage"])
|
||||
references := geminiVeoImagesFromValues(body["reference_images"], body["referenceImages"], body["reference_image"], body["referenceImage"])
|
||||
images := geminiVeoImagesFromValues(body["image"], body["images"], body["image_url"], body["imageUrl"], body["image_urls"], body["imageUrls"])
|
||||
if first.URI == "" && len(images) > 0 {
|
||||
first = images[0]
|
||||
images = images[1:]
|
||||
}
|
||||
references = append(references, images...)
|
||||
|
||||
for _, item := range contentItems(body["content"]) {
|
||||
if !strings.Contains(strings.ToLower(strings.TrimSpace(stringFromAny(item["type"]))), "image") {
|
||||
continue
|
||||
}
|
||||
image := firstGeminiVeoImage(item["image_url"], item["imageUrl"], item["url"], item["image"])
|
||||
if image.URI == "" {
|
||||
continue
|
||||
}
|
||||
if image.MimeType == "" {
|
||||
image.MimeType = firstNonEmptyString(item["mime_type"], item["mimeType"])
|
||||
}
|
||||
switch strings.ToLower(strings.TrimSpace(stringFromAny(item["role"]))) {
|
||||
case "first_frame", "firstframe":
|
||||
if first.URI == "" {
|
||||
first = image
|
||||
}
|
||||
case "last_frame", "lastframe":
|
||||
if last.URI == "" {
|
||||
last = image
|
||||
}
|
||||
default:
|
||||
if first.URI == "" {
|
||||
first = image
|
||||
} else {
|
||||
references = append(references, image)
|
||||
}
|
||||
}
|
||||
}
|
||||
return first, last, deduplicateGeminiVeoImages(references, first.URI, last.URI)
|
||||
}
|
||||
|
||||
func firstGeminiVeoImage(values ...any) geminiVeoImage {
|
||||
for _, value := range values {
|
||||
if images := geminiVeoImagesFromAny(value); len(images) > 0 {
|
||||
return images[0]
|
||||
}
|
||||
}
|
||||
return geminiVeoImage{}
|
||||
}
|
||||
|
||||
func geminiVeoImagesFromValues(values ...any) []geminiVeoImage {
|
||||
out := make([]geminiVeoImage, 0)
|
||||
for _, value := range values {
|
||||
out = append(out, geminiVeoImagesFromAny(value)...)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func geminiVeoImagesFromAny(value any) []geminiVeoImage {
|
||||
switch typed := value.(type) {
|
||||
case string:
|
||||
if uri := strings.TrimSpace(typed); uri != "" {
|
||||
return []geminiVeoImage{{URI: uri}}
|
||||
}
|
||||
case []any:
|
||||
out := make([]geminiVeoImage, 0, len(typed))
|
||||
for _, item := range typed {
|
||||
out = append(out, geminiVeoImagesFromAny(item)...)
|
||||
}
|
||||
return out
|
||||
case []string:
|
||||
out := make([]geminiVeoImage, 0, len(typed))
|
||||
for _, item := range typed {
|
||||
out = append(out, geminiVeoImagesFromAny(item)...)
|
||||
}
|
||||
return out
|
||||
case map[string]any:
|
||||
if nested := firstPresent(typed["image_url"], typed["imageUrl"], typed["image"]); nested != nil {
|
||||
if images := geminiVeoImagesFromAny(nested); len(images) > 0 {
|
||||
if images[0].MimeType == "" {
|
||||
images[0].MimeType = firstNonEmptyString(typed["mime_type"], typed["mimeType"])
|
||||
}
|
||||
return images
|
||||
}
|
||||
}
|
||||
uri := firstNonEmptyString(typed["url"], typed["uri"], typed["data"], typed["bytesBase64Encoded"], typed["bytes_base64_encoded"])
|
||||
if uri != "" {
|
||||
return []geminiVeoImage{{URI: uri, MimeType: firstNonEmptyString(typed["mime_type"], typed["mimeType"])}}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func deduplicateGeminiVeoImages(images []geminiVeoImage, exclusions ...string) []geminiVeoImage {
|
||||
seen := map[string]bool{}
|
||||
for _, exclusion := range exclusions {
|
||||
if exclusion = strings.TrimSpace(exclusion); exclusion != "" {
|
||||
seen[exclusion] = true
|
||||
}
|
||||
}
|
||||
out := make([]geminiVeoImage, 0, len(images))
|
||||
for _, image := range images {
|
||||
image.URI = strings.TrimSpace(image.URI)
|
||||
if image.URI == "" || seen[image.URI] {
|
||||
continue
|
||||
}
|
||||
seen[image.URI] = true
|
||||
out = append(out, image)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func geminiVeoInlineImage(ctx context.Context, image geminiVeoImage) (map[string]any, error) {
|
||||
parsed := geminiDataURL(image.URI)
|
||||
mimeType := strings.TrimSpace(image.MimeType)
|
||||
data := ""
|
||||
if parsed != nil {
|
||||
data = parsed.data
|
||||
mimeType = firstNonEmptyString(mimeType, parsed.mimeType)
|
||||
} else if requestAssetStringIsHTTPURL(image.URI) {
|
||||
fetchedMimeType, payload, err := fetchRemoteMediaInputPayload(ctx, image.URI, 256<<20)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
data = base64.StdEncoding.EncodeToString(payload)
|
||||
mimeType = firstNonEmptyString(mimeType, fetchedMimeType)
|
||||
} else if !requestLikeURL(image.URI) {
|
||||
data = strings.TrimSpace(image.URI)
|
||||
}
|
||||
if data == "" {
|
||||
return nil, &ClientError{
|
||||
Code: "invalid_parameter",
|
||||
Message: "gemini Veo image input must be hydrated as base64 data",
|
||||
Param: "image",
|
||||
StatusCode: http.StatusBadRequest,
|
||||
Retryable: false,
|
||||
}
|
||||
}
|
||||
if _, err := base64.StdEncoding.DecodeString(data); err != nil {
|
||||
return nil, &ClientError{Code: "invalid_parameter", Message: "gemini Veo image input contains invalid base64 data", Param: "image", StatusCode: http.StatusBadRequest, Retryable: false}
|
||||
}
|
||||
if mimeType == "" {
|
||||
mimeType = "image/png"
|
||||
}
|
||||
return map[string]any{"bytesBase64Encoded": data, "mimeType": geminiMediaMime(mimeType, "image")}, nil
|
||||
}
|
||||
|
||||
func requestAssetStringIsHTTPURL(value string) bool {
|
||||
parsed, err := url.Parse(strings.TrimSpace(value))
|
||||
return err == nil && parsed.Host != "" && (strings.EqualFold(parsed.Scheme, "http") || strings.EqualFold(parsed.Scheme, "https"))
|
||||
}
|
||||
|
||||
func requestLikeURL(value string) bool {
|
||||
normalized := strings.ToLower(strings.TrimSpace(value))
|
||||
return strings.HasPrefix(normalized, "http://") || strings.HasPrefix(normalized, "https://") || strings.HasPrefix(normalized, "/static/")
|
||||
}
|
||||
|
||||
func (c GeminiClient) geminiVeoPost(ctx context.Context, request Request, apiKey string, body map[string]any) (map[string]any, string, *WireResponse, error) {
|
||||
raw, _ := json.Marshal(body)
|
||||
endpoint := geminiVeoActionURL(request.Candidate.BaseURL, upstreamModelName(request.Candidate))
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(raw))
|
||||
if err != nil {
|
||||
return nil, "", nil, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("x-goog-api-key", apiKey)
|
||||
applyUpstreamIdempotency(req, request)
|
||||
if err := notifySubmissionStarted(request); err != nil {
|
||||
return nil, "", nil, err
|
||||
}
|
||||
resp, err := httpClient(request.HTTPClient, c.HTTPClient).Do(req)
|
||||
if err != nil {
|
||||
return nil, "", nil, transportClientError(err)
|
||||
}
|
||||
if err := notifyResponseReceived(request); err != nil {
|
||||
resp.Body.Close()
|
||||
return nil, "", nil, err
|
||||
}
|
||||
requestID := requestIDFromHTTPResponse(resp)
|
||||
result, wire, err := decodeHTTPResponseForProtocol(resp, ProtocolGeminiVeo)
|
||||
if notifyErr := notifyWireResponse(request, wire); notifyErr != nil {
|
||||
return result, requestID, wire, notifyErr
|
||||
}
|
||||
return result, requestID, wire, err
|
||||
}
|
||||
|
||||
func (c GeminiClient) geminiVeoGetOperation(ctx context.Context, request Request, apiKey string, operationName string) (map[string]any, string, *WireResponse, error) {
|
||||
endpoint, err := geminiVeoOperationURL(request.Candidate.BaseURL, operationName)
|
||||
if err != nil {
|
||||
return nil, "", nil, err
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, nil)
|
||||
if err != nil {
|
||||
return nil, "", nil, err
|
||||
}
|
||||
req.Header.Set("x-goog-api-key", apiKey)
|
||||
resp, err := httpClient(request.HTTPClient, c.HTTPClient).Do(req)
|
||||
if err != nil {
|
||||
return nil, "", nil, transportClientError(err)
|
||||
}
|
||||
requestID := requestIDFromHTTPResponse(resp)
|
||||
result, wire, err := decodeHTTPResponseForProtocol(resp, ProtocolGeminiVeo)
|
||||
return result, requestID, wire, err
|
||||
}
|
||||
|
||||
func (c GeminiClient) geminiVeoCompletedResponse(ctx context.Context, request Request, apiKey string, operationName string, operation map[string]any, wire *WireResponse, requestID string, startedAt time.Time) (Response, error) {
|
||||
finishedAt := time.Now()
|
||||
if operationError := mapFromAny(operation["error"]); len(operationError) > 0 {
|
||||
code := firstNonEmptyString(operationError["status"], operationError["code"])
|
||||
if code == "" {
|
||||
code = "gemini_veo_failed"
|
||||
}
|
||||
message := strings.TrimSpace(stringFromAny(operationError["message"]))
|
||||
if message == "" {
|
||||
message = "gemini Veo operation failed"
|
||||
}
|
||||
return Response{}, &ClientError{
|
||||
Code: strings.ToLower(code),
|
||||
Message: message,
|
||||
RequestID: requestID,
|
||||
ResponseStartedAt: startedAt,
|
||||
ResponseFinishedAt: finishedAt,
|
||||
ResponseDurationMS: responseDurationMS(startedAt, finishedAt),
|
||||
Retryable: false,
|
||||
}
|
||||
}
|
||||
items, err := c.geminiVeoResultItems(ctx, request, apiKey, operation)
|
||||
finishedAt = time.Now()
|
||||
if err != nil {
|
||||
return Response{}, annotateResponseError(err, requestID, startedAt, finishedAt)
|
||||
}
|
||||
result := map[string]any{
|
||||
"id": operationName,
|
||||
"object": "video.generation",
|
||||
"created": nowUnix(),
|
||||
"model": request.Model,
|
||||
"upstream_task_id": operationName,
|
||||
"data": items,
|
||||
}
|
||||
return Response{
|
||||
Result: result,
|
||||
RequestID: firstNonEmpty(requestID, operationName),
|
||||
Progress: providerProgress(request),
|
||||
ResponseStartedAt: startedAt,
|
||||
ResponseFinishedAt: finishedAt,
|
||||
ResponseDurationMS: responseDurationMS(startedAt, finishedAt),
|
||||
UpstreamProtocol: ProtocolGeminiVeo,
|
||||
Wire: wire,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (c GeminiClient) geminiVeoResultItems(ctx context.Context, request Request, apiKey string, operation map[string]any) ([]any, error) {
|
||||
response := mapFromAny(operation["response"])
|
||||
generateResponse := mapFromAny(firstPresent(response["generateVideoResponse"], response["generate_video_response"]))
|
||||
samples := mapListFromAny(firstPresent(generateResponse["generatedSamples"], generateResponse["generated_samples"], response["generatedVideos"], response["generated_videos"]))
|
||||
if len(samples) == 0 {
|
||||
return nil, &ClientError{Code: "invalid_response", Message: "gemini Veo operation returned no generated video", Retryable: false}
|
||||
}
|
||||
items := make([]any, 0, len(samples))
|
||||
for _, sample := range samples {
|
||||
video := mapFromAny(sample["video"])
|
||||
if len(video) == 0 {
|
||||
video = sample
|
||||
}
|
||||
mimeType := firstNonEmptyString(video["mimeType"], video["mime_type"], sample["mimeType"], sample["mime_type"])
|
||||
if mimeType == "" {
|
||||
mimeType = "video/mp4"
|
||||
}
|
||||
var payload []byte
|
||||
if encoded := firstNonEmptyString(video["videoBytes"], video["video_bytes"], video["bytesBase64Encoded"], video["bytes_base64_encoded"]); encoded != "" {
|
||||
decoded, err := base64.StdEncoding.DecodeString(encoded)
|
||||
if err != nil {
|
||||
return nil, &ClientError{Code: "invalid_response", Message: "gemini Veo returned invalid video base64 data", Retryable: false}
|
||||
}
|
||||
payload = decoded
|
||||
} else if uri := strings.TrimSpace(firstNonEmptyString(video["uri"], video["url"])); uri != "" {
|
||||
var err error
|
||||
payload, mimeType, err = c.geminiVeoDownload(ctx, request, apiKey, uri, mimeType)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
if len(payload) == 0 {
|
||||
return nil, &ClientError{Code: "invalid_response", Message: "gemini Veo generated video payload is missing", Retryable: false}
|
||||
}
|
||||
items = append(items, map[string]any{"type": "video", "video_bytes": payload, "mime_type": mimeType})
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
|
||||
func (c GeminiClient) geminiVeoDownload(ctx context.Context, request Request, apiKey string, rawURI string, fallbackMimeType string) ([]byte, string, error) {
|
||||
downloadURL, trustedHost, err := geminiVeoDownloadURL(request.Candidate.BaseURL, rawURI)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, downloadURL, nil)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
if trustedHost {
|
||||
req.Header.Set("x-goog-api-key", apiKey)
|
||||
}
|
||||
client := httpClient(request.HTTPClient, c.HTTPClient)
|
||||
redirectClient := *client
|
||||
originalRedirect := client.CheckRedirect
|
||||
redirectClient.CheckRedirect = func(next *http.Request, via []*http.Request) error {
|
||||
if originalRedirect != nil {
|
||||
if err := originalRedirect(next, via); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if len(via) > 0 && !strings.EqualFold(next.URL.Hostname(), via[0].URL.Hostname()) {
|
||||
next.Header.Del("x-goog-api-key")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
resp, err := redirectClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, "", transportClientError(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
raw, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
|
||||
return nil, "", &ClientError{Code: statusCodeName(resp.StatusCode), Message: errorMessage(raw, resp.Status), StatusCode: resp.StatusCode, RequestID: requestIDFromHTTPResponse(resp), Retryable: HTTPRetryable(resp.StatusCode)}
|
||||
}
|
||||
payload, err := io.ReadAll(io.LimitReader(resp.Body, geminiVeoMaxVideoBytes+1))
|
||||
if err != nil {
|
||||
if timeoutErr := transportClientError(err); timeoutErr.Code == "timeout" {
|
||||
timeoutErr.StatusCode = resp.StatusCode
|
||||
timeoutErr.RequestID = requestIDFromHTTPResponse(resp)
|
||||
return nil, "", timeoutErr
|
||||
}
|
||||
return nil, "", &ClientError{Code: "response_read_error", Message: err.Error(), StatusCode: resp.StatusCode, Retryable: true}
|
||||
}
|
||||
if int64(len(payload)) > geminiVeoMaxVideoBytes {
|
||||
return nil, "", &ClientError{Code: "response_too_large", Message: fmt.Sprintf("gemini Veo video exceeds %d bytes", geminiVeoMaxVideoBytes), StatusCode: resp.StatusCode, Retryable: false}
|
||||
}
|
||||
mimeType := strings.TrimSpace(strings.Split(resp.Header.Get("Content-Type"), ";")[0])
|
||||
if mimeType == "" || mimeType == "application/octet-stream" {
|
||||
mimeType = fallbackMimeType
|
||||
}
|
||||
if mimeType == "" {
|
||||
mimeType = "video/mp4"
|
||||
}
|
||||
return payload, mimeType, nil
|
||||
}
|
||||
|
||||
func geminiVeoVersionedBaseURL(baseURL string) string {
|
||||
base := strings.TrimRight(strings.TrimSpace(baseURL), "/")
|
||||
if base == "" {
|
||||
base = "https://generativelanguage.googleapis.com"
|
||||
}
|
||||
base = strings.TrimSuffix(base, "/openai")
|
||||
if !strings.HasSuffix(base, "/v1") && !strings.HasSuffix(base, "/v1beta") && !strings.HasSuffix(base, "/v1alpha") {
|
||||
base += "/v1beta"
|
||||
}
|
||||
return base
|
||||
}
|
||||
|
||||
func geminiVeoActionURL(baseURL string, model string) string {
|
||||
return fmt.Sprintf("%s/models/%s:predictLongRunning", geminiVeoVersionedBaseURL(baseURL), url.PathEscape(strings.TrimSpace(model)))
|
||||
}
|
||||
|
||||
func geminiVeoOperationURL(baseURL string, operationName string) (string, error) {
|
||||
if err := validateGeminiVeoOperationName(operationName); err != nil {
|
||||
return "", err
|
||||
}
|
||||
segments := strings.Split(strings.Trim(operationName, "/"), "/")
|
||||
for index, segment := range segments {
|
||||
segments[index] = url.PathEscape(segment)
|
||||
}
|
||||
return geminiVeoVersionedBaseURL(baseURL) + "/" + strings.Join(segments, "/"), nil
|
||||
}
|
||||
|
||||
func validateGeminiVeoOperationName(operationName string) error {
|
||||
name := strings.TrimSpace(operationName)
|
||||
if name == "" || strings.Contains(name, "..") || strings.ContainsAny(name, "?#\\") || strings.HasPrefix(name, "/") || strings.Contains(name, "://") {
|
||||
return &ClientError{Code: "invalid_response", Message: "gemini Veo operation name is invalid", Retryable: false}
|
||||
}
|
||||
if !strings.Contains(name, "/operations/") && !strings.HasPrefix(name, "operations/") {
|
||||
return &ClientError{Code: "invalid_response", Message: "gemini Veo operation name is invalid", Retryable: false}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func geminiVeoDownloadURL(baseURL string, rawURI string) (string, bool, error) {
|
||||
base, err := url.Parse(geminiVeoVersionedBaseURL(baseURL))
|
||||
if err != nil {
|
||||
return "", false, &ClientError{Code: "invalid_response", Message: "gemini Veo base URL is invalid", Retryable: false}
|
||||
}
|
||||
parsed, err := url.Parse(strings.TrimSpace(rawURI))
|
||||
if err != nil {
|
||||
return "", false, &ClientError{Code: "invalid_response", Message: "gemini Veo video URI is invalid", Retryable: false}
|
||||
}
|
||||
if !parsed.IsAbs() {
|
||||
parsed = base.ResolveReference(parsed)
|
||||
}
|
||||
if parsed.Scheme != "https" && parsed.Scheme != "http" {
|
||||
return "", false, &ClientError{Code: "invalid_response", Message: "gemini Veo video URI scheme is unsupported", Retryable: false}
|
||||
}
|
||||
hostname := strings.ToLower(strings.TrimSpace(parsed.Hostname()))
|
||||
trusted := strings.EqualFold(hostname, base.Hostname()) || hostname == "generativelanguage.googleapis.com"
|
||||
return parsed.String(), trusted, nil
|
||||
}
|
||||
|
||||
func geminiVeoOperationCheckpoint(operation map[string]any) map[string]any {
|
||||
checkpoint := map[string]any{"done": boolFromAny(operation["done"])}
|
||||
if metadata := mapFromAny(operation["metadata"]); len(metadata) > 0 {
|
||||
checkpoint["metadata"] = metadata
|
||||
}
|
||||
if operationError := mapFromAny(operation["error"]); len(operationError) > 0 {
|
||||
checkpoint["error"] = operationError
|
||||
}
|
||||
return checkpoint
|
||||
}
|
||||
|
||||
func resetGeminiVeoPollTimer(timer *time.Timer, duration time.Duration) {
|
||||
if !timer.Stop() {
|
||||
select {
|
||||
case <-timer.C:
|
||||
default:
|
||||
}
|
||||
}
|
||||
timer.Reset(duration)
|
||||
}
|
||||
|
||||
func geminiVeoRetryInterval(interval time.Duration, failures int) time.Duration {
|
||||
if failures <= 0 {
|
||||
return interval
|
||||
}
|
||||
multiplier := 1 << min(failures-1, 4)
|
||||
result := interval * time.Duration(multiplier)
|
||||
if result > 30*time.Second {
|
||||
return 30 * time.Second
|
||||
}
|
||||
return result
|
||||
}
|
||||
@@ -0,0 +1,383 @@
|
||||
package clients
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"reflect"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func TestGeminiVeoRunUsesOfficialLongRunningProtocol(t *testing.T) {
|
||||
const operationName = "models/veo-3.1-generate-preview/operations/test-operation"
|
||||
videoPayload := []byte("test mp4 payload")
|
||||
var server *httptest.Server
|
||||
var pollCount atomic.Int32
|
||||
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Header.Get("x-goog-api-key") != "test-gemini-key" {
|
||||
t.Errorf("missing Gemini API key header for %s", r.URL.Path)
|
||||
}
|
||||
switch {
|
||||
case r.Method == http.MethodPost && r.URL.Path == "/v1beta/models/veo-3.1-generate-preview:predictLongRunning":
|
||||
var body map[string]any
|
||||
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||||
t.Fatalf("decode Veo body: %v", err)
|
||||
}
|
||||
instances := body["instances"].([]any)
|
||||
if prompt := instances[0].(map[string]any)["prompt"]; prompt != "a paper boat on a lake" {
|
||||
t.Errorf("unexpected prompt: %v", prompt)
|
||||
}
|
||||
parameters := body["parameters"].(map[string]any)
|
||||
if parameters["durationSeconds"] != float64(4) || parameters["resolution"] != "720p" || parameters["aspectRatio"] != "16:9" {
|
||||
t.Errorf("unexpected parameters: %+v", parameters)
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"name": operationName})
|
||||
case r.Method == http.MethodGet && r.URL.Path == "/v1beta/"+operationName:
|
||||
if pollCount.Add(1) == 1 {
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"name": operationName, "done": false})
|
||||
return
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"name": operationName,
|
||||
"done": true,
|
||||
"response": map[string]any{"generateVideoResponse": map[string]any{"generatedSamples": []any{
|
||||
map[string]any{"video": map[string]any{"uri": server.URL + "/v1beta/files/generated-video:download", "mimeType": "video/mp4"}},
|
||||
}}},
|
||||
})
|
||||
case r.Method == http.MethodGet && r.URL.Path == "/v1beta/files/generated-video:download":
|
||||
w.Header().Set("Content-Type", "video/mp4")
|
||||
_, _ = w.Write(videoPayload)
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
var submitted string
|
||||
polled := 0
|
||||
response, err := (GeminiClient{}).Run(context.Background(), Request{
|
||||
Kind: "videos.generations",
|
||||
ModelType: "video_generate",
|
||||
Model: "veo-3.1",
|
||||
Body: map[string]any{
|
||||
"prompt": "a paper boat on a lake",
|
||||
"duration": 4,
|
||||
"resolution": "720p",
|
||||
"aspect_ratio": "16:9",
|
||||
},
|
||||
Candidate: testGeminiVeoCandidate(server.URL),
|
||||
OnRemoteTaskSubmitted: func(remoteTaskID string, _ map[string]any) error {
|
||||
submitted = remoteTaskID
|
||||
return nil
|
||||
},
|
||||
OnRemoteTaskPolled: func(remoteTaskID string, _ map[string]any) error {
|
||||
if remoteTaskID != operationName {
|
||||
t.Errorf("unexpected polled operation: %s", remoteTaskID)
|
||||
}
|
||||
polled++
|
||||
return nil
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("run Gemini Veo: %v", err)
|
||||
}
|
||||
if submitted != operationName || polled != 2 || pollCount.Load() != 2 {
|
||||
t.Fatalf("unexpected async callbacks: submitted=%q polled=%d requests=%d", submitted, polled, pollCount.Load())
|
||||
}
|
||||
if response.UpstreamProtocol != ProtocolGeminiVeo || response.RequestID != operationName {
|
||||
t.Fatalf("unexpected response metadata: %+v", response)
|
||||
}
|
||||
data := response.Result["data"].([]any)
|
||||
item := data[0].(map[string]any)
|
||||
if got := item["video_bytes"].([]byte); !reflect.DeepEqual(got, videoPayload) {
|
||||
t.Fatalf("unexpected video payload: %q", got)
|
||||
}
|
||||
if _, exposed := item["uri"]; exposed {
|
||||
t.Fatalf("upstream authenticated URI must not be exposed: %+v", item)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiVeoRunResumesOperationWithoutSubmittingAgain(t *testing.T) {
|
||||
const operationName = "models/veo-3.1-generate-preview/operations/resumed"
|
||||
var postCount atomic.Int32
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method == http.MethodPost {
|
||||
postCount.Add(1)
|
||||
http.Error(w, "unexpected submit", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"name": operationName,
|
||||
"done": true,
|
||||
"response": map[string]any{"generateVideoResponse": map[string]any{"generatedSamples": []any{
|
||||
map[string]any{"video": map[string]any{"videoBytes": base64.StdEncoding.EncodeToString([]byte("resumed video")), "mimeType": "video/mp4"}},
|
||||
}}},
|
||||
})
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
response, err := (GeminiClient{}).Run(context.Background(), Request{
|
||||
Kind: "videos.generations",
|
||||
ModelType: "video_generate",
|
||||
Model: "veo-3.1",
|
||||
Body: map[string]any{"prompt": "ignored on resume"},
|
||||
Candidate: testGeminiVeoCandidate(server.URL),
|
||||
RemoteTaskID: operationName,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("resume Gemini Veo: %v", err)
|
||||
}
|
||||
if postCount.Load() != 0 || response.Result["id"] != operationName {
|
||||
t.Fatalf("resume submitted a new operation or returned wrong ID: posts=%d result=%+v", postCount.Load(), response.Result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiVeoBodyMapsImageAndParameters(t *testing.T) {
|
||||
first := "data:image/png;base64," + base64.StdEncoding.EncodeToString([]byte("first"))
|
||||
last := "data:image/jpeg;base64," + base64.StdEncoding.EncodeToString([]byte("last"))
|
||||
reference := "data:image/webp;base64," + base64.StdEncoding.EncodeToString([]byte("reference"))
|
||||
body, err := geminiVeoBody(context.Background(), Request{Body: map[string]any{
|
||||
"prompt": "animate",
|
||||
"first_frame": first,
|
||||
"last_frame": last,
|
||||
"reference_images": []any{reference},
|
||||
"duration": 8,
|
||||
"resolution": "2160p",
|
||||
"aspect_ratio": "9:16",
|
||||
"negative_prompt": "text",
|
||||
"person_generation": "allow_adult",
|
||||
"enhance_prompt": false,
|
||||
}})
|
||||
if err != nil {
|
||||
t.Fatalf("build Gemini Veo body: %v", err)
|
||||
}
|
||||
instance := body["instances"].([]any)[0].(map[string]any)
|
||||
image := instance["image"].(map[string]any)
|
||||
lastFrame := instance["lastFrame"].(map[string]any)
|
||||
references := instance["referenceImages"].([]any)
|
||||
if image["bytesBase64Encoded"] != base64.StdEncoding.EncodeToString([]byte("first")) || image["mimeType"] != "image/png" {
|
||||
t.Fatalf("unexpected first frame: %+v", image)
|
||||
}
|
||||
if lastFrame["mimeType"] != "image/jpeg" || len(references) != 1 {
|
||||
t.Fatalf("unexpected last/reference images: last=%+v references=%+v", lastFrame, references)
|
||||
}
|
||||
parameters := body["parameters"].(map[string]any)
|
||||
if parameters["resolution"] != "4k" || parameters["durationSeconds"] != 8 || parameters["enhancePrompt"] != false {
|
||||
t.Fatalf("unexpected Gemini Veo parameters: %+v", parameters)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiVeoDownloadsImageURLWhenBuildingOfficialPayload(t *testing.T) {
|
||||
payload := []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n', 1, 2, 3, 4}
|
||||
fetchCount := 0
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
fetchCount++
|
||||
w.Header().Set("Content-Type", "image/png")
|
||||
_, _ = w.Write(payload)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
body, err := geminiVeoBody(context.Background(), Request{Body: map[string]any{
|
||||
"prompt": "animate",
|
||||
"first_frame": server.URL + "/input.png",
|
||||
"duration": 8,
|
||||
}})
|
||||
if err != nil {
|
||||
t.Fatalf("build Gemini Veo body from URL: %v", err)
|
||||
}
|
||||
instance := body["instances"].([]any)[0].(map[string]any)
|
||||
image := instance["image"].(map[string]any)
|
||||
if fetchCount != 1 || image["bytesBase64Encoded"] != base64.StdEncoding.EncodeToString(payload) || image["mimeType"] != "image/png" {
|
||||
t.Fatalf("Veo URL should be fetched during official payload construction: fetches=%d image=%+v", fetchCount, image)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiOfficialAPIClassification(t *testing.T) {
|
||||
tests := []struct {
|
||||
baseURL string
|
||||
want bool
|
||||
}{
|
||||
{baseURL: "", want: true},
|
||||
{baseURL: "https://generativelanguage.googleapis.com", want: true},
|
||||
{baseURL: "https://generativelanguage.googleapis.com/v1beta", want: true},
|
||||
{baseURL: "https://generativelanguage.googleapis.com/v1beta/openai", want: false},
|
||||
{baseURL: "https://gemini-compatible.example.com/v1beta", want: false},
|
||||
}
|
||||
for _, test := range tests {
|
||||
if got := GeminiUsesOfficialAPI(store.RuntimeModelCandidate{BaseURL: test.baseURL}); got != test.want {
|
||||
t.Fatalf("GeminiUsesOfficialAPI(%q) = %t, want %t", test.baseURL, got, test.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiCompatibleConvertsFileDataURLToInlineBase64(t *testing.T) {
|
||||
payload := []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n', 1, 2, 3, 4}
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "image/png")
|
||||
_, _ = w.Write(payload)
|
||||
}))
|
||||
defer server.Close()
|
||||
body := map[string]any{"contents": []any{map[string]any{"parts": []any{map[string]any{
|
||||
"fileData": map[string]any{"fileUri": server.URL + "/source.png", "mimeType": "image/png"},
|
||||
}}}}}
|
||||
|
||||
prepared, err := prepareCompatibleGeminiInlineImages(context.Background(), body)
|
||||
if err != nil {
|
||||
t.Fatalf("prepare compatible Gemini inline image: %v", err)
|
||||
}
|
||||
contents := prepared["contents"].([]any)
|
||||
parts := contents[0].(map[string]any)["parts"].([]any)
|
||||
part := parts[0].(map[string]any)
|
||||
inline := part["inlineData"].(map[string]any)
|
||||
if _, exists := part["fileData"]; exists || inline["mimeType"] != "image/png" || inline["data"] != base64.StdEncoding.EncodeToString(payload) {
|
||||
t.Fatalf("compatible Gemini should use inlineData: %+v", part)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiOfficialImageUploadUsesFilesAPI(t *testing.T) {
|
||||
imagePayload := []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n', 1, 2, 3, 4}
|
||||
imageServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "image/png")
|
||||
_, _ = w.Write(imagePayload)
|
||||
}))
|
||||
defer imageServer.Close()
|
||||
|
||||
var uploadServer *httptest.Server
|
||||
startCount := 0
|
||||
uploadCount := 0
|
||||
uploadServer = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/upload/v1beta/files":
|
||||
startCount++
|
||||
if r.Header.Get("x-goog-api-key") != "gemini-key" || r.Header.Get("X-Goog-Upload-Protocol") != "resumable" || r.Header.Get("X-Goog-Upload-Header-Content-Type") != "image/png" {
|
||||
t.Fatalf("unexpected Gemini Files start headers: %+v", r.Header)
|
||||
}
|
||||
w.Header().Set("X-Goog-Upload-URL", uploadServer.URL+"/upload/session")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
case "/upload/session":
|
||||
uploadCount++
|
||||
raw, err := io.ReadAll(r.Body)
|
||||
if err != nil {
|
||||
t.Fatalf("read Gemini Files upload: %v", err)
|
||||
}
|
||||
if string(raw) != string(imagePayload) || r.Header.Get("X-Goog-Upload-Command") != "upload, finalize" {
|
||||
t.Fatalf("unexpected Gemini Files upload body=%v headers=%+v", raw, r.Header)
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"file": map[string]any{
|
||||
"uri": "https://generativelanguage.googleapis.com/v1beta/files/easyai-image",
|
||||
"mimeType": "image/png",
|
||||
}})
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
defer uploadServer.Close()
|
||||
|
||||
fileData, err := uploadOfficialGeminiImageFile(context.Background(), uploadServer.Client(), store.RuntimeModelCandidate{BaseURL: uploadServer.URL + "/v1beta"}, "gemini-key", imageServer.URL+"/source.png", "image/png")
|
||||
if err != nil {
|
||||
t.Fatalf("upload official Gemini image: %v", err)
|
||||
}
|
||||
if startCount != 1 || uploadCount != 1 || fileData["fileUri"] != "https://generativelanguage.googleapis.com/v1beta/files/easyai-image" || fileData["mimeType"] != "image/png" {
|
||||
t.Fatalf("unexpected Gemini Files result: starts=%d uploads=%d fileData=%+v", startCount, uploadCount, fileData)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiVeoRejectsInvalidDurationBeforeSubmission(t *testing.T) {
|
||||
_, err := geminiVeoBody(context.Background(), Request{Body: map[string]any{"prompt": "test", "duration": 5}})
|
||||
if err == nil || ErrorCode(err) != "invalid_parameter" || ErrorParam(err) != "duration" || !strings.Contains(err.Error(), "4, 6, or 8") {
|
||||
t.Fatalf("unexpected invalid duration error: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiVeoDownloadURLOnlyTrustsProviderAndGoogleAPIs(t *testing.T) {
|
||||
for _, test := range []struct {
|
||||
uri string
|
||||
trusted bool
|
||||
}{
|
||||
{uri: "https://gemini-proxy.example.com/v1beta/files/video", trusted: true},
|
||||
{uri: "https://generativelanguage.googleapis.com/v1beta/files/video", trusted: true},
|
||||
{uri: "https://storage.googleapis.com/signed/video", trusted: false},
|
||||
{uri: "https://attacker.example.com/video", trusted: false},
|
||||
} {
|
||||
_, trusted, err := geminiVeoDownloadURL("https://gemini-proxy.example.com/v1beta", test.uri)
|
||||
if err != nil {
|
||||
t.Fatalf("parse %s: %v", test.uri, err)
|
||||
}
|
||||
if trusted != test.trusted {
|
||||
t.Fatalf("unexpected trust for %s: got %t want %t", test.uri, trusted, test.trusted)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiVeoLive(t *testing.T) {
|
||||
apiKey := strings.TrimSpace(os.Getenv("GEMINI_VEO_LIVE_API_KEY"))
|
||||
if apiKey == "" {
|
||||
t.Skip("GEMINI_VEO_LIVE_API_KEY is not configured")
|
||||
}
|
||||
baseURL := strings.TrimSpace(os.Getenv("GEMINI_VEO_LIVE_BASE_URL"))
|
||||
model := strings.TrimSpace(os.Getenv("GEMINI_VEO_LIVE_MODEL"))
|
||||
if model == "" {
|
||||
model = "veo-3.1-generate-preview"
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 12*time.Minute)
|
||||
defer cancel()
|
||||
operationName := strings.TrimSpace(os.Getenv("GEMINI_VEO_LIVE_OPERATION"))
|
||||
response, err := (GeminiClient{}).Run(ctx, Request{
|
||||
Kind: "videos.generations",
|
||||
ModelType: "video_generate",
|
||||
Model: model,
|
||||
Body: map[string]any{
|
||||
"prompt": "A small red paper boat gently floating on calm water, static camera, natural daylight.",
|
||||
"duration": 4,
|
||||
"resolution": "720p",
|
||||
"aspect_ratio": "16:9",
|
||||
},
|
||||
Candidate: store.RuntimeModelCandidate{
|
||||
Provider: "gemini",
|
||||
BaseURL: baseURL,
|
||||
ProviderModelName: model,
|
||||
Credentials: map[string]any{"apiKey": apiKey},
|
||||
PlatformConfig: map[string]any{"pollIntervalMs": 10000, "pollTimeoutMs": 660000},
|
||||
},
|
||||
RemoteTaskID: operationName,
|
||||
OnRemoteTaskSubmitted: func(remoteTaskID string, _ map[string]any) error {
|
||||
t.Logf("live Gemini Veo submitted operation %s", remoteTaskID)
|
||||
return nil
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("live Gemini Veo request failed: %v", err)
|
||||
}
|
||||
data, _ := response.Result["data"].([]any)
|
||||
if len(data) != 1 {
|
||||
t.Fatalf("live Gemini Veo returned %d videos", len(data))
|
||||
}
|
||||
item, _ := data[0].(map[string]any)
|
||||
payload, _ := item["video_bytes"].([]byte)
|
||||
if len(payload) < 1024 {
|
||||
t.Fatalf("live Gemini Veo video is unexpectedly small: %d bytes", len(payload))
|
||||
}
|
||||
if mimeType := strings.TrimSpace(stringFromAny(item["mime_type"])); !strings.HasPrefix(mimeType, "video/") {
|
||||
t.Fatalf("live Gemini Veo returned unexpected MIME type: %q", mimeType)
|
||||
}
|
||||
t.Logf("live Gemini Veo accepted operation %s and downloaded %d video bytes", response.RequestID, len(payload))
|
||||
}
|
||||
|
||||
func testGeminiVeoCandidate(baseURL string) store.RuntimeModelCandidate {
|
||||
return store.RuntimeModelCandidate{
|
||||
Provider: "gemini",
|
||||
BaseURL: baseURL,
|
||||
ProviderModelName: "veo-3.1-generate-preview",
|
||||
Credentials: map[string]any{"apiKey": "test-gemini-key"},
|
||||
PlatformConfig: map[string]any{"pollIntervalMs": 1, "pollTimeoutMs": 1000},
|
||||
}
|
||||
}
|
||||
@@ -13,6 +13,11 @@ import (
|
||||
"time"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultMaxJSONResponseBytes int64 = 16 << 20
|
||||
mediaMaxJSONResponseBytes int64 = 128 << 20
|
||||
)
|
||||
|
||||
func credential(candidate map[string]any, keys ...string) string {
|
||||
for _, key := range keys {
|
||||
if value, ok := candidate[key].(string); ok && strings.TrimSpace(value) != "" {
|
||||
@@ -44,25 +49,140 @@ func intValue(body map[string]any, key string, fallback int) int {
|
||||
}
|
||||
|
||||
func decodeHTTPResponse(resp *http.Response) (map[string]any, error) {
|
||||
result, _, err := decodeHTTPResponseForProtocol(resp, "")
|
||||
return result, err
|
||||
}
|
||||
|
||||
func decodeHTTPResponseForProtocol(resp *http.Response, protocol string) (map[string]any, *WireResponse, error) {
|
||||
defer resp.Body.Close()
|
||||
raw, _ := io.ReadAll(io.LimitReader(resp.Body, 16*1024*1024))
|
||||
maxBytes := maxJSONResponseBytesForProtocol(protocol)
|
||||
raw, readErr := io.ReadAll(io.LimitReader(resp.Body, maxBytes+1))
|
||||
tooLarge := int64(len(raw)) > maxBytes
|
||||
if tooLarge {
|
||||
raw = raw[:maxBytes]
|
||||
}
|
||||
wire := &WireResponse{
|
||||
Protocol: strings.TrimSpace(protocol),
|
||||
StatusCode: resp.StatusCode,
|
||||
Headers: compatibleResponseHeaders(resp.Header),
|
||||
// raw is already owned by this response. Reusing it avoids retaining a
|
||||
// second copy of media responses that may be as large as 128 MiB.
|
||||
RawJSON: raw,
|
||||
}
|
||||
if readErr != nil {
|
||||
if len(raw) > 0 {
|
||||
_ = json.Unmarshal(raw, &wire.Body)
|
||||
}
|
||||
if timeoutErr := transportClientError(readErr); timeoutErr.Code == "timeout" {
|
||||
timeoutErr.StatusCode = resp.StatusCode
|
||||
timeoutErr.RequestID = requestIDFromHTTPResponse(resp)
|
||||
timeoutErr.Wire = wire
|
||||
return nil, wire, timeoutErr
|
||||
}
|
||||
return nil, wire, &ClientError{
|
||||
Code: "response_read_error",
|
||||
Message: readErr.Error(),
|
||||
StatusCode: resp.StatusCode,
|
||||
RequestID: requestIDFromHTTPResponse(resp),
|
||||
Retryable: true,
|
||||
Wire: wire,
|
||||
}
|
||||
}
|
||||
if tooLarge {
|
||||
if len(raw) > 0 {
|
||||
_ = json.Unmarshal(raw, &wire.Body)
|
||||
}
|
||||
return nil, wire, &ClientError{
|
||||
Code: "response_too_large",
|
||||
Message: fmt.Sprintf("upstream JSON response exceeds %d bytes", maxBytes),
|
||||
StatusCode: resp.StatusCode,
|
||||
RequestID: requestIDFromHTTPResponse(resp),
|
||||
Retryable: false,
|
||||
Wire: wire,
|
||||
}
|
||||
}
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return nil, &ClientError{
|
||||
if len(raw) > 0 {
|
||||
_ = json.Unmarshal(raw, &wire.Body)
|
||||
}
|
||||
return nil, wire, &ClientError{
|
||||
Code: statusCodeName(resp.StatusCode),
|
||||
Message: errorMessage(raw, resp.Status),
|
||||
StatusCode: resp.StatusCode,
|
||||
RequestID: requestIDFromHTTPResponse(resp),
|
||||
Retryable: HTTPRetryable(resp.StatusCode),
|
||||
Wire: wire,
|
||||
}
|
||||
}
|
||||
var out map[string]any
|
||||
if len(raw) == 0 {
|
||||
return map[string]any{}, nil
|
||||
wire.Body = map[string]any{}
|
||||
return map[string]any{}, wire, nil
|
||||
}
|
||||
if err := json.Unmarshal(raw, &out); err != nil {
|
||||
return nil, &ClientError{Code: "invalid_response", Message: err.Error(), Retryable: false}
|
||||
return nil, wire, &ClientError{Code: "invalid_response", Message: err.Error(), StatusCode: resp.StatusCode, Retryable: false, Wire: wire}
|
||||
}
|
||||
return out, nil
|
||||
// Result and wire metadata are read-only until the runner replaces the
|
||||
// canonical result, so they can share the decoded object instead of keeping
|
||||
// another full copy of a large Base64 media response.
|
||||
wire.Body = out
|
||||
return out, wire, nil
|
||||
}
|
||||
|
||||
func maxJSONResponseBytesForProtocol(protocol string) int64 {
|
||||
switch strings.TrimSpace(protocol) {
|
||||
case ProtocolGeminiGenerateContent, ProtocolOpenAIImages:
|
||||
return mediaMaxJSONResponseBytes
|
||||
default:
|
||||
return defaultMaxJSONResponseBytes
|
||||
}
|
||||
}
|
||||
|
||||
func compatibleResponseHeaders(source http.Header) map[string][]string {
|
||||
if len(source) == 0 {
|
||||
return nil
|
||||
}
|
||||
allowed := map[string]bool{
|
||||
"content-type": true, "retry-after": true,
|
||||
"x-request-id": true, "request-id": true,
|
||||
"x-ratelimit-limit-requests": true, "x-ratelimit-remaining-requests": true,
|
||||
"x-ratelimit-reset-requests": true, "x-ratelimit-limit-tokens": true,
|
||||
"x-ratelimit-remaining-tokens": true, "x-ratelimit-reset-tokens": true,
|
||||
"x-goog-request-id": true, "x-goog-upload-status": true,
|
||||
"x-goog-upload-url": true, "x-goog-upload-size-received": true,
|
||||
}
|
||||
out := map[string][]string{}
|
||||
for key, values := range source {
|
||||
if !allowed[strings.ToLower(strings.TrimSpace(key))] {
|
||||
continue
|
||||
}
|
||||
out[http.CanonicalHeaderKey(key)] = append([]string(nil), values...)
|
||||
}
|
||||
if len(out) == 0 {
|
||||
return nil
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func notifySubmissionStarted(request Request) error {
|
||||
if request.OnUpstreamSubmissionStarted == nil {
|
||||
return nil
|
||||
}
|
||||
return request.OnUpstreamSubmissionStarted()
|
||||
}
|
||||
|
||||
func notifyResponseReceived(request Request) error {
|
||||
if request.OnUpstreamResponseReceived == nil {
|
||||
return nil
|
||||
}
|
||||
return request.OnUpstreamResponseReceived()
|
||||
}
|
||||
|
||||
func notifyWireResponse(request Request, wire *WireResponse) error {
|
||||
if wire == nil || request.OnUpstreamWireResponse == nil {
|
||||
return nil
|
||||
}
|
||||
return request.OnUpstreamWireResponse(wire)
|
||||
}
|
||||
|
||||
func decodeOpenAIStreamResponse(resp *http.Response, onDelta StreamDelta) (map[string]any, error) {
|
||||
@@ -88,6 +208,9 @@ func decodeOpenAIStreamReader(reader io.Reader, onDelta StreamDelta) (map[string
|
||||
scanner.Buffer(make([]byte, 0, 64*1024), 16*1024*1024)
|
||||
rawLines := make([]string, 0)
|
||||
parts := make([]string, 0)
|
||||
refusalParts := make([]string, 0)
|
||||
streamLogprobs := make([]any, 0)
|
||||
streamAnnotations := make([]any, 0)
|
||||
reasoningParts := make([]string, 0)
|
||||
var last map[string]any
|
||||
var usage Usage
|
||||
@@ -118,6 +241,11 @@ func decodeOpenAIStreamReader(reader io.Reader, onDelta StreamDelta) (map[string
|
||||
if reasoningText != "" {
|
||||
reasoningParts = append(reasoningParts, reasoningText)
|
||||
}
|
||||
if refusal := streamEventRefusal(event); refusal != "" {
|
||||
refusalParts = append(refusalParts, refusal)
|
||||
}
|
||||
streamLogprobs = append(streamLogprobs, streamEventLogprobs(event)...)
|
||||
streamAnnotations = append(streamAnnotations, streamEventAnnotations(event)...)
|
||||
aggregateStreamToolCalls(event, toolCalls)
|
||||
if reason := streamEventFinishReason(event); reason != "" {
|
||||
finishReason = reason
|
||||
@@ -145,7 +273,7 @@ func decodeOpenAIStreamReader(reader io.Reader, onDelta StreamDelta) (map[string
|
||||
}
|
||||
return out, true, nil
|
||||
}
|
||||
return buildOpenAIStreamResult(last, parts, reasoningParts, toolCalls, finishReason, usage), true, nil
|
||||
return buildOpenAIStreamResult(last, parts, refusalParts, streamLogprobs, streamAnnotations, reasoningParts, toolCalls, finishReason, usage), true, nil
|
||||
}
|
||||
|
||||
func decodeOpenAIStream(raw []byte) (map[string]any, bool) {
|
||||
@@ -156,8 +284,8 @@ func decodeOpenAIStream(raw []byte) (map[string]any, bool) {
|
||||
return result, ok && err == nil
|
||||
}
|
||||
|
||||
func buildOpenAIStreamResult(last map[string]any, parts []string, reasoningParts []string, toolCalls map[int]map[string]any, finishReason string, usage Usage) map[string]any {
|
||||
if len(parts) == 0 && len(reasoningParts) == 0 && len(toolCalls) == 0 {
|
||||
func buildOpenAIStreamResult(last map[string]any, parts []string, refusalParts []string, streamLogprobs []any, streamAnnotations []any, reasoningParts []string, toolCalls map[int]map[string]any, finishReason string, usage Usage) map[string]any {
|
||||
if len(parts) == 0 && len(refusalParts) == 0 && len(streamLogprobs) == 0 && len(streamAnnotations) == 0 && len(reasoningParts) == 0 && len(toolCalls) == 0 {
|
||||
return last
|
||||
}
|
||||
message := map[string]any{
|
||||
@@ -167,6 +295,12 @@ func buildOpenAIStreamResult(last map[string]any, parts []string, reasoningParts
|
||||
if len(reasoningParts) > 0 {
|
||||
message["reasoning_content"] = strings.Join(reasoningParts, "")
|
||||
}
|
||||
if len(refusalParts) > 0 {
|
||||
message["refusal"] = strings.Join(refusalParts, "")
|
||||
}
|
||||
if len(streamAnnotations) > 0 {
|
||||
message["annotations"] = streamAnnotations
|
||||
}
|
||||
if len(toolCalls) > 0 {
|
||||
message["tool_calls"] = sortedStreamToolCalls(toolCalls)
|
||||
}
|
||||
@@ -174,15 +308,17 @@ func buildOpenAIStreamResult(last map[string]any, parts []string, reasoningParts
|
||||
finishReason = "stop"
|
||||
}
|
||||
var out map[string]any
|
||||
choice := map[string]any{
|
||||
"index": 0, "message": message, "finish_reason": finishReason,
|
||||
}
|
||||
if len(streamLogprobs) > 0 {
|
||||
choice["logprobs"] = map[string]any{"content": streamLogprobs, "refusal": nil}
|
||||
}
|
||||
out = map[string]any{
|
||||
"id": stringFromAny(firstPresent(last["id"], "chatcmpl-stream")),
|
||||
"object": "chat.completion",
|
||||
"model": stringFromAny(last["model"]),
|
||||
"choices": []any{map[string]any{
|
||||
"index": 0,
|
||||
"message": message,
|
||||
"finish_reason": finishReason,
|
||||
}},
|
||||
"id": stringFromAny(firstPresent(last["id"], "chatcmpl-stream")),
|
||||
"object": "chat.completion",
|
||||
"model": stringFromAny(last["model"]),
|
||||
"choices": []any{choice},
|
||||
}
|
||||
if usage.TotalTokens > 0 {
|
||||
usageMap := map[string]any{
|
||||
@@ -200,6 +336,42 @@ func buildOpenAIStreamResult(last map[string]any, parts []string, reasoningParts
|
||||
return out
|
||||
}
|
||||
|
||||
func streamEventRefusal(event map[string]any) string {
|
||||
choices, _ := event["choices"].([]any)
|
||||
for _, rawChoice := range choices {
|
||||
choice, _ := rawChoice.(map[string]any)
|
||||
delta, _ := choice["delta"].(map[string]any)
|
||||
if refusal := stringFromAny(delta["refusal"]); refusal != "" {
|
||||
return refusal
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func streamEventLogprobs(event map[string]any) []any {
|
||||
choices, _ := event["choices"].([]any)
|
||||
out := make([]any, 0)
|
||||
for _, rawChoice := range choices {
|
||||
choice, _ := rawChoice.(map[string]any)
|
||||
logprobs, _ := choice["logprobs"].(map[string]any)
|
||||
content, _ := logprobs["content"].([]any)
|
||||
out = append(out, content...)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func streamEventAnnotations(event map[string]any) []any {
|
||||
choices, _ := event["choices"].([]any)
|
||||
out := make([]any, 0)
|
||||
for _, rawChoice := range choices {
|
||||
choice, _ := rawChoice.(map[string]any)
|
||||
delta, _ := choice["delta"].(map[string]any)
|
||||
annotations, _ := delta["annotations"].([]any)
|
||||
out = append(out, annotations...)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// NormalizeChatCompletionRequestBody 将后续请求里的工具调用上下文还原为
|
||||
// OpenAI Chat Completions 标准格式,便于再次发送给 OpenAI-compatible 上游。
|
||||
func NormalizeChatCompletionRequestBody(body map[string]any) map[string]any {
|
||||
@@ -219,7 +391,7 @@ func NormalizeChatCompletionRequestBody(body map[string]any) map[string]any {
|
||||
continue
|
||||
}
|
||||
copied := cloneMapAny(message)
|
||||
normalizeToolCallsContainer(copied, false)
|
||||
normalizeRequestToolCallsContainer(copied)
|
||||
normalizeToolMessageFields(copied)
|
||||
toolMessages, cleanContent, changed := toolResultMessagesFromContent(copied["content"])
|
||||
if changed {
|
||||
@@ -241,6 +413,67 @@ func NormalizeChatCompletionRequestBody(body map[string]any) map[string]any {
|
||||
return out
|
||||
}
|
||||
|
||||
// normalizeRequestToolCallsContainer preserves canonical OpenAI Chat
|
||||
// structures. Only known provider aliases are normalized. Legacy
|
||||
// function_call and standard custom tool calls must remain intact.
|
||||
func normalizeRequestToolCallsContainer(container map[string]any) {
|
||||
if container == nil {
|
||||
return
|
||||
}
|
||||
toolCalls := canonicalToolCalls(container["tool_calls"], false)
|
||||
for _, key := range []string{"tool_call", "toolCall", "toolCalls"} {
|
||||
if raw, ok := container[key]; ok {
|
||||
for _, normalized := range normalizeRawToolCalls(raw, len(toolCalls), false) {
|
||||
toolCalls = append(toolCalls, normalized)
|
||||
}
|
||||
delete(container, key)
|
||||
}
|
||||
}
|
||||
if contentToolCalls, cleanContent, changed := toolCallsFromContent(container["content"], len(toolCalls), false); changed {
|
||||
toolCalls = append(toolCalls, contentToolCalls...)
|
||||
setNormalizedContent(container, cleanContent, false)
|
||||
}
|
||||
if partToolCalls := toolCallsFromParts(container["parts"], len(toolCalls), false); len(partToolCalls) > 0 {
|
||||
toolCalls = append(toolCalls, partToolCalls...)
|
||||
delete(container, "parts")
|
||||
}
|
||||
if len(toolCalls) > 0 {
|
||||
container["tool_calls"] = toolCalls
|
||||
}
|
||||
if functionCall, ok := container["functionCall"]; ok {
|
||||
if _, canonical := container["function_call"]; !canonical {
|
||||
if normalized := normalizeToolCall(functionCall, 0, false); normalized != nil {
|
||||
container["function_call"] = normalized["function"]
|
||||
}
|
||||
}
|
||||
delete(container, "functionCall")
|
||||
}
|
||||
}
|
||||
|
||||
func canonicalToolCalls(value any, stream bool) []any {
|
||||
items, ok := value.([]any)
|
||||
if !ok {
|
||||
if value == nil {
|
||||
return nil
|
||||
}
|
||||
items = []any{value}
|
||||
}
|
||||
out := make([]any, 0, len(items))
|
||||
for index, raw := range items {
|
||||
toolCall, _ := raw.(map[string]any)
|
||||
typeName := stringFromAny(toolCall["type"])
|
||||
if (typeName == "custom" && len(mapFromAny(toolCall["custom"])) > 0) ||
|
||||
((typeName == "" || typeName == "function") && len(mapFromAny(toolCall["function"])) > 0) {
|
||||
out = append(out, cloneMapAny(toolCall))
|
||||
continue
|
||||
}
|
||||
if normalized := normalizeToolCall(raw, index, stream); normalized != nil {
|
||||
out = append(out, normalized)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func cloneMapAny(source map[string]any) map[string]any {
|
||||
if source == nil {
|
||||
return nil
|
||||
@@ -332,14 +565,27 @@ func normalizeToolCallsContainer(container map[string]any, stream bool) {
|
||||
if container == nil {
|
||||
return
|
||||
}
|
||||
toolCalls := make([]any, 0)
|
||||
for _, rawToolCall := range rawToolCallValues(container) {
|
||||
for _, normalized := range normalizeRawToolCalls(rawToolCall, len(toolCalls), stream) {
|
||||
toolCalls = append(toolCalls, normalized)
|
||||
toolCalls := canonicalToolCalls(container["tool_calls"], stream)
|
||||
for _, key := range []string{"tool_call", "toolCall", "toolCalls"} {
|
||||
if rawToolCall, ok := container[key]; ok {
|
||||
for _, normalized := range normalizeRawToolCalls(rawToolCall, len(toolCalls), stream) {
|
||||
toolCalls = append(toolCalls, normalized)
|
||||
}
|
||||
delete(container, key)
|
||||
}
|
||||
}
|
||||
if functionCall, ok := container["functionCall"]; ok {
|
||||
if _, canonical := container["function_call"]; !canonical {
|
||||
if normalized := normalizeToolCall(functionCall, 0, stream); normalized != nil {
|
||||
container["function_call"] = normalized["function"]
|
||||
}
|
||||
}
|
||||
delete(container, "functionCall")
|
||||
}
|
||||
if contentToolCalls, cleanContent, changed := toolCallsFromContent(container["content"], len(toolCalls), stream); changed {
|
||||
toolCalls = append(toolCalls, contentToolCalls...)
|
||||
for _, normalized := range contentToolCalls {
|
||||
toolCalls = append(toolCalls, normalized)
|
||||
}
|
||||
setNormalizedContent(container, cleanContent, stream)
|
||||
}
|
||||
if partToolCalls := toolCallsFromParts(container["parts"], len(toolCalls), stream); len(partToolCalls) > 0 {
|
||||
@@ -349,9 +595,6 @@ func normalizeToolCallsContainer(container map[string]any, stream bool) {
|
||||
if len(toolCalls) > 0 {
|
||||
container["tool_calls"] = toolCalls
|
||||
}
|
||||
for _, key := range []string{"tool_call", "toolCall", "toolCalls", "function_call", "functionCall"} {
|
||||
delete(container, key)
|
||||
}
|
||||
}
|
||||
|
||||
func normalizeToolMessageFields(message map[string]any) {
|
||||
@@ -443,6 +686,20 @@ func normalizeToolCall(value any, index int, stream bool) map[string]any {
|
||||
if len(source) == 0 {
|
||||
return nil
|
||||
}
|
||||
if stringFromAny(source["type"]) == "custom" {
|
||||
customSource := mapFromAny(source["custom"])
|
||||
if len(customSource) == 0 {
|
||||
return nil
|
||||
}
|
||||
toolCall := cloneMapAny(source)
|
||||
toolCall["custom"] = cloneMapAny(customSource)
|
||||
if stream {
|
||||
if _, ok := toolCall["index"]; !ok {
|
||||
toolCall["index"] = index
|
||||
}
|
||||
}
|
||||
return toolCall
|
||||
}
|
||||
functionSource := mapFromAny(source["function"])
|
||||
if len(functionSource) == 0 {
|
||||
functionSource = mapFromAny(firstPresent(source["function_call"], source["functionCall"]))
|
||||
@@ -862,7 +1119,7 @@ func aggregateStreamToolCalls(event map[string]any, toolCalls map[int]map[string
|
||||
for _, rawChoice := range choices {
|
||||
choice, _ := rawChoice.(map[string]any)
|
||||
delta, _ := choice["delta"].(map[string]any)
|
||||
rawToolCalls, _ := delta["tool_calls"].([]any)
|
||||
rawToolCalls := streamToolCallsFromDelta(delta)
|
||||
for _, rawToolCall := range rawToolCalls {
|
||||
incoming, _ := rawToolCall.(map[string]any)
|
||||
index := intFromAny(incoming["index"])
|
||||
@@ -876,25 +1133,39 @@ func aggregateStreamToolCalls(event map[string]any, toolCalls map[int]map[string
|
||||
current[key] = value
|
||||
}
|
||||
}
|
||||
incomingFn, _ := incoming["function"].(map[string]any)
|
||||
if len(incomingFn) == 0 {
|
||||
continue
|
||||
}
|
||||
currentFn, _ := current["function"].(map[string]any)
|
||||
if currentFn == nil {
|
||||
currentFn = map[string]any{}
|
||||
current["function"] = currentFn
|
||||
}
|
||||
if name, ok := incomingFn["name"].(string); ok && name != "" {
|
||||
currentFn["name"] = stringFromAny(currentFn["name"]) + name
|
||||
}
|
||||
if arguments, ok := incomingFn["arguments"].(string); ok && arguments != "" {
|
||||
currentFn["arguments"] = stringFromAny(currentFn["arguments"]) + arguments
|
||||
}
|
||||
aggregateStreamToolPayload(current, incoming, "function", "arguments")
|
||||
aggregateStreamToolPayload(current, incoming, "custom", "input")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func streamToolCallsFromDelta(delta map[string]any) []any {
|
||||
rawToolCalls, _ := delta["tool_calls"].([]any)
|
||||
out := append([]any(nil), rawToolCalls...)
|
||||
if functionCall, ok := delta["function_call"].(map[string]any); ok {
|
||||
out = append(out, map[string]any{"index": 0, "type": "function", "function": functionCall})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func aggregateStreamToolPayload(current map[string]any, incoming map[string]any, containerKey string, payloadKey string) {
|
||||
incomingPayload, _ := incoming[containerKey].(map[string]any)
|
||||
if len(incomingPayload) == 0 {
|
||||
return
|
||||
}
|
||||
currentPayload, _ := current[containerKey].(map[string]any)
|
||||
if currentPayload == nil {
|
||||
currentPayload = map[string]any{}
|
||||
current[containerKey] = currentPayload
|
||||
}
|
||||
if name, ok := incomingPayload["name"].(string); ok && name != "" {
|
||||
currentPayload["name"] = stringFromAny(currentPayload["name"]) + name
|
||||
}
|
||||
if payload, ok := incomingPayload[payloadKey].(string); ok && payload != "" {
|
||||
currentPayload[payloadKey] = stringFromAny(currentPayload[payloadKey]) + payload
|
||||
}
|
||||
}
|
||||
|
||||
func sortedStreamToolCalls(toolCalls map[int]map[string]any) []any {
|
||||
indices := make([]int, 0, len(toolCalls))
|
||||
for index := range toolCalls {
|
||||
|
||||
@@ -80,7 +80,18 @@ func (c KelingClient) runVideo(ctx context.Context, request Request, token strin
|
||||
}()
|
||||
|
||||
if upstreamTaskID == "" {
|
||||
submitResult, requestID, err := c.postJSON(ctx, request, prepared.Endpoint, token, prepared.Payload)
|
||||
if err := notifySubmissionStarted(request); err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
submitResult, requestID, submitWire, err := c.postJSON(ctx, request, prepared.Endpoint, token, prepared.Payload)
|
||||
if notifyErr := notifyWireResponse(request, submitWire); notifyErr != nil {
|
||||
return Response{}, notifyErr
|
||||
}
|
||||
if err == nil || ErrorResponseMetadata(err).StatusCode > 0 {
|
||||
if notifyErr := notifyResponseReceived(request); notifyErr != nil {
|
||||
return Response{}, notifyErr
|
||||
}
|
||||
}
|
||||
submitRequestID = requestID
|
||||
if err != nil {
|
||||
return Response{}, annotateResponseError(err, submitRequestID, submitStartedAt, time.Now())
|
||||
@@ -122,6 +133,11 @@ func (c KelingClient) runVideo(ctx context.Context, request Request, token strin
|
||||
return Response{}, annotateResponseError(err, requestID, pollStartedAt, pollFinishedAt)
|
||||
}
|
||||
lastResult = pollResult
|
||||
if request.OnRemoteTaskPolled != nil {
|
||||
if err := request.OnRemoteTaskPolled(upstreamTaskID, pollResult); err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
}
|
||||
|
||||
switch kelingTaskStatus(pollResult) {
|
||||
case "succeed":
|
||||
@@ -176,7 +192,18 @@ func (c KelingClient) runTaskAPIVideo(ctx context.Context, request Request, toke
|
||||
if err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
submitResult, requestID, err := c.postJSONAt(ctx, request, taskAPIBaseURL, endpoint, token, payload)
|
||||
if err := notifySubmissionStarted(request); err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
submitResult, requestID, submitWire, err := c.postJSONAt(ctx, request, taskAPIBaseURL, endpoint, token, payload, ProtocolKlingV2Omni)
|
||||
if notifyErr := notifyWireResponse(request, submitWire); notifyErr != nil {
|
||||
return Response{}, notifyErr
|
||||
}
|
||||
if err == nil || ErrorResponseMetadata(err).StatusCode > 0 {
|
||||
if notifyErr := notifyResponseReceived(request); notifyErr != nil {
|
||||
return Response{}, notifyErr
|
||||
}
|
||||
}
|
||||
submitRequestID = requestID
|
||||
if err != nil {
|
||||
return Response{}, annotateResponseError(err, submitRequestID, submitStartedAt, time.Now())
|
||||
@@ -224,6 +251,11 @@ func (c KelingClient) runTaskAPIVideo(ctx context.Context, request Request, toke
|
||||
if err != nil {
|
||||
return Response{}, annotateResponseError(err, requestID, pollStartedAt, pollFinishedAt)
|
||||
}
|
||||
if request.OnRemoteTaskPolled != nil {
|
||||
if err := request.OnRemoteTaskPolled(upstreamTaskID, pollResult); err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
}
|
||||
|
||||
task := kelingTaskAPITask(pollResult, upstreamTaskID)
|
||||
lastStatus = strings.ToLower(strings.TrimSpace(stringFromAny(task["status"])))
|
||||
@@ -567,31 +599,31 @@ func (c KelingClient) kelingOmniElementList(ctx context.Context, request Request
|
||||
return elements, createdIDs, nil
|
||||
}
|
||||
|
||||
func (c KelingClient) postJSON(ctx context.Context, request Request, path string, token string, body map[string]any) (map[string]any, string, error) {
|
||||
return c.postJSONAt(ctx, request, request.Candidate.BaseURL, path, token, body)
|
||||
func (c KelingClient) postJSON(ctx context.Context, request Request, path string, token string, body map[string]any) (map[string]any, string, *WireResponse, error) {
|
||||
return c.postJSONAt(ctx, request, request.Candidate.BaseURL, path, token, body, ProtocolKlingV1Omni)
|
||||
}
|
||||
|
||||
func (c KelingClient) postJSONAt(ctx context.Context, request Request, baseURL string, path string, token string, body map[string]any) (map[string]any, string, error) {
|
||||
func (c KelingClient) postJSONAt(ctx context.Context, request Request, baseURL string, path string, token string, body map[string]any, protocol string) (map[string]any, string, *WireResponse, error) {
|
||||
raw, _ := json.Marshal(body)
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, joinURL(baseURL, path), bytes.NewReader(raw))
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
return nil, "", nil, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
resp, err := httpClient(request.HTTPClient, c.HTTPClient).Do(req)
|
||||
if err != nil {
|
||||
return nil, "", &ClientError{Code: "network", Message: err.Error(), Retryable: true}
|
||||
return nil, "", nil, transportClientError(err)
|
||||
}
|
||||
requestID := requestIDFromHTTPResponse(resp)
|
||||
result, err := decodeHTTPResponse(resp)
|
||||
result, wire, err := decodeHTTPResponseForProtocol(resp, protocol)
|
||||
if err != nil {
|
||||
return result, requestID, err
|
||||
return result, requestID, wire, err
|
||||
}
|
||||
if code := intFromAny(result["code"]); code != 0 {
|
||||
return result, requestID, &ClientError{Code: kelingEnvelopeErrorCode(result), Message: kelingEnvelopeErrorMessage(result), RequestID: firstNonEmpty(requestID, stringFromAny(result["request_id"])), Retryable: false}
|
||||
return result, requestID, wire, &ClientError{Code: kelingEnvelopeErrorCode(result), Message: kelingEnvelopeErrorMessage(result), RequestID: firstNonEmpty(requestID, stringFromAny(result["request_id"])), Retryable: false, Wire: wire}
|
||||
}
|
||||
return result, firstNonEmpty(requestID, stringFromAny(result["request_id"])), nil
|
||||
return result, firstNonEmpty(requestID, stringFromAny(result["request_id"])), wire, nil
|
||||
}
|
||||
|
||||
func (c KelingClient) getJSON(ctx context.Context, request Request, path string, token string) (map[string]any, string, error) {
|
||||
@@ -606,7 +638,7 @@ func (c KelingClient) getJSONAt(ctx context.Context, request Request, baseURL st
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
resp, err := httpClient(request.HTTPClient, c.HTTPClient).Do(req)
|
||||
if err != nil {
|
||||
return nil, "", &ClientError{Code: "network", Message: err.Error(), Retryable: true}
|
||||
return nil, "", transportClientError(err)
|
||||
}
|
||||
requestID := requestIDFromHTTPResponse(resp)
|
||||
result, err := decodeHTTPResponse(resp)
|
||||
@@ -629,7 +661,7 @@ func (c KelingClient) createKelingElement(ctx context.Context, request Request,
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
resp, err := httpClient(request.HTTPClient, c.HTTPClient).Do(req)
|
||||
if err != nil {
|
||||
return "", &ClientError{Code: "network", Message: err.Error(), Retryable: true}
|
||||
return "", transportClientError(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
body, _ := io.ReadAll(io.LimitReader(resp.Body, 16*1024*1024))
|
||||
@@ -663,7 +695,7 @@ func (c KelingClient) cleanupKelingElements(ctx context.Context, request Request
|
||||
if id == "" {
|
||||
continue
|
||||
}
|
||||
_, _, _ = c.postJSON(ctx, request, "/general/delete-elements", token, map[string]any{"element_id": id})
|
||||
_, _, _, _ = c.postJSON(ctx, request, "/general/delete-elements", token, map[string]any{"element_id": id})
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -713,7 +745,7 @@ func kelingImageToBase64(ctx context.Context, request Request, value string) (st
|
||||
}
|
||||
resp, err := httpClient(request.HTTPClient).Do(req)
|
||||
if err != nil {
|
||||
return "", &ClientError{Code: "network", Message: err.Error(), Retryable: true}
|
||||
return "", transportClientError(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
@@ -722,7 +754,7 @@ func kelingImageToBase64(ctx context.Context, request Request, value string) (st
|
||||
}
|
||||
raw, err := io.ReadAll(io.LimitReader(resp.Body, 16*1024*1024))
|
||||
if err != nil {
|
||||
return "", &ClientError{Code: "network", Message: err.Error(), Retryable: true}
|
||||
return "", transportClientError(err)
|
||||
}
|
||||
return base64.StdEncoding.EncodeToString(raw), nil
|
||||
}
|
||||
@@ -1274,7 +1306,6 @@ func kelingVideoSuccessResult(request Request, upstreamTaskID string, raw map[st
|
||||
"status": "succeeded",
|
||||
"upstream_task_id": upstreamTaskID,
|
||||
"data": items,
|
||||
"raw": raw,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1310,7 +1341,6 @@ func kelingTaskAPIVideoSuccessResult(request Request, upstreamTaskID string, tas
|
||||
"status": "succeeded",
|
||||
"upstream_task_id": upstreamTaskID,
|
||||
"data": items,
|
||||
"raw": raw,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1343,9 +1373,10 @@ func kelingPollInterval(request Request) time.Duration {
|
||||
}
|
||||
|
||||
func kelingPollTimeout(request Request) time.Duration {
|
||||
seconds := numericValue(firstPresent(request.Candidate.PlatformConfig["kelingPollTimeoutSeconds"], request.Candidate.PlatformConfig["klingPollTimeoutSeconds"], request.Body["pollTimeoutSeconds"], request.Body["poll_timeout_seconds"]), 600)
|
||||
fallbackSeconds := ProviderRequestTimeout(request.Kind).Seconds()
|
||||
seconds := numericValue(firstPresent(request.Candidate.PlatformConfig["kelingPollTimeoutSeconds"], request.Candidate.PlatformConfig["klingPollTimeoutSeconds"], request.Body["pollTimeoutSeconds"], request.Body["poll_timeout_seconds"]), fallbackSeconds)
|
||||
if seconds < 1 {
|
||||
seconds = 600
|
||||
seconds = fallbackSeconds
|
||||
}
|
||||
return time.Duration(seconds) * time.Second
|
||||
}
|
||||
|
||||
@@ -326,17 +326,17 @@ func (c MinimaxClient) runSpeech(ctx context.Context, request Request) (Response
|
||||
if err != nil {
|
||||
return Response{}, &ClientError{Code: "invalid_response", Message: "minimax speech audio hex is invalid: " + err.Error(), RequestID: firstNonEmptyString(requestID, requestIDFromResult(result)), ResponseStartedAt: startedAt, ResponseFinishedAt: finishedAt, ResponseDurationMS: responseDurationMS(startedAt, finishedAt), Retryable: false}
|
||||
}
|
||||
normalized := cloneMapAny(result)
|
||||
normalized["status"] = "success"
|
||||
normalized["created"] = time.Now().UnixMilli()
|
||||
normalized["model"] = request.Model
|
||||
normalized["raw_data"] = cloneMapAny(result)
|
||||
normalized["data"] = []any{map[string]any{
|
||||
"type": "audio",
|
||||
"content": "data:audio/mpeg;base64," + base64.StdEncoding.EncodeToString(audioBytes),
|
||||
"mime_type": "audio/mpeg",
|
||||
"uploaded": false,
|
||||
}}
|
||||
normalized := map[string]any{
|
||||
"status": "success",
|
||||
"created": time.Now().UnixMilli(),
|
||||
"model": request.Model,
|
||||
"data": []any{map[string]any{
|
||||
"type": "audio",
|
||||
"content": "data:audio/mpeg;base64," + base64.StdEncoding.EncodeToString(audioBytes),
|
||||
"mime_type": "audio/mpeg",
|
||||
"uploaded": false,
|
||||
}},
|
||||
}
|
||||
return Response{
|
||||
Result: normalized,
|
||||
RequestID: firstNonEmptyString(requestID, requestIDFromResult(result)),
|
||||
@@ -378,12 +378,12 @@ func (c MinimaxClient) runVoiceClone(ctx context.Context, request Request) (Resp
|
||||
if isProviderTaskFailure(providerTaskSpec{Name: "minimax"}, result) {
|
||||
return Response{}, providerTaskFailure(providerTaskSpec{Name: "minimax"}, result, firstNonEmptyString(requestID, uploadRequestID, requestIDFromResult(result)), startedAt)
|
||||
}
|
||||
normalized := cloneMapAny(result)
|
||||
normalized["status"] = "success"
|
||||
normalized["created"] = time.Now().UnixMilli()
|
||||
normalized["model"] = request.Model
|
||||
normalized["voice_id"] = stringFromAny(payload["voice_id"])
|
||||
normalized["raw_data"] = cloneMapAny(result)
|
||||
normalized := map[string]any{
|
||||
"status": "success",
|
||||
"created": time.Now().UnixMilli(),
|
||||
"model": request.Model,
|
||||
"voice_id": stringFromAny(payload["voice_id"]),
|
||||
}
|
||||
if demoAudio := firstNonEmptyString(valueAtPath(result, "demo_audio"), valueAtPath(result, "data.demo_audio")); demoAudio != "" {
|
||||
normalized["demo_audio"] = demoAudio
|
||||
normalized["data"] = []any{map[string]any{"type": "audio", "url": demoAudio}}
|
||||
@@ -646,7 +646,7 @@ func providerPostMultipartFile(ctx context.Context, client *http.Client, url str
|
||||
applyProviderAuth(req, credentials, auth)
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, "", &ClientError{Code: "network", Message: err.Error(), Retryable: true}
|
||||
return nil, "", transportClientError(err)
|
||||
}
|
||||
requestID := requestIDFromHTTPResponse(resp)
|
||||
result, err := decodeHTTPResponse(resp)
|
||||
|
||||
@@ -0,0 +1,232 @@
|
||||
package clients
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func TestGeminiGenerationConfigDropsDefaultRatioTokensAndNormalizesKSize(t *testing.T) {
|
||||
config := geminiGenerationConfig(map[string]any{
|
||||
"aspect_ratio": "auto",
|
||||
"size": "2k",
|
||||
"generationConfig": map[string]any{
|
||||
"imageConfig": map[string]any{
|
||||
"aspectRatio": "adaptive",
|
||||
"imageSize": "2k",
|
||||
},
|
||||
},
|
||||
}, true)
|
||||
|
||||
imageConfig := mapFromAny(config["imageConfig"])
|
||||
if _, ok := imageConfig["aspectRatio"]; ok {
|
||||
t.Fatalf("non-ratio tokens must not be sent to Gemini: %+v", config)
|
||||
}
|
||||
if imageConfig["imageSize"] != "2K" {
|
||||
t.Fatalf("Gemini imageSize should use canonical K format: %+v", config)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeOfficialOpenAIImageSizeUsesModelSpecificWireFormat(t *testing.T) {
|
||||
legacy := map[string]any{
|
||||
"model": "gpt-image-1",
|
||||
"size": "2K",
|
||||
"width": 2048,
|
||||
"height": 1152,
|
||||
"aspect_ratio": "16:9",
|
||||
}
|
||||
normalizeOfficialOpenAIImageSize(legacy, map[string]any{
|
||||
"size": "2K",
|
||||
"aspect_ratio": "16:9",
|
||||
})
|
||||
if legacy["size"] != "1536x1024" {
|
||||
t.Fatalf("legacy GPT Image should use an official orientation size: %+v", legacy)
|
||||
}
|
||||
|
||||
flexible := map[string]any{
|
||||
"model": "gpt-image-2",
|
||||
"size": "2K",
|
||||
"resolution": "2k",
|
||||
"aspect_ratio": "16:9",
|
||||
}
|
||||
normalizeOfficialOpenAIImageSize(flexible, map[string]any{
|
||||
"size": "2K",
|
||||
"aspect_ratio": "16:9",
|
||||
})
|
||||
if flexible["size"] != "2048x1152" {
|
||||
t.Fatalf("GPT Image 2 should receive flexible pixel dimensions: %+v", flexible)
|
||||
}
|
||||
|
||||
rounded := map[string]any{
|
||||
"model": "gpt-image-2",
|
||||
"size": "2K",
|
||||
"resolution": "2K",
|
||||
"aspect_ratio": "3:2",
|
||||
}
|
||||
normalizeOfficialOpenAIImageSize(rounded, map[string]any{
|
||||
"size": "2K",
|
||||
"aspect_ratio": "3:2",
|
||||
})
|
||||
if rounded["size"] != "2048x1360" {
|
||||
t.Fatalf("GPT Image 2 dimensions should satisfy the official 16-pixel multiple: %+v", rounded)
|
||||
}
|
||||
|
||||
explicit := map[string]any{
|
||||
"model": "gpt-image-2",
|
||||
"size": "1232x768",
|
||||
"width": 1232,
|
||||
"height": 768,
|
||||
}
|
||||
normalizeOfficialOpenAIImageSize(explicit, map[string]any{
|
||||
"size": "1234x777",
|
||||
})
|
||||
if explicit["size"] != "1234x777" {
|
||||
t.Fatalf("an explicit WxH size must be sent unchanged to OpenAI: %+v", explicit)
|
||||
}
|
||||
|
||||
explicitLegacy := map[string]any{
|
||||
"model": "gpt-image-1",
|
||||
"size": "2048x1152",
|
||||
}
|
||||
normalizeOfficialOpenAIImageSize(explicitLegacy, map[string]any{
|
||||
"size": "2048x1152",
|
||||
})
|
||||
if explicitLegacy["size"] != "2048x1152" {
|
||||
t.Fatalf("explicit WxH must also win for legacy OpenAI image models: %+v", explicitLegacy)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIClientUsesOriginalWxHAndConvertsKResolution(t *testing.T) {
|
||||
received := make([]map[string]any, 0, 3)
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
var body map[string]any
|
||||
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||||
t.Fatalf("decode OpenAI image request: %v", err)
|
||||
}
|
||||
received = append(received, body)
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"data": []any{map[string]any{"b64_json": "aW1hZ2U="}},
|
||||
})
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
candidate := store.RuntimeModelCandidate{
|
||||
BaseURL: server.URL,
|
||||
Provider: "openai",
|
||||
ProviderModelName: "gpt-image-2",
|
||||
ModelType: "image_generate",
|
||||
Credentials: map[string]any{"apiKey": "openai-key"},
|
||||
}
|
||||
requests := []Request{
|
||||
{
|
||||
Kind: "images.generations",
|
||||
ModelType: "image_generate",
|
||||
Body: map[string]any{
|
||||
"prompt": "explicit dimensions",
|
||||
"size": "1232x768",
|
||||
"width": 1232,
|
||||
"height": 768,
|
||||
},
|
||||
OriginalBody: map[string]any{
|
||||
"size": "1234x777",
|
||||
},
|
||||
Candidate: candidate,
|
||||
},
|
||||
{
|
||||
Kind: "images.generations",
|
||||
ModelType: "image_generate",
|
||||
Body: map[string]any{
|
||||
"prompt": "resolution conversion",
|
||||
"size": "2048x1365",
|
||||
"resolution": "2K",
|
||||
"aspect_ratio": "3:2",
|
||||
"width": 2048,
|
||||
"height": 1365,
|
||||
},
|
||||
OriginalBody: map[string]any{
|
||||
"size": "2K",
|
||||
"aspect_ratio": "3:2",
|
||||
},
|
||||
Candidate: candidate,
|
||||
},
|
||||
{
|
||||
Kind: "images.generations",
|
||||
ModelType: "image_generate",
|
||||
Body: map[string]any{
|
||||
"prompt": "resolution field conversion",
|
||||
"size": "1024x1024",
|
||||
"resolution": "1K",
|
||||
"aspect_ratio": "1:1",
|
||||
"width": 1024,
|
||||
"height": 1024,
|
||||
},
|
||||
OriginalBody: map[string]any{
|
||||
"resolution": "1k",
|
||||
"aspect_ratio": "1:1",
|
||||
},
|
||||
Candidate: candidate,
|
||||
},
|
||||
}
|
||||
for _, request := range requests {
|
||||
if _, err := (OpenAIClient{HTTPClient: server.Client()}).Run(context.Background(), request); err != nil {
|
||||
t.Fatalf("run OpenAI image request: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
if len(received) != 3 {
|
||||
t.Fatalf("unexpected request count: %d", len(received))
|
||||
}
|
||||
if received[0]["size"] != "1234x777" {
|
||||
t.Fatalf("explicit WxH should win over preprocessed dimensions: %+v", received[0])
|
||||
}
|
||||
if received[1]["size"] != "2048x1360" {
|
||||
t.Fatalf("2K with 3:2 should use normalized OpenAI dimensions: %+v", received[1])
|
||||
}
|
||||
if received[2]["size"] != "1024x1024" {
|
||||
t.Fatalf("resolution=1K should be converted to OpenAI dimensions: %+v", received[2])
|
||||
}
|
||||
for _, body := range received {
|
||||
for _, key := range []string{"resolution", "aspect_ratio", "width", "height"} {
|
||||
if _, ok := body[key]; ok {
|
||||
t.Fatalf("generic geometry field %q must not reach OpenAI: %+v", key, body)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestVolcesImageBodySelectsConfiguredSizeFormat(t *testing.T) {
|
||||
body := volcesImageBody(Request{
|
||||
Kind: "images.generations",
|
||||
ModelType: "image_generate",
|
||||
Body: map[string]any{
|
||||
"model": "Seedream-5.0-Pro",
|
||||
"resolution": "2k",
|
||||
"size": "2048x1152",
|
||||
"width": 2048,
|
||||
"height": 1152,
|
||||
"aspect_ratio": "16:9",
|
||||
"platformId": "gateway-internal-platform",
|
||||
},
|
||||
Candidate: store.RuntimeModelCandidate{
|
||||
ProviderModelName: "doubao-seedream-5-0-pro-260628",
|
||||
ModelType: "image_generate",
|
||||
Capabilities: map[string]any{
|
||||
"image_generate": map[string]any{
|
||||
"size_param_format": "resolution",
|
||||
},
|
||||
},
|
||||
},
|
||||
})
|
||||
if body["size"] != "2K" {
|
||||
t.Fatalf("resolution-only Volces model should receive canonical K size: %+v", body)
|
||||
}
|
||||
for _, key := range []string{"resolution", "width", "height", "aspect_ratio", "platformId"} {
|
||||
if _, ok := body[key]; ok {
|
||||
t.Fatalf("generic geometry field %q must not reach Volces: %+v", key, body)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,251 @@
|
||||
package clients
|
||||
|
||||
import (
|
||||
"math"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
const (
|
||||
openAIGPTImage2MaxEdge = 3840
|
||||
openAIGPTImage2MinPixels = 655360
|
||||
openAIGPTImage2MaxPixels = 8294400
|
||||
openAIGPTImage2MaxRatio = 3
|
||||
openAIGPTImage2Multiple = 16
|
||||
)
|
||||
|
||||
type providerImageDimensions struct {
|
||||
width int
|
||||
height int
|
||||
}
|
||||
|
||||
func normalizedProviderAspectRatio(value any) string {
|
||||
compact := strings.Map(func(r rune) rune {
|
||||
switch r {
|
||||
case ' ', '\t', '\r', '\n':
|
||||
return -1
|
||||
default:
|
||||
return r
|
||||
}
|
||||
}, strings.TrimSpace(stringFromAny(value)))
|
||||
parts := strings.Split(compact, ":")
|
||||
if len(parts) != 2 {
|
||||
return ""
|
||||
}
|
||||
width, widthErr := strconv.ParseFloat(parts[0], 64)
|
||||
height, heightErr := strconv.ParseFloat(parts[1], 64)
|
||||
if widthErr != nil || heightErr != nil || width <= 0 || height <= 0 {
|
||||
return ""
|
||||
}
|
||||
return parts[0] + ":" + parts[1]
|
||||
}
|
||||
|
||||
func normalizedProviderImageResolution(value any) string {
|
||||
text := strings.ToLower(strings.TrimSpace(stringFromAny(value)))
|
||||
if len(text) < 2 || text[len(text)-1] != 'k' {
|
||||
return ""
|
||||
}
|
||||
magnitude, err := strconv.Atoi(text[:len(text)-1])
|
||||
if err != nil || magnitude <= 0 {
|
||||
return ""
|
||||
}
|
||||
return strconv.Itoa(magnitude) + "K"
|
||||
}
|
||||
|
||||
func normalizeOfficialOpenAIImageSize(body map[string]any, originalBody map[string]any) {
|
||||
if originalBody == nil {
|
||||
originalBody = body
|
||||
}
|
||||
if width, height, ok := providerPixelDimensions(stringFromAny(originalBody["size"])); ok {
|
||||
body["size"] = strconv.Itoa(width) + "x" + strconv.Itoa(height)
|
||||
return
|
||||
}
|
||||
|
||||
model := strings.ToLower(strings.TrimSpace(stringFromAny(body["model"])))
|
||||
resolution := normalizedProviderImageResolution(firstNonEmptyString(
|
||||
originalBody["resolution"],
|
||||
originalBody["size"],
|
||||
body["resolution"],
|
||||
body["size"],
|
||||
))
|
||||
if resolution != "" {
|
||||
width, height, ok := providerBodyPixelDimensions(body)
|
||||
if !ok {
|
||||
width, height, ok = providerResolutionDimensions(
|
||||
resolution,
|
||||
firstNonEmptyString(originalBody["aspect_ratio"], originalBody["aspectRatio"], originalBody["ratio"]),
|
||||
)
|
||||
}
|
||||
if !ok {
|
||||
delete(body, "size")
|
||||
return
|
||||
}
|
||||
switch {
|
||||
case strings.HasPrefix(model, "gpt-image-2"):
|
||||
dimensions := normalizeOpenAIGPTImage2Dimensions(width, height)
|
||||
body["size"] = strconv.Itoa(dimensions.width) + "x" + strconv.Itoa(dimensions.height)
|
||||
case strings.HasPrefix(model, "gpt-image-"):
|
||||
body["size"] = officialOpenAILegacyImageSize(width, height)
|
||||
default:
|
||||
body["size"] = strconv.Itoa(width) + "x" + strconv.Itoa(height)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
if !strings.HasPrefix(model, "gpt-image-") {
|
||||
return
|
||||
}
|
||||
size := strings.ToLower(strings.TrimSpace(stringFromAny(body["size"])))
|
||||
if size == "auto" {
|
||||
body["size"] = "auto"
|
||||
return
|
||||
}
|
||||
if strings.HasPrefix(model, "gpt-image-2") {
|
||||
if width, height, ok := providerPixelDimensions(size); ok {
|
||||
body["size"] = strconv.Itoa(width) + "x" + strconv.Itoa(height)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
switch size {
|
||||
case "1024x1024", "1536x1024", "1024x1536":
|
||||
body["size"] = size
|
||||
return
|
||||
}
|
||||
width, height, ok := requestedProviderDimensions(body)
|
||||
if !ok {
|
||||
body["size"] = "auto"
|
||||
return
|
||||
}
|
||||
body["size"] = officialOpenAILegacyImageSize(width, height)
|
||||
}
|
||||
|
||||
func officialOpenAILegacyImageSize(width int, height int) string {
|
||||
switch {
|
||||
case width > height:
|
||||
return "1536x1024"
|
||||
case height > width:
|
||||
return "1024x1536"
|
||||
default:
|
||||
return "1024x1024"
|
||||
}
|
||||
}
|
||||
|
||||
func providerBodyPixelDimensions(body map[string]any) (int, int, bool) {
|
||||
if width, height, ok := providerPixelDimensions(stringFromAny(body["size"])); ok {
|
||||
return width, height, true
|
||||
}
|
||||
width := intFromAny(body["width"])
|
||||
height := intFromAny(body["height"])
|
||||
return width, height, width > 0 && height > 0
|
||||
}
|
||||
|
||||
func providerResolutionDimensions(resolution string, aspectRatio any) (int, int, bool) {
|
||||
resolution = normalizedProviderImageResolution(resolution)
|
||||
if resolution == "" {
|
||||
return 0, 0, false
|
||||
}
|
||||
magnitude, err := strconv.Atoi(strings.TrimSuffix(resolution, "K"))
|
||||
if err != nil || magnitude <= 0 {
|
||||
return 0, 0, false
|
||||
}
|
||||
baseSide := magnitude * 1024
|
||||
ratio := providerAspectRatioNumber(aspectRatio)
|
||||
width, height := baseSide, baseSide
|
||||
if ratio > 1 {
|
||||
height = max(1, int(math.Round(float64(baseSide)/ratio)))
|
||||
} else if ratio > 0 && ratio < 1 {
|
||||
width = max(1, int(math.Round(float64(baseSide)*ratio)))
|
||||
}
|
||||
return width, height, true
|
||||
}
|
||||
|
||||
func requestedProviderDimensions(body map[string]any) (int, int, bool) {
|
||||
if width, height, ok := providerPixelDimensions(stringFromAny(body["size"])); ok {
|
||||
return width, height, true
|
||||
}
|
||||
width := intFromAny(body["width"])
|
||||
height := intFromAny(body["height"])
|
||||
if width > 0 && height > 0 {
|
||||
return width, height, true
|
||||
}
|
||||
ratio := providerAspectRatioNumber(firstNonEmptyString(body["aspect_ratio"], body["aspectRatio"], body["ratio"]))
|
||||
if ratio > 1 {
|
||||
return 2, 1, true
|
||||
}
|
||||
if ratio > 0 && ratio < 1 {
|
||||
return 1, 2, true
|
||||
}
|
||||
if ratio == 1 {
|
||||
return 1, 1, true
|
||||
}
|
||||
return 0, 0, false
|
||||
}
|
||||
|
||||
func providerPixelDimensions(value string) (int, int, bool) {
|
||||
parts := strings.Split(strings.ToLower(strings.TrimSpace(value)), "x")
|
||||
if len(parts) != 2 {
|
||||
return 0, 0, false
|
||||
}
|
||||
width, widthErr := strconv.Atoi(strings.TrimSpace(parts[0]))
|
||||
height, heightErr := strconv.Atoi(strings.TrimSpace(parts[1]))
|
||||
if widthErr != nil || heightErr != nil || width <= 0 || height <= 0 {
|
||||
return 0, 0, false
|
||||
}
|
||||
return width, height, true
|
||||
}
|
||||
|
||||
func providerAspectRatioNumber(value any) float64 {
|
||||
normalized := normalizedProviderAspectRatio(value)
|
||||
if normalized == "" {
|
||||
return 0
|
||||
}
|
||||
parts := strings.Split(normalized, ":")
|
||||
width, _ := strconv.ParseFloat(parts[0], 64)
|
||||
height, _ := strconv.ParseFloat(parts[1], 64)
|
||||
return width / height
|
||||
}
|
||||
|
||||
func normalizeOpenAIGPTImage2Dimensions(width int, height int) providerImageDimensions {
|
||||
if width <= 0 || height <= 0 {
|
||||
return providerImageDimensions{}
|
||||
}
|
||||
rawRatio := float64(width) / float64(height)
|
||||
ratio := rawRatio
|
||||
if ratio < 1 {
|
||||
ratio = 1 / ratio
|
||||
}
|
||||
ratio = min(ratio, float64(openAIGPTImage2MaxRatio))
|
||||
landscape := rawRatio >= 1
|
||||
longEdge := min(max(width, height), openAIGPTImage2MaxEdge)
|
||||
dimensions := buildOpenAIGPTImage2Dimensions(longEdge, ratio, landscape, false)
|
||||
for dimensions.width*dimensions.height > openAIGPTImage2MaxPixels {
|
||||
longEdge = max(openAIGPTImage2Multiple, max(dimensions.width, dimensions.height)-openAIGPTImage2Multiple)
|
||||
dimensions = buildOpenAIGPTImage2Dimensions(longEdge, ratio, landscape, false)
|
||||
}
|
||||
for dimensions.width*dimensions.height < openAIGPTImage2MinPixels && max(dimensions.width, dimensions.height) < openAIGPTImage2MaxEdge {
|
||||
longEdge = min(openAIGPTImage2MaxEdge, max(dimensions.width, dimensions.height)+openAIGPTImage2Multiple)
|
||||
dimensions = buildOpenAIGPTImage2Dimensions(longEdge, ratio, landscape, true)
|
||||
}
|
||||
return dimensions
|
||||
}
|
||||
|
||||
func buildOpenAIGPTImage2Dimensions(longEdge int, ratio float64, landscape bool, roundUp bool) providerImageDimensions {
|
||||
normalizedLongEdge := roundOpenAIGPTImage2Multiple(float64(longEdge), roundUp)
|
||||
normalizedShortEdge := roundOpenAIGPTImage2Multiple(float64(normalizedLongEdge)/ratio, roundUp)
|
||||
if float64(normalizedLongEdge)/float64(normalizedShortEdge) > openAIGPTImage2MaxRatio {
|
||||
normalizedShortEdge += openAIGPTImage2Multiple
|
||||
}
|
||||
if landscape {
|
||||
return providerImageDimensions{width: normalizedLongEdge, height: normalizedShortEdge}
|
||||
}
|
||||
return providerImageDimensions{width: normalizedShortEdge, height: normalizedLongEdge}
|
||||
}
|
||||
|
||||
func roundOpenAIGPTImage2Multiple(value float64, roundUp bool) int {
|
||||
rounded := math.Round(value / openAIGPTImage2Multiple)
|
||||
if roundUp {
|
||||
rounded = math.Ceil(value / openAIGPTImage2Multiple)
|
||||
}
|
||||
return max(openAIGPTImage2Multiple, int(rounded)*openAIGPTImage2Multiple)
|
||||
}
|
||||
@@ -0,0 +1,82 @@
|
||||
package clients
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
type timeoutResponseBody struct {
|
||||
read bool
|
||||
}
|
||||
|
||||
func (body *timeoutResponseBody) Read(buffer []byte) (int, error) {
|
||||
if body.read {
|
||||
return 0, context.DeadlineExceeded
|
||||
}
|
||||
body.read = true
|
||||
return copy(buffer, []byte(`{"partial":`)), nil
|
||||
}
|
||||
|
||||
func (*timeoutResponseBody) Close() error { return nil }
|
||||
|
||||
func TestProviderRequestTimeoutUsesMediaDefaults(t *testing.T) {
|
||||
for _, test := range []struct {
|
||||
kind string
|
||||
want time.Duration
|
||||
}{
|
||||
{kind: "images.generations", want: 20 * time.Minute},
|
||||
{kind: "images.edits", want: 20 * time.Minute},
|
||||
{kind: "videos.generations", want: 30 * time.Minute},
|
||||
{kind: "chat.completions", want: 10 * time.Minute},
|
||||
} {
|
||||
if got := ProviderRequestTimeout(test.kind); got != test.want {
|
||||
t.Fatalf("timeout for %s: got %s want %s", test.kind, got, test.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestTransportTimeoutUsesTerminalTimeoutCode(t *testing.T) {
|
||||
err := transportClientError(context.DeadlineExceeded)
|
||||
if err.Code != "timeout" || err.Retryable {
|
||||
t.Fatalf("transport timeout classification = %+v, want terminal timeout", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResponseBodyTimeoutUsesTerminalTimeoutCode(t *testing.T) {
|
||||
response := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: make(http.Header),
|
||||
Body: io.NopCloser(&timeoutResponseBody{}),
|
||||
}
|
||||
_, wire, err := decodeHTTPResponseForProtocol(response, ProtocolOpenAIImages)
|
||||
if err == nil {
|
||||
t.Fatal("expected response body timeout")
|
||||
}
|
||||
clientErr, ok := err.(*ClientError)
|
||||
if !ok || clientErr.Code != "timeout" || clientErr.Retryable {
|
||||
t.Fatalf("response body timeout classification = %#v, want terminal timeout", err)
|
||||
}
|
||||
if clientErr.StatusCode != http.StatusOK || clientErr.Wire != wire {
|
||||
t.Fatalf("response body timeout lost response metadata: %#v", clientErr)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMediaPollTimeoutsUseTaskDefaultsWithoutPlatformOverride(t *testing.T) {
|
||||
image := Request{Kind: "images.generations", Candidate: storeCandidateWithConfig("", "", nil, nil)}
|
||||
video := Request{Kind: "videos.generations", Candidate: storeCandidateWithConfig("", "", nil, nil)}
|
||||
if got := providerPollTimeout(image); got != 20*time.Minute {
|
||||
t.Fatalf("image provider poll timeout: got %s want %s", got, 20*time.Minute)
|
||||
}
|
||||
if got := providerPollTimeout(video); got != 30*time.Minute {
|
||||
t.Fatalf("video provider poll timeout: got %s want %s", got, 30*time.Minute)
|
||||
}
|
||||
if got := kelingPollTimeout(video); got != 30*time.Minute {
|
||||
t.Fatalf("Keling video poll timeout: got %s want %s", got, 30*time.Minute)
|
||||
}
|
||||
if got := volcesPollTimeout(video); got != 30*time.Minute {
|
||||
t.Fatalf("Volces video poll timeout: got %s want %s", got, 30*time.Minute)
|
||||
}
|
||||
}
|
||||
@@ -3,8 +3,14 @@ package clients
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"net/textproto"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
@@ -12,7 +18,8 @@ import (
|
||||
)
|
||||
|
||||
type OpenAIClient struct {
|
||||
HTTPClient *http.Client
|
||||
HTTPClient *http.Client
|
||||
Corrections *ParameterCorrectionCache
|
||||
}
|
||||
|
||||
func (c OpenAIClient) Run(ctx context.Context, request Request) (Response, error) {
|
||||
@@ -42,9 +49,19 @@ func (c OpenAIClient) Run(ctx context.Context, request Request) (Response, error
|
||||
}
|
||||
if endpointKind == "chat.completions" {
|
||||
body = NormalizeChatCompletionRequestBody(body)
|
||||
applyOpenAIChatReasoningParams(body, request.Candidate)
|
||||
normalizedBody, normalizeErr := normalizeOpenAIChatAudioParts(body, request.Candidate)
|
||||
if normalizeErr != nil {
|
||||
return Response{}, normalizeErr
|
||||
}
|
||||
body = normalizedBody
|
||||
if err := applyOpenAIChatReasoningParamsWithSource(body, request.Candidate, openAIAdaptationSource(request, endpointKind)); err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
body = FilterOpenAIChatRequestBody(body)
|
||||
} else if request.Kind == "responses" {
|
||||
if err := applyOpenAIResponsesReasoningParamsWithSource(body, request.Candidate, request.OriginalBody); err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
body = FilterOpenAIResponsesRequestBody(body)
|
||||
if _, hasInput := body["input"]; !hasInput {
|
||||
if messages, hasMessages := request.Body["messages"]; hasMessages {
|
||||
@@ -59,34 +76,89 @@ func (c OpenAIClient) Run(ctx context.Context, request Request) (Response, error
|
||||
}
|
||||
}
|
||||
body["model"] = upstreamModelName(request.Candidate)
|
||||
normalizeOpenAIImageRequestBody(endpointKind, body, request.OriginalBody)
|
||||
stream := openAIEndpointSupportsStream(endpointKind) && (request.Stream || boolValue(body, "stream"))
|
||||
ensureOpenAIStreamUsage(body, endpointKind, stream)
|
||||
raw, _ := json.Marshal(body)
|
||||
upstreamEndpoint := joinURL(openAIBaseURL(endpointKind, request.Candidate), endpoint)
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, upstreamEndpoint, bytes.NewReader(raw))
|
||||
if err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
responseStartedAt := time.Now()
|
||||
resp, err := httpClient(request.HTTPClient, c.HTTPClient).Do(req)
|
||||
if err != nil {
|
||||
return Response{}, &ClientError{Code: "network", Message: err.Error(), Retryable: true}
|
||||
correctionScope := newParameterCorrectionScope(request, endpointKind)
|
||||
correctionEnabled := endpointKind == "chat.completions" || endpointKind == "responses" ||
|
||||
endpointKind == "images.generations" || endpointKind == "images.edits"
|
||||
protectedCorrections := callerProtectedCorrectionParameters(request, endpointKind)
|
||||
if correctionEnabled {
|
||||
c.Corrections.apply(correctionScope, body, protectedCorrections)
|
||||
}
|
||||
provisionalRules := make([]parameterCorrectionRule, 0, 2)
|
||||
seenCorrectionErrors := make(map[string]struct{})
|
||||
var resp *http.Response
|
||||
requestClient := httpClient(request.HTTPClient, c.HTTPClient)
|
||||
for correctionAttempt := 0; ; correctionAttempt++ {
|
||||
raw, contentType, payloadErr := openAIRequestPayload(ctx, endpointKind, body, request.Candidate)
|
||||
if payloadErr != nil {
|
||||
return Response{}, payloadErr
|
||||
}
|
||||
req, requestErr := http.NewRequestWithContext(ctx, http.MethodPost, upstreamEndpoint, bytes.NewReader(raw))
|
||||
if requestErr != nil {
|
||||
return Response{}, requestErr
|
||||
}
|
||||
req.Header.Set("Content-Type", contentType)
|
||||
req.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
applyUpstreamIdempotency(req, request)
|
||||
if err := notifySubmissionStarted(request); err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
resp, requestErr = requestClient.Do(req)
|
||||
if requestErr != nil {
|
||||
return Response{}, transportClientError(requestErr)
|
||||
}
|
||||
if err := notifyResponseReceived(request); err != nil {
|
||||
_ = resp.Body.Close()
|
||||
return Response{}, err
|
||||
}
|
||||
if !correctionEnabled || (resp.StatusCode != http.StatusBadRequest && resp.StatusCode != http.StatusUnprocessableEntity) {
|
||||
break
|
||||
}
|
||||
_, _, upstreamErr := decodeHTTPResponseForProtocol(resp, openAIWireProtocol(endpointKind))
|
||||
if upstreamErr == nil {
|
||||
break
|
||||
}
|
||||
fingerprint := fmt.Sprintf("%d:%s", resp.StatusCode, strings.ToLower(upstreamErr.Error()))
|
||||
if _, repeated := seenCorrectionErrors[fingerprint]; repeated || correctionAttempt >= 2 {
|
||||
return Response{}, annotateResponseError(upstreamErr, requestIDFromParameterError(upstreamErr), responseStartedAt, time.Now())
|
||||
}
|
||||
seenCorrectionErrors[fingerprint] = struct{}{}
|
||||
rule, safe := deriveParameterCorrection(upstreamErr, body, request.Candidate, protectedCorrections)
|
||||
if !safe || !applyParameterCorrectionRule(body, rule) {
|
||||
return Response{}, annotateResponseError(upstreamErr, requestIDFromParameterError(upstreamErr), responseStartedAt, time.Now())
|
||||
}
|
||||
c.Corrections.invalidate(correctionScope, rule.Param)
|
||||
provisionalRules = append(provisionalRules, rule)
|
||||
}
|
||||
if len(provisionalRules) > 0 && resp.StatusCode >= 200 && resp.StatusCode < 300 {
|
||||
c.Corrections.commit(correctionScope, provisionalRules)
|
||||
}
|
||||
requestID := requestIDFromHTTPResponse(resp)
|
||||
var result map[string]any
|
||||
var wire *WireResponse
|
||||
var err error
|
||||
upstreamResponseID := ""
|
||||
nativeStreamDelta := openAIWireStreamDelta(request.StreamDelta, openAIWireProtocol(endpointKind), resp)
|
||||
if request.Kind == "responses" && protocol == ProtocolOpenAIResponses && stream {
|
||||
result, upstreamResponseID, err = decodeNativeResponsesStream(resp, request.StreamDelta)
|
||||
result, upstreamResponseID, err = decodeNativeResponsesStream(resp, nativeStreamDelta)
|
||||
wire = &WireResponse{Protocol: openAIWireProtocol(endpointKind), StatusCode: resp.StatusCode, Headers: compatibleResponseHeaders(resp.Header)}
|
||||
} else {
|
||||
var streamDelta StreamDelta = request.StreamDelta
|
||||
var streamDelta StreamDelta = nativeStreamDelta
|
||||
var adapter *chatResponsesStreamAdapter
|
||||
if request.Kind == "responses" && protocol == ProtocolOpenAIChatCompletions && stream {
|
||||
adapter = newChatResponsesStreamAdapter(request.PublicResponseID, request.Model)
|
||||
adapter = newChatResponsesStreamAdapter(request.PublicResponseID, request.Model, request.Body)
|
||||
streamDelta = func(event StreamDeltaEvent) error { return adapter.delta(event, request.StreamDelta) }
|
||||
}
|
||||
result, err = decodeOpenAIResponse(resp, stream, streamDelta)
|
||||
if stream {
|
||||
result, err = decodeOpenAIResponse(resp, true, streamDelta)
|
||||
wire = &WireResponse{Protocol: openAIWireProtocol(endpointKind), StatusCode: resp.StatusCode, Headers: compatibleResponseHeaders(resp.Header)}
|
||||
} else {
|
||||
result, wire, err = decodeHTTPResponseForProtocol(resp, openAIWireProtocol(endpointKind))
|
||||
}
|
||||
if err == nil && endpointKind == "chat.completions" {
|
||||
result = NormalizeChatCompletionResult(result)
|
||||
}
|
||||
@@ -103,6 +175,12 @@ func (c OpenAIClient) Run(ctx context.Context, request Request) (Response, error
|
||||
Progress: providerProgress(request), ResponseStartedAt: responseStartedAt, ResponseFinishedAt: time.Now(),
|
||||
UpstreamProtocol: protocol, UpstreamEndpoint: endpoint, UpstreamResponseID: upstreamResponseID,
|
||||
PublicResponseID: request.PublicResponseID, ResponseConverted: true,
|
||||
Wire: func() *WireResponse {
|
||||
if wire != nil {
|
||||
wire.Converted = true
|
||||
}
|
||||
return wire
|
||||
}(),
|
||||
}, nil
|
||||
}
|
||||
}
|
||||
@@ -138,9 +216,412 @@ func (c OpenAIClient) Run(ctx context.Context, request Request) (Response, error
|
||||
UpstreamEndpoint: endpoint,
|
||||
UpstreamResponseID: upstreamResponseID,
|
||||
PublicResponseID: publicResponseID,
|
||||
Wire: wire,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func openAIAdaptationSource(request Request, endpointKind string) map[string]any {
|
||||
if request.OriginalBody == nil || request.Kind != "responses" || endpointKind != "chat.completions" {
|
||||
return request.OriginalBody
|
||||
}
|
||||
out := cloneBody(request.OriginalBody)
|
||||
if reasoning, ok := request.OriginalBody["reasoning"].(map[string]any); ok {
|
||||
if effort, explicit := reasoning["effort"]; explicit {
|
||||
out["reasoning_effort"] = effort
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func normalizeOpenAIImageRequestBody(endpointKind string, body map[string]any, originalBody map[string]any) {
|
||||
if endpointKind != "images.generations" && endpointKind != "images.edits" {
|
||||
return
|
||||
}
|
||||
normalizeOfficialOpenAIImageSize(body, originalBody)
|
||||
// OpenAI image endpoints express output geometry through size. The Gateway
|
||||
// keeps the generic fields for capability validation and billing, but they
|
||||
// are not valid OpenAI wire parameters.
|
||||
for _, key := range []string{
|
||||
"aspect_ratio",
|
||||
"aspectRatio",
|
||||
"resolution",
|
||||
"width",
|
||||
"height",
|
||||
"platform_id",
|
||||
"platformId",
|
||||
"platform_model_id",
|
||||
"platformModelId",
|
||||
} {
|
||||
delete(body, key)
|
||||
}
|
||||
}
|
||||
|
||||
func openAIRequestPayload(ctx context.Context, endpointKind string, body map[string]any, candidate store.RuntimeModelCandidate) ([]byte, string, error) {
|
||||
if endpointKind != "images.edits" {
|
||||
raw, err := json.Marshal(body)
|
||||
return raw, "application/json", err
|
||||
}
|
||||
if OpenAIImageEditUsesJSONURL(candidate) {
|
||||
return openAIImageEditJSONPayload(body)
|
||||
}
|
||||
var payload bytes.Buffer
|
||||
writer := multipart.NewWriter(&payload)
|
||||
images := openAIImageEditValues(firstPresent(body["images"], body["image"]))
|
||||
if len(images) == 0 {
|
||||
return nil, "", &ClientError{
|
||||
Code: "invalid_parameter",
|
||||
Message: "image is required",
|
||||
Param: "image",
|
||||
StatusCode: http.StatusBadRequest,
|
||||
Retryable: false,
|
||||
}
|
||||
}
|
||||
imageFieldName := openAIImageEditFieldName(candidate, len(images))
|
||||
for index, value := range images {
|
||||
contentType, image, err := openAIImageEditPayload(ctx, value)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
if err := writeOpenAIImageEditFile(writer, imageFieldName, fmt.Sprintf("image-%d%s", index+1, openAIImageFileExtension(contentType)), contentType, image); err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
}
|
||||
if mask := firstPresent(body["mask"], body["mask_image"], body["maskImage"]); mask != nil {
|
||||
contentType, image, err := openAIImageEditPayload(ctx, mask)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
if err := writeOpenAIImageEditFile(writer, "mask", "mask"+openAIImageFileExtension(contentType), contentType, image); err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
}
|
||||
for key, value := range body {
|
||||
switch key {
|
||||
case "image", "images", "mask", "mask_image", "maskImage":
|
||||
continue
|
||||
}
|
||||
if strings.HasPrefix(key, "_") || value == nil {
|
||||
continue
|
||||
}
|
||||
fieldValue, err := openAIFormFieldValue(value)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
if fieldValue == "" {
|
||||
continue
|
||||
}
|
||||
if err := writer.WriteField(key, fieldValue); err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
}
|
||||
if err := writer.Close(); err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
return payload.Bytes(), writer.FormDataContentType(), nil
|
||||
}
|
||||
|
||||
// OpenAIImageEditUsesJSONURL reports whether the platform explicitly opts in to
|
||||
// JSON URL requests for image edits. Multipart remains the compatibility-safe
|
||||
// default when the setting is absent or unrecognized.
|
||||
func OpenAIImageEditUsesJSONURL(candidate store.RuntimeModelCandidate) bool {
|
||||
format := strings.ToLower(strings.TrimSpace(stringFromAny(candidate.PlatformConfig["imageEditRequestFormat"])))
|
||||
format = strings.ReplaceAll(format, "-", "_")
|
||||
format = strings.ReplaceAll(format, " ", "_")
|
||||
return format == "json" || format == "json_url"
|
||||
}
|
||||
|
||||
func openAIImageEditJSONPayload(body map[string]any) ([]byte, string, error) {
|
||||
images := openAIImageEditValues(firstPresent(body["images"], body["image"]))
|
||||
if len(images) == 0 {
|
||||
return nil, "", &ClientError{
|
||||
Code: "invalid_parameter",
|
||||
Message: "image is required",
|
||||
Param: "image",
|
||||
StatusCode: http.StatusBadRequest,
|
||||
Retryable: false,
|
||||
}
|
||||
}
|
||||
imageURLs := make([]any, 0, len(images))
|
||||
for _, image := range images {
|
||||
imageURL, err := openAIImageEditURLValue(image, "image")
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
imageURLs = append(imageURLs, imageURL)
|
||||
}
|
||||
payload := make(map[string]any, len(body))
|
||||
for key, value := range body {
|
||||
switch key {
|
||||
case "image", "images", "mask", "mask_image", "maskImage":
|
||||
continue
|
||||
}
|
||||
if strings.HasPrefix(key, "_") || value == nil {
|
||||
continue
|
||||
}
|
||||
payload[key] = value
|
||||
}
|
||||
if len(imageURLs) == 1 {
|
||||
payload["image"] = imageURLs[0]
|
||||
} else {
|
||||
payload["image"] = imageURLs
|
||||
}
|
||||
if mask := firstPresent(body["mask"], body["mask_image"], body["maskImage"]); mask != nil {
|
||||
maskURL, err := openAIImageEditURLValue(mask, "mask")
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
payload["mask"] = maskURL
|
||||
}
|
||||
raw, err := json.Marshal(payload)
|
||||
return raw, "application/json", err
|
||||
}
|
||||
|
||||
func openAIImageEditURLValue(value any, param string) (string, error) {
|
||||
switch typed := value.(type) {
|
||||
case map[string]any:
|
||||
for _, key := range []string{"url", "image_url", "imageUrl"} {
|
||||
if nested := typed[key]; nested != nil {
|
||||
return openAIImageEditURLValue(nested, param)
|
||||
}
|
||||
}
|
||||
case string:
|
||||
raw := strings.TrimSpace(typed)
|
||||
parsed, err := url.Parse(raw)
|
||||
if err == nil && parsed.Host != "" && (strings.EqualFold(parsed.Scheme, "http") || strings.EqualFold(parsed.Scheme, "https")) {
|
||||
return raw, nil
|
||||
}
|
||||
}
|
||||
return "", &ClientError{
|
||||
Code: "invalid_parameter",
|
||||
Message: "OpenAI image edit JSON mode requires " + param + " to be an HTTP(S) URL",
|
||||
Param: param,
|
||||
StatusCode: http.StatusBadRequest,
|
||||
Retryable: false,
|
||||
}
|
||||
}
|
||||
|
||||
func openAIImageEditFieldName(candidate store.RuntimeModelCandidate, imageCount int) string {
|
||||
if configured := strings.TrimSpace(stringFromAny(candidate.PlatformConfig["imageEditMultipartFieldName"])); configured == "image" || configured == "image[]" {
|
||||
return configured
|
||||
}
|
||||
if imageCount > 1 {
|
||||
return "image[]"
|
||||
}
|
||||
parsed, err := url.Parse(strings.TrimSpace(candidate.BaseURL))
|
||||
if err == nil && strings.EqualFold(parsed.Hostname(), "api.openai.com") {
|
||||
return "image[]"
|
||||
}
|
||||
return "image"
|
||||
}
|
||||
|
||||
func openAIImageEditValues(value any) []any {
|
||||
switch typed := value.(type) {
|
||||
case []any:
|
||||
return typed
|
||||
case []string:
|
||||
out := make([]any, 0, len(typed))
|
||||
for _, item := range typed {
|
||||
out = append(out, item)
|
||||
}
|
||||
return out
|
||||
case nil:
|
||||
return nil
|
||||
default:
|
||||
return []any{typed}
|
||||
}
|
||||
}
|
||||
|
||||
func openAIImageEditPayload(ctx context.Context, value any) (string, []byte, error) {
|
||||
switch typed := value.(type) {
|
||||
case map[string]any:
|
||||
for _, key := range []string{"data", "b64_json", "base64", "url"} {
|
||||
if nested := typed[key]; nested != nil {
|
||||
return openAIImageEditPayload(ctx, nested)
|
||||
}
|
||||
}
|
||||
case string:
|
||||
raw := strings.TrimSpace(typed)
|
||||
if raw == "" {
|
||||
break
|
||||
}
|
||||
contentType := ""
|
||||
encoded := raw
|
||||
if strings.HasPrefix(strings.ToLower(raw), "data:") {
|
||||
prefix, payload, ok := strings.Cut(raw, ",")
|
||||
if !ok || !strings.Contains(strings.ToLower(prefix), ";base64") {
|
||||
break
|
||||
}
|
||||
contentType = strings.TrimSpace(strings.Split(strings.TrimPrefix(prefix, "data:"), ";")[0])
|
||||
encoded = payload
|
||||
} else if parsed, parseErr := url.Parse(raw); parseErr == nil && parsed.Host != "" && (strings.EqualFold(parsed.Scheme, "http") || strings.EqualFold(parsed.Scheme, "https")) {
|
||||
return fetchRemoteMediaInputPayload(ctx, raw, 256<<20)
|
||||
}
|
||||
image, err := decodeOpenAIImageEditBase64(encoded)
|
||||
if err != nil {
|
||||
break
|
||||
}
|
||||
if contentType == "" {
|
||||
contentType = strings.TrimSpace(strings.Split(http.DetectContentType(image), ";")[0])
|
||||
}
|
||||
return contentType, image, nil
|
||||
case []byte:
|
||||
if len(typed) > 0 {
|
||||
return strings.TrimSpace(strings.Split(http.DetectContentType(typed), ";")[0]), typed, nil
|
||||
}
|
||||
}
|
||||
return "", nil, &ClientError{
|
||||
Code: "invalid_parameter",
|
||||
Message: "image must be a base64 or data URL payload",
|
||||
Param: "image",
|
||||
StatusCode: http.StatusBadRequest,
|
||||
Retryable: false,
|
||||
}
|
||||
}
|
||||
|
||||
func fetchRemoteMediaInputPayload(ctx context.Context, sourceURL string, maxBytes int64) (string, []byte, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, sourceURL, nil)
|
||||
if err != nil {
|
||||
return "", nil, &ClientError{Code: "request_asset_fetch_failed", Message: err.Error(), Retryable: false}
|
||||
}
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
return "", nil, &ClientError{Code: "request_asset_fetch_failed", Message: err.Error(), Retryable: true}
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
|
||||
return "", nil, &ClientError{
|
||||
Code: "request_asset_fetch_failed",
|
||||
Message: resp.Status,
|
||||
StatusCode: resp.StatusCode,
|
||||
Retryable: HTTPRetryable(resp.StatusCode),
|
||||
}
|
||||
}
|
||||
if maxBytes <= 0 {
|
||||
maxBytes = 256 << 20
|
||||
}
|
||||
payload, err := io.ReadAll(io.LimitReader(resp.Body, maxBytes+1))
|
||||
if err != nil {
|
||||
return "", nil, &ClientError{Code: "request_asset_fetch_failed", Message: err.Error(), Retryable: true}
|
||||
}
|
||||
if int64(len(payload)) > maxBytes {
|
||||
return "", nil, &ClientError{
|
||||
Code: "invalid_parameter",
|
||||
Message: "remote media input exceeds the download limit",
|
||||
Param: "image",
|
||||
StatusCode: http.StatusBadRequest,
|
||||
Retryable: false,
|
||||
}
|
||||
}
|
||||
contentType := strings.TrimSpace(strings.Split(resp.Header.Get("Content-Type"), ";")[0])
|
||||
if contentType == "" || contentType == "application/octet-stream" {
|
||||
contentType = strings.TrimSpace(strings.Split(http.DetectContentType(payload), ";")[0])
|
||||
}
|
||||
return contentType, payload, nil
|
||||
}
|
||||
|
||||
func decodeOpenAIImageEditBase64(value string) ([]byte, error) {
|
||||
normalized := strings.Map(func(char rune) rune {
|
||||
switch char {
|
||||
case '\n', '\r', '\t', ' ':
|
||||
return -1
|
||||
default:
|
||||
return char
|
||||
}
|
||||
}, value)
|
||||
var lastErr error
|
||||
for _, encoding := range []*base64.Encoding{
|
||||
base64.StdEncoding,
|
||||
base64.RawStdEncoding,
|
||||
base64.URLEncoding,
|
||||
base64.RawURLEncoding,
|
||||
} {
|
||||
payload, err := encoding.DecodeString(normalized)
|
||||
if err == nil && len(payload) > 0 {
|
||||
return payload, nil
|
||||
}
|
||||
lastErr = err
|
||||
}
|
||||
return nil, lastErr
|
||||
}
|
||||
|
||||
func writeOpenAIImageEditFile(writer *multipart.Writer, fieldName string, fileName string, contentType string, payload []byte) error {
|
||||
header := make(textproto.MIMEHeader)
|
||||
header.Set("Content-Disposition", fmt.Sprintf(`form-data; name="%s"; filename="%s"`, fieldName, escapeMultipartFilename(fileName)))
|
||||
if contentType != "" {
|
||||
header.Set("Content-Type", contentType)
|
||||
}
|
||||
part, err := writer.CreatePart(header)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, err = io.Copy(part, bytes.NewReader(payload))
|
||||
return err
|
||||
}
|
||||
|
||||
func openAIImageFileExtension(contentType string) string {
|
||||
switch strings.ToLower(strings.TrimSpace(strings.Split(contentType, ";")[0])) {
|
||||
case "image/jpeg":
|
||||
return ".jpg"
|
||||
case "image/webp":
|
||||
return ".webp"
|
||||
case "image/gif":
|
||||
return ".gif"
|
||||
default:
|
||||
return ".png"
|
||||
}
|
||||
}
|
||||
|
||||
func openAIFormFieldValue(value any) (string, error) {
|
||||
switch typed := value.(type) {
|
||||
case string:
|
||||
return typed, nil
|
||||
case bool:
|
||||
return fmt.Sprintf("%t", typed), nil
|
||||
case float64:
|
||||
return fmt.Sprintf("%v", typed), nil
|
||||
case float32:
|
||||
return fmt.Sprintf("%v", typed), nil
|
||||
case int:
|
||||
return fmt.Sprintf("%d", typed), nil
|
||||
case int64:
|
||||
return fmt.Sprintf("%d", typed), nil
|
||||
case json.Number:
|
||||
return typed.String(), nil
|
||||
default:
|
||||
encoded, err := json.Marshal(value)
|
||||
return string(encoded), err
|
||||
}
|
||||
}
|
||||
|
||||
func openAIWireStreamDelta(next StreamDelta, protocol string, response *http.Response) StreamDelta {
|
||||
if next == nil {
|
||||
return nil
|
||||
}
|
||||
return func(event StreamDeltaEvent) error {
|
||||
event.WireProtocol = protocol
|
||||
if response != nil {
|
||||
event.WireStatusCode = response.StatusCode
|
||||
event.WireHeaders = compatibleResponseHeaders(response.Header)
|
||||
}
|
||||
return next(event)
|
||||
}
|
||||
}
|
||||
|
||||
func openAIWireProtocol(kind string) string {
|
||||
switch kind {
|
||||
case "chat.completions":
|
||||
return ProtocolOpenAIChatCompletions
|
||||
case "responses":
|
||||
return ProtocolOpenAIResponses
|
||||
case "embeddings":
|
||||
return ProtocolOpenAIEmbeddings
|
||||
case "images.generations", "images.edits":
|
||||
return ProtocolOpenAIImages
|
||||
default:
|
||||
return "openai_" + strings.ReplaceAll(kind, ".", "_")
|
||||
}
|
||||
}
|
||||
|
||||
func decodeOpenAIResponse(resp *http.Response, stream bool, onDelta StreamDelta) (map[string]any, error) {
|
||||
if stream {
|
||||
result, err := decodeOpenAIStreamResponse(resp, onDelta)
|
||||
@@ -202,12 +683,18 @@ func ensureOpenAIStreamUsage(body map[string]any, kind string, stream bool) {
|
||||
return
|
||||
}
|
||||
streamOptions := map[string]any{}
|
||||
if existing, ok := body["stream_options"].(map[string]any); ok {
|
||||
if raw, explicit := body["stream_options"]; explicit {
|
||||
existing, ok := raw.(map[string]any)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
for key, value := range existing {
|
||||
streamOptions[key] = value
|
||||
}
|
||||
}
|
||||
streamOptions["include_usage"] = true
|
||||
if _, explicit := streamOptions["include_usage"]; !explicit {
|
||||
streamOptions["include_usage"] = true
|
||||
}
|
||||
body["stream_options"] = streamOptions
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,169 @@
|
||||
package clients
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
var openAIChatAudioFormatByMIME = map[string]string{
|
||||
"audio/aac": "aac",
|
||||
"audio/flac": "flac",
|
||||
"audio/m4a": "m4a",
|
||||
"audio/mp4": "m4a",
|
||||
"audio/mpeg": "mp3",
|
||||
"audio/ogg": "ogg",
|
||||
"audio/opus": "opus",
|
||||
"audio/wav": "wav",
|
||||
"audio/webm": "webm",
|
||||
"audio/x-m4a": "m4a",
|
||||
"audio/x-wav": "wav",
|
||||
}
|
||||
|
||||
var openAIChatAudioFormatByExtension = map[string]string{
|
||||
"aac": "aac",
|
||||
"flac": "flac",
|
||||
"m4a": "m4a",
|
||||
"mp3": "mp3",
|
||||
"mp4": "m4a",
|
||||
"oga": "ogg",
|
||||
"ogg": "ogg",
|
||||
"opus": "opus",
|
||||
"wav": "wav",
|
||||
"webm": "webm",
|
||||
}
|
||||
|
||||
func normalizeOpenAIChatAudioParts(body map[string]any, candidate store.RuntimeModelCandidate) (map[string]any, error) {
|
||||
if !isAliyunBailianOpenAI(candidate) {
|
||||
return body, nil
|
||||
}
|
||||
messages, ok := body["messages"].([]any)
|
||||
if !ok {
|
||||
return body, nil
|
||||
}
|
||||
|
||||
normalizedMessages := make([]any, 0, len(messages))
|
||||
for _, rawMessage := range messages {
|
||||
message, ok := rawMessage.(map[string]any)
|
||||
if !ok {
|
||||
normalizedMessages = append(normalizedMessages, rawMessage)
|
||||
continue
|
||||
}
|
||||
content, ok := message["content"].([]any)
|
||||
if !ok {
|
||||
normalizedMessages = append(normalizedMessages, rawMessage)
|
||||
continue
|
||||
}
|
||||
|
||||
normalizedContent := make([]any, 0, len(content))
|
||||
for _, rawPart := range content {
|
||||
part, ok := rawPart.(map[string]any)
|
||||
if !ok {
|
||||
normalizedContent = append(normalizedContent, rawPart)
|
||||
continue
|
||||
}
|
||||
normalizedPart, err := normalizeAliyunOpenAIChatAudioPart(part)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
normalizedContent = append(normalizedContent, normalizedPart)
|
||||
}
|
||||
copiedMessage := cloneMapAny(message)
|
||||
copiedMessage["content"] = normalizedContent
|
||||
normalizedMessages = append(normalizedMessages, copiedMessage)
|
||||
}
|
||||
|
||||
out := cloneMapAny(body)
|
||||
out["messages"] = normalizedMessages
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func normalizeAliyunOpenAIChatAudioPart(part map[string]any) (map[string]any, error) {
|
||||
switch strings.TrimSpace(stringFromAny(part["type"])) {
|
||||
case "audio_url":
|
||||
audioURL, _ := part["audio_url"].(map[string]any)
|
||||
data := strings.TrimSpace(stringFromAny(audioURL["url"]))
|
||||
format := resolveOpenAIChatAudioFormat(data, audioURL["format"], audioURL["mime_type"], audioURL["mimeType"])
|
||||
if data == "" || format == "" {
|
||||
return nil, invalidOpenAIChatAudioError()
|
||||
}
|
||||
return map[string]any{
|
||||
"type": "input_audio",
|
||||
"input_audio": map[string]any{
|
||||
"data": data,
|
||||
"format": format,
|
||||
},
|
||||
}, nil
|
||||
case "input_audio":
|
||||
inputAudio, _ := part["input_audio"].(map[string]any)
|
||||
data := firstNonEmptyString(inputAudio["data"], inputAudio["url"])
|
||||
format := resolveOpenAIChatAudioFormat(data, inputAudio["format"], inputAudio["mime_type"], inputAudio["mimeType"])
|
||||
if data == "" || format == "" {
|
||||
return nil, invalidOpenAIChatAudioError()
|
||||
}
|
||||
copiedInput := cloneMapAny(inputAudio)
|
||||
copiedInput["data"] = data
|
||||
copiedInput["format"] = format
|
||||
delete(copiedInput, "url")
|
||||
copiedPart := cloneMapAny(part)
|
||||
copiedPart["input_audio"] = copiedInput
|
||||
return copiedPart, nil
|
||||
default:
|
||||
return part, nil
|
||||
}
|
||||
}
|
||||
|
||||
func resolveOpenAIChatAudioFormat(source string, values ...any) string {
|
||||
if explicit := strings.ToLower(strings.TrimSpace(firstNonEmptyString(values...))); explicit != "" {
|
||||
if normalized := openAIChatAudioFormatByExtension[explicit]; normalized != "" {
|
||||
return normalized
|
||||
}
|
||||
if normalized := openAIChatAudioFormatByMIME[strings.TrimSpace(strings.Split(explicit, ";")[0])]; normalized != "" {
|
||||
return normalized
|
||||
}
|
||||
return explicit
|
||||
}
|
||||
if dataMIME := openAIChatAudioDataMIME(source); dataMIME != "" {
|
||||
if normalized := openAIChatAudioFormatByMIME[dataMIME]; normalized != "" {
|
||||
return normalized
|
||||
}
|
||||
}
|
||||
if extension := openAIChatAudioSourceExtension(source); extension != "" {
|
||||
return openAIChatAudioFormatByExtension[extension]
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func openAIChatAudioDataMIME(source string) string {
|
||||
if !strings.HasPrefix(strings.ToLower(source), "data:") {
|
||||
return ""
|
||||
}
|
||||
header := strings.TrimPrefix(source[:strings.Index(source+",", ",")], "data:")
|
||||
return strings.ToLower(strings.TrimSpace(strings.Split(header, ";")[0]))
|
||||
}
|
||||
|
||||
func openAIChatAudioSourceExtension(source string) string {
|
||||
if source == "" || strings.HasPrefix(strings.ToLower(source), "data:") {
|
||||
return ""
|
||||
}
|
||||
if parsed, err := url.Parse(source); err == nil && parsed.Path != "" {
|
||||
source = parsed.Path
|
||||
}
|
||||
source = strings.TrimSuffix(strings.Split(strings.Split(source, "?")[0], "#")[0], "/")
|
||||
if index := strings.LastIndex(source, "."); index >= 0 && index+1 < len(source) {
|
||||
return strings.ToLower(source[index+1:])
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func invalidOpenAIChatAudioError() error {
|
||||
return &ClientError{
|
||||
Code: "invalid_parameter",
|
||||
Message: "input_audio requires data and a resolvable format",
|
||||
Param: "messages.content.input_audio",
|
||||
StatusCode: http.StatusBadRequest,
|
||||
Retryable: false,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,92 @@
|
||||
package clients
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func TestOpenAIImageEndpointsRemoveRejectedResponseFormatAndCacheCorrection(t *testing.T) {
|
||||
for _, test := range []struct {
|
||||
name string
|
||||
kind string
|
||||
body map[string]any
|
||||
}{
|
||||
{
|
||||
name: "generation JSON",
|
||||
kind: "images.generations",
|
||||
body: map[string]any{"prompt": "test", "response_format": "url"},
|
||||
},
|
||||
{
|
||||
name: "edit multipart",
|
||||
kind: "images.edits",
|
||||
body: map[string]any{"prompt": "test", "image": "aW1hZ2U=", "response_format": "url"},
|
||||
},
|
||||
} {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
var requests atomic.Int32
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
requestNo := requests.Add(1)
|
||||
responseFormat := ""
|
||||
if test.kind == "images.edits" {
|
||||
if err := r.ParseMultipartForm(1 << 20); err != nil {
|
||||
t.Fatalf("parse multipart image request: %v", err)
|
||||
}
|
||||
responseFormat = r.FormValue("response_format")
|
||||
} else {
|
||||
var body map[string]any
|
||||
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||||
t.Fatalf("decode image request: %v", err)
|
||||
}
|
||||
responseFormat, _ = body["response_format"].(string)
|
||||
}
|
||||
if requestNo == 1 {
|
||||
if responseFormat != "url" {
|
||||
t.Fatalf("first request lost caller response_format: %q", responseFormat)
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusBadRequest)
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"error": map[string]any{
|
||||
"message": "Unknown parameter: 'response_format'.",
|
||||
"type": "invalid_request_error", "param": "response_format", "code": "unknown_parameter",
|
||||
}})
|
||||
return
|
||||
}
|
||||
if responseFormat != "" {
|
||||
t.Fatalf("corrected request still contains response_format: %q", responseFormat)
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"data": []any{map[string]any{"b64_json": "aW1hZ2U="}}})
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
cache := NewParameterCorrectionCache()
|
||||
request := Request{
|
||||
Kind: test.kind, Model: "gpt-image-2", Body: test.body, OriginalBody: test.body,
|
||||
UpstreamIdempotencyKey: "gateway-task-id",
|
||||
Candidate: store.RuntimeModelCandidate{
|
||||
Provider: "openai", BaseURL: server.URL, ProviderModelName: "gpt-image-2",
|
||||
Credentials: map[string]any{"apiKey": "test-key"},
|
||||
},
|
||||
}
|
||||
client := OpenAIClient{HTTPClient: server.Client(), Corrections: cache}
|
||||
if _, err := client.Run(context.Background(), request); err != nil {
|
||||
t.Fatalf("correct rejected response_format: %v", err)
|
||||
}
|
||||
if requests.Load() != 2 {
|
||||
t.Fatalf("first corrected call requests=%d, want 2", requests.Load())
|
||||
}
|
||||
if _, err := client.Run(context.Background(), request); err != nil {
|
||||
t.Fatalf("reuse cached response_format correction: %v", err)
|
||||
}
|
||||
if requests.Load() != 3 {
|
||||
t.Fatalf("cached correction should avoid another 400, requests=%d", requests.Load())
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,154 @@
|
||||
package clients
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"reflect"
|
||||
"testing"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func TestOpenAIChatNativeDeepPassthroughPreservesOfficialAndFutureFields(t *testing.T) {
|
||||
requestBody := completeChatPassthroughBody()
|
||||
var captured map[string]any
|
||||
var idempotencyKey string
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
idempotencyKey = r.Header.Get("Idempotency-Key")
|
||||
if err := json.NewDecoder(r.Body).Decode(&captured); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"id": "chatcmpl-native", "object": "chat.completion", "model": "provider-chat",
|
||||
"choices": []any{map[string]any{"message": map[string]any{"role": "assistant", "content": "ok"}}},
|
||||
})
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
_, err := (OpenAIClient{HTTPClient: server.Client()}).Run(context.Background(), Request{
|
||||
Kind: "chat.completions", Model: "public-chat", Body: requestBody, OriginalBody: requestBody,
|
||||
UpstreamIdempotencyKey: "gateway-task-id",
|
||||
Candidate: store.RuntimeModelCandidate{
|
||||
Provider: "openai", BaseURL: server.URL, ProviderModelName: "provider-chat",
|
||||
Credentials: map[string]any{"apiKey": "test-key"},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
want := jsonRoundTripMap(t, requestBody)
|
||||
delete(want, "request_id")
|
||||
want["model"] = "provider-chat"
|
||||
if !reflect.DeepEqual(captured, want) {
|
||||
t.Fatalf("native Chat request changed\n got: %#v\nwant: %#v", captured, want)
|
||||
}
|
||||
if idempotencyKey != "gateway-task-id" {
|
||||
t.Fatalf("missing Gateway task idempotency key: %q", idempotencyKey)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesNativeDeepPassthroughPreservesOfficialAndFutureFields(t *testing.T) {
|
||||
requestBody := completeResponsesPassthroughBody()
|
||||
var captured map[string]any
|
||||
var idempotencyKey string
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
idempotencyKey = r.Header.Get("Idempotency-Key")
|
||||
if err := json.NewDecoder(r.Body).Decode(&captured); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"id": "resp-native", "object": "response", "status": "completed", "output": []any{},
|
||||
})
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
_, err := (OpenAIClient{HTTPClient: server.Client()}).Run(context.Background(), Request{
|
||||
Kind: "responses", Model: "public-responses", Body: requestBody, OriginalBody: requestBody,
|
||||
UpstreamProtocol: ProtocolOpenAIResponses, UpstreamPreviousResponseID: "resp_upstream_parent",
|
||||
UpstreamIdempotencyKey: "gateway-task-id",
|
||||
Candidate: store.RuntimeModelCandidate{
|
||||
Provider: "openai", BaseURL: server.URL, ProviderModelName: "provider-responses",
|
||||
Credentials: map[string]any{"apiKey": "test-key"},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
want := jsonRoundTripMap(t, requestBody)
|
||||
delete(want, "request_id")
|
||||
want["model"] = "provider-responses"
|
||||
want["previous_response_id"] = "resp_upstream_parent"
|
||||
if !reflect.DeepEqual(captured, want) {
|
||||
t.Fatalf("native Responses request changed\n got: %#v\nwant: %#v", captured, want)
|
||||
}
|
||||
if idempotencyKey != "gateway-task-id" {
|
||||
t.Fatalf("missing Gateway task idempotency key: %q", idempotencyKey)
|
||||
}
|
||||
}
|
||||
|
||||
func completeChatPassthroughBody() map[string]any {
|
||||
return map[string]any{
|
||||
"model": "caller-chat", "messages": []any{
|
||||
map[string]any{"role": "assistant", "content": nil, "function_call": map[string]any{"name": "legacy", "arguments": "{\"x\":1}"}},
|
||||
map[string]any{"role": "assistant", "content": nil, "tool_calls": []any{
|
||||
map[string]any{"id": "call_function", "type": "function", "function": map[string]any{"name": "lookup", "arguments": "{\"q\":1}"}},
|
||||
map[string]any{"id": "call_custom", "type": "custom", "custom": map[string]any{"name": "shell", "input": "pwd"}},
|
||||
}},
|
||||
},
|
||||
"audio": map[string]any{"format": "wav", "voice": "alloy"}, "frequency_penalty": 0.1,
|
||||
"function_call": map[string]any{"name": "legacy"}, "functions": []any{map[string]any{"name": "legacy", "parameters": map[string]any{"type": "object"}}},
|
||||
"logit_bias": map[string]any{"1": 2}, "logprobs": true, "max_completion_tokens": 101, "max_tokens": 102,
|
||||
"metadata": map[string]any{"trace": "1"}, "modalities": []any{"text"}, "moderation": map[string]any{"type": "auto"}, "n": 1,
|
||||
"parallel_tool_calls": true, "prediction": map[string]any{"type": "content", "content": "answer"}, "presence_penalty": 0.2,
|
||||
"prompt_cache_key": "cache", "prompt_cache_options": map[string]any{"type": "ephemeral"}, "prompt_cache_retention": "in_memory",
|
||||
"reasoning_effort": "low", "response_format": map[string]any{"type": "json_object"}, "safety_identifier": "safe", "seed": 7,
|
||||
"service_tier": "default", "stop": []any{"END"}, "store": false, "stream": false,
|
||||
"stream_options": map[string]any{"include_usage": true}, "temperature": 0.7,
|
||||
"tool_choice": map[string]any{"type": "custom", "custom": map[string]any{"name": "shell"}},
|
||||
"tools": []any{
|
||||
map[string]any{"type": "function", "function": map[string]any{"name": "lookup", "parameters": map[string]any{"type": "object"}}},
|
||||
map[string]any{"type": "custom", "custom": map[string]any{"name": "shell", "format": map[string]any{"type": "text"}}},
|
||||
},
|
||||
"top_logprobs": 2, "top_p": 0.9, "user": "user-1", "verbosity": "low",
|
||||
"web_search_options": map[string]any{"search_context_size": "low"},
|
||||
"future_official_field": map[string]any{"nested": []any{map[string]any{"keep": true}}},
|
||||
"request_id": "gateway-only",
|
||||
}
|
||||
}
|
||||
|
||||
func completeResponsesPassthroughBody() map[string]any {
|
||||
return map[string]any{
|
||||
"background": false, "context_management": map[string]any{"type": "compaction", "compact_threshold": 2000},
|
||||
"conversation": "conv_1", "include": []any{"message.output_text.logprobs"},
|
||||
"input": []any{map[string]any{"type": "message", "role": "user", "content": []any{
|
||||
map[string]any{"type": "input_text", "text": "hello", "prompt_cache_breakpoint": map[string]any{"type": "ephemeral"}},
|
||||
}}},
|
||||
"instructions": "be concise", "max_output_tokens": 200, "max_tool_calls": 3, "metadata": map[string]any{"trace": "1"},
|
||||
"model": "caller-responses", "moderation": map[string]any{"type": "auto"}, "parallel_tool_calls": true,
|
||||
"previous_response_id": "resp_caller_parent", "prompt": map[string]any{"id": "pmpt_1", "variables": map[string]any{"x": "y"}},
|
||||
"prompt_cache_key": "cache", "prompt_cache_options": map[string]any{"type": "ephemeral"}, "prompt_cache_retention": "24h",
|
||||
"reasoning": map[string]any{"effort": "low", "summary": "auto"}, "safety_identifier": "safe", "service_tier": "default",
|
||||
"store": true, "stream": false, "stream_options": map[string]any{"include_obfuscation": true}, "temperature": 0.7,
|
||||
"text": map[string]any{"format": map[string]any{"type": "json_schema", "name": "answer", "schema": map[string]any{"type": "object"}}, "verbosity": "low"},
|
||||
"tool_choice": map[string]any{"type": "custom", "name": "shell"},
|
||||
"tools": []any{map[string]any{"type": "custom", "name": "shell", "description": "run", "format": map[string]any{"type": "text"}}},
|
||||
"top_logprobs": 2, "top_p": 0.9, "truncation": "auto", "user": "user-1",
|
||||
"future_official_field": map[string]any{"nested": []any{map[string]any{"keep": true}}},
|
||||
"request_id": "gateway-only",
|
||||
}
|
||||
}
|
||||
|
||||
func jsonRoundTripMap(t *testing.T, value map[string]any) map[string]any {
|
||||
t.Helper()
|
||||
raw, err := json.Marshal(value)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var out map[string]any
|
||||
if err := json.Unmarshal(raw, &out); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -1,15 +1,9 @@
|
||||
package clients
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"sort"
|
||||
)
|
||||
|
||||
// Keep these lists aligned with openai-node 6.47.0 and the public OpenAI API
|
||||
// reference. The Gateway accepts a small, explicit set of routing extensions at
|
||||
// ingress, but only protocol fields (plus controlled provider adaptations) are
|
||||
// allowed across the upstream boundary.
|
||||
// reference. They document and test the currently known surface, but are not an
|
||||
// upstream allowlist: unknown fields may be introduced by OpenAI after a
|
||||
// Gateway release and must remain transparent.
|
||||
var openAIChatRequestParameters = stringSet(
|
||||
"messages", "model", "audio", "frequency_penalty", "function_call", "functions",
|
||||
"logit_bias", "logprobs", "max_completion_tokens", "max_tokens", "metadata",
|
||||
@@ -32,7 +26,7 @@ var openAIResponsesRequestParameters = stringSet(
|
||||
var gatewayOpenAIRequestExtensions = stringSet(
|
||||
"runMode", "run_mode", "conversationId", "conversation_id", "sessionId", "session_id",
|
||||
"requestId", "request_id", "signal", "userMessage", "user_message", "platformId",
|
||||
"platform_id", "options", "enable_thinking", "thinking_budget_tokens", "enable_web_search",
|
||||
"platform_id", "options", "enable_thinking", "thinking_budget", "thinking_budget_tokens", "enable_web_search",
|
||||
"modelType", "model_type", "capability", "capabilityType", "mode", "simulation", "testMode",
|
||||
"cacheAffinityKey", "cache_affinity_key", "simulationDurationMs", "simulationDurationSeconds",
|
||||
"simulationMinDurationMs", "simulationMaxDurationMs", "simulationMinDurationSeconds",
|
||||
@@ -50,56 +44,36 @@ var controlledOpenAIChatProviderParameters = stringSet(
|
||||
var controlledOpenAIResponsesProviderParameters = stringSet("presence_penalty", "frequency_penalty")
|
||||
|
||||
func ValidateOpenAIRequestParameters(kind string, body map[string]any) error {
|
||||
allowed := openAIChatRequestParameters
|
||||
if kind == "responses" {
|
||||
allowed = openAIResponsesRequestParameters
|
||||
}
|
||||
unknown := make([]string, 0)
|
||||
for key := range body {
|
||||
if _, ok := allowed[key]; ok {
|
||||
continue
|
||||
}
|
||||
if _, ok := gatewayOpenAIRequestExtensions[key]; ok {
|
||||
continue
|
||||
}
|
||||
if kind == "responses" {
|
||||
if _, ok := gatewayResponsesRequestExtensions[key]; ok {
|
||||
continue
|
||||
}
|
||||
}
|
||||
unknown = append(unknown, key)
|
||||
}
|
||||
if len(unknown) == 0 {
|
||||
return nil
|
||||
}
|
||||
sort.Strings(unknown)
|
||||
return &ClientError{
|
||||
Code: "invalid_parameter",
|
||||
Message: fmt.Sprintf("Unknown parameter: %s", unknown[0]),
|
||||
Param: unknown[0],
|
||||
StatusCode: http.StatusBadRequest,
|
||||
Retryable: false,
|
||||
}
|
||||
// OpenAI-compatible public endpoints intentionally accept future official
|
||||
// fields. Provider validation remains authoritative for fields the selected
|
||||
// upstream does not support.
|
||||
_ = kind
|
||||
_ = body
|
||||
return nil
|
||||
}
|
||||
|
||||
func FilterOpenAIChatRequestBody(body map[string]any) map[string]any {
|
||||
return filterOpenAIRequestBody(body, openAIChatRequestParameters, controlledOpenAIChatProviderParameters)
|
||||
return filterOpenAIRequestBody(body, controlledOpenAIChatProviderParameters, nil)
|
||||
}
|
||||
|
||||
func FilterOpenAIResponsesRequestBody(body map[string]any) map[string]any {
|
||||
return filterOpenAIRequestBody(body, openAIResponsesRequestParameters, controlledOpenAIResponsesProviderParameters)
|
||||
return filterOpenAIRequestBody(body, controlledOpenAIResponsesProviderParameters, gatewayResponsesRequestExtensions)
|
||||
}
|
||||
|
||||
func filterOpenAIRequestBody(body map[string]any, allowed map[string]struct{}, extensions map[string]struct{}) map[string]any {
|
||||
func filterOpenAIRequestBody(body map[string]any, controlled map[string]struct{}, protocolInternal map[string]struct{}) map[string]any {
|
||||
out := make(map[string]any, len(body))
|
||||
for key, value := range body {
|
||||
if _, ok := allowed[key]; ok {
|
||||
if _, ok := controlled[key]; ok {
|
||||
out[key] = value
|
||||
continue
|
||||
}
|
||||
if _, ok := extensions[key]; ok {
|
||||
out[key] = value
|
||||
if _, internal := gatewayOpenAIRequestExtensions[key]; internal {
|
||||
continue
|
||||
}
|
||||
if _, internal := protocolInternal[key]; internal {
|
||||
continue
|
||||
}
|
||||
out[key] = value
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
@@ -1,9 +1,6 @@
|
||||
package clients
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
import "testing"
|
||||
|
||||
func TestOpenAIChatOfficialParametersSurviveBoundary(t *testing.T) {
|
||||
body := map[string]any{}
|
||||
@@ -11,7 +8,7 @@ func TestOpenAIChatOfficialParametersSurviveBoundary(t *testing.T) {
|
||||
body[key] = "sentinel-" + key
|
||||
}
|
||||
body["conversationId"] = "internal"
|
||||
body["unknown"] = "must-not-leak"
|
||||
body["future_official_field"] = map[string]any{"nested": []any{"must-survive"}}
|
||||
|
||||
filtered := FilterOpenAIChatRequestBody(body)
|
||||
for key := range openAIChatRequestParameters {
|
||||
@@ -19,7 +16,10 @@ func TestOpenAIChatOfficialParametersSurviveBoundary(t *testing.T) {
|
||||
t.Fatalf("official Chat parameter %q was removed", key)
|
||||
}
|
||||
}
|
||||
for _, key := range []string{"conversationId", "unknown"} {
|
||||
if _, ok := filtered["future_official_field"]; !ok {
|
||||
t.Fatal("future Chat field was removed at the upstream boundary")
|
||||
}
|
||||
for _, key := range []string{"conversationId"} {
|
||||
if _, ok := filtered[key]; ok {
|
||||
t.Fatalf("internal/unknown parameter %q leaked upstream", key)
|
||||
}
|
||||
@@ -32,7 +32,7 @@ func TestOpenAIResponsesOfficialParametersSurviveBoundary(t *testing.T) {
|
||||
body[key] = "sentinel-" + key
|
||||
}
|
||||
body["request_id"] = "internal"
|
||||
body["unknown"] = "must-not-leak"
|
||||
body["future_official_field"] = map[string]any{"nested": []any{"must-survive"}}
|
||||
|
||||
filtered := FilterOpenAIResponsesRequestBody(body)
|
||||
for key := range openAIResponsesRequestParameters {
|
||||
@@ -40,20 +40,19 @@ func TestOpenAIResponsesOfficialParametersSurviveBoundary(t *testing.T) {
|
||||
t.Fatalf("official Responses parameter %q was removed", key)
|
||||
}
|
||||
}
|
||||
for _, key := range []string{"request_id", "unknown"} {
|
||||
if _, ok := filtered["future_official_field"]; !ok {
|
||||
t.Fatal("future Responses field was removed at the upstream boundary")
|
||||
}
|
||||
for _, key := range []string{"request_id"} {
|
||||
if _, ok := filtered[key]; ok {
|
||||
t.Fatalf("internal/unknown parameter %q leaked upstream", key)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateOpenAIRequestParametersRejectsUnknownTopLevelField(t *testing.T) {
|
||||
err := ValidateOpenAIRequestParameters("responses", map[string]any{"model": "demo", "input": "hello", "rogue": true})
|
||||
if err == nil || ErrorCode(err) != "invalid_parameter" || !strings.Contains(err.Error(), "rogue") {
|
||||
t.Fatalf("expected OpenAI-style invalid_parameter for rogue field, got %v", err)
|
||||
}
|
||||
if ErrorParam(err) != "rogue" {
|
||||
t.Fatalf("expected rogue parameter attribution, got %q", ErrorParam(err))
|
||||
func TestValidateOpenAIRequestParametersAcceptsFutureTopLevelField(t *testing.T) {
|
||||
if err := ValidateOpenAIRequestParameters("responses", map[string]any{"model": "demo", "input": "hello", "future_official_field": true}); err != nil {
|
||||
t.Fatalf("future Responses fields must remain forward compatible, got %v", err)
|
||||
}
|
||||
if err := ValidateOpenAIRequestParameters("responses", map[string]any{"model": "demo", "input": "hello", "messages": []any{}, "request_id": "internal"}); err != nil {
|
||||
t.Fatalf("expected controlled Responses extensions to remain accepted, got %v", err)
|
||||
@@ -104,3 +103,34 @@ func TestResponsesFallbackMapsEquivalentCurrentParameters(t *testing.T) {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeChatRequestPreservesCustomAndLegacyFunctionCall(t *testing.T) {
|
||||
legacy := map[string]any{"name": "legacy", "arguments": "{\"x\":1}"}
|
||||
custom := map[string]any{
|
||||
"id": "call_custom", "type": "custom",
|
||||
"custom": map[string]any{"name": "shell", "input": "pwd"},
|
||||
}
|
||||
body := NormalizeChatCompletionRequestBody(map[string]any{"messages": []any{map[string]any{
|
||||
"role": "assistant", "content": nil, "function_call": legacy, "tool_calls": []any{custom},
|
||||
}}})
|
||||
messages, _ := body["messages"].([]any)
|
||||
message, _ := messages[0].(map[string]any)
|
||||
if got, ok := message["function_call"].(map[string]any); !ok || got["name"] != "legacy" || got["arguments"] != "{\"x\":1}" {
|
||||
t.Fatalf("legacy function_call changed: %+v", message)
|
||||
}
|
||||
toolCalls, _ := message["tool_calls"].([]any)
|
||||
got, _ := toolCalls[0].(map[string]any)
|
||||
gotCustom, _ := got["custom"].(map[string]any)
|
||||
if got["type"] != "custom" || got["id"] != "call_custom" || gotCustom["name"] != "shell" || gotCustom["input"] != "pwd" || got["function"] != nil {
|
||||
t.Fatalf("standard custom tool call changed: %+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEnsureOpenAIStreamUsageDoesNotOverrideCallerChoice(t *testing.T) {
|
||||
body := map[string]any{"stream_options": map[string]any{"include_usage": false, "include_obfuscation": false}}
|
||||
ensureOpenAIStreamUsage(body, "chat.completions", true)
|
||||
options, _ := body["stream_options"].(map[string]any)
|
||||
if options["include_usage"] != false || options["include_obfuscation"] != false {
|
||||
t.Fatalf("caller stream options were changed: %+v", options)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,531 @@
|
||||
package clients
|
||||
|
||||
import (
|
||||
"container/list"
|
||||
"errors"
|
||||
"fmt"
|
||||
"math"
|
||||
"net/url"
|
||||
"regexp"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
const parameterCorrectionCacheCapacity = 512
|
||||
|
||||
var safeUpstreamCorrectionParameters = stringSet(
|
||||
"reasoning_effort", "reasoning.effort", "enable_thinking", "thinking_budget",
|
||||
"temperature", "top_p", "frequency_penalty", "presence_penalty", "seed", "stop", "n",
|
||||
"max_tokens", "max_completion_tokens", "max_output_tokens",
|
||||
"logprobs", "top_logprobs", "service_tier", "verbosity",
|
||||
"response_format",
|
||||
)
|
||||
|
||||
var (
|
||||
anchoredParameterPattern = regexp.MustCompile(`(?i)^(?:unknown|unsupported|invalid) parameter\s*:?\s*[` + "`" + `'\"]?([a-z][a-z0-9_.]*)`)
|
||||
restrictedParameterPattern = regexp.MustCompile(`(?i)^the value of the ([a-z][a-z0-9_.]*) parameter is restricted to\s+([^\s,.;]+)`)
|
||||
fixedValuePattern = regexp.MustCompile(`(?i)(?:restricted to|must be(?: set)?(?: to)?|only supports?(?: the)? value(?: of)?)\s*[` + "`" + `'\"]?([^\s,.;` + "`" + `'\"]+)`)
|
||||
betweenPattern = regexp.MustCompile(`(?i)\bbetween\s+(-?[0-9]+(?:\.[0-9]+)?)\s+and\s+(-?[0-9]+(?:\.[0-9]+)?)\b`)
|
||||
maximumPattern = regexp.MustCompile(`(?i)(?:less than or equal to|at most|maximum(?: value)?(?: is| of)?|<=)\s+(-?[0-9]+(?:\.[0-9]+)?)`)
|
||||
minimumPattern = regexp.MustCompile(`(?i)(?:greater than or equal to|at least|minimum(?: value)?(?: is| of)?|>=)\s+(-?[0-9]+(?:\.[0-9]+)?)`)
|
||||
)
|
||||
|
||||
type parameterCorrectionAction string
|
||||
|
||||
const (
|
||||
parameterCorrectionRemove parameterCorrectionAction = "remove"
|
||||
parameterCorrectionSet parameterCorrectionAction = "set"
|
||||
parameterCorrectionClamp parameterCorrectionAction = "clamp"
|
||||
)
|
||||
|
||||
type parameterCorrectionRule struct {
|
||||
Param string
|
||||
Action parameterCorrectionAction
|
||||
Value any
|
||||
Min *float64
|
||||
Max *float64
|
||||
}
|
||||
|
||||
type parameterCorrectionScope struct {
|
||||
Provider string
|
||||
BaseURL string
|
||||
Protocol string
|
||||
Kind string
|
||||
Model string
|
||||
}
|
||||
|
||||
func newParameterCorrectionScope(request Request, endpointKind string) parameterCorrectionScope {
|
||||
return parameterCorrectionScope{
|
||||
Provider: normalizedReasoningString(request.Candidate.Provider),
|
||||
BaseURL: normalizedCorrectionBaseURL(request.Candidate.BaseURL),
|
||||
Protocol: strings.TrimSpace(request.UpstreamProtocol),
|
||||
Kind: strings.TrimSpace(endpointKind),
|
||||
Model: normalizedReasoningString(upstreamModelName(request.Candidate)),
|
||||
}
|
||||
}
|
||||
|
||||
func normalizedCorrectionBaseURL(raw string) string {
|
||||
trimmed := strings.TrimRight(strings.TrimSpace(raw), "/")
|
||||
parsed, err := url.Parse(trimmed)
|
||||
if err != nil || parsed.Host == "" {
|
||||
return strings.ToLower(trimmed)
|
||||
}
|
||||
parsed.Scheme = strings.ToLower(parsed.Scheme)
|
||||
parsed.Host = strings.ToLower(parsed.Host)
|
||||
parsed.RawQuery = ""
|
||||
parsed.Fragment = ""
|
||||
return strings.TrimRight(parsed.String(), "/")
|
||||
}
|
||||
|
||||
func (scope parameterCorrectionScope) prefix() string {
|
||||
return strings.Join([]string{scope.Provider, scope.BaseURL, scope.Protocol, scope.Kind, scope.Model}, "\x1f") + "\x1f"
|
||||
}
|
||||
|
||||
func (scope parameterCorrectionScope) key(param string) string {
|
||||
return scope.prefix() + param
|
||||
}
|
||||
|
||||
type parameterCorrectionCacheEntry struct {
|
||||
key string
|
||||
rule parameterCorrectionRule
|
||||
}
|
||||
|
||||
type ParameterCorrectionCache struct {
|
||||
mu sync.Mutex
|
||||
capacity int
|
||||
entries map[string]*list.Element
|
||||
lru *list.List
|
||||
}
|
||||
|
||||
func NewParameterCorrectionCache() *ParameterCorrectionCache {
|
||||
return &ParameterCorrectionCache{
|
||||
capacity: parameterCorrectionCacheCapacity,
|
||||
entries: make(map[string]*list.Element),
|
||||
lru: list.New(),
|
||||
}
|
||||
}
|
||||
|
||||
func (cache *ParameterCorrectionCache) apply(scope parameterCorrectionScope, body map[string]any, protected ...map[string]struct{}) []string {
|
||||
if cache == nil {
|
||||
return nil
|
||||
}
|
||||
cache.mu.Lock()
|
||||
defer cache.mu.Unlock()
|
||||
prefix := scope.prefix()
|
||||
elements := make([]*list.Element, 0)
|
||||
for key, element := range cache.entries {
|
||||
if strings.HasPrefix(key, prefix) {
|
||||
elements = append(elements, element)
|
||||
}
|
||||
}
|
||||
sort.Slice(elements, func(i, j int) bool {
|
||||
return elements[i].Value.(parameterCorrectionCacheEntry).rule.Param < elements[j].Value.(parameterCorrectionCacheEntry).rule.Param
|
||||
})
|
||||
applied := make([]string, 0, len(elements))
|
||||
for _, element := range elements {
|
||||
rule := element.Value.(parameterCorrectionCacheEntry).rule
|
||||
if correctionParamProtected(rule.Param, protected...) {
|
||||
continue
|
||||
}
|
||||
if applyParameterCorrectionRule(body, rule) {
|
||||
applied = append(applied, rule.Param)
|
||||
cache.lru.MoveToFront(element)
|
||||
}
|
||||
}
|
||||
return applied
|
||||
}
|
||||
|
||||
func (cache *ParameterCorrectionCache) commit(scope parameterCorrectionScope, rules []parameterCorrectionRule) {
|
||||
if cache == nil || len(rules) == 0 {
|
||||
return
|
||||
}
|
||||
cache.mu.Lock()
|
||||
defer cache.mu.Unlock()
|
||||
for _, rule := range rules {
|
||||
key := scope.key(rule.Param)
|
||||
if existing := cache.entries[key]; existing != nil {
|
||||
existing.Value = parameterCorrectionCacheEntry{key: key, rule: rule}
|
||||
cache.lru.MoveToFront(existing)
|
||||
continue
|
||||
}
|
||||
element := cache.lru.PushFront(parameterCorrectionCacheEntry{key: key, rule: rule})
|
||||
cache.entries[key] = element
|
||||
for cache.lru.Len() > cache.capacity {
|
||||
oldest := cache.lru.Back()
|
||||
if oldest == nil {
|
||||
break
|
||||
}
|
||||
delete(cache.entries, oldest.Value.(parameterCorrectionCacheEntry).key)
|
||||
cache.lru.Remove(oldest)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (cache *ParameterCorrectionCache) invalidate(scope parameterCorrectionScope, param string) {
|
||||
if cache == nil {
|
||||
return
|
||||
}
|
||||
cache.mu.Lock()
|
||||
defer cache.mu.Unlock()
|
||||
key := scope.key(param)
|
||||
if element := cache.entries[key]; element != nil {
|
||||
delete(cache.entries, key)
|
||||
cache.lru.Remove(element)
|
||||
}
|
||||
}
|
||||
|
||||
func (cache *ParameterCorrectionCache) size() int {
|
||||
if cache == nil {
|
||||
return 0
|
||||
}
|
||||
cache.mu.Lock()
|
||||
defer cache.mu.Unlock()
|
||||
return cache.lru.Len()
|
||||
}
|
||||
|
||||
func deriveParameterCorrection(err error, body map[string]any, candidate store.RuntimeModelCandidate, protected map[string]struct{}) (parameterCorrectionRule, bool) {
|
||||
var clientErr *ClientError
|
||||
if !errors.As(err, &clientErr) || (clientErr.StatusCode != 400 && clientErr.StatusCode != 422) {
|
||||
return parameterCorrectionRule{}, false
|
||||
}
|
||||
param, code, message := structuredParameterError(clientErr)
|
||||
if param == "" {
|
||||
param = anchoredParameterFromMessage(message)
|
||||
}
|
||||
param = normalizeCorrectionParam(param)
|
||||
if !isSafeCorrectionParam(param) {
|
||||
return parameterCorrectionRule{}, false
|
||||
}
|
||||
if correctionParamProtected(param, protected) {
|
||||
return parameterCorrectionRule{}, false
|
||||
}
|
||||
lowerMessage := strings.ToLower(message)
|
||||
lowerCode := strings.ToLower(code)
|
||||
if conflictRule, ok := deriveConflictCorrection(param, lowerMessage); ok {
|
||||
if correctionParamProtected(conflictRule.Param, protected) {
|
||||
return parameterCorrectionRule{}, false
|
||||
}
|
||||
return conflictRule, true
|
||||
}
|
||||
if strings.Contains(lowerMessage, "unknown parameter") ||
|
||||
strings.Contains(lowerMessage, "unsupported parameter") ||
|
||||
strings.Contains(lowerMessage, "not supported") ||
|
||||
strings.Contains(lowerMessage, "does not support") ||
|
||||
strings.Contains(lowerCode, "unknown_parameter") ||
|
||||
strings.Contains(lowerCode, "unsupported_parameter") {
|
||||
return parameterCorrectionRule{Param: param, Action: parameterCorrectionRemove}, true
|
||||
}
|
||||
if param == "reasoning_effort" || param == "reasoning.effort" {
|
||||
if effort := nearestErrorSupportedReasoningEffort(message, normalizedReasoningString(valueAtCorrectionPath(body, param)), candidate); effort != "" {
|
||||
return parameterCorrectionRule{Param: param, Action: parameterCorrectionSet, Value: effort}, true
|
||||
}
|
||||
}
|
||||
if match := fixedValuePattern.FindStringSubmatch(message); len(match) == 2 {
|
||||
if value, ok := parseCorrectionScalar(match[1]); ok {
|
||||
return parameterCorrectionRule{Param: param, Action: parameterCorrectionSet, Value: value}, true
|
||||
}
|
||||
}
|
||||
if min, max, ok := explicitNumericBounds(message); ok {
|
||||
if isOutputLimitCorrectionParam(param) {
|
||||
min = nil
|
||||
}
|
||||
return parameterCorrectionRule{Param: param, Action: parameterCorrectionClamp, Min: min, Max: max}, true
|
||||
}
|
||||
return parameterCorrectionRule{}, false
|
||||
}
|
||||
|
||||
func correctionParamProtected(param string, protected ...map[string]struct{}) bool {
|
||||
for _, values := range protected {
|
||||
if _, ok := values[normalizeCorrectionParam(param)]; ok {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func callerProtectedCorrectionParameters(request Request, endpointKind string) map[string]struct{} {
|
||||
if request.OriginalBody == nil {
|
||||
return nil
|
||||
}
|
||||
protected := make(map[string]struct{})
|
||||
protect := func(param string) {
|
||||
param = normalizeCorrectionParam(param)
|
||||
if isSafeCorrectionParam(param) {
|
||||
protected[param] = struct{}{}
|
||||
}
|
||||
}
|
||||
for key := range request.OriginalBody {
|
||||
protect(key)
|
||||
}
|
||||
if reasoning, ok := request.OriginalBody["reasoning"].(map[string]any); ok {
|
||||
if _, explicit := reasoning["effort"]; explicit {
|
||||
if endpointKind == "chat.completions" {
|
||||
protect("reasoning_effort")
|
||||
} else {
|
||||
protect("reasoning.effort")
|
||||
}
|
||||
}
|
||||
}
|
||||
if request.Kind == "responses" && endpointKind == "chat.completions" {
|
||||
if _, explicit := request.OriginalBody["max_output_tokens"]; explicit {
|
||||
protect("max_completion_tokens")
|
||||
}
|
||||
if _, explicit := request.OriginalBody["top_logprobs"]; explicit {
|
||||
protect("top_logprobs")
|
||||
protect("logprobs")
|
||||
}
|
||||
if responseIncludeContains(request.OriginalBody["include"], "message.output_text.logprobs") {
|
||||
protect("logprobs")
|
||||
}
|
||||
if text, ok := request.OriginalBody["text"].(map[string]any); ok {
|
||||
if _, explicit := text["verbosity"]; explicit {
|
||||
protect("verbosity")
|
||||
}
|
||||
}
|
||||
}
|
||||
if endpointKind == "images.generations" || endpointKind == "images.edits" {
|
||||
// The Gateway canonicalizes image output independently of the provider's
|
||||
// URL/Base64 preference. Allow an explicitly supplied response_format to
|
||||
// be removed only after the selected upstream rejects it as unsupported.
|
||||
delete(protected, "response_format")
|
||||
}
|
||||
return protected
|
||||
}
|
||||
|
||||
func responseIncludeContains(value any, target string) bool {
|
||||
switch values := value.(type) {
|
||||
case []any:
|
||||
for _, value := range values {
|
||||
if stringFromAny(value) == target {
|
||||
return true
|
||||
}
|
||||
}
|
||||
case []string:
|
||||
for _, value := range values {
|
||||
if value == target {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func requestIDFromParameterError(err error) string {
|
||||
var clientErr *ClientError
|
||||
if errors.As(err, &clientErr) {
|
||||
return clientErr.RequestID
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func structuredParameterError(clientErr *ClientError) (string, string, string) {
|
||||
param := strings.TrimSpace(clientErr.Param)
|
||||
code := strings.TrimSpace(clientErr.Code)
|
||||
message := strings.TrimSpace(clientErr.Message)
|
||||
if clientErr.Wire == nil {
|
||||
return param, code, message
|
||||
}
|
||||
errorObject, _ := clientErr.Wire.Body["error"].(map[string]any)
|
||||
if errorObject == nil {
|
||||
return param, code, message
|
||||
}
|
||||
if value := strings.TrimSpace(stringFromAny(errorObject["param"])); value != "" {
|
||||
param = value
|
||||
}
|
||||
if value := strings.TrimSpace(stringFromAny(errorObject["code"])); value != "" {
|
||||
code = value
|
||||
}
|
||||
if value := strings.TrimSpace(stringFromAny(errorObject["message"])); value != "" {
|
||||
message = value
|
||||
}
|
||||
return param, code, message
|
||||
}
|
||||
|
||||
func anchoredParameterFromMessage(message string) string {
|
||||
trimmed := strings.TrimSpace(message)
|
||||
for _, pattern := range []*regexp.Regexp{anchoredParameterPattern, restrictedParameterPattern} {
|
||||
if match := pattern.FindStringSubmatch(trimmed); len(match) >= 2 {
|
||||
return match[1]
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func normalizeCorrectionParam(param string) string {
|
||||
return strings.ToLower(strings.Trim(strings.TrimSpace(param), "`'\"[]"))
|
||||
}
|
||||
|
||||
func isSafeCorrectionParam(param string) bool {
|
||||
_, ok := safeUpstreamCorrectionParameters[param]
|
||||
return ok
|
||||
}
|
||||
|
||||
func deriveConflictCorrection(param string, message string) (parameterCorrectionRule, bool) {
|
||||
if !strings.Contains(message, "cannot be used together") &&
|
||||
!strings.Contains(message, "mutually exclusive") &&
|
||||
!strings.Contains(message, "not allowed when") {
|
||||
return parameterCorrectionRule{}, false
|
||||
}
|
||||
if strings.Contains(message, "thinking_budget") && strings.Contains(message, "reasoning_effort") {
|
||||
return parameterCorrectionRule{Param: "reasoning_effort", Action: parameterCorrectionRemove}, true
|
||||
}
|
||||
return parameterCorrectionRule{Param: param, Action: parameterCorrectionRemove}, true
|
||||
}
|
||||
|
||||
func nearestErrorSupportedReasoningEffort(message string, requested string, candidate store.RuntimeModelCandidate) string {
|
||||
allowed := make(map[string]struct{})
|
||||
lowerMessage := strings.ToLower(message)
|
||||
for _, effort := range chatReasoningEffortOrder {
|
||||
if regexp.MustCompile(`\b` + regexp.QuoteMeta(effort) + `\b`).MatchString(lowerMessage) {
|
||||
allowed[effort] = struct{}{}
|
||||
}
|
||||
}
|
||||
if len(allowed) == 0 {
|
||||
if resolved, state := resolveCandidateReasoningEffort(requested, candidate); state == reasoningCapabilitySupported {
|
||||
return resolved
|
||||
}
|
||||
return ""
|
||||
}
|
||||
requestedIndex := reasoningEffortIndex(requested)
|
||||
best := ""
|
||||
bestDistance := len(chatReasoningEffortOrder) + 1
|
||||
for index, effort := range chatReasoningEffortOrder {
|
||||
if _, ok := allowed[effort]; !ok {
|
||||
continue
|
||||
}
|
||||
distance := int(math.Abs(float64(index - requestedIndex)))
|
||||
if distance < bestDistance {
|
||||
best = effort
|
||||
bestDistance = distance
|
||||
}
|
||||
}
|
||||
return best
|
||||
}
|
||||
|
||||
func parseCorrectionScalar(raw string) (any, bool) {
|
||||
trimmed := strings.Trim(strings.TrimSpace(raw), "`'\"")
|
||||
switch strings.ToLower(trimmed) {
|
||||
case "true":
|
||||
return true, true
|
||||
case "false":
|
||||
return false, true
|
||||
}
|
||||
if number, err := strconv.ParseFloat(trimmed, 64); err == nil {
|
||||
if math.Trunc(number) == number {
|
||||
return int(number), true
|
||||
}
|
||||
return number, true
|
||||
}
|
||||
if isOpenAIReasoningEffort(strings.ToLower(trimmed)) {
|
||||
return strings.ToLower(trimmed), true
|
||||
}
|
||||
return nil, false
|
||||
}
|
||||
|
||||
func explicitNumericBounds(message string) (*float64, *float64, bool) {
|
||||
if match := betweenPattern.FindStringSubmatch(message); len(match) == 3 {
|
||||
minimum, minErr := strconv.ParseFloat(match[1], 64)
|
||||
maximum, maxErr := strconv.ParseFloat(match[2], 64)
|
||||
if minErr == nil && maxErr == nil {
|
||||
return &minimum, &maximum, true
|
||||
}
|
||||
}
|
||||
var minimum *float64
|
||||
var maximum *float64
|
||||
if match := minimumPattern.FindStringSubmatch(message); len(match) == 2 {
|
||||
if value, err := strconv.ParseFloat(match[1], 64); err == nil {
|
||||
minimum = &value
|
||||
}
|
||||
}
|
||||
if match := maximumPattern.FindStringSubmatch(message); len(match) == 2 {
|
||||
if value, err := strconv.ParseFloat(match[1], 64); err == nil {
|
||||
maximum = &value
|
||||
}
|
||||
}
|
||||
return minimum, maximum, minimum != nil || maximum != nil
|
||||
}
|
||||
|
||||
func isOutputLimitCorrectionParam(param string) bool {
|
||||
return param == "max_tokens" || param == "max_completion_tokens" || param == "max_output_tokens"
|
||||
}
|
||||
|
||||
func applyParameterCorrectionRule(body map[string]any, rule parameterCorrectionRule) bool {
|
||||
current := valueAtCorrectionPath(body, rule.Param)
|
||||
switch rule.Action {
|
||||
case parameterCorrectionRemove:
|
||||
if current == nil {
|
||||
return false
|
||||
}
|
||||
deleteAtCorrectionPath(body, rule.Param)
|
||||
return true
|
||||
case parameterCorrectionSet:
|
||||
if current == nil {
|
||||
return false
|
||||
}
|
||||
if fmt.Sprint(current) == fmt.Sprint(rule.Value) {
|
||||
return false
|
||||
}
|
||||
setAtCorrectionPath(body, rule.Param, rule.Value)
|
||||
return true
|
||||
case parameterCorrectionClamp:
|
||||
number, ok := finiteFloatFromAny(current)
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
corrected := number
|
||||
if rule.Min != nil && corrected < *rule.Min {
|
||||
corrected = *rule.Min
|
||||
}
|
||||
if rule.Max != nil && corrected > *rule.Max {
|
||||
corrected = *rule.Max
|
||||
}
|
||||
if corrected == number {
|
||||
return false
|
||||
}
|
||||
if math.Trunc(corrected) == corrected {
|
||||
setAtCorrectionPath(body, rule.Param, int(corrected))
|
||||
} else {
|
||||
setAtCorrectionPath(body, rule.Param, corrected)
|
||||
}
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func valueAtCorrectionPath(body map[string]any, param string) any {
|
||||
if param != "reasoning.effort" {
|
||||
return body[param]
|
||||
}
|
||||
reasoning, _ := body["reasoning"].(map[string]any)
|
||||
return reasoning["effort"]
|
||||
}
|
||||
|
||||
func setAtCorrectionPath(body map[string]any, param string, value any) {
|
||||
if param != "reasoning.effort" {
|
||||
body[param] = value
|
||||
return
|
||||
}
|
||||
reasoning, _ := body["reasoning"].(map[string]any)
|
||||
if reasoning == nil {
|
||||
reasoning = map[string]any{}
|
||||
body["reasoning"] = reasoning
|
||||
}
|
||||
reasoning["effort"] = value
|
||||
}
|
||||
|
||||
func deleteAtCorrectionPath(body map[string]any, param string) {
|
||||
if param != "reasoning.effort" {
|
||||
delete(body, param)
|
||||
return
|
||||
}
|
||||
reasoning, _ := body["reasoning"].(map[string]any)
|
||||
delete(reasoning, "effort")
|
||||
if len(reasoning) == 0 {
|
||||
delete(body, "reasoning")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,444 @@
|
||||
package clients
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func correctionTestRequest(server *httptest.Server, body map[string]any) Request {
|
||||
return Request{
|
||||
Kind: "chat.completions",
|
||||
Model: "Qwen3.8-Max-Preview",
|
||||
Body: body,
|
||||
Candidate: store.RuntimeModelCandidate{
|
||||
Provider: "aliyun-bailian-openai",
|
||||
BaseURL: server.URL,
|
||||
ProviderModelName: "qwen3.8-max-preview",
|
||||
Credentials: map[string]any{"apiKey": "test-key"},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func writeCorrectionSuccess(w http.ResponseWriter, model any) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"id": "chatcmpl-corrected", "object": "chat.completion", "model": model,
|
||||
"choices": []any{map[string]any{"message": map[string]any{"role": "assistant", "content": "ok"}}},
|
||||
"usage": map[string]any{"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
|
||||
})
|
||||
}
|
||||
|
||||
func writeParameterError(w http.ResponseWriter, param string, message string) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusBadRequest)
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"error": map[string]any{"code": "invalid_parameter", "param": param, "message": message},
|
||||
})
|
||||
}
|
||||
|
||||
func TestOpenAIParameterCorrectionLearnsOnlyAfterSuccess(t *testing.T) {
|
||||
var requests atomic.Int32
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
requests.Add(1)
|
||||
var body map[string]any
|
||||
_ = json.NewDecoder(r.Body).Decode(&body)
|
||||
if body["n"] != float64(1) {
|
||||
writeParameterError(w, "n", "Parameter n must be set to 1.")
|
||||
return
|
||||
}
|
||||
writeCorrectionSuccess(w, body["model"])
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
cache := NewParameterCorrectionCache()
|
||||
client := OpenAIClient{HTTPClient: server.Client(), Corrections: cache}
|
||||
request := correctionTestRequest(server, map[string]any{
|
||||
"messages": []any{map[string]any{"role": "user", "content": "hello"}},
|
||||
"n": 2,
|
||||
})
|
||||
if _, err := client.Run(context.Background(), request); err != nil {
|
||||
t.Fatalf("first corrected request failed: %v", err)
|
||||
}
|
||||
if got := requests.Load(); got != 2 {
|
||||
t.Fatalf("first logical request should send twice, got %d", got)
|
||||
}
|
||||
if cache.size() != 1 {
|
||||
t.Fatalf("successful correction should be cached, size=%d", cache.size())
|
||||
}
|
||||
if _, err := client.Run(context.Background(), request); err != nil {
|
||||
t.Fatalf("cached request failed: %v", err)
|
||||
}
|
||||
if got := requests.Load(); got != 3 {
|
||||
t.Fatalf("second logical request should hit cache on first send, got total %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIParameterCorrectionLearnsDashScopeThinkingNConstraint(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
var body map[string]any
|
||||
_ = json.NewDecoder(r.Body).Decode(&body)
|
||||
if body["n"] != float64(1) {
|
||||
writeParameterError(w, "n", "The n parameter must be 1 when enable_thinking is true")
|
||||
return
|
||||
}
|
||||
writeCorrectionSuccess(w, body["model"])
|
||||
}))
|
||||
defer server.Close()
|
||||
cache := NewParameterCorrectionCache()
|
||||
client := OpenAIClient{HTTPClient: server.Client(), Corrections: cache}
|
||||
_, err := client.Run(context.Background(), correctionTestRequest(server, map[string]any{
|
||||
"messages": []any{}, "n": 2,
|
||||
}))
|
||||
if err != nil {
|
||||
t.Fatalf("DashScope n constraint should be corrected: %v", err)
|
||||
}
|
||||
if cache.size() != 1 {
|
||||
t.Fatalf("successful DashScope n correction should be cached, size=%d", cache.size())
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIParameterCorrectionFailureIsNotCached(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
writeParameterError(w, "n", "Parameter n must be set to 1.")
|
||||
}))
|
||||
defer server.Close()
|
||||
cache := NewParameterCorrectionCache()
|
||||
client := OpenAIClient{HTTPClient: server.Client(), Corrections: cache}
|
||||
_, err := client.Run(context.Background(), correctionTestRequest(server, map[string]any{
|
||||
"messages": []any{}, "n": 2,
|
||||
}))
|
||||
if err == nil {
|
||||
t.Fatal("expected upstream error")
|
||||
}
|
||||
if cache.size() != 0 {
|
||||
t.Fatalf("failed correction must not be cached, size=%d", cache.size())
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIParameterCorrectionStopsAfterTwoRetries(t *testing.T) {
|
||||
var requests atomic.Int32
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
requests.Add(1)
|
||||
var body map[string]any
|
||||
_ = json.NewDecoder(r.Body).Decode(&body)
|
||||
if body["n"] != float64(1) {
|
||||
writeParameterError(w, "n", "Parameter n must be set to 1.")
|
||||
return
|
||||
}
|
||||
if body["temperature"] != float64(1) {
|
||||
writeParameterError(w, "temperature", "Parameter temperature must be set to 1.")
|
||||
return
|
||||
}
|
||||
writeParameterError(w, "top_p", "Parameter top_p must be set to 1.")
|
||||
}))
|
||||
defer server.Close()
|
||||
client := OpenAIClient{HTTPClient: server.Client(), Corrections: NewParameterCorrectionCache()}
|
||||
_, err := client.Run(context.Background(), correctionTestRequest(server, map[string]any{
|
||||
"messages": []any{}, "n": 2, "temperature": 2, "top_p": 2,
|
||||
}))
|
||||
if err == nil {
|
||||
t.Fatal("expected third parameter error")
|
||||
}
|
||||
if got := requests.Load(); got != 3 {
|
||||
t.Fatalf("expected initial request plus two correction retries, got %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIParameterCorrectionEvictsConflictingRule(t *testing.T) {
|
||||
cache := NewParameterCorrectionCache()
|
||||
var seen []float64
|
||||
var mu sync.Mutex
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
var body map[string]any
|
||||
_ = json.NewDecoder(r.Body).Decode(&body)
|
||||
n, _ := body["n"].(float64)
|
||||
mu.Lock()
|
||||
seen = append(seen, n)
|
||||
mu.Unlock()
|
||||
if n != 2 {
|
||||
writeParameterError(w, "n", "Parameter n must be set to 2.")
|
||||
return
|
||||
}
|
||||
writeCorrectionSuccess(w, body["model"])
|
||||
}))
|
||||
defer server.Close()
|
||||
request := correctionTestRequest(server, map[string]any{"messages": []any{}, "n": 9})
|
||||
scope := newParameterCorrectionScope(request, "chat.completions")
|
||||
cache.commit(scope, []parameterCorrectionRule{{Param: "n", Action: parameterCorrectionSet, Value: 1}})
|
||||
client := OpenAIClient{HTTPClient: server.Client(), Corrections: cache}
|
||||
if _, err := client.Run(context.Background(), request); err != nil {
|
||||
t.Fatalf("conflicting cache correction failed: %v", err)
|
||||
}
|
||||
if _, err := client.Run(context.Background(), request); err != nil {
|
||||
t.Fatalf("relearned cache request failed: %v", err)
|
||||
}
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if fmt.Sprint(seen) != "[1 2 2]" {
|
||||
t.Fatalf("expected old rule, relearn retry, then new cache hit; got %v", seen)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIParameterCorrectionNeverTouchesDangerousFields(t *testing.T) {
|
||||
var requests atomic.Int32
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
requests.Add(1)
|
||||
writeParameterError(w, "model", "Unsupported parameter: model")
|
||||
}))
|
||||
defer server.Close()
|
||||
messages := []any{map[string]any{"role": "user", "content": "secret text"}}
|
||||
request := correctionTestRequest(server, map[string]any{"messages": messages})
|
||||
client := OpenAIClient{HTTPClient: server.Client(), Corrections: NewParameterCorrectionCache()}
|
||||
_, err := client.Run(context.Background(), request)
|
||||
if err == nil {
|
||||
t.Fatal("expected unsafe parameter error")
|
||||
}
|
||||
if requests.Load() != 1 {
|
||||
t.Fatalf("unsafe parameter must not retry, got %d requests", requests.Load())
|
||||
}
|
||||
if fmt.Sprint(request.Body["messages"]) != fmt.Sprint(messages) {
|
||||
t.Fatal("messages were unexpectedly modified")
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIParameterCorrectionWorksBeforeStreamingOutput(t *testing.T) {
|
||||
var requests atomic.Int32
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
requests.Add(1)
|
||||
var body map[string]any
|
||||
_ = json.NewDecoder(r.Body).Decode(&body)
|
||||
if body["n"] != float64(1) {
|
||||
writeParameterError(w, "n", "Parameter n must be set to 1.")
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
_, _ = fmt.Fprint(w, "data: {\"id\":\"chatcmpl-stream\",\"choices\":[{\"delta\":{\"content\":\"ok\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n")
|
||||
}))
|
||||
defer server.Close()
|
||||
var deltas []string
|
||||
request := correctionTestRequest(server, map[string]any{"messages": []any{}, "n": 2, "stream": true})
|
||||
request.Stream = true
|
||||
request.StreamDelta = func(event StreamDeltaEvent) error {
|
||||
deltas = append(deltas, event.Text)
|
||||
return nil
|
||||
}
|
||||
client := OpenAIClient{HTTPClient: server.Client(), Corrections: NewParameterCorrectionCache()}
|
||||
if _, err := client.Run(context.Background(), request); err != nil {
|
||||
t.Fatalf("stream correction failed: %v", err)
|
||||
}
|
||||
if requests.Load() != 2 || fmt.Sprint(deltas) != "[ok]" {
|
||||
t.Fatalf("unexpected stream correction requests=%d deltas=%v", requests.Load(), deltas)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesParameterCorrectionLearnsAfterSuccess(t *testing.T) {
|
||||
var requests atomic.Int32
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
requests.Add(1)
|
||||
var body map[string]any
|
||||
_ = json.NewDecoder(r.Body).Decode(&body)
|
||||
if body["top_p"] != float64(1) {
|
||||
writeParameterError(w, "top_p", "Parameter top_p must be set to 1.")
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"id": "resp-corrected", "object": "response", "status": "completed",
|
||||
"output": []any{}, "usage": map[string]any{"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
|
||||
})
|
||||
}))
|
||||
defer server.Close()
|
||||
cache := NewParameterCorrectionCache()
|
||||
client := OpenAIClient{HTTPClient: server.Client(), Corrections: cache}
|
||||
request := Request{
|
||||
Kind: "responses",
|
||||
Model: "Qwen3.8-Max-Preview",
|
||||
UpstreamProtocol: ProtocolOpenAIResponses,
|
||||
Body: map[string]any{"input": "hello", "top_p": 2, "reasoning": map[string]any{"effort": "none"}},
|
||||
Candidate: store.RuntimeModelCandidate{
|
||||
Provider: "aliyun-bailian-openai", BaseURL: server.URL,
|
||||
ProviderModelName: "qwen3.8-max-preview", Credentials: map[string]any{"apiKey": "test-key"},
|
||||
},
|
||||
}
|
||||
if _, err := client.Run(context.Background(), request); err != nil {
|
||||
t.Fatalf("first corrected Responses request failed: %v", err)
|
||||
}
|
||||
if _, err := client.Run(context.Background(), request); err != nil {
|
||||
t.Fatalf("cached Responses request failed: %v", err)
|
||||
}
|
||||
if requests.Load() != 3 || cache.size() != 1 {
|
||||
t.Fatalf("expected retry then Responses cache hit, requests=%d cache=%d", requests.Load(), cache.size())
|
||||
}
|
||||
}
|
||||
|
||||
func TestParameterCorrectionCacheIsConcurrentAndBounded(t *testing.T) {
|
||||
cache := NewParameterCorrectionCache()
|
||||
var group sync.WaitGroup
|
||||
for index := 0; index < 1024; index++ {
|
||||
group.Add(1)
|
||||
go func(index int) {
|
||||
defer group.Done()
|
||||
scope := parameterCorrectionScope{Provider: "p", BaseURL: "https://example.com", Protocol: "chat", Kind: "chat", Model: fmt.Sprint(index)}
|
||||
cache.commit(scope, []parameterCorrectionRule{{Param: "n", Action: parameterCorrectionSet, Value: 1}})
|
||||
cache.apply(scope, map[string]any{"n": 2})
|
||||
}(index)
|
||||
}
|
||||
group.Wait()
|
||||
if cache.size() > parameterCorrectionCacheCapacity {
|
||||
t.Fatalf("cache exceeded capacity: %d", cache.size())
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIParameterCorrectionNeverChangesCallerExplicitParameterColdOrHot(t *testing.T) {
|
||||
t.Run("cold", func(t *testing.T) {
|
||||
var requests atomic.Int32
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
requests.Add(1)
|
||||
writeParameterError(w, "n", "Parameter n must be set to 1.")
|
||||
}))
|
||||
defer server.Close()
|
||||
body := map[string]any{"messages": []any{}, "n": 2}
|
||||
request := correctionTestRequest(server, body)
|
||||
request.OriginalBody = body
|
||||
cache := NewParameterCorrectionCache()
|
||||
_, err := (OpenAIClient{HTTPClient: server.Client(), Corrections: cache}).Run(context.Background(), request)
|
||||
if err == nil || requests.Load() != 1 || cache.size() != 0 {
|
||||
t.Fatalf("explicit cold parameter was corrected: requests=%d cache=%d err=%v", requests.Load(), cache.size(), err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("hot", func(t *testing.T) {
|
||||
var captured float64
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
var body map[string]any
|
||||
_ = json.NewDecoder(r.Body).Decode(&body)
|
||||
captured, _ = body["n"].(float64)
|
||||
writeCorrectionSuccess(w, body["model"])
|
||||
}))
|
||||
defer server.Close()
|
||||
body := map[string]any{"messages": []any{}, "n": 2}
|
||||
request := correctionTestRequest(server, body)
|
||||
request.OriginalBody = body
|
||||
cache := NewParameterCorrectionCache()
|
||||
cache.commit(newParameterCorrectionScope(request, "chat.completions"), []parameterCorrectionRule{{Param: "n", Action: parameterCorrectionSet, Value: 1}})
|
||||
if _, err := (OpenAIClient{HTTPClient: server.Client(), Corrections: cache}).Run(context.Background(), request); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if captured != 2 {
|
||||
t.Fatalf("cached correction changed explicit n: %v", captured)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestOpenAIParameterCorrectionMayChangeGatewayInjectedParameter(t *testing.T) {
|
||||
var requests atomic.Int32
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
requests.Add(1)
|
||||
var body map[string]any
|
||||
_ = json.NewDecoder(r.Body).Decode(&body)
|
||||
if body["n"] != float64(1) {
|
||||
writeParameterError(w, "n", "Parameter n must be set to 1.")
|
||||
return
|
||||
}
|
||||
writeCorrectionSuccess(w, body["model"])
|
||||
}))
|
||||
defer server.Close()
|
||||
request := correctionTestRequest(server, map[string]any{"messages": []any{}, "n": 2})
|
||||
request.OriginalBody = map[string]any{"messages": []any{}}
|
||||
if _, err := (OpenAIClient{HTTPClient: server.Client(), Corrections: NewParameterCorrectionCache()}).Run(context.Background(), request); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if requests.Load() != 2 {
|
||||
t.Fatalf("Gateway-injected parameter was not corrected: requests=%d", requests.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIProviderAdaptationRejectsExplicitSemanticChange(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
|
||||
t.Fatal("semantically incompatible explicit parameter must fail before upstream")
|
||||
}))
|
||||
defer server.Close()
|
||||
body := map[string]any{"messages": []any{}, "reasoning_effort": "none", "temperature": 0.2}
|
||||
request := correctionTestRequest(server, body)
|
||||
request.OriginalBody = body
|
||||
_, err := (OpenAIClient{HTTPClient: server.Client()}).Run(context.Background(), request)
|
||||
if ErrorCode(err) != "invalid_parameter" || ErrorParam(err) != "reasoning_effort" {
|
||||
t.Fatalf("expected exact reasoning incompatibility, got %v param=%q", err, ErrorParam(err))
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIProviderAdaptationRejectsInexactDeepSeekEffort(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
|
||||
t.Fatal("inexact DeepSeek effort must fail before upstream")
|
||||
}))
|
||||
defer server.Close()
|
||||
body := map[string]any{"messages": []any{}, "reasoning_effort": "low"}
|
||||
request := Request{
|
||||
Kind: "chat.completions", Model: "deepseek-v4-pro", Body: body, OriginalBody: body,
|
||||
Candidate: store.RuntimeModelCandidate{
|
||||
Provider: "deepseek-openai", BaseURL: server.URL, ProviderModelName: "deepseek-v4-pro", Credentials: map[string]any{"apiKey": "test-key"},
|
||||
},
|
||||
}
|
||||
_, err := (OpenAIClient{HTTPClient: server.Client()}).Run(context.Background(), request)
|
||||
if ErrorCode(err) != "invalid_parameter" || ErrorParam(err) != "reasoning_effort" {
|
||||
t.Fatalf("expected exact DeepSeek reasoning incompatibility, got %v param=%q", err, ErrorParam(err))
|
||||
}
|
||||
}
|
||||
|
||||
func TestResponsesFallbackProtectsExplicitMappedOutputLimitFromCorrection(t *testing.T) {
|
||||
var requests atomic.Int32
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
requests.Add(1)
|
||||
var body map[string]any
|
||||
_ = json.NewDecoder(r.Body).Decode(&body)
|
||||
if body["max_completion_tokens"] != float64(128) || body["max_tokens"] != nil {
|
||||
t.Fatalf("Responses output limit mapped incorrectly: %+v", body)
|
||||
}
|
||||
writeParameterError(w, "max_completion_tokens", "Unsupported parameter: max_completion_tokens")
|
||||
}))
|
||||
defer server.Close()
|
||||
body := map[string]any{"input": "hello", "max_output_tokens": 128}
|
||||
request := Request{
|
||||
Kind: "responses", Model: "demo", Body: body, OriginalBody: body, UpstreamProtocol: ProtocolOpenAIChatCompletions,
|
||||
Candidate: store.RuntimeModelCandidate{
|
||||
Provider: "openai", BaseURL: server.URL, ProviderModelName: "demo", Credentials: map[string]any{"apiKey": "test-key"},
|
||||
},
|
||||
}
|
||||
_, err := (OpenAIClient{HTTPClient: server.Client(), Corrections: NewParameterCorrectionCache()}).Run(context.Background(), request)
|
||||
if err == nil || requests.Load() != 1 {
|
||||
t.Fatalf("explicit mapped output limit was silently corrected: requests=%d err=%v", requests.Load(), err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResponsesFallbackProtectsIncludeMappedLogprobsFromCorrection(t *testing.T) {
|
||||
var requests atomic.Int32
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
requests.Add(1)
|
||||
var body map[string]any
|
||||
_ = json.NewDecoder(r.Body).Decode(&body)
|
||||
if body["logprobs"] != true {
|
||||
t.Fatalf("Responses include did not enable Chat logprobs: %+v", body)
|
||||
}
|
||||
writeParameterError(w, "logprobs", "Unsupported parameter: logprobs")
|
||||
}))
|
||||
defer server.Close()
|
||||
body := map[string]any{"input": "hello", "include": []any{"message.output_text.logprobs"}}
|
||||
request := Request{
|
||||
Kind: "responses", Model: "demo", Body: body, OriginalBody: body, UpstreamProtocol: ProtocolOpenAIChatCompletions,
|
||||
Candidate: store.RuntimeModelCandidate{
|
||||
Provider: "openai", BaseURL: server.URL, ProviderModelName: "demo", Credentials: map[string]any{"apiKey": "test-key"},
|
||||
},
|
||||
}
|
||||
_, err := (OpenAIClient{HTTPClient: server.Client(), Corrections: NewParameterCorrectionCache()}).Run(context.Background(), request)
|
||||
if err == nil || requests.Load() != 1 {
|
||||
t.Fatalf("explicit include logprobs was silently corrected: requests=%d err=%v", requests.Load(), err)
|
||||
}
|
||||
}
|
||||
@@ -45,7 +45,15 @@ func (c providerTaskClient) Run(ctx context.Context, request Request) (Response,
|
||||
requestID := upstreamTaskID
|
||||
var submitResult map[string]any
|
||||
if upstreamTaskID == "" {
|
||||
if err := notifySubmissionStarted(request); err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
result, id, err := c.submit(ctx, request, payload)
|
||||
if err == nil || ErrorResponseMetadata(err).StatusCode > 0 {
|
||||
if notifyErr := notifyResponseReceived(request); notifyErr != nil {
|
||||
return Response{}, notifyErr
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
return Response{}, annotateResponseError(err, id, startedAt, time.Now())
|
||||
}
|
||||
@@ -182,7 +190,7 @@ func providerPostJSON(ctx context.Context, client *http.Client, url string, body
|
||||
applyProviderAuth(req, credentials, auth)
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, "", &ClientError{Code: "network", Message: err.Error(), Retryable: true}
|
||||
return nil, "", transportClientError(err)
|
||||
}
|
||||
requestID := requestIDFromHTTPResponse(resp)
|
||||
result, err := decodeHTTPResponse(resp)
|
||||
@@ -197,7 +205,7 @@ func providerGetJSON(ctx context.Context, client *http.Client, url string, crede
|
||||
applyProviderAuth(req, credentials, auth)
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, "", &ClientError{Code: "network", Message: err.Error(), Retryable: true}
|
||||
return nil, "", transportClientError(err)
|
||||
}
|
||||
requestID := requestIDFromHTTPResponse(resp)
|
||||
result, err := decodeHTTPResponse(resp)
|
||||
@@ -332,24 +340,20 @@ func hasProviderTaskResult(result map[string]any) bool {
|
||||
return result["data"] != nil || valueAtPath(result, "data.result") != nil || valueAtPath(result, "data.audio") != nil || valueAtPath(result, "output.image_urls") != nil || valueAtPath(result, "output.video_url") != nil || valueAtPath(result, "Response.ResultVideoUrl") != nil || valueAtPath(result, "Response.ResultImages") != nil || result["audio_url"] != nil || result["urls"] != nil
|
||||
}
|
||||
|
||||
func normalizeProviderTaskResult(request Request, spec providerTaskSpec, result map[string]any, upstreamTaskID string) map[string]any {
|
||||
out := cloneMapAny(result)
|
||||
out["status"] = "success"
|
||||
func normalizeProviderTaskResult(request Request, _ providerTaskSpec, result map[string]any, upstreamTaskID string) map[string]any {
|
||||
data, ok := result["data"].([]any)
|
||||
if !ok {
|
||||
data = providerTaskData(request, result)
|
||||
}
|
||||
out := map[string]any{
|
||||
"status": "success",
|
||||
"created": time.Now().UnixMilli(),
|
||||
"model": request.Model,
|
||||
"data": data,
|
||||
}
|
||||
if upstreamTaskID != "" {
|
||||
out["upstream_task_id"] = upstreamTaskID
|
||||
}
|
||||
if out["created"] == nil {
|
||||
out["created"] = time.Now().UnixMilli()
|
||||
}
|
||||
if out["model"] == nil {
|
||||
out["model"] = request.Model
|
||||
}
|
||||
if _, ok := out["data"].([]any); !ok {
|
||||
if out["data"] != nil {
|
||||
out["raw_data"] = out["data"]
|
||||
}
|
||||
out["data"] = providerTaskData(request, result)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
@@ -456,7 +460,7 @@ func providerPollInterval(request Request) time.Duration {
|
||||
}
|
||||
|
||||
func providerPollTimeout(request Request) time.Duration {
|
||||
return durationFromConfig(request.Candidate.PlatformConfig, 10*time.Minute, "pollTimeoutMs", "poll_timeout_ms", "timeoutMs")
|
||||
return durationFromConfig(request.Candidate.PlatformConfig, ProviderRequestTimeout(request.Kind), "pollTimeoutMs", "poll_timeout_ms", "timeoutMs")
|
||||
}
|
||||
|
||||
func durationFromConfig(config map[string]any, fallback time.Duration, keys ...string) time.Duration {
|
||||
|
||||
@@ -0,0 +1,312 @@
|
||||
package clients
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"os"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
type qwen38LiveProxy struct {
|
||||
server *httptest.Server
|
||||
upstreamBase *url.URL
|
||||
mu sync.Mutex
|
||||
shapes []map[string]any
|
||||
}
|
||||
|
||||
func newQwen38LiveProxy(t *testing.T, upstream string) *qwen38LiveProxy {
|
||||
t.Helper()
|
||||
upstreamBase, err := url.Parse(strings.TrimRight(upstream, "/"))
|
||||
if err != nil {
|
||||
t.Fatalf("parse Qwen3.8 live upstream: %v", err)
|
||||
}
|
||||
proxy := &qwen38LiveProxy{upstreamBase: upstreamBase}
|
||||
proxy.server = httptest.NewServer(http.HandlerFunc(proxy.forward))
|
||||
t.Cleanup(proxy.server.Close)
|
||||
return proxy
|
||||
}
|
||||
|
||||
func (p *qwen38LiveProxy) forward(w http.ResponseWriter, request *http.Request) {
|
||||
body, err := io.ReadAll(request.Body)
|
||||
if err != nil {
|
||||
http.Error(w, "read request", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
var decoded map[string]any
|
||||
if json.Unmarshal(body, &decoded) == nil {
|
||||
p.mu.Lock()
|
||||
p.shapes = append(p.shapes, qwen38SafeRequestShape(decoded))
|
||||
p.mu.Unlock()
|
||||
}
|
||||
target := *p.upstreamBase
|
||||
target.Path = strings.TrimRight(target.Path, "/") + request.URL.Path
|
||||
upstreamRequest, err := http.NewRequestWithContext(request.Context(), request.Method, target.String(), bytes.NewReader(body))
|
||||
if err != nil {
|
||||
http.Error(w, "create upstream request", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
upstreamRequest.Header.Set("Authorization", request.Header.Get("Authorization"))
|
||||
upstreamRequest.Header.Set("Content-Type", "application/json")
|
||||
response, err := http.DefaultClient.Do(upstreamRequest)
|
||||
if err != nil {
|
||||
http.Error(w, "upstream request failed", http.StatusBadGateway)
|
||||
return
|
||||
}
|
||||
defer response.Body.Close()
|
||||
responseBody, err := io.ReadAll(response.Body)
|
||||
if err != nil {
|
||||
http.Error(w, "read upstream response", http.StatusBadGateway)
|
||||
return
|
||||
}
|
||||
status := response.StatusCode
|
||||
if status == http.StatusInternalServerError && strings.Contains(string(responseBody), "The n parameter must be 1 when enable_thinking is true") {
|
||||
// The deployed server-main currently wraps DashScope's 400 as 500. The
|
||||
// adapter only restores the original parameter-error status and shape so
|
||||
// this live test can exercise the local client's correction loop.
|
||||
status = http.StatusBadRequest
|
||||
responseBody, _ = json.Marshal(map[string]any{
|
||||
"error": map[string]any{
|
||||
"code": "invalid_parameter_error",
|
||||
"param": "n",
|
||||
"message": "The n parameter must be 1 when enable_thinking is true",
|
||||
},
|
||||
})
|
||||
}
|
||||
if contentType := response.Header.Get("Content-Type"); contentType != "" {
|
||||
w.Header().Set("Content-Type", contentType)
|
||||
}
|
||||
w.WriteHeader(status)
|
||||
_, _ = w.Write(responseBody)
|
||||
}
|
||||
|
||||
func (p *qwen38LiveProxy) requestShapes() []map[string]any {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
out := make([]map[string]any, len(p.shapes))
|
||||
copy(out, p.shapes)
|
||||
return out
|
||||
}
|
||||
|
||||
func qwen38SafeRequestShape(body map[string]any) map[string]any {
|
||||
shape := map[string]any{}
|
||||
for _, key := range []string{"model", "stream", "reasoning_effort", "enable_thinking", "thinking_budget", "temperature", "n"} {
|
||||
if value, ok := body[key]; ok {
|
||||
shape[key] = value
|
||||
}
|
||||
}
|
||||
if reasoning, ok := body["reasoning"].(map[string]any); ok {
|
||||
safeReasoning := map[string]any{}
|
||||
for _, key := range []string{"effort", "summary"} {
|
||||
if value, exists := reasoning[key]; exists {
|
||||
safeReasoning[key] = value
|
||||
}
|
||||
}
|
||||
shape["reasoning"] = safeReasoning
|
||||
}
|
||||
return shape
|
||||
}
|
||||
|
||||
func qwen38LiveRequest(proxy *qwen38LiveProxy, apiKey string, body map[string]any, stream bool) Request {
|
||||
body = cloneBody(body)
|
||||
body["stream"] = stream
|
||||
return Request{
|
||||
Kind: "chat.completions",
|
||||
Model: "Qwen3.8-Max-Preview",
|
||||
Body: body,
|
||||
Stream: stream,
|
||||
UpstreamProtocol: ProtocolOpenAIChatCompletions,
|
||||
Candidate: store.RuntimeModelCandidate{
|
||||
Provider: "aliyun-bailian-openai",
|
||||
BaseURL: proxy.server.URL,
|
||||
ProviderModelName: "Qwen3.8-Max-Preview",
|
||||
Credentials: map[string]any{"apiKey": apiKey},
|
||||
Capabilities: map[string]any{
|
||||
"text_generate": map[string]any{
|
||||
"supportThinking": true,
|
||||
"supportThinkingModeSwitch": false,
|
||||
"thinkingEffortLevels": []any{"low", "medium", "xhigh"},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func qwen38LiveBody(extra map[string]any) map[string]any {
|
||||
body := map[string]any{
|
||||
"messages": []any{map[string]any{"role": "user", "content": "Reply only OK"}},
|
||||
"max_tokens": 256,
|
||||
}
|
||||
for key, value := range extra {
|
||||
body[key] = value
|
||||
}
|
||||
return body
|
||||
}
|
||||
|
||||
func qwen38ResponseHasReasoning(response Response, streamedReasoning string) bool {
|
||||
if strings.TrimSpace(streamedReasoning) != "" {
|
||||
return true
|
||||
}
|
||||
choices, _ := response.Result["choices"].([]any)
|
||||
if len(choices) == 0 {
|
||||
return false
|
||||
}
|
||||
choice, _ := choices[0].(map[string]any)
|
||||
message, _ := choice["message"].(map[string]any)
|
||||
reasoning, _ := message["reasoning_content"].(string)
|
||||
return strings.TrimSpace(reasoning) != ""
|
||||
}
|
||||
|
||||
func qwen38LiveConfig(t *testing.T) (string, string) {
|
||||
t.Helper()
|
||||
baseURL := strings.TrimSpace(os.Getenv("QWEN38_LIVE_BASE_URL"))
|
||||
apiKey := strings.TrimSpace(os.Getenv("QWEN38_LIVE_API_KEY"))
|
||||
if baseURL == "" || apiKey == "" {
|
||||
t.Skip("set QWEN38_LIVE_BASE_URL and QWEN38_LIVE_API_KEY to run real-model acceptance")
|
||||
}
|
||||
return baseURL, apiKey
|
||||
}
|
||||
|
||||
func TestQwen38LiveStaticNormalizationAndReasoning(t *testing.T) {
|
||||
baseURL, apiKey := qwen38LiveConfig(t)
|
||||
proxy := newQwen38LiveProxy(t, baseURL)
|
||||
client := OpenAIClient{Corrections: NewParameterCorrectionCache()}
|
||||
tests := []struct {
|
||||
name string
|
||||
input map[string]any
|
||||
stream bool
|
||||
expected map[string]any
|
||||
absent []string
|
||||
allowCompatibilityRetry bool
|
||||
}{
|
||||
{name: "none", input: map[string]any{"reasoning_effort": "none"}, expected: map[string]any{"reasoning_effort": "low", "enable_thinking": true}},
|
||||
{name: "minimal stream", input: map[string]any{"reasoning_effort": "minimal"}, stream: true, expected: map[string]any{"reasoning_effort": "low", "enable_thinking": true}},
|
||||
{name: "medium", input: map[string]any{"reasoning_effort": "medium"}, expected: map[string]any{"reasoning_effort": "medium", "enable_thinking": true}},
|
||||
{name: "high stream", input: map[string]any{"reasoning_effort": "high"}, stream: true, expected: map[string]any{"reasoning_effort": "xhigh", "enable_thinking": true}},
|
||||
{name: "max", input: map[string]any{"reasoning_effort": "max"}, expected: map[string]any{"reasoning_effort": "xhigh", "enable_thinking": true}},
|
||||
{name: "forced thinking stream", input: map[string]any{"enable_thinking": false}, stream: true, expected: map[string]any{"reasoning_effort": "low", "enable_thinking": true}},
|
||||
{name: "budget wins", input: map[string]any{"reasoning_effort": "high", "thinking_budget": 128}, expected: map[string]any{"thinking_budget": float64(128), "enable_thinking": true}, absent: []string{"reasoning_effort"}, allowCompatibilityRetry: true},
|
||||
{name: "temperature floor stream", input: map[string]any{"temperature": 0.1}, stream: true, expected: map[string]any{"temperature": 0.6, "enable_thinking": true}},
|
||||
}
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
before := len(proxy.requestShapes())
|
||||
streamedReasoning := ""
|
||||
request := qwen38LiveRequest(proxy, apiKey, qwen38LiveBody(test.input), test.stream)
|
||||
if test.stream {
|
||||
request.StreamDelta = func(event StreamDeltaEvent) error {
|
||||
streamedReasoning += event.ReasoningContent
|
||||
return nil
|
||||
}
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Minute)
|
||||
defer cancel()
|
||||
response, err := client.Run(ctx, request)
|
||||
if err != nil {
|
||||
shapes := proxy.requestShapes()
|
||||
t.Fatalf("real Qwen3.8 request failed: %v; safe request shapes=%+v", err, shapes[before:])
|
||||
}
|
||||
if !qwen38ResponseHasReasoning(response, streamedReasoning) {
|
||||
t.Fatal("real Qwen3.8 response did not contain reasoning_content")
|
||||
}
|
||||
shapes := proxy.requestShapes()
|
||||
submissions := len(shapes) - before
|
||||
if submissions != 1 && !(test.allowCompatibilityRetry && submissions == 2) {
|
||||
t.Fatalf("unexpected submission count, before=%d after=%d safe request shapes=%+v", before, len(shapes), shapes[before:])
|
||||
}
|
||||
shape := shapes[before]
|
||||
for key, expected := range test.expected {
|
||||
if normalizedJSONScalar(shape[key]) != normalizedJSONScalar(expected) {
|
||||
t.Fatalf("%s expected %#v, got %#v in shape %+v", key, expected, shape[key], shape)
|
||||
}
|
||||
}
|
||||
for _, key := range test.absent {
|
||||
if _, ok := shape[key]; ok {
|
||||
t.Fatalf("%s must be absent in shape %+v", key, shape)
|
||||
}
|
||||
}
|
||||
t.Logf("case=%s response_id=%s correction=%t cache_hit=false", test.name, response.RequestID, submissions > 1)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestQwen38LiveCorrectionCacheAndRestart(t *testing.T) {
|
||||
baseURL, apiKey := qwen38LiveConfig(t)
|
||||
proxy := newQwen38LiveProxy(t, baseURL)
|
||||
request := qwen38LiveRequest(proxy, apiKey, qwen38LiveBody(map[string]any{"reasoning_effort": "none", "n": 2}), false)
|
||||
cache := NewParameterCorrectionCache()
|
||||
client := OpenAIClient{Corrections: cache}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Minute)
|
||||
defer cancel()
|
||||
first, err := client.Run(ctx, request)
|
||||
if err != nil || !qwen38ResponseHasReasoning(first, "") {
|
||||
t.Fatalf("cold corrected request failed or lacked reasoning: %v", err)
|
||||
}
|
||||
second, err := client.Run(ctx, request)
|
||||
if err != nil || !qwen38ResponseHasReasoning(second, "") {
|
||||
t.Fatalf("cached corrected request failed or lacked reasoning: %v", err)
|
||||
}
|
||||
shapes := proxy.requestShapes()
|
||||
if len(shapes) != 3 || normalizedJSONScalar(shapes[0]["n"]) != float64(2) || normalizedJSONScalar(shapes[1]["n"]) != float64(1) || normalizedJSONScalar(shapes[2]["n"]) != float64(1) {
|
||||
t.Fatalf("expected cold n=2 -> corrected n=1 -> cached n=1, got %+v", shapes)
|
||||
}
|
||||
|
||||
restarted := OpenAIClient{Corrections: NewParameterCorrectionCache()}
|
||||
third, err := restarted.Run(ctx, request)
|
||||
if err != nil || !qwen38ResponseHasReasoning(third, "") {
|
||||
t.Fatalf("post-restart corrected request failed or lacked reasoning: %v", err)
|
||||
}
|
||||
shapes = proxy.requestShapes()
|
||||
if len(shapes) != 5 || normalizedJSONScalar(shapes[3]["n"]) != float64(2) || normalizedJSONScalar(shapes[4]["n"]) != float64(1) {
|
||||
t.Fatalf("new process cache should relearn n rule, got %+v", shapes)
|
||||
}
|
||||
t.Logf("response_ids=%s,%s,%s correction=n:set:1 cache_hit_sequence=false,true,false", first.RequestID, second.RequestID, third.RequestID)
|
||||
}
|
||||
|
||||
func TestQwen38LiveNativeResponsesNormalization(t *testing.T) {
|
||||
baseURL, apiKey := qwen38LiveConfig(t)
|
||||
proxy := newQwen38LiveProxy(t, baseURL)
|
||||
client := OpenAIClient{Corrections: NewParameterCorrectionCache()}
|
||||
request := Request{
|
||||
Kind: "responses",
|
||||
Model: "Qwen3.8-Max-Preview",
|
||||
UpstreamProtocol: ProtocolOpenAIResponses,
|
||||
Body: map[string]any{
|
||||
"input": "Reply only OK", "max_output_tokens": 256,
|
||||
"reasoning": map[string]any{"effort": "none"},
|
||||
},
|
||||
Candidate: store.RuntimeModelCandidate{
|
||||
Provider: "aliyun-bailian-openai", BaseURL: proxy.server.URL,
|
||||
ProviderModelName: "Qwen3.8-Max-Preview", Credentials: map[string]any{"apiKey": apiKey},
|
||||
Capabilities: map[string]any{"text_generate": map[string]any{
|
||||
"supportThinking": true, "supportThinkingModeSwitch": false,
|
||||
"thinkingEffortLevels": []any{"low", "medium", "xhigh"},
|
||||
}},
|
||||
},
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Minute)
|
||||
defer cancel()
|
||||
response, err := client.Run(ctx, request)
|
||||
if err != nil {
|
||||
t.Fatalf("real Qwen3.8 Responses request failed: %v; safe shapes=%+v", err, proxy.requestShapes())
|
||||
}
|
||||
shapes := proxy.requestShapes()
|
||||
if len(shapes) != 1 {
|
||||
t.Fatalf("expected one Responses submission, got %+v", shapes)
|
||||
}
|
||||
reasoning, _ := shapes[0]["reasoning"].(map[string]any)
|
||||
if reasoning["effort"] != "low" {
|
||||
t.Fatalf("Responses none was not normalized to low: %+v", shapes)
|
||||
}
|
||||
t.Logf("case=responses-none response_id=%s correction=false cache_hit=false", response.RequestID)
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -74,7 +74,7 @@ func TestOpenAIResponsesChatFallbackPreservesHistoryToolsUsageAndReasoningIntern
|
||||
if len(messages) != 4 {
|
||||
t.Fatalf("expected prior user/assistant plus current system/user/tool messages, got %+v", messages)
|
||||
}
|
||||
if body["max_tokens"] != float64(128) || body["reasoning_effort"] != "high" {
|
||||
if body["max_completion_tokens"] != float64(128) || body["reasoning_effort"] != "high" {
|
||||
t.Fatalf("expected mapped max/reasoning fields: %+v", body)
|
||||
}
|
||||
tools, _ := body["tools"].([]any)
|
||||
@@ -165,6 +165,27 @@ func TestOpenAIResponsesNativeStreamForwardsEventsWithoutChatAggregation(t *test
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesNativeStreamAcceptsIncompleteTerminalEvent(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
_, _ = w.Write([]byte("event: response.incomplete\ndata: {\"type\":\"response.incomplete\",\"response\":{\"id\":\"resp_incomplete\",\"object\":\"response\",\"status\":\"incomplete\",\"output\":[],\"incomplete_details\":{\"reason\":\"max_output_tokens\"}}}\n\n"))
|
||||
}))
|
||||
defer server.Close()
|
||||
events := make([]StreamDeltaEvent, 0, 1)
|
||||
response, err := (OpenAIClient{}).Run(context.Background(), Request{
|
||||
Kind: "responses", Model: "Demo", Body: map[string]any{"input": "hello", "stream": true}, Stream: true,
|
||||
Candidate: store.RuntimeModelCandidate{BaseURL: server.URL, ProviderModelName: "demo", Credentials: map[string]any{"apiKey": "secret"}},
|
||||
UpstreamProtocol: ProtocolOpenAIResponses,
|
||||
StreamDelta: func(event StreamDeltaEvent) error { events = append(events, event); return nil },
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(events) != 1 || events[0].Event["type"] != "response.incomplete" || response.Result["status"] != "incomplete" || response.UpstreamResponseID != "resp_incomplete" {
|
||||
t.Fatalf("native incomplete terminal event was not preserved: response=%+v events=%+v", response, events)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesChatFallbackStreamsFunctionArgumentFragments(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/chat/completions" {
|
||||
@@ -220,6 +241,54 @@ func TestOpenAIResponsesChatFallbackStreamsFunctionArgumentFragments(t *testing.
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesChatFallbackStreamsCustomRefusalAndLogprobsEndToEnd(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
_, _ = w.Write([]byte("data: {\"id\":\"chatcmpl-custom\",\"object\":\"chat.completion.chunk\",\"model\":\"demo\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"hi\",\"refusal\":\"no\",\"annotations\":[{\"type\":\"url_citation\"}],\"tool_calls\":[{\"index\":0,\"id\":\"call_custom\",\"type\":\"custom\",\"custom\":{\"name\":\"shell\",\"input\":\"pw\"}}]},\"logprobs\":{\"content\":[{\"token\":\"hi\",\"logprob\":-0.1}]},\"finish_reason\":null}]}\n\n"))
|
||||
_, _ = w.Write([]byte("data: {\"id\":\"chatcmpl-custom\",\"object\":\"chat.completion.chunk\",\"model\":\"demo\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,\"type\":\"custom\",\"custom\":{\"input\":\"d\"}}]},\"finish_reason\":\"tool_calls\"}],\"usage\":{\"prompt_tokens\":2,\"completion_tokens\":3,\"total_tokens\":5}}\n\n"))
|
||||
_, _ = w.Write([]byte("data: [DONE]\n\n"))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
events := make([]StreamDeltaEvent, 0)
|
||||
response, err := (OpenAIClient{}).Run(context.Background(), Request{
|
||||
Kind: "responses", Model: "Demo", Body: map[string]any{
|
||||
"input": "call it", "stream": true, "include": []any{"message.output_text.logprobs"},
|
||||
"tools": []any{map[string]any{"type": "custom", "name": "shell", "format": map[string]any{"type": "text"}}},
|
||||
}, Stream: true,
|
||||
Candidate: store.RuntimeModelCandidate{BaseURL: server.URL, ProviderModelName: "demo", Credentials: map[string]any{"apiKey": "secret"}},
|
||||
UpstreamProtocol: ProtocolOpenAIChatCompletions, PublicResponseID: "resp_12345678901234567890123456789012",
|
||||
StreamDelta: func(event StreamDeltaEvent) error { events = append(events, event); return nil },
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
output, _ := response.Result["output"].([]any)
|
||||
if len(output) != 2 {
|
||||
t.Fatalf("expected message and custom tool call: %+v", response.Result)
|
||||
}
|
||||
message, _ := output[0].(map[string]any)
|
||||
parts, _ := message["content"].([]any)
|
||||
text, _ := parts[0].(map[string]any)
|
||||
refusal, _ := parts[1].(map[string]any)
|
||||
custom, _ := output[1].(map[string]any)
|
||||
if custom["type"] != "custom_tool_call" || custom["call_id"] != "call_custom" || custom["input"] != "pwd" {
|
||||
t.Fatalf("custom tool stream was not aggregated: %+v", output)
|
||||
}
|
||||
if refusal["refusal"] != "no" || len(text["logprobs"].([]any)) != 1 || len(text["annotations"].([]any)) != 1 {
|
||||
t.Fatalf("refusal/logprobs were not aggregated: %+v", message)
|
||||
}
|
||||
types := make([]string, 0, len(events))
|
||||
for _, event := range events {
|
||||
types = append(types, stringFromAny(event.Event["type"]))
|
||||
}
|
||||
for _, required := range []string{"response.custom_tool_call_input.delta", "response.custom_tool_call_input.done", "response.refusal.delta", "response.refusal.done", "response.completed"} {
|
||||
if !containsTestString(types, required) {
|
||||
t.Fatalf("missing %s in converted stream: %v", required, types)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatResultToResponseMapsIncompleteFinishReason(t *testing.T) {
|
||||
response := ChatResultToResponse(map[string]any{
|
||||
"choices": []any{map[string]any{"finish_reason": "length", "message": map[string]any{"role": "assistant", "content": "partial"}}},
|
||||
@@ -233,6 +302,23 @@ func TestChatResultToResponseMapsIncompleteFinishReason(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatResultToResponseConvertsLegacyFunctionCall(t *testing.T) {
|
||||
response := ChatResultToResponse(map[string]any{
|
||||
"choices": []any{map[string]any{"finish_reason": "function_call", "message": map[string]any{
|
||||
"role": "assistant", "content": nil,
|
||||
"function_call": map[string]any{"name": "legacy_lookup", "arguments": "{\"city\":\"Paris\"}"},
|
||||
}}},
|
||||
}, "resp_12345678901234567890123456789012", "demo", map[string]any{})
|
||||
output, _ := response["output"].([]any)
|
||||
if len(output) != 1 {
|
||||
t.Fatalf("legacy function_call was lost: %+v", response)
|
||||
}
|
||||
call, _ := output[0].(map[string]any)
|
||||
if call["type"] != "function_call" || call["name"] != "legacy_lookup" || call["arguments"] != "{\"city\":\"Paris\"}" {
|
||||
t.Fatalf("legacy function_call was not converted: %+v", call)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNativeResponsesStreamSupportsMultilineDataFrames(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
@@ -277,6 +363,7 @@ func TestChatResponsesStreamAdapterEmitsTextItemLifecycle(t *testing.T) {
|
||||
want := []string{
|
||||
"response.created", "response.in_progress", "response.output_item.added", "response.content_part.added",
|
||||
"response.output_text.delta", "response.output_text.done", "response.content_part.done", "response.output_item.done",
|
||||
"response.completed",
|
||||
}
|
||||
if strings.Join(events, ",") != strings.Join(want, ",") {
|
||||
t.Fatalf("unexpected text event lifecycle got=%v want=%v", events, want)
|
||||
@@ -285,8 +372,8 @@ func TestChatResponsesStreamAdapterEmitsTextItemLifecycle(t *testing.T) {
|
||||
|
||||
func TestResponsesChatFallbackRejectsBuiltInToolsAndUnknownParameters(t *testing.T) {
|
||||
_, err := ResponsesRequestToChat(map[string]any{"input": "hello", "tools": []any{map[string]any{"type": "web_search_preview"}}}, nil)
|
||||
if ErrorCode(err) != "unsupported_response_tool" {
|
||||
t.Fatalf("expected unsupported_response_tool, got %v", err)
|
||||
if ErrorCode(err) != "unsupported_response_parameter" || ErrorParam(err) != "tools[0].type" {
|
||||
t.Fatalf("expected precise unsupported_response_parameter, got %v param=%q", err, ErrorParam(err))
|
||||
}
|
||||
_, err = ResponsesRequestToChat(map[string]any{"input": "hello", "conversation": "conv_1"}, nil)
|
||||
if ErrorCode(err) != "unsupported_response_parameter" {
|
||||
@@ -317,6 +404,223 @@ func TestResponsesChatFallbackKeepsClientManagedStateAuthoritative(t *testing.T)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResponsesChatFallbackMapsMultimodalCustomToolsAndCallHistory(t *testing.T) {
|
||||
body, err := ResponsesRequestToChat(map[string]any{
|
||||
"input": []any{
|
||||
map[string]any{"type": "message", "role": "user", "content": []any{
|
||||
map[string]any{"type": "input_text", "text": "hello", "prompt_cache_breakpoint": map[string]any{"type": "ephemeral"}},
|
||||
map[string]any{"type": "input_image", "image_url": "https://example.com/a.png", "detail": "high"},
|
||||
map[string]any{"type": "input_file", "file_id": "file_1", "filename": "a.pdf"},
|
||||
map[string]any{"type": "input_audio", "input_audio": map[string]any{"data": "AAAA", "format": "wav"}},
|
||||
}},
|
||||
map[string]any{"type": "function_call", "call_id": "call_fn", "name": "lookup", "arguments": "{\"x\":1}"},
|
||||
map[string]any{"type": "function_call_output", "call_id": "call_fn", "output": map[string]any{"ok": true}},
|
||||
map[string]any{"type": "custom_tool_call", "call_id": "call_custom", "name": "shell", "input": "pwd"},
|
||||
map[string]any{"type": "custom_tool_call_output", "call_id": "call_custom", "output": "done"},
|
||||
},
|
||||
"tools": []any{
|
||||
map[string]any{"type": "function", "name": "lookup", "parameters": map[string]any{"type": "object"}, "strict": true},
|
||||
map[string]any{"type": "custom", "name": "shell", "description": "run", "format": map[string]any{"type": "text"}},
|
||||
},
|
||||
"tool_choice": map[string]any{"type": "custom", "name": "shell"},
|
||||
"include": []any{"message.output_text.logprobs"}, "max_output_tokens": 256,
|
||||
}, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if body["max_completion_tokens"] != 256 || body["logprobs"] != true {
|
||||
t.Fatalf("token/logprobs mapping is incomplete: %+v", body)
|
||||
}
|
||||
messages, _ := body["messages"].([]any)
|
||||
if len(messages) != 5 {
|
||||
t.Fatalf("call history conversion changed item ownership: %+v", messages)
|
||||
}
|
||||
first, _ := messages[0].(map[string]any)
|
||||
parts, _ := first["content"].([]any)
|
||||
if len(parts) != 4 {
|
||||
t.Fatalf("multimodal content was dropped: %+v", parts)
|
||||
}
|
||||
text, _ := parts[0].(map[string]any)
|
||||
image, _ := parts[1].(map[string]any)
|
||||
imageURL, _ := image["image_url"].(map[string]any)
|
||||
file, _ := parts[2].(map[string]any)
|
||||
fileValue, _ := file["file"].(map[string]any)
|
||||
if text["prompt_cache_breakpoint"] == nil || imageURL["detail"] != "high" || fileValue["file_id"] != "file_1" || fileValue["filename"] != "a.pdf" {
|
||||
t.Fatalf("multimodal nested fields changed: %+v", parts)
|
||||
}
|
||||
customMessage, _ := messages[3].(map[string]any)
|
||||
customCalls, _ := customMessage["tool_calls"].([]any)
|
||||
customCall, _ := customCalls[0].(map[string]any)
|
||||
custom, _ := customCall["custom"].(map[string]any)
|
||||
if customCall["type"] != "custom" || custom["name"] != "shell" || custom["input"] != "pwd" {
|
||||
t.Fatalf("custom call history changed: %+v", customMessage)
|
||||
}
|
||||
tools, _ := body["tools"].([]any)
|
||||
customTool, _ := tools[1].(map[string]any)
|
||||
if customTool["type"] != "custom" || customTool["custom"] == nil {
|
||||
t.Fatalf("custom tool definition changed: %+v", tools)
|
||||
}
|
||||
choice, _ := body["tool_choice"].(map[string]any)
|
||||
if choice["type"] != "custom" || choice["custom"] == nil {
|
||||
t.Fatalf("custom tool choice changed: %+v", choice)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResponsesChatFallbackMapsTextToolOutputAndRejectsNativeOnlyToolOutputContent(t *testing.T) {
|
||||
body, err := ResponsesRequestToChat(map[string]any{"input": []any{map[string]any{
|
||||
"type": "function_call_output", "call_id": "call_1", "output": []any{map[string]any{
|
||||
"type": "input_text", "text": "done", "prompt_cache_breakpoint": map[string]any{"mode": "explicit"},
|
||||
}},
|
||||
}}}, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
messages, _ := body["messages"].([]any)
|
||||
message, _ := messages[0].(map[string]any)
|
||||
parts, _ := message["content"].([]any)
|
||||
part, _ := parts[0].(map[string]any)
|
||||
if part["type"] != "text" || part["text"] != "done" || part["prompt_cache_breakpoint"] == nil {
|
||||
t.Fatalf("text tool output was not preserved: %+v", message)
|
||||
}
|
||||
|
||||
err = ValidateResponsesChatFallback(map[string]any{"input": []any{map[string]any{
|
||||
"type": "function_call_output", "call_id": "call_1", "output": []any{map[string]any{
|
||||
"type": "input_image", "image_url": "https://example.com/a.png", "detail": "high",
|
||||
}},
|
||||
}}})
|
||||
if ErrorCode(err) != "unsupported_response_parameter" || ErrorParam(err) != "input[0].output[0].type" {
|
||||
t.Fatalf("native-only tool output must be rejected precisely, got %v param=%q", err, ErrorParam(err))
|
||||
}
|
||||
}
|
||||
|
||||
func TestResponsesChatFallbackAcceptsOnlyNoopNativeFields(t *testing.T) {
|
||||
if err := ValidateResponsesChatFallback(map[string]any{
|
||||
"input": "hello", "background": false, "context_management": map[string]any{}, "conversation": "",
|
||||
"prompt": nil, "truncation": "disabled", "reasoning": map[string]any{"summary": nil},
|
||||
}); err != nil {
|
||||
t.Fatalf("no-op native fields should be accepted: %v", err)
|
||||
}
|
||||
for param, value := range map[string]any{
|
||||
"background": true, "conversation": "conv_1", "max_tool_calls": 2, "prompt": map[string]any{"id": "pmpt_1"}, "truncation": "auto",
|
||||
} {
|
||||
err := ValidateResponsesChatFallback(map[string]any{"input": "hello", param: value})
|
||||
if ErrorCode(err) != "unsupported_response_parameter" || ErrorParam(err) != param {
|
||||
t.Fatalf("expected precise rejection for %s, got %v param=%q", param, err, ErrorParam(err))
|
||||
}
|
||||
}
|
||||
err := ValidateResponsesChatFallback(map[string]any{"input": []any{map[string]any{
|
||||
"type": "custom_tool_call", "call_id": "call_1", "name": "tool", "input": "x", "namespace": "native-only",
|
||||
}}})
|
||||
if ErrorCode(err) != "unsupported_response_parameter" || ErrorParam(err) != "input[0].namespace" {
|
||||
t.Fatalf("custom call namespace must be rejected precisely, got %v param=%q", err, ErrorParam(err))
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatResultToResponseBuildsCompleteResponseWithRefusalLogprobsAndCustomTool(t *testing.T) {
|
||||
request := map[string]any{
|
||||
"background": false, "conversation": nil, "instructions": "be concise", "max_output_tokens": 300,
|
||||
"max_tool_calls": nil, "metadata": map[string]any{"trace": "1"}, "moderation": map[string]any{"type": "auto"},
|
||||
"parallel_tool_calls": true, "previous_response_id": "resp_parent", "prompt": nil, "prompt_cache_key": "cache",
|
||||
"prompt_cache_options": map[string]any{"type": "ephemeral"}, "prompt_cache_retention": "in_memory",
|
||||
"reasoning": map[string]any{"effort": "low"}, "safety_identifier": "safe", "service_tier": "default",
|
||||
"temperature": 0.4, "text": map[string]any{"verbosity": "low"}, "tool_choice": "auto", "tools": []any{},
|
||||
"top_logprobs": 2, "top_p": 0.9, "truncation": "disabled", "user": "user-1",
|
||||
}
|
||||
response := ChatResultToResponse(map[string]any{
|
||||
"created": 1710000000, "service_tier": "priority", "moderation": map[string]any{"flagged": false},
|
||||
"choices": []any{map[string]any{
|
||||
"finish_reason": "stop", "logprobs": map[string]any{"content": []any{map[string]any{"token": "hello", "logprob": -0.1}}},
|
||||
"message": map[string]any{
|
||||
"role": "assistant", "content": "hello", "refusal": "cannot continue", "annotations": []any{map[string]any{"type": "url_citation"}},
|
||||
"tool_calls": []any{
|
||||
map[string]any{"id": "call_fn", "type": "function", "function": map[string]any{"name": "lookup", "arguments": "{\"x\":1}"}},
|
||||
map[string]any{"id": "call_custom", "type": "custom", "custom": map[string]any{"name": "shell", "input": "pwd"}},
|
||||
},
|
||||
},
|
||||
}},
|
||||
"usage": map[string]any{"prompt_tokens": 2, "completion_tokens": 3, "total_tokens": 5},
|
||||
}, "resp_12345678901234567890123456789012", "demo", request)
|
||||
for _, key := range []string{
|
||||
"id", "object", "created_at", "status", "completed_at", "error", "incomplete_details", "instructions", "metadata", "model", "output",
|
||||
"parallel_tool_calls", "temperature", "tool_choice", "tools", "top_p", "background", "conversation", "max_output_tokens", "max_tool_calls",
|
||||
"moderation", "output_text", "previous_response_id", "prompt", "prompt_cache_key", "prompt_cache_options", "prompt_cache_retention", "reasoning",
|
||||
"safety_identifier", "service_tier", "text", "top_logprobs", "truncation", "usage", "user",
|
||||
} {
|
||||
if _, ok := response[key]; !ok {
|
||||
t.Fatalf("complete fallback Response omitted %q: %+v", key, response)
|
||||
}
|
||||
}
|
||||
output, _ := response["output"].([]any)
|
||||
if len(output) != 3 {
|
||||
t.Fatalf("expected message, function, and custom output items: %+v", output)
|
||||
}
|
||||
message, _ := output[0].(map[string]any)
|
||||
parts, _ := message["content"].([]any)
|
||||
text, _ := parts[0].(map[string]any)
|
||||
refusal, _ := parts[1].(map[string]any)
|
||||
custom, _ := output[2].(map[string]any)
|
||||
if len(text["logprobs"].([]any)) != 1 || len(text["annotations"].([]any)) != 1 || refusal["type"] != "refusal" || custom["type"] != "custom_tool_call" || custom["input"] != "pwd" {
|
||||
t.Fatalf("Response content/tool details changed: %+v", output)
|
||||
}
|
||||
if response["service_tier"] != "priority" {
|
||||
t.Fatalf("upstream response service tier was not preserved: %+v", response)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatResponsesStreamAdapterEmitsRefusalLogprobsAndCustomToolLifecycle(t *testing.T) {
|
||||
adapter := newChatResponsesStreamAdapter("resp_12345678901234567890123456789012", "demo")
|
||||
events := make([]map[string]any, 0)
|
||||
onDelta := func(event StreamDeltaEvent) error {
|
||||
events = append(events, event.Event)
|
||||
return nil
|
||||
}
|
||||
if err := adapter.delta(StreamDeltaEvent{Event: map[string]any{"choices": []any{map[string]any{
|
||||
"logprobs": map[string]any{"content": []any{map[string]any{"token": "hi", "logprob": -0.1}}},
|
||||
"delta": map[string]any{
|
||||
"content": "hi", "refusal": "no",
|
||||
"tool_calls": []any{map[string]any{"index": 0, "id": "call_custom", "type": "custom", "custom": map[string]any{"name": "shell", "input": "pw"}}},
|
||||
},
|
||||
}}}}, onDelta); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := adapter.delta(StreamDeltaEvent{Event: map[string]any{"choices": []any{map[string]any{
|
||||
"delta": map[string]any{"tool_calls": []any{map[string]any{"index": 0, "type": "custom", "custom": map[string]any{"input": "d"}}}},
|
||||
}}}}, onDelta); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
result := ChatResultToResponse(map[string]any{"choices": []any{map[string]any{
|
||||
"message": map[string]any{"role": "assistant", "content": "hi", "refusal": "no", "tool_calls": []any{
|
||||
map[string]any{"id": "call_custom", "type": "custom", "custom": map[string]any{"name": "shell", "input": "pwd"}},
|
||||
}},
|
||||
}}}, adapter.publicID, "demo", map[string]any{})
|
||||
if err := adapter.done(result, onDelta); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
types := make([]string, 0, len(events))
|
||||
var sawLogprobs bool
|
||||
for index, event := range events {
|
||||
types = append(types, stringFromAny(event["type"]))
|
||||
if intFromAny(event["sequence_number"]) != index {
|
||||
t.Fatalf("non-contiguous sequence at %d: %+v", index, event)
|
||||
}
|
||||
if event["type"] == "response.output_text.delta" {
|
||||
logprobs, _ := event["logprobs"].([]any)
|
||||
sawLogprobs = len(logprobs) == 1
|
||||
}
|
||||
}
|
||||
for _, required := range []string{
|
||||
"response.refusal.delta", "response.refusal.done", "response.custom_tool_call_input.delta",
|
||||
"response.custom_tool_call_input.done", "response.output_item.done", "response.completed",
|
||||
} {
|
||||
if !containsTestString(types, required) {
|
||||
t.Fatalf("missing %s lifecycle event: %v", required, types)
|
||||
}
|
||||
}
|
||||
if !sawLogprobs {
|
||||
t.Fatalf("streamed output_text delta lost logprobs: %+v", events)
|
||||
}
|
||||
}
|
||||
|
||||
func mustJSON(t *testing.T, value any) []byte {
|
||||
t.Helper()
|
||||
data, err := json.Marshal(value)
|
||||
|
||||
@@ -0,0 +1,51 @@
|
||||
package clients
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func TestLocalUnsupportedKindDoesNotMarkUpstreamSubmissionStarted(t *testing.T) {
|
||||
started, responseReceived := false, false
|
||||
_, err := (VolcesClient{}).Run(context.Background(), Request{
|
||||
Kind: "speech.generations",
|
||||
Candidate: store.RuntimeModelCandidate{
|
||||
Provider: "volces", Credentials: map[string]any{"apiKey": "test"},
|
||||
},
|
||||
OnUpstreamSubmissionStarted: func() error { started = true; return nil },
|
||||
OnUpstreamResponseReceived: func() error { responseReceived = true; return nil },
|
||||
})
|
||||
if err == nil || ErrorCode(err) != "unsupported_kind" {
|
||||
t.Fatalf("expected local unsupported_kind, got %v", err)
|
||||
}
|
||||
if started || responseReceived {
|
||||
t.Fatalf("local validation must remain not_submitted: started=%v response=%v", started, responseReceived)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNetworkDisconnectMarksSubmittingWithoutResponseReceived(t *testing.T) {
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}))
|
||||
baseURL := upstream.URL
|
||||
upstream.Close()
|
||||
started, responseReceived := false, false
|
||||
_, err := (VolcesClient{HTTPClient: upstream.Client()}).Run(context.Background(), Request{
|
||||
Kind: "images.generations",
|
||||
Body: map[string]any{"prompt": "test"},
|
||||
Candidate: store.RuntimeModelCandidate{
|
||||
Provider: "volces", BaseURL: baseURL, ModelName: "seedream-test",
|
||||
Credentials: map[string]any{"apiKey": "test"},
|
||||
},
|
||||
OnUpstreamSubmissionStarted: func() error { started = true; return nil },
|
||||
OnUpstreamResponseReceived: func() error { responseReceived = true; return nil },
|
||||
})
|
||||
if err == nil || ErrorCode(err) != "network" {
|
||||
t.Fatalf("expected network error, got %v", err)
|
||||
}
|
||||
if !started || responseReceived {
|
||||
t.Fatalf("network disconnect must remain submitting: started=%v response=%v", started, responseReceived)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,10 @@
|
||||
[
|
||||
{ "name": "none-to-low", "input": { "reasoning_effort": "none" }, "expected": { "enable_thinking": true, "reasoning_effort": "low" } },
|
||||
{ "name": "minimal-to-low", "input": { "reasoning_effort": "minimal" }, "expected": { "enable_thinking": true, "reasoning_effort": "low" } },
|
||||
{ "name": "medium", "input": { "reasoning_effort": "medium" }, "expected": { "enable_thinking": true, "reasoning_effort": "medium" } },
|
||||
{ "name": "high-to-xhigh", "input": { "reasoning_effort": "high" }, "expected": { "enable_thinking": true, "reasoning_effort": "xhigh" } },
|
||||
{ "name": "max-to-xhigh", "input": { "reasoning_effort": "max" }, "expected": { "enable_thinking": true, "reasoning_effort": "xhigh" } },
|
||||
{ "name": "force-thinking", "input": { "enable_thinking": false }, "expected": { "enable_thinking": true, "reasoning_effort": "low" } },
|
||||
{ "name": "budget-wins", "input": { "reasoning_effort": "high", "thinking_budget": 300000 }, "expected": { "enable_thinking": true, "thinking_budget": 262144 } },
|
||||
{ "name": "temperature-floor", "input": { "temperature": 0.1 }, "expected": { "enable_thinking": true, "temperature": 0.6 } }
|
||||
]
|
||||
@@ -0,0 +1,617 @@
|
||||
package clients
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/md5"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"math"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
const (
|
||||
topazDefaultPollInterval = 15 * time.Second
|
||||
topazDefaultPollTimeout = 60 * time.Minute
|
||||
topazDefaultMaxInput = int64(2 << 30)
|
||||
)
|
||||
|
||||
type TopazClient struct {
|
||||
HTTPClient *http.Client
|
||||
Probe func(context.Context, string) (TopazVideoMetadata, error)
|
||||
}
|
||||
|
||||
type TopazVideoMetadata struct {
|
||||
Width int
|
||||
Height int
|
||||
Duration float64
|
||||
FrameRate float64
|
||||
FrameCount int
|
||||
HasAudio bool
|
||||
}
|
||||
|
||||
type topazSource struct {
|
||||
Path string
|
||||
Container string
|
||||
Size int64
|
||||
MD5 string
|
||||
Resolution map[string]any
|
||||
Duration float64
|
||||
FrameRate float64
|
||||
FrameCount int
|
||||
}
|
||||
|
||||
func (c TopazClient) Run(ctx context.Context, request Request) (Response, error) {
|
||||
startedAt := time.Now()
|
||||
apiKey := credential(request.Candidate.Credentials, "apiKey", "api_key", "key", "token")
|
||||
if apiKey == "" {
|
||||
return Response{}, &ClientError{Code: "missing_credentials", Message: "Topaz API key is required", Retryable: false}
|
||||
}
|
||||
requestID := strings.TrimSpace(request.RemoteTaskID)
|
||||
payload := cloneMapAny(request.RemoteTaskPayload)
|
||||
var source *topazSource
|
||||
var target map[string]int
|
||||
if requestID == "" || strings.TrimSpace(firstNonEmptyString(payload["phase"])) != "uploaded" {
|
||||
videoURL := strings.TrimSpace(firstNonEmptyString(request.Body["video_url"], request.Body["videoUrl"]))
|
||||
if videoURL == "" {
|
||||
return Response{}, &ClientError{Code: "invalid_parameter", Message: "video_url is required", Param: "video_url", StatusCode: http.StatusBadRequest}
|
||||
}
|
||||
prepared, err := c.prepareSource(ctx, request, videoURL)
|
||||
if err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
source = &prepared
|
||||
defer os.Remove(prepared.Path)
|
||||
target = topazTargetResolution(request.Body, prepared.Resolution)
|
||||
if requestID == "" {
|
||||
created, err := c.createRequest(ctx, request, apiKey, prepared, target)
|
||||
if err != nil {
|
||||
return Response{}, annotateResponseError(err, "", startedAt, time.Now())
|
||||
}
|
||||
requestID = strings.TrimSpace(firstNonEmptyString(created["requestId"], created["request_id"], created["id"]))
|
||||
if requestID == "" {
|
||||
return Response{}, &ClientError{Code: "invalid_response", Message: "Topaz create response is missing requestId", Retryable: false}
|
||||
}
|
||||
payload = map[string]any{"phase": "created", "targetResolution": target}
|
||||
if request.OnRemoteTaskSubmitted != nil {
|
||||
if err := request.OnRemoteTaskSubmitted(requestID, payload); err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
}
|
||||
}
|
||||
if err := c.uploadSource(ctx, request, apiKey, requestID, prepared); err != nil {
|
||||
return Response{}, annotateResponseError(err, requestID, startedAt, time.Now())
|
||||
}
|
||||
payload = map[string]any{"phase": "uploaded", "targetResolution": target}
|
||||
if request.OnRemoteTaskPolled != nil {
|
||||
if err := request.OnRemoteTaskPolled(requestID, payload); err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
}
|
||||
}
|
||||
if target == nil {
|
||||
target = topazTargetFromPayload(payload, request.Body)
|
||||
}
|
||||
completed, err := c.poll(ctx, request, apiKey, requestID, target)
|
||||
if err != nil {
|
||||
return Response{}, annotateResponseError(err, requestID, startedAt, time.Now())
|
||||
}
|
||||
outputURL := topazDownloadURL(completed)
|
||||
if outputURL == "" {
|
||||
return Response{}, &ClientError{Code: "invalid_response", Message: "Topaz completed without a download URL", RequestID: requestID, Retryable: false}
|
||||
}
|
||||
metadata, err := c.probe(ctx, outputURL)
|
||||
if err != nil {
|
||||
return Response{}, &ClientError{Code: "invalid_response", Message: "cannot validate Topaz output metadata: " + err.Error(), RequestID: requestID, Retryable: true}
|
||||
}
|
||||
if target["width"] > 0 && target["height"] > 0 && (metadata.Width+4 < target["width"] || metadata.Height+4 < target["height"]) {
|
||||
return Response{}, &ClientError{Code: "invalid_response", Message: fmt.Sprintf("Topaz output resolution %dx%d does not reach target %dx%d", metadata.Width, metadata.Height, target["width"], target["height"]), RequestID: requestID, Retryable: false}
|
||||
}
|
||||
finishedAt := time.Now()
|
||||
resultItem := map[string]any{
|
||||
"type": "video",
|
||||
"url": outputURL,
|
||||
"video_url": outputURL,
|
||||
"width": metadata.Width,
|
||||
"height": metadata.Height,
|
||||
"target_resolution": fmt.Sprintf("%dx%d", target["width"], target["height"]),
|
||||
"duration": metadata.Duration,
|
||||
"target_frame_rate": metadata.FrameRate,
|
||||
"slow_motion_rate": numericValue(request.Body["slow_motion_rate"], 1),
|
||||
}
|
||||
if source != nil {
|
||||
resultItem["source_resolution"] = fmt.Sprintf("%vx%v", source.Resolution["width"], source.Resolution["height"])
|
||||
resultItem["source_frame_rate"] = source.FrameRate
|
||||
resultItem["duration"] = source.Duration
|
||||
}
|
||||
return Response{
|
||||
Result: map[string]any{
|
||||
"status": "success",
|
||||
"model": request.Model,
|
||||
"task_id": requestID,
|
||||
"upstream_task_id": requestID,
|
||||
"data": []any{resultItem},
|
||||
},
|
||||
RequestID: requestID,
|
||||
Progress: append(providerProgress(request), Progress{Phase: "polling", Progress: 0.9, Message: "Topaz upscale completed", Payload: map[string]any{"upstreamTaskId": requestID}}),
|
||||
ResponseStartedAt: startedAt,
|
||||
ResponseFinishedAt: finishedAt,
|
||||
ResponseDurationMS: responseDurationMS(startedAt, finishedAt),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (c TopazClient) prepareSource(ctx context.Context, request Request, rawURL string) (topazSource, error) {
|
||||
parsed, err := url.Parse(rawURL)
|
||||
if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") {
|
||||
return topazSource{}, &ClientError{Code: "invalid_parameter", Message: "video_url must be an http(s) URL", Param: "video_url", StatusCode: http.StatusBadRequest}
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil)
|
||||
if err != nil {
|
||||
return topazSource{}, err
|
||||
}
|
||||
downloadClient := topazSourceHTTPClient(httpClient(request.HTTPClient, c.HTTPClient), boolishDefault(request.Candidate.PlatformConfig["allowPrivateSourceDownloads"], false))
|
||||
resp, err := downloadClient.Do(req)
|
||||
if err != nil {
|
||||
return topazSource{}, transportClientError(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return topazSource{}, &ClientError{Code: "invalid_parameter", Message: "video_url download failed: " + resp.Status, Param: "video_url", StatusCode: http.StatusBadRequest}
|
||||
}
|
||||
limit := int64(numericValue(firstPresent(request.Candidate.PlatformConfig["maxInputBytes"], request.Candidate.PlatformConfig["max_input_bytes"]), float64(topazDefaultMaxInput)))
|
||||
if limit <= 0 {
|
||||
limit = topazDefaultMaxInput
|
||||
}
|
||||
tmp, err := os.CreateTemp("", "easyai-topaz-*"+topazContainerExtension(parsed.Path))
|
||||
if err != nil {
|
||||
return topazSource{}, err
|
||||
}
|
||||
path := tmp.Name()
|
||||
cleanup := func() {
|
||||
_ = tmp.Close()
|
||||
_ = os.Remove(path)
|
||||
}
|
||||
hash := md5.New()
|
||||
written, err := io.Copy(io.MultiWriter(tmp, hash), io.LimitReader(resp.Body, limit+1))
|
||||
closeErr := tmp.Close()
|
||||
if err != nil || closeErr != nil || written <= 0 || written > limit {
|
||||
cleanup()
|
||||
if written > limit {
|
||||
return topazSource{}, &ClientError{Code: "invalid_parameter", Message: "video source exceeds configured size limit", Param: "video_url", StatusCode: http.StatusBadRequest}
|
||||
}
|
||||
return topazSource{}, firstError(err, closeErr)
|
||||
}
|
||||
metadata, err := c.probe(ctx, path)
|
||||
if err != nil || metadata.Width <= 0 || metadata.Height <= 0 {
|
||||
cleanup()
|
||||
return topazSource{}, &ClientError{Code: "invalid_parameter", Message: "cannot probe source video metadata", Param: "video_url", StatusCode: http.StatusBadRequest}
|
||||
}
|
||||
frameRate := metadata.FrameRate
|
||||
if frameRate <= 0 {
|
||||
frameRate = numericValue(request.Candidate.PlatformConfig["frameRate"], 30)
|
||||
}
|
||||
duration := math.Max(1, metadata.Duration)
|
||||
frameCount := metadata.FrameCount
|
||||
if frameCount <= 0 {
|
||||
frameCount = int(math.Round(duration * frameRate))
|
||||
}
|
||||
return topazSource{
|
||||
Path: path, Container: strings.TrimPrefix(strings.ToLower(filepath.Ext(path)), "."), Size: written,
|
||||
MD5: hex.EncodeToString(hash.Sum(nil)), Resolution: map[string]any{"width": metadata.Width, "height": metadata.Height},
|
||||
Duration: duration, FrameRate: frameRate, FrameCount: frameCount,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (c TopazClient) createRequest(ctx context.Context, request Request, apiKey string, source topazSource, target map[string]int) (map[string]any, error) {
|
||||
model := topazModelName(firstNonEmptyString(request.Candidate.ProviderModelName, request.Model))
|
||||
scale := math.Max(float64(target["width"])/numericValue(source.Resolution["width"], 1), float64(target["height"])/numericValue(source.Resolution["height"], 1))
|
||||
filter := map[string]any{"model": model}
|
||||
if scale > 1 {
|
||||
key := "scale"
|
||||
if model == "slf-2" || model == "slhq-1" || model == "slm-1" || model == "slp-2.5" {
|
||||
key = "upscaling_factor"
|
||||
}
|
||||
filter[key] = math.Round(scale*1000) / 1000
|
||||
}
|
||||
preserveAudio := boolishDefault(request.Body["preserve_audio"], true)
|
||||
filters := []any{filter}
|
||||
targetFrameRate := numericValue(firstPresent(request.Body["target_frame_rate"], request.Body["output_frame_rate"]), source.FrameRate)
|
||||
slowMotionRate := numericValue(request.Body["slow_motion_rate"], 1)
|
||||
if targetFrameRate <= 0 {
|
||||
targetFrameRate = source.FrameRate
|
||||
}
|
||||
if slowMotionRate <= 0 {
|
||||
slowMotionRate = 1
|
||||
}
|
||||
if targetFrameRate != source.FrameRate || slowMotionRate > 1 {
|
||||
interpolationModel := strings.TrimSpace(firstNonEmptyString(request.Body["frame_interpolation_model"]))
|
||||
if interpolationModel == "" {
|
||||
return nil, &ClientError{Code: "invalid_parameter", Message: "frame_interpolation_model is required when output frame rate or slow motion changes", Param: "frame_interpolation_model", StatusCode: http.StatusBadRequest}
|
||||
}
|
||||
interpolation := map[string]any{"model": topazModelName(interpolationModel)}
|
||||
if targetFrameRate > 0 {
|
||||
interpolation["fps"] = targetFrameRate
|
||||
}
|
||||
if slowMotionRate > 1 {
|
||||
interpolation["slowmo"] = slowMotionRate
|
||||
}
|
||||
filters = append(filters, interpolation)
|
||||
}
|
||||
body := map[string]any{
|
||||
"source": map[string]any{"container": source.Container, "size": source.Size, "duration": source.Duration, "frameCount": source.FrameCount, "frameRate": source.FrameRate, "resolution": source.Resolution},
|
||||
"filters": filters,
|
||||
"output": map[string]any{"resolution": target, "frameRate": targetFrameRate, "audioCodec": "AAC", "audioTransfer": map[bool]string{true: "Copy", false: "None"}[preserveAudio], "dynamicCompressionLevel": "High", "videoEncoder": "H265", "videoProfile": "Main", "container": "mp4"},
|
||||
}
|
||||
return c.topazJSON(ctx, request, apiKey, http.MethodPost, "/video/", body)
|
||||
}
|
||||
|
||||
func (c TopazClient) uploadSource(ctx context.Context, request Request, apiKey, requestID string, source topazSource) error {
|
||||
accepted, err := c.topazJSON(ctx, request, apiKey, http.MethodPatch, "/video/"+url.PathEscape(requestID)+"/accept", nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
urls := stringList(accepted["urls"])
|
||||
if len(urls) == 0 {
|
||||
return &ClientError{Code: "invalid_response", Message: "Topaz accept response has no upload URLs", RequestID: requestID}
|
||||
}
|
||||
file, err := os.Open(source.Path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer file.Close()
|
||||
chunkSize := (source.Size + int64(len(urls)) - 1) / int64(len(urls))
|
||||
results := make([]any, 0, len(urls))
|
||||
for index, uploadURL := range urls {
|
||||
remaining := source.Size - int64(index)*chunkSize
|
||||
if remaining <= 0 {
|
||||
break
|
||||
}
|
||||
length := minInt64(chunkSize, remaining)
|
||||
section := io.NewSectionReader(file, int64(index)*chunkSize, length)
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPut, uploadURL, section)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
req.ContentLength = length
|
||||
req.Header.Set("Content-Type", topazContainerContentType(source.Container))
|
||||
resp, err := httpClient(request.HTTPClient, c.HTTPClient).Do(req)
|
||||
if err != nil {
|
||||
return transportClientError(err)
|
||||
}
|
||||
_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 1<<20))
|
||||
resp.Body.Close()
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return &ClientError{Code: "provider_failed", Message: "Topaz part upload failed: " + resp.Status, StatusCode: resp.StatusCode, Retryable: resp.StatusCode >= 500}
|
||||
}
|
||||
etag := strings.Trim(resp.Header.Get("ETag"), `"`)
|
||||
if etag == "" {
|
||||
return &ClientError{Code: "invalid_response", Message: "Topaz part upload is missing ETag", Retryable: false}
|
||||
}
|
||||
results = append(results, map[string]any{"partNum": index + 1, "eTag": etag})
|
||||
}
|
||||
_, err = c.topazJSON(ctx, request, apiKey, http.MethodPatch, "/video/"+url.PathEscape(requestID)+"/complete-upload/", map[string]any{"md5Hash": source.MD5, "uploadResults": results})
|
||||
return err
|
||||
}
|
||||
|
||||
func (c TopazClient) poll(ctx context.Context, request Request, apiKey, requestID string, target map[string]int) (map[string]any, error) {
|
||||
interval := universalDurationConfig(request.Candidate.PlatformConfig, topazDefaultPollInterval, "pollIntervalMs", "poll_interval_ms")
|
||||
timeout := universalDurationConfig(request.Candidate.PlatformConfig, topazDefaultPollTimeout, "pollTimeoutMs", "poll_timeout_ms", "timeoutMs")
|
||||
deadline := time.NewTimer(timeout)
|
||||
defer deadline.Stop()
|
||||
for {
|
||||
status, err := c.topazJSON(ctx, request, apiKey, http.MethodGet, "/video/"+url.PathEscape(requestID)+"/status", nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
state := strings.ToLower(strings.TrimSpace(firstNonEmptyString(status["status"], status["state"])))
|
||||
if request.OnRemoteTaskPolled != nil {
|
||||
if err := request.OnRemoteTaskPolled(requestID, map[string]any{"phase": "uploaded", "status": state, "targetResolution": target}); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
switch state {
|
||||
case "complete", "completed", "success", "succeeded":
|
||||
return status, nil
|
||||
case "failed", "failure", "cancelled", "canceled", "error":
|
||||
return nil, &ClientError{Code: "provider_failed", Message: firstNonEmptyString(status["message"], "Topaz task failed"), RequestID: requestID, Retryable: false}
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, &ClientError{Code: "cancelled", Message: ctx.Err().Error(), RequestID: requestID, Retryable: true}
|
||||
case <-deadline.C:
|
||||
return nil, &ClientError{Code: "timeout", Message: "Topaz task polling timed out", RequestID: requestID, Retryable: true}
|
||||
case <-time.After(interval):
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c TopazClient) topazJSON(ctx context.Context, request Request, apiKey, method, path string, body map[string]any) (map[string]any, error) {
|
||||
var reader io.Reader
|
||||
if body != nil {
|
||||
raw, _ := json.Marshal(body)
|
||||
reader = bytes.NewReader(raw)
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, method, providerURL(request.Candidate.BaseURL, path), reader)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("X-API-Key", apiKey)
|
||||
req.Header.Set("Accept", "application/json")
|
||||
if body != nil {
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
}
|
||||
resp, err := httpClient(request.HTTPClient, c.HTTPClient).Do(req)
|
||||
if err != nil {
|
||||
return nil, transportClientError(err)
|
||||
}
|
||||
result, decodeErr := decodeHTTPResponse(resp)
|
||||
if decodeErr != nil {
|
||||
return nil, decodeErr
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (c TopazClient) probe(ctx context.Context, target string) (TopazVideoMetadata, error) {
|
||||
if c.Probe != nil {
|
||||
return c.Probe(ctx, target)
|
||||
}
|
||||
return probeTopazVideo(ctx, target)
|
||||
}
|
||||
|
||||
func probeTopazVideo(ctx context.Context, target string) (TopazVideoMetadata, error) {
|
||||
if _, err := exec.LookPath("ffprobe"); err != nil {
|
||||
return TopazVideoMetadata{}, err
|
||||
}
|
||||
probeCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
|
||||
defer cancel()
|
||||
output, err := exec.CommandContext(probeCtx, "ffprobe", "-v", "error", "-show_entries", "format=duration:stream=codec_type,width,height,avg_frame_rate,nb_frames", "-of", "json", target).Output()
|
||||
if err != nil {
|
||||
return TopazVideoMetadata{}, err
|
||||
}
|
||||
var decoded struct {
|
||||
Format struct {
|
||||
Duration string `json:"duration"`
|
||||
} `json:"format"`
|
||||
Streams []struct {
|
||||
CodecType string `json:"codec_type"`
|
||||
Width int `json:"width"`
|
||||
Height int `json:"height"`
|
||||
AverageRate string `json:"avg_frame_rate"`
|
||||
FrameCount string `json:"nb_frames"`
|
||||
} `json:"streams"`
|
||||
}
|
||||
if err := json.Unmarshal(output, &decoded); err != nil {
|
||||
return TopazVideoMetadata{}, err
|
||||
}
|
||||
metadata := TopazVideoMetadata{}
|
||||
metadata.Duration, _ = strconv.ParseFloat(decoded.Format.Duration, 64)
|
||||
for _, stream := range decoded.Streams {
|
||||
if stream.CodecType == "audio" {
|
||||
metadata.HasAudio = true
|
||||
}
|
||||
if stream.CodecType != "video" || stream.Width <= 0 || stream.Height <= 0 {
|
||||
continue
|
||||
}
|
||||
metadata.Width, metadata.Height = stream.Width, stream.Height
|
||||
metadata.FrameRate = parseTopazRate(stream.AverageRate)
|
||||
metadata.FrameCount, _ = strconv.Atoi(stream.FrameCount)
|
||||
}
|
||||
if metadata.Width <= 0 || metadata.Height <= 0 {
|
||||
return TopazVideoMetadata{}, errors.New("video stream metadata is missing")
|
||||
}
|
||||
return metadata, nil
|
||||
}
|
||||
|
||||
func topazTargetResolution(body map[string]any, source map[string]any) map[string]int {
|
||||
width := int(numericValue(body["output_width"], 0))
|
||||
height := int(numericValue(body["output_height"], 0))
|
||||
if width > 0 && height > 0 {
|
||||
return map[string]int{"width": width, "height": height}
|
||||
}
|
||||
value := strings.ToLower(strings.TrimSpace(firstNonEmptyString(body["target_resolution"], body["output_resolution"], "1080p")))
|
||||
if parsedWidth, parsedHeight, ok := parseTopazSize(value); ok {
|
||||
return map[string]int{"width": parsedWidth, "height": parsedHeight}
|
||||
}
|
||||
base := map[string][2]int{"480p": {854, 480}, "720p": {1280, 720}, "1080p": {1920, 1080}, "1440p": {2560, 1440}, "2k": {2560, 1440}, "2160p": {3840, 2160}, "4k": {3840, 2160}}
|
||||
value = strings.TrimSuffix(value, "_upscale")
|
||||
target, ok := base[value]
|
||||
if !ok {
|
||||
target = base["1080p"]
|
||||
}
|
||||
sourceWidth := numericValue(source["width"], 1)
|
||||
sourceHeight := numericValue(source["height"], 1)
|
||||
if sourceWidth == sourceHeight {
|
||||
side := minInt(target[0], target[1])
|
||||
return map[string]int{"width": side, "height": side}
|
||||
}
|
||||
if sourceWidth > sourceHeight {
|
||||
return map[string]int{"width": int(math.Round(sourceWidth / sourceHeight * float64(target[1]))), "height": target[1]}
|
||||
}
|
||||
return map[string]int{"width": target[1], "height": int(math.Round(sourceHeight / sourceWidth * float64(target[1])))}
|
||||
}
|
||||
|
||||
func topazTargetFromPayload(payload map[string]any, body map[string]any) map[string]int {
|
||||
if raw, ok := payload["targetResolution"].(map[string]any); ok {
|
||||
return map[string]int{"width": int(numericValue(raw["width"], 0)), "height": int(numericValue(raw["height"], 0))}
|
||||
}
|
||||
return topazTargetResolution(body, map[string]any{"width": 16, "height": 9})
|
||||
}
|
||||
|
||||
func topazDownloadURL(result map[string]any) string {
|
||||
if download, ok := result["download"].(map[string]any); ok {
|
||||
return firstNonEmptyString(download["url"], download["downloadUrl"])
|
||||
}
|
||||
return firstNonEmptyString(result["download_url"], result["downloadUrl"], result["url"])
|
||||
}
|
||||
|
||||
func topazModelName(model string) string {
|
||||
switch strings.TrimSpace(model) {
|
||||
case "easy-proteus-standard-4":
|
||||
return "prob-4"
|
||||
case "easy-starlight-fast-2":
|
||||
return "slf-2"
|
||||
case "easy-starlight-hq-1":
|
||||
return "slhq-1"
|
||||
case "easy-starlight-mini-1":
|
||||
return "slm-1"
|
||||
default:
|
||||
return strings.TrimSpace(model)
|
||||
}
|
||||
}
|
||||
|
||||
func topazContainerExtension(path string) string {
|
||||
ext := strings.ToLower(filepath.Ext(strings.Split(path, "?")[0]))
|
||||
switch ext {
|
||||
case ".mov", ".mkv", ".webm", ".mp4":
|
||||
return ext
|
||||
default:
|
||||
return ".mp4"
|
||||
}
|
||||
}
|
||||
|
||||
func topazContainerContentType(container string) string {
|
||||
switch strings.ToLower(container) {
|
||||
case "mov":
|
||||
return "video/quicktime"
|
||||
case "mkv":
|
||||
return "video/x-matroska"
|
||||
case "webm":
|
||||
return "video/webm"
|
||||
default:
|
||||
return "video/mp4"
|
||||
}
|
||||
}
|
||||
|
||||
func parseTopazRate(value string) float64 {
|
||||
parts := strings.Split(value, "/")
|
||||
if len(parts) == 2 {
|
||||
numerator, _ := strconv.ParseFloat(parts[0], 64)
|
||||
denominator, _ := strconv.ParseFloat(parts[1], 64)
|
||||
if denominator > 0 {
|
||||
return numerator / denominator
|
||||
}
|
||||
}
|
||||
parsed, _ := strconv.ParseFloat(value, 64)
|
||||
return parsed
|
||||
}
|
||||
|
||||
func parseTopazSize(value string) (int, int, bool) {
|
||||
parts := strings.Split(strings.ReplaceAll(value, " ", ""), "x")
|
||||
if len(parts) != 2 {
|
||||
return 0, 0, false
|
||||
}
|
||||
width, errWidth := strconv.Atoi(parts[0])
|
||||
height, errHeight := strconv.Atoi(parts[1])
|
||||
return width, height, errWidth == nil && errHeight == nil && width > 0 && height > 0
|
||||
}
|
||||
|
||||
func stringList(value any) []string {
|
||||
items, ok := value.([]any)
|
||||
if !ok {
|
||||
if values, ok := value.([]string); ok {
|
||||
return values
|
||||
}
|
||||
return nil
|
||||
}
|
||||
out := make([]string, 0, len(items))
|
||||
for _, item := range items {
|
||||
if text := strings.TrimSpace(fmt.Sprint(item)); text != "" {
|
||||
out = append(out, text)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func boolishDefault(value any, fallback bool) bool {
|
||||
if value == nil {
|
||||
return fallback
|
||||
}
|
||||
switch typed := value.(type) {
|
||||
case bool:
|
||||
return typed
|
||||
case string:
|
||||
parsed, err := strconv.ParseBool(strings.TrimSpace(typed))
|
||||
if err == nil {
|
||||
return parsed
|
||||
}
|
||||
}
|
||||
return fallback
|
||||
}
|
||||
|
||||
func firstError(values ...error) error {
|
||||
for _, value := range values {
|
||||
if value != nil {
|
||||
return value
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func minInt64(left, right int64) int64 {
|
||||
if left < right {
|
||||
return left
|
||||
}
|
||||
return right
|
||||
}
|
||||
|
||||
func minInt(left, right int) int {
|
||||
if left < right {
|
||||
return left
|
||||
}
|
||||
return right
|
||||
}
|
||||
|
||||
func topazSourceHTTPClient(base *http.Client, allowPrivate bool) *http.Client {
|
||||
client := *base
|
||||
client.CheckRedirect = func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }
|
||||
if client.Timeout <= 0 {
|
||||
client.Timeout = 10 * time.Minute
|
||||
}
|
||||
if allowPrivate {
|
||||
return &client
|
||||
}
|
||||
transport := http.DefaultTransport.(*http.Transport).Clone()
|
||||
transport.Proxy = nil
|
||||
transport.DialContext = func(ctx context.Context, network, address string) (net.Conn, error) {
|
||||
host, port, err := net.SplitHostPort(address)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
addresses, err := net.DefaultResolver.LookupIPAddr(ctx, host)
|
||||
if err != nil || len(addresses) == 0 {
|
||||
return nil, errors.New("video source DNS resolution failed")
|
||||
}
|
||||
for _, address := range addresses {
|
||||
if topazBlockedAddress(address.IP) {
|
||||
return nil, errors.New("video source resolved to a blocked network")
|
||||
}
|
||||
}
|
||||
dialer := &net.Dialer{Timeout: 10 * time.Second}
|
||||
var attempts []error
|
||||
for _, resolved := range addresses {
|
||||
connection, dialErr := dialer.DialContext(ctx, network, net.JoinHostPort(resolved.IP.String(), port))
|
||||
if dialErr == nil {
|
||||
return connection, nil
|
||||
}
|
||||
attempts = append(attempts, dialErr)
|
||||
}
|
||||
return nil, errors.Join(attempts...)
|
||||
}
|
||||
client.Transport = transport
|
||||
return &client
|
||||
}
|
||||
|
||||
func topazBlockedAddress(ip net.IP) bool {
|
||||
return ip == nil || ip.IsLoopback() || ip.IsPrivate() || ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() || ip.IsUnspecified() || ip.IsMulticast()
|
||||
}
|
||||
@@ -3,6 +3,7 @@ package clients
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
@@ -10,24 +11,85 @@ import (
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultProviderRequestTimeout = 10 * time.Minute
|
||||
imageProviderRequestTimeout = 20 * time.Minute
|
||||
videoProviderRequestTimeout = 30 * time.Minute
|
||||
)
|
||||
|
||||
// ProviderRequestTimeout returns the default end-to-end provider request or
|
||||
// polling budget for a task kind. Explicit platform polling configuration may
|
||||
// still override these defaults for providers with a documented requirement.
|
||||
func ProviderRequestTimeout(kind string) time.Duration {
|
||||
switch {
|
||||
case strings.HasPrefix(strings.TrimSpace(kind), "images."):
|
||||
return imageProviderRequestTimeout
|
||||
case strings.HasPrefix(strings.TrimSpace(kind), "videos."):
|
||||
return videoProviderRequestTimeout
|
||||
default:
|
||||
return defaultProviderRequestTimeout
|
||||
}
|
||||
}
|
||||
|
||||
func transportClientError(err error) *ClientError {
|
||||
if err == nil {
|
||||
return nil
|
||||
}
|
||||
var netErr net.Error
|
||||
if errors.Is(err, context.DeadlineExceeded) || (errors.As(err, &netErr) && netErr.Timeout()) {
|
||||
return &ClientError{
|
||||
Code: "timeout",
|
||||
Message: "upstream request timed out: " + err.Error(),
|
||||
Retryable: false,
|
||||
}
|
||||
}
|
||||
return &ClientError{Code: "network", Message: err.Error(), Retryable: true}
|
||||
}
|
||||
|
||||
type Request struct {
|
||||
Kind string
|
||||
ModelType string
|
||||
Model string
|
||||
Body map[string]any
|
||||
Candidate store.RuntimeModelCandidate
|
||||
HTTPClient *http.Client
|
||||
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
|
||||
PublicResponseID string
|
||||
PublicPreviousResponseID string
|
||||
UpstreamPreviousResponseID string
|
||||
PreviousResponseTurns []ResponseTurn
|
||||
Kind string
|
||||
ModelType string
|
||||
Model string
|
||||
Body map[string]any
|
||||
OriginalBody map[string]any
|
||||
Candidate store.RuntimeModelCandidate
|
||||
HTTPClient *http.Client
|
||||
RemoteTaskID string
|
||||
RemoteTaskPayload map[string]any
|
||||
UpstreamIdempotencyKey string
|
||||
OnRemoteTaskSubmitted func(remoteTaskID string, payload map[string]any) error
|
||||
OnRemoteTaskPolled func(remoteTaskID string, payload map[string]any) error
|
||||
OnUpstreamSubmissionStarted func() error
|
||||
OnUpstreamResponseReceived func() error
|
||||
OnUpstreamWireResponse func(*WireResponse) error
|
||||
Stream bool
|
||||
StreamDelta StreamDelta
|
||||
UpstreamProtocol string
|
||||
PublicResponseID string
|
||||
PublicPreviousResponseID string
|
||||
UpstreamPreviousResponseID string
|
||||
PreviousResponseTurns []ResponseTurn
|
||||
}
|
||||
|
||||
func applyUpstreamIdempotency(request *http.Request, input Request) {
|
||||
if request == nil {
|
||||
return
|
||||
}
|
||||
if key := strings.TrimSpace(input.UpstreamIdempotencyKey); key != "" {
|
||||
request.Header.Set("Idempotency-Key", key)
|
||||
}
|
||||
}
|
||||
|
||||
// WireResponse preserves the provider-facing response independently from the
|
||||
// canonical result used by billing, retries, and cross-protocol conversion.
|
||||
// Headers are filtered before this value leaves the clients package.
|
||||
type WireResponse struct {
|
||||
Protocol string `json:"protocol,omitempty"`
|
||||
StatusCode int `json:"statusCode,omitempty"`
|
||||
Headers map[string][]string `json:"headers,omitempty"`
|
||||
Body map[string]any `json:"body,omitempty"`
|
||||
RawJSON []byte `json:"-"`
|
||||
Converted bool `json:"converted,omitempty"`
|
||||
}
|
||||
|
||||
type ResponseTurn struct {
|
||||
@@ -54,6 +116,7 @@ type Response struct {
|
||||
ResponseChainDepth int
|
||||
ResponseConverted bool
|
||||
InternalResult map[string]any
|
||||
Wire *WireResponse
|
||||
}
|
||||
|
||||
type Usage struct {
|
||||
@@ -75,6 +138,10 @@ type StreamDeltaEvent struct {
|
||||
Text string
|
||||
ReasoningContent string
|
||||
Event map[string]any
|
||||
WireProtocol string
|
||||
WireConverted bool
|
||||
WireStatusCode int
|
||||
WireHeaders map[string][]string
|
||||
}
|
||||
|
||||
type StreamDelta func(event StreamDeltaEvent) error
|
||||
@@ -106,12 +173,22 @@ type ClientError struct {
|
||||
Code string
|
||||
Message string
|
||||
Param string
|
||||
Details map[string]any
|
||||
StatusCode int
|
||||
RequestID string
|
||||
ResponseStartedAt time.Time
|
||||
ResponseFinishedAt time.Time
|
||||
ResponseDurationMS int64
|
||||
Retryable bool
|
||||
Wire *WireResponse
|
||||
}
|
||||
|
||||
func ErrorWireResponse(err error) *WireResponse {
|
||||
var clientErr *ClientError
|
||||
if errors.As(err, &clientErr) {
|
||||
return clientErr.Wire
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func ErrorParam(err error) string {
|
||||
@@ -122,6 +199,18 @@ func ErrorParam(err error) string {
|
||||
return ""
|
||||
}
|
||||
|
||||
func ErrorDetails(err error) map[string]any {
|
||||
var clientErr *ClientError
|
||||
if !errors.As(err, &clientErr) || len(clientErr.Details) == 0 {
|
||||
return nil
|
||||
}
|
||||
details := make(map[string]any, len(clientErr.Details))
|
||||
for key, value := range clientErr.Details {
|
||||
details[key] = value
|
||||
}
|
||||
return details
|
||||
}
|
||||
|
||||
func (e *ClientError) Error() string {
|
||||
if e.Message != "" {
|
||||
return e.Message
|
||||
@@ -142,6 +231,10 @@ func upstreamModelName(candidate store.RuntimeModelCandidate) string {
|
||||
}
|
||||
|
||||
func ErrorCode(err error) string {
|
||||
var coded interface{ ErrorCode() string }
|
||||
if errors.As(err, &coded) && strings.TrimSpace(coded.ErrorCode()) != "" {
|
||||
return strings.TrimSpace(coded.ErrorCode())
|
||||
}
|
||||
var clientErr *ClientError
|
||||
if errors.As(err, &clientErr) && clientErr.Code != "" {
|
||||
return clientErr.Code
|
||||
|
||||
@@ -40,7 +40,15 @@ func (c UniversalClient) Run(ctx context.Context, request Request) (Response, er
|
||||
if err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
if err := notifySubmissionStarted(request); err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
submitResult, submitRequestID, err = c.universalSubmit(ctx, executor, request, modelType, payload)
|
||||
if err == nil || ErrorResponseMetadata(err).StatusCode > 0 {
|
||||
if notifyErr := notifyResponseReceived(request); notifyErr != nil {
|
||||
return Response{}, notifyErr
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
return Response{}, annotateResponseError(err, submitRequestID, startedAt, time.Now())
|
||||
}
|
||||
@@ -145,7 +153,7 @@ func (c UniversalClient) universalSubmit(ctx context.Context, executor *scripten
|
||||
|
||||
func (c UniversalClient) universalPollUntilDone(ctx context.Context, executor *scriptengine.Executor, request Request, modelType string, upstreamTaskID string, payload map[string]any, requestID string, startedAt time.Time) (map[string]any, string, error) {
|
||||
interval := universalDurationConfig(request.Candidate.PlatformConfig, 2*time.Second, "pollIntervalMs", "poll_interval_ms")
|
||||
timeout := universalDurationConfig(request.Candidate.PlatformConfig, 10*time.Minute, "pollTimeoutMs", "poll_timeout_ms", "timeoutMs")
|
||||
timeout := universalDurationConfig(request.Candidate.PlatformConfig, ProviderRequestTimeout(request.Kind), "pollTimeoutMs", "poll_timeout_ms", "timeoutMs")
|
||||
deadline := time.NewTimer(timeout)
|
||||
defer deadline.Stop()
|
||||
ticker := time.NewTicker(interval)
|
||||
@@ -226,7 +234,7 @@ func universalScriptContext(request Request, modelType string, payload map[strin
|
||||
"platformModelId": request.Candidate.PlatformModelID,
|
||||
"canonicalModelKey": request.Candidate.CanonicalModelKey,
|
||||
"modelType": modelType,
|
||||
"timeout": universalDurationConfig(request.Candidate.PlatformConfig, 10*time.Minute, "pollTimeoutMs", "poll_timeout_ms").Milliseconds(),
|
||||
"timeout": universalDurationConfig(request.Candidate.PlatformConfig, ProviderRequestTimeout(request.Kind), "pollTimeoutMs", "poll_timeout_ms").Milliseconds(),
|
||||
},
|
||||
"env": cloneMapAny(request.Candidate.PlatformConfig),
|
||||
"candidate": universalCandidateSnapshot(request),
|
||||
@@ -310,7 +318,7 @@ func universalPostJSON(ctx context.Context, client *http.Client, baseURL string,
|
||||
}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, "", &ClientError{Code: "network", Message: err.Error(), Retryable: true}
|
||||
return nil, "", transportClientError(err)
|
||||
}
|
||||
requestID := requestIDFromHTTPResponse(resp)
|
||||
result, err := decodeHTTPResponse(resp)
|
||||
@@ -327,7 +335,7 @@ func universalGetJSON(ctx context.Context, client *http.Client, url string, cred
|
||||
}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, "", &ClientError{Code: "network", Message: err.Error(), Retryable: true}
|
||||
return nil, "", transportClientError(err)
|
||||
}
|
||||
requestID := requestIDFromHTTPResponse(resp)
|
||||
result, err := decodeHTTPResponse(resp)
|
||||
|
||||
@@ -0,0 +1,271 @@
|
||||
package clients
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"io"
|
||||
"mime/multipart"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/textproto"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
const vectorizerMaxResponseBytes = 128 << 20
|
||||
|
||||
// VectorizerClient implements the Vectorizer.AI binary vectorization API.
|
||||
// Gateway async semantics are supplied by the outer River-backed task runner.
|
||||
type VectorizerClient struct {
|
||||
HTTPClient *http.Client
|
||||
LookupIP func(context.Context, string) ([]net.IPAddr, error)
|
||||
}
|
||||
|
||||
func (c VectorizerClient) Run(ctx context.Context, request Request) (Response, error) {
|
||||
startedAt := time.Now()
|
||||
format := strings.ToLower(strings.TrimSpace(firstNonEmptyString(request.Body["format"], request.Body["output_format"])))
|
||||
if format == "" {
|
||||
format = "svg"
|
||||
}
|
||||
if !vectorizerFormatAllowed(format) {
|
||||
return Response{}, &ClientError{Code: "invalid_parameter", Message: "vectorizer format must be svg, eps, pdf, dxf, or png", Param: "format", StatusCode: http.StatusBadRequest}
|
||||
}
|
||||
|
||||
imageToken := strings.TrimSpace(firstNonEmptyString(request.Body["_vectorizer_image_token"]))
|
||||
receipt := strings.TrimSpace(firstNonEmptyString(request.Body["_vectorizer_receipt"]))
|
||||
endpoint := "vectorize"
|
||||
fields := map[string]string{"output.file_format": format}
|
||||
if imageToken != "" {
|
||||
endpoint = "download"
|
||||
fields["image.token"] = imageToken
|
||||
if receipt != "" {
|
||||
fields["receipt"] = receipt
|
||||
}
|
||||
} else {
|
||||
imageURL := vectorizerImageURL(request.Body)
|
||||
if imageURL == "" {
|
||||
return Response{}, &ClientError{Code: "invalid_parameter", Message: "vectorizer source image URL is required", Param: "source.url", StatusCode: http.StatusBadRequest}
|
||||
}
|
||||
if err := c.validateSourceURL(ctx, imageURL, boolishDefault(request.Candidate.PlatformConfig["allowPrivateSourceDownloads"], false)); err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
fields["image.url"] = imageURL
|
||||
fields["mode"] = firstNonEmptyString(request.Candidate.PlatformConfig["mode"], "production")
|
||||
fields["policy.retention_days"] = firstNonEmptyString(request.Candidate.PlatformConfig["retentionDays"], request.Candidate.PlatformConfig["retention_days"], "7")
|
||||
fields["processing.shapes.min_area_px"] = vectorizerCleanupMinArea(request.Body)
|
||||
if maxColors := vectorizerMaxColors(request.Body); maxColors > 0 {
|
||||
fields["processing.max_colors"] = fmt.Sprint(maxColors)
|
||||
}
|
||||
}
|
||||
appendVectorizerOutputFields(fields, format)
|
||||
|
||||
payload, contentType, err := vectorizerMultipartBody(fields)
|
||||
if err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
url := providerURL(request.Candidate.BaseURL, endpoint)
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, payload)
|
||||
if err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
req.Header.Set("Content-Type", contentType)
|
||||
applyVectorizerAuth(req, request.Candidate.Credentials)
|
||||
client := httpClient(request.HTTPClient, c.HTTPClient)
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return Response{}, transportClientError(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
requestID := requestIDFromHTTPResponse(resp)
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
result, decodeErr := decodeHTTPResponse(resp)
|
||||
if decodeErr != nil {
|
||||
return Response{}, annotateResponseError(decodeErr, requestID, startedAt, time.Now())
|
||||
}
|
||||
return Response{}, &ClientError{Code: "provider_failed", Message: firstNonEmptyString(result["message"], result["error"], resp.Status), RequestID: requestID, StatusCode: resp.StatusCode, Retryable: resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode >= 500}
|
||||
}
|
||||
raw, err := io.ReadAll(io.LimitReader(resp.Body, vectorizerMaxResponseBytes+1))
|
||||
if err != nil {
|
||||
return Response{}, &ClientError{Code: "invalid_response", Message: err.Error(), RequestID: requestID, Retryable: true}
|
||||
}
|
||||
if len(raw) == 0 || len(raw) > vectorizerMaxResponseBytes {
|
||||
return Response{}, &ClientError{Code: "invalid_response", Message: "vectorizer response is empty or too large", RequestID: requestID, Retryable: false}
|
||||
}
|
||||
returnedToken := strings.TrimSpace(resp.Header.Get("X-Image-Token"))
|
||||
returnedReceipt := strings.TrimSpace(resp.Header.Get("X-Receipt"))
|
||||
if request.OnRemoteTaskSubmitted != nil && returnedToken != "" {
|
||||
digest := sha256.Sum256([]byte(returnedToken))
|
||||
if err := request.OnRemoteTaskSubmitted("vectorizer-"+hex.EncodeToString(digest[:8]), map[string]any{
|
||||
"imageToken": returnedToken,
|
||||
"receipt": returnedReceipt,
|
||||
}); err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
}
|
||||
finishedAt := time.Now()
|
||||
mimeType := strings.TrimSpace(strings.Split(resp.Header.Get("Content-Type"), ";")[0])
|
||||
if mimeType == "" || mimeType == "application/octet-stream" {
|
||||
mimeType = vectorizerContentType(format)
|
||||
}
|
||||
return Response{
|
||||
Result: map[string]any{
|
||||
"status": "success",
|
||||
"model": request.Model,
|
||||
"data": []any{map[string]any{
|
||||
"type": vectorizerOutputKind(format),
|
||||
"b64_json": base64.StdEncoding.EncodeToString(raw),
|
||||
"mime_type": mimeType,
|
||||
"format": format,
|
||||
}},
|
||||
"vectorizer": map[string]any{
|
||||
"format": format,
|
||||
"creditsCharged": numericHeader(resp.Header.Get("X-Credits-Charged")),
|
||||
"creditsCalculated": numericHeader(resp.Header.Get("X-Credits-Calculated")),
|
||||
},
|
||||
},
|
||||
RequestID: requestID,
|
||||
Progress: append(providerProgress(request), Progress{Phase: "uploading", Progress: 0.9, Message: "vectorizer result received"}),
|
||||
ResponseStartedAt: startedAt,
|
||||
ResponseFinishedAt: finishedAt,
|
||||
ResponseDurationMS: responseDurationMS(startedAt, finishedAt),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (c VectorizerClient) validateSourceURL(ctx context.Context, rawURL string, allowPrivate bool) error {
|
||||
parsed, err := url.Parse(rawURL)
|
||||
if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") || strings.TrimSpace(parsed.Hostname()) == "" || parsed.User != nil {
|
||||
return &ClientError{Code: "invalid_parameter", Message: "vectorizer source URL must be a public http(s) URL without userinfo", Param: "source.url", StatusCode: http.StatusBadRequest}
|
||||
}
|
||||
if allowPrivate {
|
||||
return nil
|
||||
}
|
||||
lookup := c.LookupIP
|
||||
if lookup == nil {
|
||||
lookup = net.DefaultResolver.LookupIPAddr
|
||||
}
|
||||
addresses, err := lookup(ctx, parsed.Hostname())
|
||||
if err != nil || len(addresses) == 0 {
|
||||
return &ClientError{Code: "invalid_parameter", Message: "vectorizer source DNS resolution failed", Param: "source.url", StatusCode: http.StatusBadRequest}
|
||||
}
|
||||
for _, address := range addresses {
|
||||
if topazBlockedAddress(address.IP) {
|
||||
return &ClientError{Code: "invalid_parameter", Message: "vectorizer source resolved to a blocked network", Param: "source.url", StatusCode: http.StatusBadRequest}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func vectorizerImageURL(body map[string]any) string {
|
||||
if source, ok := body["source"].(map[string]any); ok {
|
||||
return firstNonEmptyString(source["url"], source["image_url"], source["imageUrl"])
|
||||
}
|
||||
return firstNonEmptyString(body["image_url"], body["imageUrl"])
|
||||
}
|
||||
|
||||
func vectorizerMultipartBody(fields map[string]string) (*bytes.Buffer, string, error) {
|
||||
var payload bytes.Buffer
|
||||
writer := multipart.NewWriter(&payload)
|
||||
for key, value := range fields {
|
||||
header := make(textproto.MIMEHeader)
|
||||
header.Set("Content-Disposition", fmt.Sprintf(`form-data; name="%s"`, strings.ReplaceAll(key, `"`, `\"`)))
|
||||
part, err := writer.CreatePart(header)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
if _, err := io.WriteString(part, value); err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
}
|
||||
if err := writer.Close(); err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
return &payload, writer.FormDataContentType(), nil
|
||||
}
|
||||
|
||||
func applyVectorizerAuth(req *http.Request, credentials map[string]any) {
|
||||
username := credential(credentials, "username", "accessKey", "access_key", "apiId", "api_id", "id")
|
||||
password := credential(credentials, "password", "secretKey", "secret_key", "apiSecret", "api_secret", "secret")
|
||||
if username != "" || password != "" {
|
||||
req.SetBasicAuth(username, password)
|
||||
return
|
||||
}
|
||||
if apiKey := credential(credentials, "apiKey", "api_key", "token"); apiKey != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
}
|
||||
}
|
||||
|
||||
func vectorizerFormatAllowed(format string) bool {
|
||||
switch format {
|
||||
case "svg", "eps", "pdf", "dxf", "png":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func vectorizerCleanupMinArea(body map[string]any) string {
|
||||
cleanup := strings.ToLower(strings.TrimSpace(firstNonEmptyString(body["cleanupLevel"], body["cleanup_level"])))
|
||||
switch cleanup {
|
||||
case "low":
|
||||
return "0"
|
||||
case "strong":
|
||||
return "1"
|
||||
default:
|
||||
return "0.125"
|
||||
}
|
||||
}
|
||||
|
||||
func vectorizerMaxColors(body map[string]any) int {
|
||||
for _, key := range []string{"maxColors", "max_colors"} {
|
||||
if value := int(numericValue(body[key], 0)); value == 0 || value == 2 || value == 4 || value == 8 || value == 16 || value == 32 {
|
||||
return value
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func appendVectorizerOutputFields(fields map[string]string, format string) {
|
||||
if format == "svg" {
|
||||
fields["output.svg.version"] = "svg_1_1"
|
||||
fields["output.svg.fixed_size"] = "false"
|
||||
fields["output.svg.adobe_compatibility_mode"] = "true"
|
||||
}
|
||||
if format == "dxf" {
|
||||
fields["output.dxf.compatibility_level"] = "lines_and_arcs"
|
||||
}
|
||||
}
|
||||
|
||||
func vectorizerContentType(format string) string {
|
||||
switch format {
|
||||
case "svg":
|
||||
return "image/svg+xml"
|
||||
case "eps":
|
||||
return "application/postscript"
|
||||
case "pdf":
|
||||
return "application/pdf"
|
||||
case "dxf":
|
||||
return "application/dxf"
|
||||
default:
|
||||
return "image/png"
|
||||
}
|
||||
}
|
||||
|
||||
func vectorizerOutputKind(format string) string {
|
||||
if format == "svg" || format == "png" {
|
||||
return "image"
|
||||
}
|
||||
return "file"
|
||||
}
|
||||
|
||||
func numericHeader(value string) any {
|
||||
value = strings.TrimSpace(value)
|
||||
if value == "" {
|
||||
return nil
|
||||
}
|
||||
return numericValue(value, 0)
|
||||
}
|
||||
@@ -44,14 +44,21 @@ func (c VolcesClient) runImage(ctx context.Context, request Request, apiKey stri
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
applyUpstreamIdempotency(req, request)
|
||||
|
||||
responseStartedAt := time.Now()
|
||||
if err := notifySubmissionStarted(request); err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
resp, err := httpClient(request.HTTPClient, c.HTTPClient).Do(req)
|
||||
if err != nil {
|
||||
return Response{}, &ClientError{Code: "network", Message: err.Error(), Retryable: true}
|
||||
return Response{}, transportClientError(err)
|
||||
}
|
||||
if err := notifyResponseReceived(request); err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
requestID := requestIDFromHTTPResponse(resp)
|
||||
result, err := decodeHTTPResponse(resp)
|
||||
result, wire, err := decodeHTTPResponseForProtocol(resp, ProtocolOpenAIImages)
|
||||
responseFinishedAt := time.Now()
|
||||
if err != nil {
|
||||
return Response{}, annotateResponseError(err, requestID, responseStartedAt, responseFinishedAt)
|
||||
@@ -67,6 +74,8 @@ func (c VolcesClient) runImage(ctx context.Context, request Request, apiKey stri
|
||||
ResponseStartedAt: responseStartedAt,
|
||||
ResponseFinishedAt: responseFinishedAt,
|
||||
ResponseDurationMS: responseDurationMS(responseStartedAt, responseFinishedAt),
|
||||
UpstreamProtocol: ProtocolOpenAIImages,
|
||||
Wire: wire,
|
||||
}, nil
|
||||
}
|
||||
|
||||
@@ -80,7 +89,7 @@ func (c VolcesClient) runVideo(ctx context.Context, request Request, apiKey stri
|
||||
if err := validateVolcesVideoTaskBody(body); err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
submitResult, requestID, err := c.postJSON(ctx, request, request.Candidate.BaseURL, taskPath, apiKey, body)
|
||||
submitResult, requestID, _, err := c.postJSON(ctx, request, request.Candidate.BaseURL, taskPath, apiKey, body)
|
||||
submitRequestID = requestID
|
||||
if err != nil {
|
||||
return Response{}, annotateResponseError(err, submitRequestID, submitStartedAt, time.Now())
|
||||
@@ -119,7 +128,7 @@ func (c VolcesClient) runVideo(ctx context.Context, request Request, apiKey stri
|
||||
}
|
||||
case <-nextPoll.C:
|
||||
pollStartedAt := time.Now()
|
||||
pollResult, pollRequestID, err := c.getJSON(ctx, request, request.Candidate.BaseURL, taskPath+"/"+upstreamTaskID, apiKey)
|
||||
pollResult, pollRequestID, pollWire, err := c.getJSON(ctx, request, request.Candidate.BaseURL, taskPath+"/"+upstreamTaskID, apiKey)
|
||||
pollFinishedAt := time.Now()
|
||||
requestID := firstNonEmpty(pollRequestID, submitRequestID, upstreamTaskID)
|
||||
lastRequestID = requestID
|
||||
@@ -151,6 +160,8 @@ func (c VolcesClient) runVideo(ctx context.Context, request Request, apiKey stri
|
||||
ResponseStartedAt: submitStartedAt,
|
||||
ResponseFinishedAt: pollFinishedAt,
|
||||
ResponseDurationMS: responseDurationMS(submitStartedAt, pollFinishedAt),
|
||||
UpstreamProtocol: ProtocolVolcesContents,
|
||||
Wire: pollWire,
|
||||
}, nil
|
||||
case "failed", "cancelled":
|
||||
return Response{}, &ClientError{
|
||||
@@ -188,7 +199,7 @@ func (c VolcesClient) DeleteVideoTask(ctx context.Context, request Request) (map
|
||||
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}
|
||||
return nil, "", transportClientError(err)
|
||||
}
|
||||
requestID := requestIDFromHTTPResponse(response)
|
||||
result, err := decodeHTTPResponse(response)
|
||||
@@ -214,44 +225,74 @@ func volcesVideoTaskPath(request Request) string {
|
||||
return strings.TrimRight(path, "/")
|
||||
}
|
||||
|
||||
func (c VolcesClient) postJSON(ctx context.Context, request Request, baseURL string, path string, apiKey string, body map[string]any) (map[string]any, string, error) {
|
||||
func (c VolcesClient) postJSON(ctx context.Context, request Request, baseURL string, path string, apiKey string, body map[string]any) (map[string]any, string, *WireResponse, error) {
|
||||
raw, _ := json.Marshal(body)
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, joinURL(baseURL, path), bytes.NewReader(raw))
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
return nil, "", nil, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
applyUpstreamIdempotency(req, request)
|
||||
if err := notifySubmissionStarted(request); err != nil {
|
||||
return nil, "", nil, err
|
||||
}
|
||||
resp, err := httpClient(request.HTTPClient, c.HTTPClient).Do(req)
|
||||
if err != nil {
|
||||
return nil, "", &ClientError{Code: "network", Message: err.Error(), Retryable: true}
|
||||
return nil, "", nil, transportClientError(err)
|
||||
}
|
||||
if err := notifyResponseReceived(request); err != nil {
|
||||
return nil, "", nil, err
|
||||
}
|
||||
requestID := requestIDFromHTTPResponse(resp)
|
||||
result, err := decodeHTTPResponse(resp)
|
||||
result, wire, err := decodeHTTPResponseForProtocol(resp, ProtocolVolcesContents)
|
||||
if err != nil {
|
||||
return result, requestID, err
|
||||
if notifyErr := notifyWireResponse(request, wire); notifyErr != nil {
|
||||
return result, requestID, wire, notifyErr
|
||||
}
|
||||
return result, requestID, wire, err
|
||||
}
|
||||
original := result
|
||||
result, envelopeRequestID, err := normalizeVolcesCompatibleResult(result)
|
||||
return result, firstNonEmpty(requestID, envelopeRequestID), err
|
||||
if wire != nil {
|
||||
wire.Converted = volcesCompatibleResultConverted(original)
|
||||
}
|
||||
if notifyErr := notifyWireResponse(request, wire); notifyErr != nil {
|
||||
return result, firstNonEmpty(requestID, envelopeRequestID), wire, notifyErr
|
||||
}
|
||||
return result, firstNonEmpty(requestID, envelopeRequestID), wire, err
|
||||
}
|
||||
|
||||
func (c VolcesClient) getJSON(ctx context.Context, request Request, baseURL string, path string, apiKey string) (map[string]any, string, error) {
|
||||
func (c VolcesClient) getJSON(ctx context.Context, request Request, baseURL string, path string, apiKey string) (map[string]any, string, *WireResponse, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, joinURL(baseURL, path), nil)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
return nil, "", nil, err
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
resp, err := httpClient(request.HTTPClient, c.HTTPClient).Do(req)
|
||||
if err != nil {
|
||||
return nil, "", &ClientError{Code: "network", Message: err.Error(), Retryable: true}
|
||||
return nil, "", nil, transportClientError(err)
|
||||
}
|
||||
requestID := requestIDFromHTTPResponse(resp)
|
||||
result, err := decodeHTTPResponse(resp)
|
||||
result, wire, err := decodeHTTPResponseForProtocol(resp, ProtocolVolcesContents)
|
||||
if err != nil {
|
||||
return result, requestID, err
|
||||
return result, requestID, wire, err
|
||||
}
|
||||
original := result
|
||||
result, envelopeRequestID, err := normalizeVolcesCompatibleResult(result)
|
||||
return result, firstNonEmpty(requestID, envelopeRequestID), err
|
||||
if wire != nil {
|
||||
wire.Converted = volcesCompatibleResultConverted(original)
|
||||
}
|
||||
return result, firstNonEmpty(requestID, envelopeRequestID), wire, err
|
||||
}
|
||||
|
||||
func volcesCompatibleResultConverted(result map[string]any) bool {
|
||||
if _, ok := result["error"].(map[string]any); ok {
|
||||
return true
|
||||
}
|
||||
_, hasCode := result["code"]
|
||||
_, hasData := result["data"].(map[string]any)
|
||||
return hasCode && hasData
|
||||
}
|
||||
|
||||
func normalizeVolcesCompatibleResult(result map[string]any) (map[string]any, string, error) {
|
||||
@@ -318,16 +359,33 @@ func volcesImageBody(request Request) map[string]any {
|
||||
body["seed"] = -1
|
||||
}
|
||||
}
|
||||
if resolution := strings.TrimSpace(stringFromAny(body["resolution"])); resolution != "" {
|
||||
resolution := normalizedProviderImageResolution(body["resolution"])
|
||||
size := widthHeightSize(body)
|
||||
if volcesImageUsesResolutionSize(request) && resolution != "" {
|
||||
body["size"] = resolution
|
||||
} else if size != "" {
|
||||
body["size"] = size
|
||||
} else if resolution != "" {
|
||||
body["size"] = resolution
|
||||
}
|
||||
if size := widthHeightSize(body); size != "" {
|
||||
body["size"] = size
|
||||
for _, key := range []string{"aspect_ratio", "aspectRatio", "ratio", "resolution", "width", "height"} {
|
||||
delete(body, key)
|
||||
}
|
||||
normalizeVolcesSequentialImageGeneration(body, request)
|
||||
return body
|
||||
}
|
||||
|
||||
func volcesImageUsesResolutionSize(request Request) bool {
|
||||
modelType := firstNonEmpty(request.ModelType, request.Candidate.ModelType)
|
||||
if capability, ok := request.Candidate.Capabilities[modelType].(map[string]any); ok {
|
||||
if strings.EqualFold(strings.TrimSpace(stringFromAny(capability["size_param_format"])), "resolution") {
|
||||
return true
|
||||
}
|
||||
}
|
||||
model := strings.ToLower(strings.TrimSpace(upstreamModelName(request.Candidate)))
|
||||
return strings.HasPrefix(model, "doubao-seedream-5-0")
|
||||
}
|
||||
|
||||
func volcesVideoBody(request Request) map[string]any {
|
||||
body := cleanProviderBody(request.Body)
|
||||
body["model"] = upstreamModelName(request.Candidate)
|
||||
@@ -354,6 +412,16 @@ func cleanProviderBody(body map[string]any) map[string]any {
|
||||
"poll_interval_ms",
|
||||
"pollTimeoutSeconds",
|
||||
"poll_timeout_seconds",
|
||||
"_gateway_compatibility",
|
||||
"_gateway_target_protocol",
|
||||
"_compat_provider",
|
||||
"_kling_compat_version",
|
||||
"platform_id",
|
||||
"platformId",
|
||||
"platform_model_id",
|
||||
"platformModelId",
|
||||
"modelType",
|
||||
"model_type",
|
||||
} {
|
||||
delete(out, key)
|
||||
}
|
||||
@@ -1057,9 +1125,11 @@ func volcesTaskErrorMessage(result map[string]any) string {
|
||||
}
|
||||
|
||||
func volcesVideoSuccessResult(request Request, upstreamTaskID string, raw map[string]any) map[string]any {
|
||||
result := cloneMapAny(raw)
|
||||
if result == nil {
|
||||
result = map[string]any{}
|
||||
result := map[string]any{}
|
||||
for _, key := range []string{"seed", "resolution", "ratio", "duration", "frames", "framespersecond"} {
|
||||
if raw[key] != nil {
|
||||
result[key] = raw[key]
|
||||
}
|
||||
}
|
||||
content, _ := raw["content"].(map[string]any)
|
||||
videoURL := strings.TrimSpace(stringFromAny(content["video_url"]))
|
||||
@@ -1069,18 +1139,18 @@ func volcesVideoSuccessResult(request Request, upstreamTaskID string, raw map[st
|
||||
}
|
||||
data := []any{}
|
||||
if videoURL != "" {
|
||||
data = append(data, map[string]any{"url": videoURL, "type": "video"})
|
||||
item := map[string]any{"url": videoURL, "type": "video"}
|
||||
if lastFrameURL := strings.TrimSpace(stringFromAny(content["last_frame_url"])); lastFrameURL != "" {
|
||||
item["last_frame_url"] = lastFrameURL
|
||||
}
|
||||
data = append(data, item)
|
||||
}
|
||||
result["id"] = firstNonEmpty(stringFromAny(raw["id"]), upstreamTaskID)
|
||||
if strings.TrimSpace(stringFromAny(result["model"])) == "" {
|
||||
result["model"] = upstreamModelName(request.Candidate)
|
||||
}
|
||||
result["model"] = firstNonEmpty(stringFromAny(raw["model"]), upstreamModelName(request.Candidate))
|
||||
result["status"] = "succeeded"
|
||||
result["object"] = "video.generation"
|
||||
result["created"] = created
|
||||
result["upstream_task_id"] = upstreamTaskID
|
||||
result["data"] = data
|
||||
result["raw"] = cloneMapAny(raw)
|
||||
return result
|
||||
}
|
||||
|
||||
@@ -1114,9 +1184,10 @@ func volcesPollInterval(request Request) time.Duration {
|
||||
}
|
||||
|
||||
func volcesPollTimeout(request Request) time.Duration {
|
||||
seconds := numericValue(firstPresent(request.Candidate.PlatformConfig["volcesPollTimeoutSeconds"], request.Body["pollTimeoutSeconds"], request.Body["poll_timeout_seconds"]), 600)
|
||||
fallbackSeconds := ProviderRequestTimeout(request.Kind).Seconds()
|
||||
seconds := numericValue(firstPresent(request.Candidate.PlatformConfig["volcesPollTimeoutSeconds"], request.Body["pollTimeoutSeconds"], request.Body["poll_timeout_seconds"]), fallbackSeconds)
|
||||
if seconds < 1 {
|
||||
seconds = 600
|
||||
seconds = fallbackSeconds
|
||||
}
|
||||
return time.Duration(seconds) * time.Second
|
||||
}
|
||||
|
||||
@@ -106,7 +106,7 @@ func (c VolcesAssetClient) call(ctx context.Context, credentials VolcesAssetCred
|
||||
|
||||
response, err := httpClient(nil, c.HTTPClient).Do(req)
|
||||
if err != nil {
|
||||
return "", &ClientError{Code: "network", Message: err.Error(), Retryable: true}
|
||||
return "", transportClientError(err)
|
||||
}
|
||||
defer response.Body.Close()
|
||||
var envelope struct {
|
||||
|
||||
@@ -79,7 +79,7 @@ func TestVolcesClientSupportsDeyunEnvelope(t *testing.T) {
|
||||
}
|
||||
data, _ := response.Result["data"].([]any)
|
||||
item, _ := data[0].(map[string]any)
|
||||
if response.Result["upstream_task_id"] != "deyun-task-1" || item["url"] != "https://example.com/deyun.mp4" {
|
||||
if response.Result["upstream_task_id"] != nil || item["url"] != "https://example.com/deyun.mp4" {
|
||||
t.Fatalf("unexpected response: %+v", response.Result)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,125 @@
|
||||
package clients
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func TestDecodeHTTPResponseForProtocolCapturesOfficialErrorWire(t *testing.T) {
|
||||
response := &http.Response{
|
||||
StatusCode: http.StatusTooManyRequests,
|
||||
Status: "429 Too Many Requests",
|
||||
Header: http.Header{
|
||||
"Content-Type": {"application/json"},
|
||||
"X-Request-Id": {"req_1"},
|
||||
"X-Ratelimit-Reset-Requests": {"1s"},
|
||||
"Set-Cookie": {"secret=must-not-leak"},
|
||||
"Connection": {"keep-alive"},
|
||||
},
|
||||
Body: io.NopCloser(strings.NewReader(`{"error":{"message":"slow down","future_field":true}}`)),
|
||||
}
|
||||
_, wire, err := decodeHTTPResponseForProtocol(response, ProtocolOpenAIResponses)
|
||||
if err == nil {
|
||||
t.Fatal("expected upstream error")
|
||||
}
|
||||
if wire == nil || wire.StatusCode != http.StatusTooManyRequests || wire.Protocol != ProtocolOpenAIResponses {
|
||||
t.Fatalf("unexpected wire response: %+v", wire)
|
||||
}
|
||||
if wire.Headers["X-Request-Id"][0] != "req_1" || wire.Headers["X-Ratelimit-Reset-Requests"][0] != "1s" {
|
||||
t.Fatalf("official response headers were lost: %+v", wire.Headers)
|
||||
}
|
||||
if _, ok := wire.Headers["Set-Cookie"]; ok {
|
||||
t.Fatalf("sensitive header leaked: %+v", wire.Headers)
|
||||
}
|
||||
if _, ok := wire.Headers["Connection"]; ok {
|
||||
t.Fatalf("hop-by-hop header leaked: %+v", wire.Headers)
|
||||
}
|
||||
if ErrorWireResponse(err) != wire {
|
||||
t.Fatal("client error did not retain its wire response")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDecodeHTTPResponseForProtocolAllowsLargeImagePayload(t *testing.T) {
|
||||
payload := `{"data":"` + strings.Repeat("a", int(defaultMaxJSONResponseBytes)) + `"}`
|
||||
response := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Status: "200 OK",
|
||||
Header: http.Header{"Content-Type": {"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(payload)),
|
||||
}
|
||||
|
||||
result, wire, err := decodeHTTPResponseForProtocol(response, ProtocolGeminiGenerateContent)
|
||||
if err != nil {
|
||||
t.Fatalf("large Gemini image response failed: %v", err)
|
||||
}
|
||||
if got, _ := result["data"].(string); len(got) != int(defaultMaxJSONResponseBytes) {
|
||||
t.Fatalf("large Gemini image payload length = %d", len(got))
|
||||
}
|
||||
if wire == nil || len(wire.RawJSON) != len(payload) {
|
||||
t.Fatalf("large Gemini wire payload was truncated: %+v", wire)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDecodeHTTPResponseForProtocolRejectsOversizedDefaultPayload(t *testing.T) {
|
||||
payload := `{"data":"` + strings.Repeat("a", int(defaultMaxJSONResponseBytes)) + `"}`
|
||||
response := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Status: "200 OK",
|
||||
Header: http.Header{"Content-Type": {"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(payload)),
|
||||
}
|
||||
|
||||
_, _, err := decodeHTTPResponseForProtocol(response, ProtocolOpenAIResponses)
|
||||
if err == nil {
|
||||
t.Fatal("expected oversized default response to fail")
|
||||
}
|
||||
clientErr, ok := err.(*ClientError)
|
||||
if !ok || clientErr.Code != "response_too_large" {
|
||||
t.Fatalf("unexpected oversized response error: %T %v", err, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGeminiNativeStreamPreservesOfficialEvents(t *testing.T) {
|
||||
var requestedPath string
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
requestedPath = r.URL.RequestURI()
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.Header().Set("X-Goog-Request-Id", "goog-stream-1")
|
||||
_, _ = io.WriteString(w, "data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"hello\"}]}}],\"futureOfficialField\":true}\n\n")
|
||||
}))
|
||||
defer upstream.Close()
|
||||
|
||||
var events []StreamDeltaEvent
|
||||
response, err := (GeminiClient{HTTPClient: upstream.Client()}).Run(context.Background(), Request{
|
||||
Kind: "chat.completions", ModelType: "text", Model: "gemini-test", Stream: true,
|
||||
Body: map[string]any{"prompt": "hello"},
|
||||
Candidate: store.RuntimeModelCandidate{
|
||||
Provider: "gemini", BaseURL: upstream.URL, ProviderModelName: "gemini-test",
|
||||
Credentials: map[string]any{"apiKey": "test-key"},
|
||||
},
|
||||
StreamDelta: func(event StreamDeltaEvent) error {
|
||||
events = append(events, event)
|
||||
return nil
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Gemini stream failed: %v", err)
|
||||
}
|
||||
if !strings.Contains(requestedPath, ":streamGenerateContent") || !strings.Contains(requestedPath, "alt=sse") {
|
||||
t.Fatalf("unexpected Gemini stream endpoint: %s", requestedPath)
|
||||
}
|
||||
if len(events) != 1 || events[0].WireProtocol != ProtocolGeminiGenerateContent || events[0].Event["futureOfficialField"] != true {
|
||||
raw, _ := json.Marshal(events)
|
||||
t.Fatalf("official Gemini event was not preserved: %s", raw)
|
||||
}
|
||||
if response.Wire == nil || response.Wire.Headers["X-Goog-Request-Id"][0] != "goog-stream-1" {
|
||||
t.Fatalf("Gemini stream wire metadata was lost: %+v", response.Wire)
|
||||
}
|
||||
}
|
||||
@@ -1,6 +1,7 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"net/url"
|
||||
@@ -39,15 +40,81 @@ type Config struct {
|
||||
LocalGeneratedStorageDir string
|
||||
LocalUploadedStorageDir string
|
||||
LocalTempAssetTTLHours int
|
||||
LocalResultTTLHours int
|
||||
LocalResultMinFreeBytes int64
|
||||
LocalResultMaxBytes int64
|
||||
LocalResultMaxTaskBytes int64
|
||||
TaskProgressCallbackEnabled bool
|
||||
TaskProgressCallbackURL string
|
||||
TaskProgressCallbackTimeoutMS string
|
||||
TaskProgressCallbackMaxAttempts string
|
||||
TaskProgressCallbackTimeoutMS int
|
||||
TaskProgressCallbackMaxAttempts int
|
||||
TaskCleanupEnabled bool
|
||||
TaskRetentionDays int
|
||||
TaskAnalysisRetentionDays int
|
||||
TaskCleanupIntervalSeconds int
|
||||
TaskCleanupBatchSize int
|
||||
CORSAllowedOrigin string
|
||||
GlobalHTTPProxy string
|
||||
GlobalHTTPProxySource string
|
||||
PlatformProxyBypassIDs string
|
||||
LogLevel slog.Level
|
||||
BillingEngineMode string
|
||||
ProcessRole string
|
||||
DatabaseMaxConns int
|
||||
DatabaseCriticalMaxConns int
|
||||
DatabaseRiverMaxConns int
|
||||
DatabaseMinIdleConns int
|
||||
DatabaseMaxConnIdleSeconds int
|
||||
DatabaseIdleInTransactionTimeoutSeconds int
|
||||
DatabaseLockTimeoutSeconds int
|
||||
MediaRequestConcurrency int
|
||||
MediaMaterializationConcurrency int
|
||||
MediaImageNormalizationConcurrency int
|
||||
MediaOSSDirectEnabled bool
|
||||
MediaOSSEndpoint string
|
||||
MediaOSSBucket string
|
||||
MediaOSSAccessKeyID string
|
||||
MediaOSSAccessKeySecret string
|
||||
MediaOSSPublicBaseURL string
|
||||
MediaOSSObjectPrefix string
|
||||
AsyncQueueWorkerEnabled bool
|
||||
AsyncWorkerHardLimit int
|
||||
AsyncWorkerInstanceHardLimit int
|
||||
AsyncWorkerRefreshIntervalSeconds int
|
||||
AsyncWorkerLoadMode string
|
||||
AsyncAdmissionMicrobatchSize int
|
||||
AsyncAdmissionDispatcherEnabled bool
|
||||
AsyncAdmissionDispatcherConfigured bool
|
||||
RoutingMode string
|
||||
ExecutionPoolID string
|
||||
ExecutionPoolLabels string
|
||||
WorkerID string
|
||||
WorkerAdvertiseEndpoint string
|
||||
WorkerOrchestratorInstanceRef string
|
||||
WorkerEndpointAllowedSuffixes string
|
||||
WorkerEndpointAllowPrivate bool
|
||||
WorkerExecutionSecret string
|
||||
WorkerExecutionCAFile string
|
||||
WorkerExecutionCertFile string
|
||||
WorkerExecutionKeyFile string
|
||||
RouteProbeEnabled bool
|
||||
RouteProbeHotIntervalSeconds int
|
||||
RouteProbeColdIntervalSeconds int
|
||||
RouteProbeTimeoutMS int
|
||||
RouteProbeConcurrency int
|
||||
WorkerAutoscalingEnabled bool
|
||||
CapacityOrchestratorAdapter string
|
||||
CapacityPoolsJSON string
|
||||
WorkerTargetOutstandingPerReplica int
|
||||
WorkerScaleUpWindowSeconds int
|
||||
WorkerScaleDownStabilizationSeconds int
|
||||
WorkerDrainTimeoutSeconds int
|
||||
NodeMemoryTargetPercent int
|
||||
NodeMemoryHardPercent int
|
||||
NodeCPUTargetPercent int
|
||||
PostgresConnectionBudget int
|
||||
PostgresNonWorkerConnectionBudget int
|
||||
WorkerDatabaseMaxConns int
|
||||
}
|
||||
|
||||
func Load() Config {
|
||||
@@ -81,26 +148,291 @@ func Load() Config {
|
||||
LocalGeneratedStorageDir: env("AI_GATEWAY_GENERATED_STORAGE_DIR", env("LOCAL_GENERATED_STORAGE_DIR", env("AI_GATEWAY_STATIC_STORAGE_DIR", DefaultLocalGeneratedStorageDir))),
|
||||
LocalUploadedStorageDir: env("AI_GATEWAY_UPLOADED_STORAGE_DIR", env("LOCAL_UPLOADED_STORAGE_DIR", DefaultLocalUploadedStorageDir)),
|
||||
LocalTempAssetTTLHours: envInt("AI_GATEWAY_LOCAL_TEMP_ASSET_TTL_HOURS", 24),
|
||||
LocalResultTTLHours: envIntValidated("AI_GATEWAY_LOCAL_RESULT_TTL_HOURS", 24),
|
||||
LocalResultMinFreeBytes: envInt64Validated("AI_GATEWAY_LOCAL_RESULT_MIN_FREE_BYTES", 10*1024*1024*1024),
|
||||
LocalResultMaxBytes: envInt64Validated("AI_GATEWAY_LOCAL_RESULT_MAX_BYTES", 256*1024*1024),
|
||||
LocalResultMaxTaskBytes: envInt64Validated("AI_GATEWAY_LOCAL_RESULT_MAX_TASK_BYTES", 512*1024*1024),
|
||||
TaskProgressCallbackEnabled: env("TASK_PROGRESS_CALLBACK_ENABLED", "true") == "true",
|
||||
TaskProgressCallbackURL: env("TASK_PROGRESS_CALLBACK_URL",
|
||||
strings.TrimRight(env("SERVER_MAIN_BASE_URL", "http://localhost:3000"), "/")+"/internal/platform/task-progress-callbacks",
|
||||
),
|
||||
TaskProgressCallbackTimeoutMS: env("TASK_PROGRESS_CALLBACK_TIMEOUT_MS", "5000"),
|
||||
TaskProgressCallbackMaxAttempts: env("TASK_PROGRESS_CALLBACK_MAX_ATTEMPTS", "10"),
|
||||
TaskProgressCallbackTimeoutMS: envIntValidated("TASK_PROGRESS_CALLBACK_TIMEOUT_MS", 5000),
|
||||
TaskProgressCallbackMaxAttempts: envIntValidated("TASK_PROGRESS_CALLBACK_MAX_ATTEMPTS", 10),
|
||||
TaskCleanupEnabled: env("AI_GATEWAY_TASK_CLEANUP_ENABLED", "false") == "true",
|
||||
TaskRetentionDays: envIntValidated("AI_GATEWAY_TASK_RETENTION_DAYS", 30),
|
||||
TaskAnalysisRetentionDays: envIntValidated("AI_GATEWAY_TASK_ANALYSIS_RETENTION_DAYS", 7),
|
||||
TaskCleanupIntervalSeconds: envIntValidated("AI_GATEWAY_TASK_CLEANUP_INTERVAL_SECONDS", 300),
|
||||
TaskCleanupBatchSize: envIntValidated("AI_GATEWAY_TASK_CLEANUP_BATCH_SIZE", 1000),
|
||||
CORSAllowedOrigin: env("CORS_ALLOWED_ORIGIN", "http://localhost:5178,http://127.0.0.1:5178"),
|
||||
GlobalHTTPProxy: globalProxy.HTTPProxy,
|
||||
GlobalHTTPProxySource: globalProxy.Source,
|
||||
PlatformProxyBypassIDs: env("AI_GATEWAY_PLATFORM_PROXY_BYPASS_IDS", ""),
|
||||
LogLevel: logLevel(env("LOG_LEVEL", "info")),
|
||||
BillingEngineMode: strings.ToLower(env("BILLING_ENGINE_MODE", "observe")),
|
||||
ProcessRole: strings.ToLower(strings.TrimSpace(env("AI_GATEWAY_PROCESS_ROLE", "all"))),
|
||||
DatabaseMaxConns: envInt("AI_GATEWAY_DATABASE_MAX_CONNS", 0),
|
||||
DatabaseCriticalMaxConns: envOptionalIntValidated("AI_GATEWAY_DATABASE_CRITICAL_MAX_CONNS", 0),
|
||||
DatabaseRiverMaxConns: envOptionalIntValidated("AI_GATEWAY_DATABASE_RIVER_MAX_CONNS", 0),
|
||||
DatabaseMinIdleConns: envOptionalIntValidated("AI_GATEWAY_DATABASE_MIN_IDLE_CONNS", 0),
|
||||
DatabaseMaxConnIdleSeconds: envOptionalIntValidated("AI_GATEWAY_DATABASE_MAX_CONN_IDLE_SECONDS", 0),
|
||||
|
||||
DatabaseIdleInTransactionTimeoutSeconds: envOptionalIntValidated(
|
||||
"AI_GATEWAY_DATABASE_IDLE_IN_TRANSACTION_TIMEOUT_SECONDS",
|
||||
60,
|
||||
),
|
||||
DatabaseLockTimeoutSeconds: envOptionalIntValidated("AI_GATEWAY_DATABASE_LOCK_TIMEOUT_SECONDS", 30),
|
||||
|
||||
MediaRequestConcurrency: envIntValidated("AI_GATEWAY_MEDIA_REQUEST_CONCURRENCY", 16),
|
||||
MediaMaterializationConcurrency: envIntValidated("AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY", 8),
|
||||
MediaImageNormalizationConcurrency: envIntValidated(
|
||||
"AI_GATEWAY_MEDIA_IMAGE_NORMALIZATION_CONCURRENCY",
|
||||
2,
|
||||
),
|
||||
MediaOSSDirectEnabled: env("AI_GATEWAY_MEDIA_OSS_DIRECT_ENABLED", "false") == "true",
|
||||
MediaOSSEndpoint: strings.TrimRight(env("AI_GATEWAY_MEDIA_OSS_ENDPOINT", ""), "/"),
|
||||
MediaOSSBucket: strings.TrimSpace(env("AI_GATEWAY_MEDIA_OSS_BUCKET", "")),
|
||||
MediaOSSAccessKeyID: strings.TrimSpace(env("AI_GATEWAY_MEDIA_OSS_ACCESS_KEY_ID", "")),
|
||||
MediaOSSAccessKeySecret: strings.TrimSpace(env("AI_GATEWAY_MEDIA_OSS_ACCESS_KEY_SECRET", "")),
|
||||
MediaOSSPublicBaseURL: strings.TrimRight(env("AI_GATEWAY_MEDIA_OSS_PUBLIC_BASE_URL", ""), "/"),
|
||||
MediaOSSObjectPrefix: strings.Trim(env("AI_GATEWAY_MEDIA_OSS_OBJECT_PREFIX", "easyai-ai-gateway/media"), "/"),
|
||||
AsyncQueueWorkerEnabled: env("AI_GATEWAY_ASYNC_QUEUE_WORKER_ENABLED", "true") == "true",
|
||||
AsyncWorkerHardLimit: envIntValidated(
|
||||
"AI_GATEWAY_ASYNC_WORKER_GLOBAL_HARD_LIMIT",
|
||||
envOptionalIntValidated("AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT", 2048),
|
||||
),
|
||||
AsyncWorkerInstanceHardLimit: envIntValidated("AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT", 32),
|
||||
AsyncWorkerRefreshIntervalSeconds: envIntValidated("AI_GATEWAY_ASYNC_WORKER_REFRESH_INTERVAL_SECONDS", 5),
|
||||
AsyncWorkerLoadMode: strings.ToLower(strings.TrimSpace(env("AI_GATEWAY_WORKER_LOAD_MODE", "adaptive"))),
|
||||
AsyncAdmissionMicrobatchSize: envIntValidated("AI_GATEWAY_ASYNC_ADMISSION_MICROBATCH_SIZE", 8),
|
||||
AsyncAdmissionDispatcherEnabled: envValue("AI_GATEWAY_ASYNC_ADMISSION_DISPATCHER_ENABLED") == "true",
|
||||
AsyncAdmissionDispatcherConfigured: envValue(
|
||||
"AI_GATEWAY_ASYNC_ADMISSION_DISPATCHER_ENABLED",
|
||||
) != "",
|
||||
RoutingMode: strings.ToLower(strings.TrimSpace(env("AI_GATEWAY_ROUTING_MODE", "legacy"))),
|
||||
ExecutionPoolID: strings.TrimSpace(env("AI_GATEWAY_EXECUTION_POOL_ID", "legacy-default")),
|
||||
ExecutionPoolLabels: strings.TrimSpace(env("AI_GATEWAY_EXECUTION_POOL_LABELS", "{}")),
|
||||
WorkerID: strings.TrimSpace(env("AI_GATEWAY_WORKER_ID", "")),
|
||||
WorkerAdvertiseEndpoint: workerAdvertiseEndpoint(),
|
||||
WorkerOrchestratorInstanceRef: strings.TrimSpace(env("AI_GATEWAY_ORCHESTRATOR_INSTANCE_REF", "")),
|
||||
WorkerEndpointAllowedSuffixes: strings.TrimSpace(env("AI_GATEWAY_WORKER_ENDPOINT_ALLOWED_SUFFIXES", "svc,cluster.local")),
|
||||
WorkerEndpointAllowPrivate: env("AI_GATEWAY_WORKER_ENDPOINT_ALLOW_PRIVATE", "true") == "true",
|
||||
WorkerExecutionSecret: env("AI_GATEWAY_WORKER_EXECUTION_SECRET", env("SERVER_MAIN_INTERNAL_SECRET", env("SERVER_MAIN_INTERNAL_TOKEN", ""))),
|
||||
WorkerExecutionCAFile: strings.TrimSpace(env("AI_GATEWAY_WORKER_EXECUTION_CA_FILE", "")),
|
||||
WorkerExecutionCertFile: strings.TrimSpace(env("AI_GATEWAY_WORKER_EXECUTION_CERT_FILE", "")),
|
||||
WorkerExecutionKeyFile: strings.TrimSpace(env("AI_GATEWAY_WORKER_EXECUTION_KEY_FILE", "")),
|
||||
RouteProbeEnabled: env("AI_GATEWAY_ROUTE_PROBE_ENABLED", "true") == "true",
|
||||
RouteProbeHotIntervalSeconds: envIntValidated("AI_GATEWAY_ROUTE_PROBE_HOT_INTERVAL_SECONDS", 15),
|
||||
RouteProbeColdIntervalSeconds: envIntValidated("AI_GATEWAY_ROUTE_PROBE_COLD_INTERVAL_SECONDS", 60),
|
||||
RouteProbeTimeoutMS: envIntValidated("AI_GATEWAY_ROUTE_PROBE_TIMEOUT_MS", 3000),
|
||||
RouteProbeConcurrency: envIntValidated("AI_GATEWAY_ROUTE_PROBE_CONCURRENCY", 8),
|
||||
WorkerAutoscalingEnabled: env("AI_GATEWAY_WORKER_AUTOSCALING_ENABLED", "false") == "true",
|
||||
CapacityOrchestratorAdapter: strings.ToLower(strings.TrimSpace(env("AI_GATEWAY_CAPACITY_ORCHESTRATOR_ADAPTER", "kubernetes"))),
|
||||
CapacityPoolsJSON: strings.TrimSpace(env("AI_GATEWAY_CAPACITY_POOLS", "")),
|
||||
WorkerTargetOutstandingPerReplica: envOptionalIntValidated("AI_GATEWAY_WORKER_TARGET_OUTSTANDING_PER_REPLICA", 0),
|
||||
WorkerScaleUpWindowSeconds: envIntValidated("AI_GATEWAY_WORKER_SCALE_UP_WINDOW_SECONDS", 20),
|
||||
WorkerScaleDownStabilizationSeconds: envIntValidated(
|
||||
"AI_GATEWAY_WORKER_SCALE_DOWN_STABILIZATION_SECONDS",
|
||||
600,
|
||||
),
|
||||
WorkerDrainTimeoutSeconds: envIntValidated("AI_GATEWAY_WORKER_DRAIN_TIMEOUT_SECONDS", 600),
|
||||
NodeMemoryTargetPercent: envIntValidated("AI_GATEWAY_NODE_MEMORY_TARGET_PERCENT", 75),
|
||||
NodeMemoryHardPercent: envIntValidated("AI_GATEWAY_NODE_MEMORY_HARD_PERCENT", 85),
|
||||
NodeCPUTargetPercent: envIntValidated("AI_GATEWAY_NODE_CPU_TARGET_PERCENT", 70),
|
||||
PostgresConnectionBudget: envIntValidated("AI_GATEWAY_POSTGRES_CONNECTION_BUDGET", 150),
|
||||
PostgresNonWorkerConnectionBudget: envIntValidated(
|
||||
"AI_GATEWAY_POSTGRES_NON_WORKER_CONNECTION_BUDGET",
|
||||
72,
|
||||
),
|
||||
WorkerDatabaseMaxConns: envIntValidated("AI_GATEWAY_WORKER_DATABASE_MAX_CONNS", 32),
|
||||
}
|
||||
}
|
||||
|
||||
func (c Config) Validate() error {
|
||||
switch strings.ToLower(strings.TrimSpace(c.ProcessRole)) {
|
||||
case "", "all", "api", "worker", "capacity-controller":
|
||||
default:
|
||||
return errors.New("AI_GATEWAY_PROCESS_ROLE must be all, api, worker, or capacity-controller")
|
||||
}
|
||||
if c.DatabaseMaxConns < 0 || c.DatabaseMaxConns > 1000 {
|
||||
return errors.New("AI_GATEWAY_DATABASE_MAX_CONNS must be between 1 and 1000 when configured")
|
||||
}
|
||||
if c.DatabaseMinIdleConns < 0 || c.DatabaseMinIdleConns > 1000 ||
|
||||
(c.DatabaseMaxConns > 0 && c.DatabaseMinIdleConns > c.DatabaseMaxConns) {
|
||||
return errors.New("AI_GATEWAY_DATABASE_MIN_IDLE_CONNS must be between 0 and AI_GATEWAY_DATABASE_MAX_CONNS")
|
||||
}
|
||||
if c.DatabaseCriticalMaxConns < 0 || c.DatabaseCriticalMaxConns > 64 ||
|
||||
(c.DatabaseMaxConns > 0 && c.DatabaseCriticalMaxConns >= c.DatabaseMaxConns) {
|
||||
return errors.New("AI_GATEWAY_DATABASE_CRITICAL_MAX_CONNS must be between 0 and less than AI_GATEWAY_DATABASE_MAX_CONNS")
|
||||
}
|
||||
if c.DatabaseRiverMaxConns < 0 || c.DatabaseRiverMaxConns > 256 ||
|
||||
(c.DatabaseMaxConns > 0 &&
|
||||
c.DatabaseCriticalMaxConns+c.DatabaseRiverMaxConns >= c.DatabaseMaxConns) {
|
||||
return errors.New("critical and River database pools must leave at least one execution connection")
|
||||
}
|
||||
if c.DatabaseMaxConns > 0 &&
|
||||
c.DatabaseMinIdleConns > c.DatabaseMaxConns-c.DatabaseCriticalMaxConns-c.DatabaseRiverMaxConns {
|
||||
return errors.New("AI_GATEWAY_DATABASE_MIN_IDLE_CONNS must not exceed the execution pool after reserving critical and River connections")
|
||||
}
|
||||
if c.DatabaseMaxConnIdleSeconds < 0 || c.DatabaseMaxConnIdleSeconds > 3600 {
|
||||
return errors.New("AI_GATEWAY_DATABASE_MAX_CONN_IDLE_SECONDS must be between 0 and 3600")
|
||||
}
|
||||
if c.DatabaseIdleInTransactionTimeoutSeconds < 0 || c.DatabaseIdleInTransactionTimeoutSeconds > 3600 {
|
||||
return errors.New("AI_GATEWAY_DATABASE_IDLE_IN_TRANSACTION_TIMEOUT_SECONDS must be between 0 and 3600")
|
||||
}
|
||||
if c.DatabaseLockTimeoutSeconds < 0 || c.DatabaseLockTimeoutSeconds > 3600 {
|
||||
return errors.New("AI_GATEWAY_DATABASE_LOCK_TIMEOUT_SECONDS must be between 0 and 3600")
|
||||
}
|
||||
if c.MediaMaterializationConcurrency != 0 && (c.MediaMaterializationConcurrency < 1 || c.MediaMaterializationConcurrency > 256) {
|
||||
return errors.New("AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY must be between 1 and 256")
|
||||
}
|
||||
if c.MediaRequestConcurrency != 0 && (c.MediaRequestConcurrency < 1 || c.MediaRequestConcurrency > 1024) {
|
||||
return errors.New("AI_GATEWAY_MEDIA_REQUEST_CONCURRENCY must be between 1 and 1024")
|
||||
}
|
||||
if c.MediaImageNormalizationConcurrency != 0 && (c.MediaImageNormalizationConcurrency < 1 || c.MediaImageNormalizationConcurrency > 64) {
|
||||
return errors.New("AI_GATEWAY_MEDIA_IMAGE_NORMALIZATION_CONCURRENCY must be between 1 and 64")
|
||||
}
|
||||
if c.MediaOSSDirectEnabled && !c.RunsCapacityController() {
|
||||
if strings.TrimSpace(c.MediaOSSAccessKeyID) == "" || strings.TrimSpace(c.MediaOSSAccessKeySecret) == "" {
|
||||
return errors.New("AI_GATEWAY_MEDIA_OSS_ACCESS_KEY_ID and AI_GATEWAY_MEDIA_OSS_ACCESS_KEY_SECRET are required when direct media OSS is enabled")
|
||||
}
|
||||
if !validMediaOSSBucket(c.MediaOSSBucket) {
|
||||
return errors.New("AI_GATEWAY_MEDIA_OSS_BUCKET is invalid")
|
||||
}
|
||||
if !validMediaOSSURL(c.MediaOSSEndpoint, c.AppEnv) {
|
||||
return errors.New("AI_GATEWAY_MEDIA_OSS_ENDPOINT must be an HTTPS URL")
|
||||
}
|
||||
if !validMediaOSSURL(c.MediaOSSPublicBaseURL, c.AppEnv) {
|
||||
return errors.New("AI_GATEWAY_MEDIA_OSS_PUBLIC_BASE_URL must be an HTTPS URL")
|
||||
}
|
||||
if !validMediaOSSObjectPrefix(c.MediaOSSObjectPrefix) {
|
||||
return errors.New("AI_GATEWAY_MEDIA_OSS_OBJECT_PREFIX is invalid")
|
||||
}
|
||||
}
|
||||
switch strings.ToLower(strings.TrimSpace(c.BillingEngineMode)) {
|
||||
case "", "observe", "enforce", "hold":
|
||||
default:
|
||||
return errors.New("BILLING_ENGINE_MODE must be observe, enforce, or hold")
|
||||
}
|
||||
if c.AsyncWorkerHardLimit < 1 || c.AsyncWorkerHardLimit > 10000 {
|
||||
return errors.New("AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT must be between 1 and 10000")
|
||||
}
|
||||
if c.AsyncWorkerInstanceHardLimit < 1 || c.AsyncWorkerInstanceHardLimit > 10000 {
|
||||
return errors.New("AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT must be between 1 and 10000")
|
||||
}
|
||||
if c.AsyncWorkerRefreshIntervalSeconds < 1 {
|
||||
return errors.New("AI_GATEWAY_ASYNC_WORKER_REFRESH_INTERVAL_SECONDS must be positive")
|
||||
}
|
||||
switch strings.ToLower(strings.TrimSpace(c.AsyncWorkerLoadMode)) {
|
||||
case "", "adaptive", "legacy":
|
||||
default:
|
||||
return errors.New("AI_GATEWAY_WORKER_LOAD_MODE must be adaptive or legacy")
|
||||
}
|
||||
switch strings.ToLower(strings.TrimSpace(c.RoutingMode)) {
|
||||
case "", "legacy", "shadow", "enforced":
|
||||
default:
|
||||
return errors.New("AI_GATEWAY_ROUTING_MODE must be legacy, shadow, or enforced")
|
||||
}
|
||||
if strings.TrimSpace(c.ExecutionPoolID) != "" && !validExecutionPoolID(c.ExecutionPoolID) {
|
||||
return errors.New("AI_GATEWAY_EXECUTION_POOL_ID is invalid")
|
||||
}
|
||||
if strings.TrimSpace(c.ExecutionPoolLabels) != "" && !jsonObject(c.ExecutionPoolLabels) {
|
||||
return errors.New("AI_GATEWAY_EXECUTION_POOL_LABELS must be a JSON object")
|
||||
}
|
||||
if c.RouteProbeHotIntervalSeconds != 0 && (c.RouteProbeHotIntervalSeconds < 5 || c.RouteProbeHotIntervalSeconds > 300) {
|
||||
return errors.New("AI_GATEWAY_ROUTE_PROBE_HOT_INTERVAL_SECONDS must be between 5 and 300")
|
||||
}
|
||||
if c.RouteProbeColdIntervalSeconds != 0 && (c.RouteProbeColdIntervalSeconds < c.RouteProbeHotIntervalSeconds || c.RouteProbeColdIntervalSeconds > 3600) {
|
||||
return errors.New("AI_GATEWAY_ROUTE_PROBE_COLD_INTERVAL_SECONDS must be between the hot interval and 3600")
|
||||
}
|
||||
if c.RouteProbeTimeoutMS != 0 && (c.RouteProbeTimeoutMS < 250 || c.RouteProbeTimeoutMS > 10000) {
|
||||
return errors.New("AI_GATEWAY_ROUTE_PROBE_TIMEOUT_MS must be between 250 and 10000")
|
||||
}
|
||||
if c.RouteProbeConcurrency != 0 && (c.RouteProbeConcurrency < 1 || c.RouteProbeConcurrency > 64) {
|
||||
return errors.New("AI_GATEWAY_ROUTE_PROBE_CONCURRENCY must be between 1 and 64")
|
||||
}
|
||||
if strings.EqualFold(c.RoutingMode, "enforced") && len(c.WorkerExecutionSecret) < 32 {
|
||||
return errors.New("AI_GATEWAY_WORKER_EXECUTION_SECRET must be at least 32 bytes in enforced routing mode")
|
||||
}
|
||||
if (c.WorkerExecutionCertFile == "") != (c.WorkerExecutionKeyFile == "") {
|
||||
return errors.New("AI_GATEWAY_WORKER_EXECUTION_CERT_FILE and AI_GATEWAY_WORKER_EXECUTION_KEY_FILE must be configured together")
|
||||
}
|
||||
if c.AsyncAdmissionMicrobatchSize < 1 || c.AsyncAdmissionMicrobatchSize > 32 {
|
||||
return errors.New("AI_GATEWAY_ASYNC_ADMISSION_MICROBATCH_SIZE must be between 1 and 32")
|
||||
}
|
||||
pools, err := c.CapacityPools()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if c.RunsCapacityController() && len(pools) == 0 {
|
||||
return errors.New("AI_GATEWAY_CAPACITY_POOLS must contain at least one pool for the capacity-controller role")
|
||||
}
|
||||
switch strings.ToLower(strings.TrimSpace(c.CapacityOrchestratorAdapter)) {
|
||||
case "", "kubernetes", "static":
|
||||
default:
|
||||
return errors.New("AI_GATEWAY_CAPACITY_ORCHESTRATOR_ADAPTER must be kubernetes or static")
|
||||
}
|
||||
if c.WorkerTargetOutstandingPerReplica < 0 || c.WorkerTargetOutstandingPerReplica > 100000 {
|
||||
return errors.New("AI_GATEWAY_WORKER_TARGET_OUTSTANDING_PER_REPLICA must be between 0 and 100000")
|
||||
}
|
||||
if c.WorkerScaleUpWindowSeconds != 0 && (c.WorkerScaleUpWindowSeconds < 5 || c.WorkerScaleUpWindowSeconds > 600) {
|
||||
return errors.New("AI_GATEWAY_WORKER_SCALE_UP_WINDOW_SECONDS must be between 5 and 600")
|
||||
}
|
||||
if c.WorkerScaleDownStabilizationSeconds != 0 &&
|
||||
(c.WorkerScaleDownStabilizationSeconds < 60 || c.WorkerScaleDownStabilizationSeconds > 86400) {
|
||||
return errors.New("AI_GATEWAY_WORKER_SCALE_DOWN_STABILIZATION_SECONDS must be between 60 and 86400")
|
||||
}
|
||||
if c.WorkerDrainTimeoutSeconds != 0 && (c.WorkerDrainTimeoutSeconds < 60 || c.WorkerDrainTimeoutSeconds > 86400) {
|
||||
return errors.New("AI_GATEWAY_WORKER_DRAIN_TIMEOUT_SECONDS must be between 60 and 86400")
|
||||
}
|
||||
if (c.NodeMemoryTargetPercent != 0 || c.NodeMemoryHardPercent != 0) &&
|
||||
(c.NodeMemoryTargetPercent < 50 || c.NodeMemoryTargetPercent > 90 ||
|
||||
c.NodeMemoryHardPercent <= c.NodeMemoryTargetPercent || c.NodeMemoryHardPercent > 95) {
|
||||
return errors.New("Worker node memory target/hard percentages must be ordered within 50..95")
|
||||
}
|
||||
if c.NodeCPUTargetPercent != 0 && (c.NodeCPUTargetPercent < 40 || c.NodeCPUTargetPercent > 90) {
|
||||
return errors.New("AI_GATEWAY_NODE_CPU_TARGET_PERCENT must be between 40 and 90")
|
||||
}
|
||||
if c.PostgresConnectionBudget < 0 || c.PostgresConnectionBudget > 1000 {
|
||||
return errors.New("AI_GATEWAY_POSTGRES_CONNECTION_BUDGET must be between 1 and 1000")
|
||||
}
|
||||
if c.PostgresNonWorkerConnectionBudget < 0 ||
|
||||
(c.PostgresConnectionBudget > 0 &&
|
||||
c.PostgresNonWorkerConnectionBudget >= c.PostgresConnectionBudget) {
|
||||
return errors.New("AI_GATEWAY_POSTGRES_NON_WORKER_CONNECTION_BUDGET must be below the PostgreSQL connection budget")
|
||||
}
|
||||
if c.WorkerDatabaseMaxConns < 0 || c.WorkerDatabaseMaxConns > 256 {
|
||||
return errors.New("AI_GATEWAY_WORKER_DATABASE_MAX_CONNS must be between 1 and 256")
|
||||
}
|
||||
if c.TaskProgressCallbackTimeoutMS != 0 && (c.TaskProgressCallbackTimeoutMS < 100 || c.TaskProgressCallbackTimeoutMS > 60000) {
|
||||
return errors.New("TASK_PROGRESS_CALLBACK_TIMEOUT_MS must be between 100 and 60000")
|
||||
}
|
||||
if c.TaskProgressCallbackMaxAttempts != 0 && (c.TaskProgressCallbackMaxAttempts < 1 || c.TaskProgressCallbackMaxAttempts > 100) {
|
||||
return errors.New("TASK_PROGRESS_CALLBACK_MAX_ATTEMPTS must be between 1 and 100")
|
||||
}
|
||||
if c.TaskRetentionDays != 0 && (c.TaskRetentionDays < 1 || c.TaskRetentionDays > 3650) {
|
||||
return errors.New("AI_GATEWAY_TASK_RETENTION_DAYS must be between 1 and 3650")
|
||||
}
|
||||
if c.TaskAnalysisRetentionDays != 0 && (c.TaskAnalysisRetentionDays < 1 || (c.TaskRetentionDays != 0 && c.TaskAnalysisRetentionDays > c.TaskRetentionDays)) {
|
||||
return errors.New("AI_GATEWAY_TASK_ANALYSIS_RETENTION_DAYS must be between 1 and AI_GATEWAY_TASK_RETENTION_DAYS")
|
||||
}
|
||||
if c.TaskCleanupIntervalSeconds != 0 && c.TaskCleanupIntervalSeconds < 60 {
|
||||
return errors.New("AI_GATEWAY_TASK_CLEANUP_INTERVAL_SECONDS must be at least 60")
|
||||
}
|
||||
if c.TaskCleanupBatchSize != 0 && (c.TaskCleanupBatchSize < 100 || c.TaskCleanupBatchSize > 5000) {
|
||||
return errors.New("AI_GATEWAY_TASK_CLEANUP_BATCH_SIZE must be between 100 and 5000")
|
||||
}
|
||||
if c.LocalResultTTLHours != 0 && (c.LocalResultTTLHours < 1 || c.LocalResultTTLHours > 24*30) {
|
||||
return errors.New("AI_GATEWAY_LOCAL_RESULT_TTL_HOURS must be between 1 and 720")
|
||||
}
|
||||
if c.LocalResultMinFreeBytes < 0 {
|
||||
return errors.New("AI_GATEWAY_LOCAL_RESULT_MIN_FREE_BYTES must not be negative")
|
||||
}
|
||||
if c.LocalResultMaxBytes != 0 && c.LocalResultMaxBytes < 1024 {
|
||||
return errors.New("AI_GATEWAY_LOCAL_RESULT_MAX_BYTES must be at least 1024")
|
||||
}
|
||||
if c.LocalResultMaxTaskBytes != 0 && c.LocalResultMaxTaskBytes < c.LocalResultMaxBytes {
|
||||
return errors.New("AI_GATEWAY_LOCAL_RESULT_MAX_TASK_BYTES must be at least AI_GATEWAY_LOCAL_RESULT_MAX_BYTES")
|
||||
}
|
||||
switch strings.ToLower(strings.TrimSpace(c.IdentitySecretStore)) {
|
||||
case "":
|
||||
case "file":
|
||||
@@ -124,11 +456,174 @@ func (c Config) Validate() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func validMediaOSSBucket(value string) bool {
|
||||
value = strings.TrimSpace(value)
|
||||
if len(value) < 3 || len(value) > 63 {
|
||||
return false
|
||||
}
|
||||
for index, item := range value {
|
||||
if (item >= 'a' && item <= 'z') || (item >= '0' && item <= '9') || (item == '-' && index > 0 && index < len(value)-1) {
|
||||
continue
|
||||
}
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func validMediaOSSURL(value string, appEnv string) bool {
|
||||
parsed, err := url.Parse(strings.TrimSpace(value))
|
||||
if err != nil || parsed.Host == "" || parsed.RawQuery != "" || parsed.Fragment != "" || parsed.User != nil {
|
||||
return false
|
||||
}
|
||||
if parsed.Scheme == "https" {
|
||||
return true
|
||||
}
|
||||
host := strings.ToLower(parsed.Hostname())
|
||||
return !strings.EqualFold(strings.TrimSpace(appEnv), "production") &&
|
||||
parsed.Scheme == "http" && (host == "127.0.0.1" || host == "localhost" || host == "::1")
|
||||
}
|
||||
|
||||
func validMediaOSSObjectPrefix(value string) bool {
|
||||
value = strings.Trim(value, "/")
|
||||
if value == "" || len(value) > 256 || strings.Contains(value, "..") {
|
||||
return false
|
||||
}
|
||||
for _, item := range value {
|
||||
if (item >= 'a' && item <= 'z') || (item >= 'A' && item <= 'Z') || (item >= '0' && item <= '9') ||
|
||||
item == '-' || item == '_' || item == '/' {
|
||||
continue
|
||||
}
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (c Config) EffectiveProcessRole() string {
|
||||
role := strings.ToLower(strings.TrimSpace(c.ProcessRole))
|
||||
if role == "" {
|
||||
return "all"
|
||||
}
|
||||
return role
|
||||
}
|
||||
|
||||
func (c Config) RunsPublicHTTP() bool {
|
||||
switch c.EffectiveProcessRole() {
|
||||
case "worker", "capacity-controller":
|
||||
return false
|
||||
default:
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
func (c Config) RunsAsyncExecutionWorker() bool {
|
||||
switch c.EffectiveProcessRole() {
|
||||
case "api", "capacity-controller":
|
||||
return false
|
||||
case "worker":
|
||||
return true
|
||||
default:
|
||||
return c.AsyncQueueWorkerEnabled
|
||||
}
|
||||
}
|
||||
|
||||
// RunsAsyncAdmissionDispatcher keeps the legacy all-in-one/Worker behaviour
|
||||
// unless a deployment explicitly separates admission coordination from task
|
||||
// execution. Production enables this on both API sites so the instance nearest
|
||||
// the current PostgreSQL primary wins the existing database scope locks, while
|
||||
// remote Worker replicas only execute River jobs.
|
||||
func (c Config) RunsAsyncAdmissionDispatcher() bool {
|
||||
if c.AsyncAdmissionDispatcherConfigured {
|
||||
return c.AsyncAdmissionDispatcherEnabled
|
||||
}
|
||||
return c.RunsAsyncExecutionWorker()
|
||||
}
|
||||
|
||||
func (c Config) RunsBackgroundWorkers() bool {
|
||||
switch c.EffectiveProcessRole() {
|
||||
case "api", "capacity-controller":
|
||||
return false
|
||||
default:
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
func (c Config) RunsCapacityController() bool {
|
||||
return c.EffectiveProcessRole() == "capacity-controller"
|
||||
}
|
||||
|
||||
func (c Config) RunsRouteProber() bool {
|
||||
return c.RunsAsyncExecutionWorker() && c.RouteProbeEnabled && !strings.EqualFold(c.RoutingMode, "legacy")
|
||||
}
|
||||
|
||||
func workerAdvertiseEndpoint() string {
|
||||
return strings.TrimRight(strings.TrimSpace(os.Getenv("AI_GATEWAY_WORKER_ADVERTISE_ENDPOINT")), "/")
|
||||
}
|
||||
|
||||
func validExecutionPoolID(value string) bool {
|
||||
value = strings.TrimSpace(value)
|
||||
if value == "" || len(value) > 128 {
|
||||
return false
|
||||
}
|
||||
for index, character := range value {
|
||||
valid := (character >= 'a' && character <= 'z') ||
|
||||
(character >= 'A' && character <= 'Z') ||
|
||||
(character >= '0' && character <= '9') ||
|
||||
character == '.' || character == '_' || character == ':' || character == '/' || character == '-'
|
||||
if !valid || (index == 0 && !((character >= 'a' && character <= 'z') || (character >= 'A' && character <= 'Z') || (character >= '0' && character <= '9'))) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func jsonObject(value string) bool {
|
||||
var object map[string]any
|
||||
return json.Unmarshal([]byte(value), &object) == nil && object != nil
|
||||
}
|
||||
|
||||
type GlobalHTTPProxyStatus struct {
|
||||
HTTPProxy string
|
||||
Source string
|
||||
}
|
||||
|
||||
type ExecutionPoolCapacityConfig struct {
|
||||
ID string `json:"id"`
|
||||
AdapterRef string `json:"adapterRef"`
|
||||
BootstrapReplicas int `json:"bootstrapReplicas"`
|
||||
MinReplicas int `json:"minReplicas"`
|
||||
MaxReplicas int `json:"maxReplicas"`
|
||||
}
|
||||
|
||||
func (c Config) CapacityPools() ([]ExecutionPoolCapacityConfig, error) {
|
||||
raw := strings.TrimSpace(c.CapacityPoolsJSON)
|
||||
if raw == "" {
|
||||
return nil, nil
|
||||
}
|
||||
var pools []ExecutionPoolCapacityConfig
|
||||
if err := json.Unmarshal([]byte(raw), &pools); err != nil {
|
||||
return nil, errors.New("AI_GATEWAY_CAPACITY_POOLS must be a JSON array")
|
||||
}
|
||||
if len(pools) == 0 {
|
||||
return nil, errors.New("AI_GATEWAY_CAPACITY_POOLS must contain at least one pool")
|
||||
}
|
||||
seen := make(map[string]struct{}, len(pools))
|
||||
for index := range pools {
|
||||
pools[index].ID = strings.TrimSpace(pools[index].ID)
|
||||
pools[index].AdapterRef = strings.TrimSpace(pools[index].AdapterRef)
|
||||
if pools[index].AdapterRef == "" {
|
||||
pools[index].AdapterRef = pools[index].ID
|
||||
}
|
||||
if !validExecutionPoolID(pools[index].ID) || pools[index].MinReplicas < 0 || pools[index].MaxReplicas < pools[index].MinReplicas || pools[index].MaxReplicas > 64 || pools[index].BootstrapReplicas < pools[index].MinReplicas || pools[index].BootstrapReplicas > pools[index].MaxReplicas {
|
||||
return nil, errors.New("AI_GATEWAY_CAPACITY_POOLS contains an invalid pool definition")
|
||||
}
|
||||
if _, duplicate := seen[pools[index].ID]; duplicate {
|
||||
return nil, errors.New("AI_GATEWAY_CAPACITY_POOLS contains duplicate pool IDs")
|
||||
}
|
||||
seen[pools[index].ID] = struct{}{}
|
||||
}
|
||||
return pools, nil
|
||||
}
|
||||
|
||||
func LoadGlobalHTTPProxyStatus() GlobalHTTPProxyStatus {
|
||||
for _, key := range []string{
|
||||
"AI_GATEWAY_GLOBAL_HTTP_PROXY",
|
||||
@@ -210,6 +705,42 @@ func envInt(key string, fallback int) int {
|
||||
return parsed
|
||||
}
|
||||
|
||||
func envIntValidated(key string, fallback int) int {
|
||||
value := envValue(key)
|
||||
if value == "" {
|
||||
return fallback
|
||||
}
|
||||
parsed, err := strconv.Atoi(value)
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return parsed
|
||||
}
|
||||
|
||||
func envOptionalIntValidated(key string, fallback int) int {
|
||||
value := envValue(key)
|
||||
if value == "" {
|
||||
return fallback
|
||||
}
|
||||
parsed, err := strconv.Atoi(value)
|
||||
if err != nil {
|
||||
return -1
|
||||
}
|
||||
return parsed
|
||||
}
|
||||
|
||||
func envInt64Validated(key string, fallback int64) int64 {
|
||||
value := envValue(key)
|
||||
if value == "" {
|
||||
return fallback
|
||||
}
|
||||
parsed, err := strconv.ParseInt(value, 10, 64)
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return parsed
|
||||
}
|
||||
|
||||
func logLevel(value string) slog.Level {
|
||||
switch strings.ToLower(value) {
|
||||
case "debug":
|
||||
|
||||
@@ -26,7 +26,13 @@ func TestLoadIdentitySecretStoreUsesNewEnvironmentNamesAndIgnoresLegacyBusinessV
|
||||
}
|
||||
|
||||
func TestValidateIdentityFileSecretStoreRequiresDirectory(t *testing.T) {
|
||||
cfg := Config{IdentitySecretStore: "file"}
|
||||
cfg := Config{
|
||||
IdentitySecretStore: "file",
|
||||
AsyncWorkerHardLimit: 2048,
|
||||
AsyncWorkerInstanceHardLimit: 32,
|
||||
AsyncWorkerRefreshIntervalSeconds: 5,
|
||||
AsyncAdmissionMicrobatchSize: 8,
|
||||
}
|
||||
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "IDENTITY_SECRET_DIR") {
|
||||
t.Fatalf("Validate() error = %v, want missing identity secret directory", err)
|
||||
}
|
||||
@@ -37,7 +43,14 @@ func TestValidateIdentityFileSecretStoreRequiresDirectory(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestValidateIdentityKubernetesSecretStore(t *testing.T) {
|
||||
cfg := Config{IdentitySecretStore: "kubernetes", IdentityKubernetesSecretName: "easyai-gateway-identity"}
|
||||
cfg := Config{
|
||||
IdentitySecretStore: "kubernetes",
|
||||
IdentityKubernetesSecretName: "easyai-gateway-identity",
|
||||
AsyncWorkerHardLimit: 2048,
|
||||
AsyncWorkerInstanceHardLimit: 32,
|
||||
AsyncWorkerRefreshIntervalSeconds: 5,
|
||||
AsyncAdmissionMicrobatchSize: 8,
|
||||
}
|
||||
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "namespace") {
|
||||
t.Fatalf("Validate() error = %v, want missing namespace", err)
|
||||
}
|
||||
@@ -52,8 +65,341 @@ func TestValidateIdentitySecurityEventTiming(t *testing.T) {
|
||||
IdentitySecurityEventHeartbeatIntervalSeconds: 60,
|
||||
IdentitySecurityEventStaleAfterSeconds: 60,
|
||||
IdentitySecurityEventClockSkewSeconds: 60,
|
||||
AsyncWorkerHardLimit: 2048,
|
||||
AsyncWorkerInstanceHardLimit: 32,
|
||||
AsyncWorkerRefreshIntervalSeconds: 5,
|
||||
AsyncAdmissionMicrobatchSize: 8,
|
||||
}
|
||||
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "heartbeat") {
|
||||
t.Fatalf("Validate() error = %v, want invalid stale threshold", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateAsyncWorkerSettings(t *testing.T) {
|
||||
cfg := Config{
|
||||
AsyncWorkerHardLimit: 10001,
|
||||
AsyncWorkerInstanceHardLimit: 32,
|
||||
AsyncWorkerRefreshIntervalSeconds: 5,
|
||||
AsyncAdmissionMicrobatchSize: 8,
|
||||
}
|
||||
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "HARD_LIMIT") {
|
||||
t.Fatalf("Validate() error = %v, want invalid hard limit", err)
|
||||
}
|
||||
cfg.AsyncWorkerHardLimit = 2048
|
||||
cfg.AsyncWorkerInstanceHardLimit = 10001
|
||||
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "INSTANCE_HARD_LIMIT") {
|
||||
t.Fatalf("Validate() error = %v, want invalid instance hard limit", err)
|
||||
}
|
||||
cfg.AsyncWorkerInstanceHardLimit = 32
|
||||
cfg.AsyncWorkerRefreshIntervalSeconds = 0
|
||||
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "REFRESH_INTERVAL") {
|
||||
t.Fatalf("Validate() error = %v, want invalid refresh interval", err)
|
||||
}
|
||||
cfg.AsyncWorkerRefreshIntervalSeconds = 5
|
||||
cfg.AsyncWorkerLoadMode = "invalid"
|
||||
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "LOAD_MODE") {
|
||||
t.Fatalf("Validate() error = %v, want invalid worker load mode", err)
|
||||
}
|
||||
t.Setenv("AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT", "not-an-integer")
|
||||
loaded := Load()
|
||||
if err := loaded.Validate(); err == nil || !strings.Contains(err.Error(), "HARD_LIMIT") {
|
||||
t.Fatalf("Validate() error = %v, want invalid non-integer hard limit", err)
|
||||
}
|
||||
t.Setenv("AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT", "2048")
|
||||
t.Setenv("AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT", "not-an-integer")
|
||||
loaded = Load()
|
||||
if err := loaded.Validate(); err == nil || !strings.Contains(err.Error(), "INSTANCE_HARD_LIMIT") {
|
||||
t.Fatalf("Validate() error = %v, want invalid non-integer instance hard limit", err)
|
||||
}
|
||||
t.Setenv("AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT", "32")
|
||||
t.Setenv("AI_GATEWAY_ASYNC_ADMISSION_MICROBATCH_SIZE", "33")
|
||||
loaded = Load()
|
||||
if err := loaded.Validate(); err == nil || !strings.Contains(err.Error(), "MICROBATCH_SIZE") {
|
||||
t.Fatalf("Validate() error = %v, want invalid admission microbatch size", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadAsyncAdmissionMicrobatchSize(t *testing.T) {
|
||||
cfg := Load()
|
||||
if cfg.AsyncAdmissionMicrobatchSize != 8 {
|
||||
t.Fatalf("admission microbatch size=%d, want 8", cfg.AsyncAdmissionMicrobatchSize)
|
||||
}
|
||||
t.Setenv("AI_GATEWAY_ASYNC_ADMISSION_MICROBATCH_SIZE", "4")
|
||||
cfg = Load()
|
||||
if err := cfg.Validate(); err != nil {
|
||||
t.Fatalf("valid admission microbatch size was rejected: %v", err)
|
||||
}
|
||||
if cfg.AsyncAdmissionMicrobatchSize != 4 {
|
||||
t.Fatalf("admission microbatch size=%d, want 4", cfg.AsyncAdmissionMicrobatchSize)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadAsyncQueueWorkerEnabled(t *testing.T) {
|
||||
cfg := Load()
|
||||
if !cfg.AsyncQueueWorkerEnabled {
|
||||
t.Fatal("async queue worker must be enabled by default")
|
||||
}
|
||||
t.Setenv("AI_GATEWAY_ASYNC_QUEUE_WORKER_ENABLED", "false")
|
||||
cfg = Load()
|
||||
if cfg.AsyncQueueWorkerEnabled {
|
||||
t.Fatal("async queue worker remained enabled after explicit disable")
|
||||
}
|
||||
}
|
||||
|
||||
func TestProcessRolePrecedenceAndCompatibility(t *testing.T) {
|
||||
t.Setenv("AI_GATEWAY_ASYNC_QUEUE_WORKER_ENABLED", "false")
|
||||
cfg := Load()
|
||||
if got := cfg.EffectiveProcessRole(); got != "all" {
|
||||
t.Fatalf("legacy configuration changed process role to %q", got)
|
||||
}
|
||||
if cfg.RunsAsyncExecutionWorker() {
|
||||
t.Fatal("legacy async worker disable was ignored")
|
||||
}
|
||||
|
||||
t.Setenv("AI_GATEWAY_PROCESS_ROLE", "worker")
|
||||
cfg = Load()
|
||||
if got := cfg.EffectiveProcessRole(); got != "worker" {
|
||||
t.Fatalf("effective process role = %q, want worker", got)
|
||||
}
|
||||
if cfg.RunsPublicHTTP() || !cfg.RunsAsyncExecutionWorker() || !cfg.RunsBackgroundWorkers() {
|
||||
t.Fatalf("worker role capabilities are inconsistent: %+v", cfg)
|
||||
}
|
||||
if !cfg.RunsAsyncAdmissionDispatcher() {
|
||||
t.Fatal("legacy Worker role no longer runs the admission dispatcher")
|
||||
}
|
||||
|
||||
t.Setenv("AI_GATEWAY_PROCESS_ROLE", "api")
|
||||
cfg = Load()
|
||||
if !cfg.RunsPublicHTTP() || cfg.RunsAsyncExecutionWorker() || cfg.RunsBackgroundWorkers() {
|
||||
t.Fatalf("api role capabilities are inconsistent: %+v", cfg)
|
||||
}
|
||||
if cfg.RunsAsyncAdmissionDispatcher() {
|
||||
t.Fatal("API role unexpectedly runs the admission dispatcher without explicit configuration")
|
||||
}
|
||||
|
||||
t.Setenv("AI_GATEWAY_ASYNC_ADMISSION_DISPATCHER_ENABLED", "true")
|
||||
cfg = Load()
|
||||
if !cfg.RunsAsyncAdmissionDispatcher() || cfg.RunsAsyncExecutionWorker() {
|
||||
t.Fatalf("API admission-only role is inconsistent: %+v", cfg)
|
||||
}
|
||||
t.Setenv("AI_GATEWAY_ASYNC_ADMISSION_DISPATCHER_ENABLED", "false")
|
||||
t.Setenv("AI_GATEWAY_PROCESS_ROLE", "worker")
|
||||
cfg = Load()
|
||||
if cfg.RunsAsyncAdmissionDispatcher() || !cfg.RunsAsyncExecutionWorker() {
|
||||
t.Fatalf("Worker execution-only role is inconsistent: %+v", cfg)
|
||||
}
|
||||
|
||||
t.Setenv("AI_GATEWAY_PROCESS_ROLE", "capacity-controller")
|
||||
cfg = Load()
|
||||
if cfg.RunsPublicHTTP() || cfg.RunsAsyncExecutionWorker() || cfg.RunsBackgroundWorkers() || !cfg.RunsCapacityController() {
|
||||
t.Fatalf("capacity controller role capabilities are inconsistent: %+v", cfg)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWorkerAutoscalingConfiguration(t *testing.T) {
|
||||
t.Setenv("AI_GATEWAY_WORKER_AUTOSCALING_ENABLED", "true")
|
||||
t.Setenv("AI_GATEWAY_CAPACITY_POOLS", `[{"id":"pool-a","bootstrapReplicas":2,"minReplicas":1,"maxReplicas":3}]`)
|
||||
cfg := Load()
|
||||
if err := cfg.Validate(); err != nil {
|
||||
t.Fatalf("valid autoscaling configuration was rejected: %v", err)
|
||||
}
|
||||
pools, err := cfg.CapacityPools()
|
||||
if err != nil || !cfg.WorkerAutoscalingEnabled || len(pools) != 1 || pools[0].BootstrapReplicas != 2 || pools[0].MaxReplicas != 3 {
|
||||
t.Fatalf("autoscaling configuration=%+v", cfg)
|
||||
}
|
||||
cfg.CapacityPoolsJSON = `[{"id":"pool-a","bootstrapReplicas":4,"minReplicas":1,"maxReplicas":3}]`
|
||||
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "invalid pool") {
|
||||
t.Fatalf("Validate() error=%v, want bootstrap bound failure", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWorkerTopologyHasNoBuiltInSites(t *testing.T) {
|
||||
t.Setenv("AI_GATEWAY_CAPACITY_POOLS", "")
|
||||
cfg := Load()
|
||||
pools, err := cfg.CapacityPools()
|
||||
if err != nil || len(pools) != 0 {
|
||||
t.Fatalf("default pools=%+v error=%v, want no built-in sites", pools, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadPrefersWorkerGlobalHardLimitAlias(t *testing.T) {
|
||||
t.Setenv("AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT", "48")
|
||||
t.Setenv("AI_GATEWAY_ASYNC_WORKER_GLOBAL_HARD_LIMIT", "96")
|
||||
cfg := Load()
|
||||
if cfg.AsyncWorkerHardLimit != 96 {
|
||||
t.Fatalf("global worker hard limit=%d, want 96", cfg.AsyncWorkerHardLimit)
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidatePostgresReplicaBudgets(t *testing.T) {
|
||||
cfg := Load()
|
||||
if cfg.PostgresConnectionBudget != 150 || cfg.PostgresNonWorkerConnectionBudget != 72 {
|
||||
t.Fatalf(
|
||||
"PostgreSQL budgets=%d/%d, want 150/72",
|
||||
cfg.PostgresConnectionBudget,
|
||||
cfg.PostgresNonWorkerConnectionBudget,
|
||||
)
|
||||
}
|
||||
cfg.PostgresNonWorkerConnectionBudget = cfg.PostgresConnectionBudget
|
||||
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "NON_WORKER") {
|
||||
t.Fatalf("Validate() error=%v, want non-Worker budget failure", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateProcessRoleAndDatabasePool(t *testing.T) {
|
||||
cfg := Load()
|
||||
if cfg.DatabaseMinIdleConns != 0 ||
|
||||
cfg.DatabaseMaxConnIdleSeconds != 0 ||
|
||||
cfg.DatabaseIdleInTransactionTimeoutSeconds != 60 ||
|
||||
cfg.DatabaseLockTimeoutSeconds != 30 {
|
||||
t.Fatalf(
|
||||
"database pool minimum/max-idle/transaction timeouts = %d/%d/%d/%d, want 0/0/60/30",
|
||||
cfg.DatabaseMinIdleConns,
|
||||
cfg.DatabaseMaxConnIdleSeconds,
|
||||
cfg.DatabaseIdleInTransactionTimeoutSeconds,
|
||||
cfg.DatabaseLockTimeoutSeconds,
|
||||
)
|
||||
}
|
||||
cfg.ProcessRole = "invalid"
|
||||
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "PROCESS_ROLE") {
|
||||
t.Fatalf("Validate() error = %v, want invalid process role", err)
|
||||
}
|
||||
cfg.ProcessRole = "api"
|
||||
cfg.DatabaseMaxConns = -1
|
||||
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "DATABASE_MAX_CONNS") {
|
||||
t.Fatalf("Validate() error = %v, want invalid database max conns", err)
|
||||
}
|
||||
cfg.DatabaseMaxConns = 16
|
||||
cfg.DatabaseMinIdleConns = 17
|
||||
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "DATABASE_MIN_IDLE_CONNS") {
|
||||
t.Fatalf("Validate() error = %v, want invalid database min idle conns", err)
|
||||
}
|
||||
cfg.DatabaseMinIdleConns = 4
|
||||
cfg.DatabaseCriticalMaxConns = 4
|
||||
cfg.DatabaseRiverMaxConns = 12
|
||||
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "River") {
|
||||
t.Fatalf("Validate() error = %v, want database pool budget failure", err)
|
||||
}
|
||||
cfg.DatabaseRiverMaxConns = 4
|
||||
t.Setenv("AI_GATEWAY_DATABASE_MIN_IDLE_CONNS", "not-an-integer")
|
||||
if err := Load().Validate(); err == nil || !strings.Contains(err.Error(), "DATABASE_MIN_IDLE_CONNS") {
|
||||
t.Fatalf("Validate() error = %v, want invalid non-integer database min idle conns", err)
|
||||
}
|
||||
t.Setenv("AI_GATEWAY_DATABASE_MIN_IDLE_CONNS", "4")
|
||||
cfg.DatabaseMaxConnIdleSeconds = 3601
|
||||
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "DATABASE_MAX_CONN_IDLE_SECONDS") {
|
||||
t.Fatalf("Validate() error = %v, want invalid database max connection idle seconds", err)
|
||||
}
|
||||
cfg.DatabaseMaxConnIdleSeconds = 30
|
||||
cfg.DatabaseIdleInTransactionTimeoutSeconds = 3601
|
||||
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "IDLE_IN_TRANSACTION") {
|
||||
t.Fatalf("Validate() error = %v, want invalid idle transaction timeout", err)
|
||||
}
|
||||
cfg.DatabaseIdleInTransactionTimeoutSeconds = 60
|
||||
cfg.DatabaseLockTimeoutSeconds = 3601
|
||||
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "LOCK_TIMEOUT") {
|
||||
t.Fatalf("Validate() error = %v, want invalid lock timeout", err)
|
||||
}
|
||||
cfg.DatabaseLockTimeoutSeconds = 30
|
||||
t.Setenv("AI_GATEWAY_DATABASE_LOCK_TIMEOUT_SECONDS", "not-an-integer")
|
||||
if err := Load().Validate(); err == nil || !strings.Contains(err.Error(), "LOCK_TIMEOUT") {
|
||||
t.Fatalf("Validate() error = %v, want invalid non-integer lock timeout", err)
|
||||
}
|
||||
t.Setenv("AI_GATEWAY_DATABASE_LOCK_TIMEOUT_SECONDS", "30")
|
||||
cfg.MediaMaterializationConcurrency = 257
|
||||
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "MEDIA_MATERIALIZATION_CONCURRENCY") {
|
||||
t.Fatalf("Validate() error = %v, want invalid media materialization concurrency", err)
|
||||
}
|
||||
t.Setenv("AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY", "12")
|
||||
if got := Load().MediaMaterializationConcurrency; got != 12 {
|
||||
t.Fatalf("media materialization concurrency = %d, want 12", got)
|
||||
}
|
||||
cfg = Load()
|
||||
cfg.MediaRequestConcurrency = 1025
|
||||
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "MEDIA_REQUEST_CONCURRENCY") {
|
||||
t.Fatalf("Validate() error = %v, want invalid media request concurrency", err)
|
||||
}
|
||||
t.Setenv("AI_GATEWAY_MEDIA_REQUEST_CONCURRENCY", "96")
|
||||
if got := Load().MediaRequestConcurrency; got != 96 {
|
||||
t.Fatalf("media request concurrency = %d, want 96", got)
|
||||
}
|
||||
cfg = Load()
|
||||
cfg.MediaImageNormalizationConcurrency = 65
|
||||
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "MEDIA_IMAGE_NORMALIZATION_CONCURRENCY") {
|
||||
t.Fatalf("Validate() error = %v, want invalid image normalization concurrency", err)
|
||||
}
|
||||
t.Setenv("AI_GATEWAY_MEDIA_IMAGE_NORMALIZATION_CONCURRENCY", "3")
|
||||
if got := Load().MediaImageNormalizationConcurrency; got != 3 {
|
||||
t.Fatalf("image normalization concurrency = %d, want 3", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateTaskHistorySettings(t *testing.T) {
|
||||
cfg := Load()
|
||||
cfg.TaskRetentionDays = 30
|
||||
cfg.TaskAnalysisRetentionDays = 31
|
||||
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "ANALYSIS_RETENTION") {
|
||||
t.Fatalf("Validate() error = %v, want invalid analysis retention", err)
|
||||
}
|
||||
cfg.TaskAnalysisRetentionDays = 7
|
||||
cfg.TaskCleanupIntervalSeconds = 59
|
||||
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "CLEANUP_INTERVAL") {
|
||||
t.Fatalf("Validate() error = %v, want invalid cleanup interval", err)
|
||||
}
|
||||
cfg.TaskCleanupIntervalSeconds = 300
|
||||
cfg.TaskCleanupBatchSize = 5001
|
||||
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "CLEANUP_BATCH") {
|
||||
t.Fatalf("Validate() error = %v, want invalid cleanup batch", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadAndValidateDirectMediaOSS(t *testing.T) {
|
||||
t.Setenv("APP_ENV", "production")
|
||||
t.Setenv("AI_GATEWAY_MEDIA_OSS_DIRECT_ENABLED", "true")
|
||||
t.Setenv("AI_GATEWAY_MEDIA_OSS_ENDPOINT", "https://media-bucket.oss-cn-shanghai.aliyuncs.com")
|
||||
t.Setenv("AI_GATEWAY_MEDIA_OSS_BUCKET", "media-bucket")
|
||||
t.Setenv("AI_GATEWAY_MEDIA_OSS_ACCESS_KEY_ID", "test-access-key")
|
||||
t.Setenv("AI_GATEWAY_MEDIA_OSS_ACCESS_KEY_SECRET", "test-access-secret")
|
||||
t.Setenv("AI_GATEWAY_MEDIA_OSS_PUBLIC_BASE_URL", "https://cdn.example.com")
|
||||
t.Setenv("AI_GATEWAY_MEDIA_OSS_OBJECT_PREFIX", "easyai-ai-gateway/production/media")
|
||||
|
||||
cfg := Load()
|
||||
if !cfg.MediaOSSDirectEnabled || cfg.MediaOSSBucket != "media-bucket" ||
|
||||
cfg.MediaOSSObjectPrefix != "easyai-ai-gateway/production/media" {
|
||||
t.Fatalf("unexpected direct media OSS config: %+v", cfg)
|
||||
}
|
||||
if err := cfg.Validate(); err != nil {
|
||||
t.Fatalf("validate direct media OSS: %v", err)
|
||||
}
|
||||
cfg.MediaOSSAccessKeySecret = ""
|
||||
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "ACCESS_KEY_SECRET") {
|
||||
t.Fatalf("Validate() error = %v, want missing direct OSS secret", err)
|
||||
}
|
||||
cfg.MediaOSSAccessKeySecret = "test-access-secret"
|
||||
cfg.MediaOSSEndpoint = "http://media-bucket.example.com"
|
||||
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "HTTPS URL") {
|
||||
t.Fatalf("Validate() error = %v, want insecure production endpoint", err)
|
||||
}
|
||||
cfg.MediaOSSEndpoint = "https://media-bucket.example.com"
|
||||
cfg.MediaOSSObjectPrefix = "../outside"
|
||||
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "OBJECT_PREFIX") {
|
||||
t.Fatalf("Validate() error = %v, want invalid object prefix", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCapacityControllerDoesNotRequireMediaOSSCredentials(t *testing.T) {
|
||||
cfg := Load()
|
||||
cfg.ProcessRole = "capacity-controller"
|
||||
cfg.CapacityPoolsJSON = `[{"id":"pool-a","bootstrapReplicas":1,"minReplicas":1,"maxReplicas":1}]`
|
||||
cfg.MediaOSSDirectEnabled = true
|
||||
cfg.MediaOSSAccessKeyID = ""
|
||||
cfg.MediaOSSAccessKeySecret = ""
|
||||
if err := cfg.Validate(); err != nil {
|
||||
t.Fatalf("capacity controller rejected unused media OSS credentials: %v", err)
|
||||
}
|
||||
|
||||
cfg.ProcessRole = "worker"
|
||||
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "ACCESS_KEY") {
|
||||
t.Fatalf("worker Validate() error = %v, want missing direct OSS credentials", err)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,27 @@
|
||||
package config
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestCapacityPoolsAreDataDriven(t *testing.T) {
|
||||
cfg := Config{CapacityPoolsJSON: `[
|
||||
{"id":"pool-a","adapterRef":"deployment-a","bootstrapReplicas":1,"minReplicas":1,"maxReplicas":3},
|
||||
{"id":"pool-b","adapterRef":"deployment-b","bootstrapReplicas":0,"minReplicas":0,"maxReplicas":2}
|
||||
]`}
|
||||
pools, err := cfg.CapacityPools()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(pools) != 2 || pools[0].ID != "pool-a" || pools[1].AdapterRef != "deployment-b" {
|
||||
t.Fatalf("pools=%+v", pools)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCapacityPoolsRejectDuplicates(t *testing.T) {
|
||||
cfg := Config{CapacityPoolsJSON: `[
|
||||
{"id":"pool-a","bootstrapReplicas":1,"minReplicas":1,"maxReplicas":1},
|
||||
{"id":"pool-a","bootstrapReplicas":1,"minReplicas":1,"maxReplicas":1}
|
||||
]`}
|
||||
if _, err := cfg.CapacityPools(); err == nil {
|
||||
t.Fatal("expected duplicate pool rejection")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
package executionpool
|
||||
|
||||
import (
|
||||
"crypto/hmac"
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
func QueueName(poolID string) string {
|
||||
poolID = strings.TrimSpace(poolID)
|
||||
if poolID == "" {
|
||||
return "gateway_tasks"
|
||||
}
|
||||
digest := sha256.Sum256([]byte(poolID))
|
||||
return "gateway_pool_" + fmt.Sprintf("%x", digest[:8])
|
||||
}
|
||||
|
||||
func RouteProfileKey(provider, protocol, endpointHost, proxyMode, configRevision string) string {
|
||||
value := strings.Join([]string{
|
||||
strings.ToLower(strings.TrimSpace(provider)),
|
||||
strings.ToLower(strings.TrimSpace(protocol)),
|
||||
strings.ToLower(strings.TrimSpace(endpointHost)),
|
||||
strings.ToLower(strings.TrimSpace(proxyMode)),
|
||||
strings.TrimSpace(configRevision),
|
||||
}, "\x00")
|
||||
digest := sha256.Sum256([]byte(value))
|
||||
return fmt.Sprintf("route_%x", digest[:16])
|
||||
}
|
||||
|
||||
func EndpointHost(rawURL string) string {
|
||||
parsed, err := url.Parse(strings.TrimSpace(rawURL))
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
return strings.ToLower(parsed.Hostname())
|
||||
}
|
||||
|
||||
func ValidateAdvertisedEndpoint(rawURL string, allowedSuffixes []string, allowPrivate bool) error {
|
||||
parsed, err := url.Parse(strings.TrimSpace(rawURL))
|
||||
if err != nil || parsed.Host == "" || parsed.User != nil {
|
||||
return errors.New("worker endpoint must be an absolute URL without user info")
|
||||
}
|
||||
if parsed.Scheme != "https" && parsed.Scheme != "http" {
|
||||
return errors.New("worker endpoint scheme must be http or https")
|
||||
}
|
||||
host := strings.ToLower(parsed.Hostname())
|
||||
if host == "" {
|
||||
return errors.New("worker endpoint host is required")
|
||||
}
|
||||
if ip := net.ParseIP(host); ip != nil {
|
||||
if allowPrivate && (ip.IsPrivate() || ip.IsLoopback()) {
|
||||
return nil
|
||||
}
|
||||
return errors.New("worker endpoint IP is outside the configured trust boundary")
|
||||
}
|
||||
for _, suffix := range allowedSuffixes {
|
||||
suffix = strings.ToLower(strings.TrimSpace(suffix))
|
||||
if suffix != "" && (host == strings.TrimPrefix(suffix, ".") || strings.HasSuffix(host, "."+strings.TrimPrefix(suffix, "."))) {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
return errors.New("worker endpoint host is outside the configured trust boundary")
|
||||
}
|
||||
|
||||
type ExecutionClaims struct {
|
||||
Audience string `json:"aud"`
|
||||
TaskID string `json:"task_id"`
|
||||
PoolID string `json:"pool_id"`
|
||||
WorkerID string `json:"worker_id"`
|
||||
Nonce string `json:"nonce"`
|
||||
ExpiresAt int64 `json:"exp"`
|
||||
}
|
||||
|
||||
type TokenSigner struct {
|
||||
Secret []byte
|
||||
Now func() time.Time
|
||||
}
|
||||
|
||||
func (s TokenSigner) Sign(claims ExecutionClaims) (string, error) {
|
||||
if len(s.Secret) < 32 {
|
||||
return "", errors.New("execution token secret must be at least 32 bytes")
|
||||
}
|
||||
if claims.Audience == "" || claims.TaskID == "" || claims.PoolID == "" || claims.WorkerID == "" || claims.Nonce == "" || claims.ExpiresAt <= 0 {
|
||||
return "", errors.New("execution token claims are incomplete")
|
||||
}
|
||||
payload, err := json.Marshal(claims)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
encoded := base64.RawURLEncoding.EncodeToString(payload)
|
||||
mac := hmac.New(sha256.New, s.Secret)
|
||||
_, _ = mac.Write([]byte(encoded))
|
||||
signature := base64.RawURLEncoding.EncodeToString(mac.Sum(nil))
|
||||
return encoded + "." + signature, nil
|
||||
}
|
||||
|
||||
func (s TokenSigner) Verify(token, audience string) (ExecutionClaims, error) {
|
||||
var claims ExecutionClaims
|
||||
if len(s.Secret) < 32 {
|
||||
return claims, errors.New("execution token secret must be at least 32 bytes")
|
||||
}
|
||||
parts := strings.Split(token, ".")
|
||||
if len(parts) != 2 {
|
||||
return claims, errors.New("invalid execution token")
|
||||
}
|
||||
signature, err := base64.RawURLEncoding.DecodeString(parts[1])
|
||||
if err != nil {
|
||||
return claims, errors.New("invalid execution token signature")
|
||||
}
|
||||
mac := hmac.New(sha256.New, s.Secret)
|
||||
_, _ = mac.Write([]byte(parts[0]))
|
||||
if !hmac.Equal(signature, mac.Sum(nil)) {
|
||||
return claims, errors.New("invalid execution token signature")
|
||||
}
|
||||
payload, err := base64.RawURLEncoding.DecodeString(parts[0])
|
||||
if err != nil || json.Unmarshal(payload, &claims) != nil {
|
||||
return ExecutionClaims{}, errors.New("invalid execution token payload")
|
||||
}
|
||||
now := time.Now()
|
||||
if s.Now != nil {
|
||||
now = s.Now()
|
||||
}
|
||||
if claims.ExpiresAt <= now.Unix() {
|
||||
return ExecutionClaims{}, errors.New("execution token expired")
|
||||
}
|
||||
if claims.Audience != audience {
|
||||
return ExecutionClaims{}, errors.New("execution token audience mismatch")
|
||||
}
|
||||
return claims, nil
|
||||
}
|
||||
@@ -0,0 +1,44 @@
|
||||
package executionpool
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestQueueNameDoesNotExposePoolID(t *testing.T) {
|
||||
name := QueueName("region-sensitive-name")
|
||||
if strings.Contains(name, "sensitive") || name == QueueName("other") {
|
||||
t.Fatalf("unexpected queue name %q", name)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTokenSigner(t *testing.T) {
|
||||
now := time.Unix(1000, 0)
|
||||
signer := TokenSigner{Secret: []byte("01234567890123456789012345678901"), Now: func() time.Time { return now }}
|
||||
token, err := signer.Sign(ExecutionClaims{
|
||||
Audience: "worker", TaskID: "task", PoolID: "pool", WorkerID: "worker", Nonce: "nonce", ExpiresAt: now.Add(30 * time.Second).Unix(),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
claims, err := signer.Verify(token, "worker")
|
||||
if err != nil || claims.TaskID != "task" {
|
||||
t.Fatalf("claims=%+v err=%v", claims, err)
|
||||
}
|
||||
if _, err := signer.Verify(token, "other"); err == nil {
|
||||
t.Fatal("expected audience mismatch")
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateAdvertisedEndpoint(t *testing.T) {
|
||||
if err := ValidateAdvertisedEndpoint("http://10.0.0.2:8088", nil, true); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := ValidateAdvertisedEndpoint("https://worker.internal.example", []string{"internal.example"}, false); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := ValidateAdvertisedEndpoint("https://public.example", []string{"internal.example"}, false); err == nil {
|
||||
t.Fatal("expected untrusted endpoint rejection")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
package executionpool
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
type RoutePreference struct {
|
||||
RouteProfileKey string
|
||||
CurrentPoolID string
|
||||
CurrentSince time.Time
|
||||
ChallengerPoolID string
|
||||
ChallengerWins int
|
||||
UpdatedAt time.Time
|
||||
}
|
||||
|
||||
func AdvanceRoutePreference(
|
||||
previous RoutePreference,
|
||||
proposedPoolID string,
|
||||
scores map[string]float64,
|
||||
now time.Time,
|
||||
) RoutePreference {
|
||||
if now.IsZero() {
|
||||
now = time.Now()
|
||||
}
|
||||
proposedPoolID = strings.TrimSpace(proposedPoolID)
|
||||
next := previous
|
||||
next.UpdatedAt = now
|
||||
if proposedPoolID == "" {
|
||||
return next
|
||||
}
|
||||
if _, eligible := scores[next.CurrentPoolID]; next.CurrentPoolID == "" || !eligible {
|
||||
next.CurrentPoolID = proposedPoolID
|
||||
next.CurrentSince = now
|
||||
next.ChallengerPoolID = ""
|
||||
next.ChallengerWins = 0
|
||||
return next
|
||||
}
|
||||
if proposedPoolID == next.CurrentPoolID {
|
||||
next.ChallengerPoolID = ""
|
||||
next.ChallengerWins = 0
|
||||
return next
|
||||
}
|
||||
if relativeImprovement(scores[proposedPoolID], scores[next.CurrentPoolID]) < defaultSwitchImprovement {
|
||||
next.ChallengerPoolID = ""
|
||||
next.ChallengerWins = 0
|
||||
return next
|
||||
}
|
||||
if next.ChallengerPoolID == proposedPoolID {
|
||||
next.ChallengerWins++
|
||||
} else {
|
||||
next.ChallengerPoolID = proposedPoolID
|
||||
next.ChallengerWins = 1
|
||||
}
|
||||
if next.ChallengerWins >= defaultBetterWindowCount &&
|
||||
!next.CurrentSince.IsZero() && now.Sub(next.CurrentSince) >= defaultMinimumDwell {
|
||||
next.CurrentPoolID = proposedPoolID
|
||||
next.CurrentSince = now
|
||||
next.ChallengerPoolID = ""
|
||||
next.ChallengerWins = 0
|
||||
}
|
||||
return next
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
package executionpool
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestAdvanceRoutePreferenceRequiresThreeWinsAndMinimumDwell(t *testing.T) {
|
||||
now := time.Date(2026, 8, 5, 12, 0, 0, 0, time.UTC)
|
||||
state := RoutePreference{CurrentPoolID: "slow", CurrentSince: now}
|
||||
scores := map[string]float64{"slow": 0.5, "fast": 0.7}
|
||||
for window := 1; window <= 3; window++ {
|
||||
state = AdvanceRoutePreference(state, "fast", scores, now.Add(time.Duration(window)*10*time.Second))
|
||||
if state.CurrentPoolID != "slow" {
|
||||
t.Fatalf("switched before minimum dwell in window %d", window)
|
||||
}
|
||||
}
|
||||
state = AdvanceRoutePreference(state, "fast", scores, now.Add(2*time.Minute))
|
||||
if state.CurrentPoolID != "fast" {
|
||||
t.Fatalf("current pool=%q, want fast", state.CurrentPoolID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdvanceRoutePreferenceImmediatelyLeavesIneligibleCurrentPool(t *testing.T) {
|
||||
now := time.Now()
|
||||
state := AdvanceRoutePreference(RoutePreference{
|
||||
CurrentPoolID: "offline", CurrentSince: now.Add(-time.Minute),
|
||||
}, "healthy", map[string]float64{"healthy": 0.6}, now)
|
||||
if state.CurrentPoolID != "healthy" || state.ChallengerWins != 0 {
|
||||
t.Fatalf("unexpected preference: %#v", state)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,101 @@
|
||||
package executionpool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"net/http"
|
||||
"net/http/httptrace"
|
||||
"time"
|
||||
)
|
||||
|
||||
type ProbeTarget struct {
|
||||
RouteProfile RouteProfile
|
||||
URL string
|
||||
Method string
|
||||
Timeout time.Duration
|
||||
}
|
||||
|
||||
type ProbeResult struct {
|
||||
Reachable bool
|
||||
StatusCode int
|
||||
DNSDuration time.Duration
|
||||
TCPDuration time.Duration
|
||||
TLSDuration time.Duration
|
||||
FirstByte time.Duration
|
||||
SampledAt time.Time
|
||||
ErrorClass string
|
||||
}
|
||||
|
||||
type NetworkProbe interface {
|
||||
Probe(context.Context, ProbeTarget) ProbeResult
|
||||
}
|
||||
|
||||
type HTTPProbe struct {
|
||||
Client *http.Client
|
||||
Now func() time.Time
|
||||
}
|
||||
|
||||
func (p HTTPProbe) Probe(ctx context.Context, target ProbeTarget) ProbeResult {
|
||||
result := ProbeResult{}
|
||||
now := time.Now
|
||||
if p.Now != nil {
|
||||
now = p.Now
|
||||
}
|
||||
result.SampledAt = now()
|
||||
method := target.Method
|
||||
if method == "" {
|
||||
method = http.MethodHead
|
||||
}
|
||||
timeout := target.Timeout
|
||||
if timeout <= 0 {
|
||||
timeout = 3 * time.Second
|
||||
}
|
||||
probeCtx, cancel := context.WithTimeout(ctx, timeout)
|
||||
defer cancel()
|
||||
var dnsStart, tcpStart, tlsStart time.Time
|
||||
requestStart := now()
|
||||
trace := &httptrace.ClientTrace{
|
||||
DNSStart: func(httptrace.DNSStartInfo) { dnsStart = now() },
|
||||
DNSDone: func(httptrace.DNSDoneInfo) {
|
||||
if !dnsStart.IsZero() {
|
||||
result.DNSDuration = now().Sub(dnsStart)
|
||||
}
|
||||
},
|
||||
ConnectStart: func(_, _ string) { tcpStart = now() },
|
||||
ConnectDone: func(_, _ string, _ error) {
|
||||
if !tcpStart.IsZero() {
|
||||
result.TCPDuration = now().Sub(tcpStart)
|
||||
}
|
||||
},
|
||||
TLSHandshakeStart: func() { tlsStart = now() },
|
||||
TLSHandshakeDone: func(tls.ConnectionState, error) {
|
||||
if !tlsStart.IsZero() {
|
||||
result.TLSDuration = now().Sub(tlsStart)
|
||||
}
|
||||
},
|
||||
GotFirstResponseByte: func() { result.FirstByte = now().Sub(requestStart) },
|
||||
}
|
||||
request, err := http.NewRequestWithContext(httptrace.WithClientTrace(probeCtx, trace), method, target.URL, nil)
|
||||
if err != nil {
|
||||
result.ErrorClass = "invalid_target"
|
||||
return result
|
||||
}
|
||||
client := p.Client
|
||||
if client == nil {
|
||||
client = http.DefaultClient
|
||||
}
|
||||
response, err := client.Do(request)
|
||||
if err != nil {
|
||||
if probeCtx.Err() != nil {
|
||||
result.ErrorClass = "timeout"
|
||||
} else {
|
||||
result.ErrorClass = "network"
|
||||
}
|
||||
return result
|
||||
}
|
||||
defer response.Body.Close()
|
||||
result.StatusCode = response.StatusCode
|
||||
// Any HTTP response, including 401/403/404/405, proves transport reachability.
|
||||
result.Reachable = true
|
||||
return result
|
||||
}
|
||||
@@ -0,0 +1,202 @@
|
||||
package executionpool
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"math"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultHealthMaxAge = 45 * time.Second
|
||||
defaultMinimumDwell = 2 * time.Minute
|
||||
defaultSwitchImprovement = 0.20
|
||||
defaultBetterWindowCount = 3
|
||||
)
|
||||
|
||||
var ErrNoEligiblePool = errors.New("no eligible execution pool")
|
||||
|
||||
type SelectionCandidate struct {
|
||||
Pool ExecutionPool
|
||||
Health RouteHealth
|
||||
Capacity CapacitySnapshot
|
||||
CapabilityMatched bool
|
||||
BetterWindows int
|
||||
}
|
||||
|
||||
type SelectionRequest struct {
|
||||
Candidates []SelectionCandidate
|
||||
CurrentPoolID string
|
||||
CurrentPoolSince time.Time
|
||||
Now time.Time
|
||||
}
|
||||
|
||||
type Decision struct {
|
||||
PoolID string
|
||||
Score float64
|
||||
Reason string
|
||||
Components map[string]float64
|
||||
Rejected map[string]string
|
||||
}
|
||||
|
||||
type Selector struct {
|
||||
HealthMaxAge time.Duration
|
||||
MinimumDwell time.Duration
|
||||
SwitchImprovement float64
|
||||
BetterWindows int
|
||||
}
|
||||
|
||||
func NewSelector() Selector {
|
||||
return Selector{
|
||||
HealthMaxAge: defaultHealthMaxAge, MinimumDwell: defaultMinimumDwell,
|
||||
SwitchImprovement: defaultSwitchImprovement, BetterWindows: defaultBetterWindowCount,
|
||||
}
|
||||
}
|
||||
|
||||
type scoredCandidate struct {
|
||||
candidate SelectionCandidate
|
||||
score float64
|
||||
components map[string]float64
|
||||
}
|
||||
|
||||
func (s Selector) Select(request SelectionRequest) (Decision, error) {
|
||||
now := request.Now
|
||||
if now.IsZero() {
|
||||
now = time.Now()
|
||||
}
|
||||
if s.HealthMaxAge <= 0 {
|
||||
s.HealthMaxAge = defaultHealthMaxAge
|
||||
}
|
||||
if s.MinimumDwell <= 0 {
|
||||
s.MinimumDwell = defaultMinimumDwell
|
||||
}
|
||||
if s.SwitchImprovement <= 0 {
|
||||
s.SwitchImprovement = defaultSwitchImprovement
|
||||
}
|
||||
if s.BetterWindows <= 0 {
|
||||
s.BetterWindows = defaultBetterWindowCount
|
||||
}
|
||||
rejected := make(map[string]string)
|
||||
scored := make([]scoredCandidate, 0, len(request.Candidates))
|
||||
for _, candidate := range request.Candidates {
|
||||
if reason := rejectionReason(candidate, now, s.HealthMaxAge); reason != "" {
|
||||
rejected[candidate.Pool.ID] = reason
|
||||
continue
|
||||
}
|
||||
network := routeQuality(candidate.Health)
|
||||
queue := queueQuality(candidate.Capacity)
|
||||
resources := clamp01(candidate.Capacity.ResourceHeadroom)
|
||||
stability := clamp01(candidate.Capacity.Stability)
|
||||
components := map[string]float64{
|
||||
"network": network, "queue": queue, "resources": resources, "stability": stability,
|
||||
}
|
||||
scored = append(scored, scoredCandidate{
|
||||
candidate: candidate,
|
||||
score: 0.55*network + 0.25*queue + 0.15*resources + 0.05*stability,
|
||||
components: components,
|
||||
})
|
||||
}
|
||||
if len(scored) == 0 {
|
||||
return Decision{Rejected: rejected}, ErrNoEligiblePool
|
||||
}
|
||||
sort.SliceStable(scored, func(i, j int) bool {
|
||||
if math.Abs(scored[i].score-scored[j].score) > 0.000001 {
|
||||
return scored[i].score > scored[j].score
|
||||
}
|
||||
return scored[i].candidate.Pool.ID < scored[j].candidate.Pool.ID
|
||||
})
|
||||
selected := scored[0]
|
||||
if currentIndex := indexOfPool(scored, strings.TrimSpace(request.CurrentPoolID)); currentIndex >= 0 && currentIndex != 0 {
|
||||
current := scored[currentIndex]
|
||||
withinDwell := !request.CurrentPoolSince.IsZero() && now.Sub(request.CurrentPoolSince) < s.MinimumDwell
|
||||
improvement := relativeImprovement(selected.score, current.score)
|
||||
if withinDwell || improvement < s.SwitchImprovement || selected.candidate.BetterWindows < s.BetterWindows {
|
||||
selected = current
|
||||
}
|
||||
}
|
||||
return Decision{
|
||||
PoolID: selected.candidate.Pool.ID,
|
||||
Score: selected.score,
|
||||
Reason: "network_quality_then_safe_capacity",
|
||||
Components: selected.components,
|
||||
Rejected: rejected,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func rejectionReason(candidate SelectionCandidate, now time.Time, maxAge time.Duration) string {
|
||||
if strings.TrimSpace(candidate.Pool.ID) == "" {
|
||||
return "invalid_pool"
|
||||
}
|
||||
if candidate.Pool.State != PoolActive {
|
||||
return "pool_not_active"
|
||||
}
|
||||
if !candidate.CapabilityMatched {
|
||||
return "capability_mismatch"
|
||||
}
|
||||
if candidate.Health.State == RouteUnreachable {
|
||||
return "route_unreachable"
|
||||
}
|
||||
if candidate.Health.State == RouteUnknown || candidate.Health.State == "" {
|
||||
return "route_unknown"
|
||||
}
|
||||
if candidate.Health.SampledAt.IsZero() || now.Sub(candidate.Health.SampledAt) > maxAge || (!candidate.Health.ExpiresAt.IsZero() && !now.Before(candidate.Health.ExpiresAt)) {
|
||||
return "route_stale"
|
||||
}
|
||||
if candidate.Capacity.Critical {
|
||||
return "resource_critical"
|
||||
}
|
||||
if candidate.Capacity.WorkerCount < 1 {
|
||||
return "no_worker"
|
||||
}
|
||||
if candidate.Capacity.SafeCapacity-candidate.Capacity.ActiveTasks < 1 {
|
||||
return "capacity_exhausted"
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func routeQuality(health RouteHealth) float64 {
|
||||
success := clamp01(health.SuccessRate)
|
||||
connect := latencyQuality(health.ConnectTLSP95, 100*time.Millisecond)
|
||||
firstByte := latencyQuality(health.FirstByteP95, 250*time.Millisecond)
|
||||
throughput := clamp01(health.UploadBytesPerSecond / (10 * 1024 * 1024))
|
||||
jitter := latencyQuality(health.JitterP95, 100*time.Millisecond)
|
||||
return 0.35*success + 0.25*connect + 0.25*firstByte + 0.10*throughput + 0.05*jitter
|
||||
}
|
||||
|
||||
func queueQuality(capacity CapacitySnapshot) float64 {
|
||||
available := capacity.SafeCapacity - capacity.ActiveTasks
|
||||
if available <= 0 {
|
||||
return 0
|
||||
}
|
||||
headroom := float64(available) / float64(max(capacity.SafeCapacity, 1))
|
||||
waitPenalty := 1 / (1 + capacity.EstimatedQueueWait.Seconds())
|
||||
return clamp01(0.6*headroom + 0.4*waitPenalty)
|
||||
}
|
||||
|
||||
func latencyQuality(value, target time.Duration) float64 {
|
||||
if value <= 0 {
|
||||
return 0
|
||||
}
|
||||
return clamp01(1 / (1 + float64(value)/float64(target)))
|
||||
}
|
||||
|
||||
func relativeImprovement(candidate, current float64) float64 {
|
||||
if current <= 0 {
|
||||
return 1
|
||||
}
|
||||
return (candidate - current) / current
|
||||
}
|
||||
|
||||
func indexOfPool(candidates []scoredCandidate, poolID string) int {
|
||||
for index := range candidates {
|
||||
if candidates[index].candidate.Pool.ID == poolID {
|
||||
return index
|
||||
}
|
||||
}
|
||||
return -1
|
||||
}
|
||||
|
||||
func clamp01(value float64) float64 {
|
||||
return math.Max(0, math.Min(1, value))
|
||||
}
|
||||
@@ -0,0 +1,86 @@
|
||||
package executionpool
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestSelectorRejectsUnreachableAndPrefersNetwork(t *testing.T) {
|
||||
now := time.Now()
|
||||
selector := NewSelector()
|
||||
decision, err := selector.Select(SelectionRequest{Now: now, Candidates: []SelectionCandidate{
|
||||
candidate("slow", RouteHealthy, now, 120*time.Millisecond, 8, 1),
|
||||
candidate("fast", RouteHealthy, now, 20*time.Millisecond, 8, 1),
|
||||
candidate("down", RouteUnreachable, now, time.Millisecond, 32, 0),
|
||||
}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if decision.PoolID != "fast" {
|
||||
t.Fatalf("pool=%s, want fast", decision.PoolID)
|
||||
}
|
||||
if decision.Rejected["down"] != "route_unreachable" {
|
||||
t.Fatalf("down rejection=%q", decision.Rejected["down"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestSelectorUsesCapacityAsHardGate(t *testing.T) {
|
||||
now := time.Now()
|
||||
selector := NewSelector()
|
||||
decision, err := selector.Select(SelectionRequest{Now: now, Candidates: []SelectionCandidate{
|
||||
candidate("fast-full", RouteHealthy, now, 10*time.Millisecond, 2, 2),
|
||||
candidate("slower-free", RouteHealthy, now, 80*time.Millisecond, 8, 1),
|
||||
}})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if decision.PoolID != "slower-free" {
|
||||
t.Fatalf("pool=%s, want slower-free", decision.PoolID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSelectorRejectsUnknownAndStale(t *testing.T) {
|
||||
now := time.Now()
|
||||
selector := NewSelector()
|
||||
_, err := selector.Select(SelectionRequest{Now: now, Candidates: []SelectionCandidate{
|
||||
candidate("unknown", RouteUnknown, now, time.Millisecond, 8, 0),
|
||||
candidate("stale", RouteHealthy, now.Add(-time.Minute), time.Millisecond, 8, 0),
|
||||
}})
|
||||
if !errors.Is(err, ErrNoEligiblePool) {
|
||||
t.Fatalf("err=%v, want ErrNoEligiblePool", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSelectorHonorsHysteresis(t *testing.T) {
|
||||
now := time.Now()
|
||||
selector := NewSelector()
|
||||
fast := candidate("fast", RouteHealthy, now, 20*time.Millisecond, 8, 1)
|
||||
fast.BetterWindows = 3
|
||||
current := candidate("current", RouteHealthy, now, 80*time.Millisecond, 8, 1)
|
||||
decision, err := selector.Select(SelectionRequest{
|
||||
Now: now, Candidates: []SelectionCandidate{fast, current},
|
||||
CurrentPoolID: "current", CurrentPoolSince: now.Add(-time.Minute),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if decision.PoolID != "current" {
|
||||
t.Fatalf("pool=%s, want current during dwell", decision.PoolID)
|
||||
}
|
||||
}
|
||||
|
||||
func candidate(id string, state RouteState, sampled time.Time, latency time.Duration, capacity, active int) SelectionCandidate {
|
||||
return SelectionCandidate{
|
||||
Pool: ExecutionPool{ID: id, State: PoolActive}, CapabilityMatched: true,
|
||||
Health: RouteHealth{
|
||||
PoolID: id, State: state, SuccessRate: 1, ConnectTLSP95: latency,
|
||||
FirstByteP95: latency, UploadBytesPerSecond: 10 * 1024 * 1024,
|
||||
JitterP95: latency / 10, SampledAt: sampled, ExpiresAt: sampled.Add(45 * time.Second),
|
||||
},
|
||||
Capacity: CapacitySnapshot{
|
||||
PoolID: id, WorkerCount: 1, SafeCapacity: capacity, ActiveTasks: active,
|
||||
ResourceHeadroom: 0.8, Stability: 1,
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,169 @@
|
||||
package executionpool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"time"
|
||||
)
|
||||
|
||||
const ProtocolVersion = "v1"
|
||||
|
||||
type PoolState string
|
||||
|
||||
const (
|
||||
PoolActive PoolState = "active"
|
||||
PoolDraining PoolState = "draining"
|
||||
PoolDisabled PoolState = "disabled"
|
||||
)
|
||||
|
||||
type RouteState string
|
||||
|
||||
const (
|
||||
RouteHealthy RouteState = "healthy"
|
||||
RouteDegraded RouteState = "degraded"
|
||||
RouteUnreachable RouteState = "unreachable"
|
||||
RouteUnknown RouteState = "unknown"
|
||||
)
|
||||
|
||||
type ExecutionPool struct {
|
||||
ID string
|
||||
Labels map[string]string
|
||||
Capabilities map[string]any
|
||||
State PoolState
|
||||
UpdatedAt time.Time
|
||||
}
|
||||
|
||||
type WorkerDescriptor struct {
|
||||
WorkerID string
|
||||
InstanceID string
|
||||
PoolID string
|
||||
Endpoint string
|
||||
ProtocolVersion string
|
||||
Revision string
|
||||
Capabilities map[string]any
|
||||
Allocated int
|
||||
SafeCapacity int
|
||||
HeavyCapacity int
|
||||
ActiveTasks int
|
||||
PressureState string
|
||||
HeartbeatAt time.Time
|
||||
LoadSampledAt time.Time
|
||||
}
|
||||
|
||||
func (w WorkerDescriptor) AvailableCapacity() int {
|
||||
limit := w.Allocated
|
||||
if w.SafeCapacity < limit {
|
||||
limit = w.SafeCapacity
|
||||
}
|
||||
if limit < 0 {
|
||||
return 0
|
||||
}
|
||||
available := limit - w.ActiveTasks
|
||||
if available < 0 {
|
||||
return 0
|
||||
}
|
||||
return available
|
||||
}
|
||||
|
||||
type RouteProfile struct {
|
||||
Key string
|
||||
Provider string
|
||||
Protocol string
|
||||
EndpointHost string
|
||||
ProxyMode string
|
||||
ConfigRevision string
|
||||
}
|
||||
|
||||
type RouteHealth struct {
|
||||
PoolID string
|
||||
RouteProfileKey string
|
||||
State RouteState
|
||||
SuccessRate float64
|
||||
ConnectTLSP95 time.Duration
|
||||
FirstByteP95 time.Duration
|
||||
UploadBytesPerSecond float64
|
||||
JitterP95 time.Duration
|
||||
ConsecutiveFailures int
|
||||
ConsecutiveSuccesses int
|
||||
SampleCount int
|
||||
SampledAt time.Time
|
||||
ExpiresAt time.Time
|
||||
}
|
||||
|
||||
// RouteObservation contains only transport-phase facts. Model generation time
|
||||
// is deliberately excluded from this structure and therefore from routing
|
||||
// quality updates.
|
||||
type RouteObservation struct {
|
||||
PoolID string
|
||||
RouteProfileKey string
|
||||
SampleCount int
|
||||
SuccessCount int
|
||||
ConnectTLSP95 time.Duration
|
||||
UploadBytesPerSecond float64
|
||||
}
|
||||
|
||||
type CapacitySnapshot struct {
|
||||
PoolID string
|
||||
WorkerCount int
|
||||
Allocated int
|
||||
SafeCapacity int
|
||||
ActiveTasks int
|
||||
HeavyCapacity int
|
||||
HeavyTasks int
|
||||
QueueDepth int
|
||||
EstimatedQueueWait time.Duration
|
||||
ResourceHeadroom float64
|
||||
Stability float64
|
||||
Critical bool
|
||||
SampledAt time.Time
|
||||
}
|
||||
|
||||
type DesiredCapacity struct {
|
||||
PoolID string
|
||||
Desired int
|
||||
Reason string
|
||||
ValidUntil time.Time
|
||||
}
|
||||
|
||||
// WorkerDirectory is the platform-neutral discovery boundary. Implementations
|
||||
// may use PostgreSQL, Consul, etcd, or another registry.
|
||||
type WorkerDirectory interface {
|
||||
ListWorkers(context.Context, time.Time) ([]WorkerDescriptor, error)
|
||||
}
|
||||
|
||||
// ExecutionBroker hides the durable queue implementation from routing code.
|
||||
type ExecutionBroker interface {
|
||||
Publish(context.Context, string, string, time.Time) (int64, error)
|
||||
}
|
||||
|
||||
// ExecutionTransport hides HTTP, gRPC, or service-mesh transport details.
|
||||
type ExecutionTransport interface {
|
||||
Execute(context.Context, WorkerDescriptor, ExecutionRequest) (ExecutionResponse, error)
|
||||
}
|
||||
|
||||
type ExecutionRequest struct {
|
||||
TaskID string
|
||||
PoolID string
|
||||
WorkerID string
|
||||
LeaseID string
|
||||
AuthorizationToken string
|
||||
Deadline time.Time
|
||||
Stream bool
|
||||
}
|
||||
|
||||
type ExecutionResponse struct {
|
||||
StatusCode int
|
||||
Headers map[string][]string
|
||||
Body io.ReadCloser
|
||||
}
|
||||
|
||||
type RouteHealthRepository interface {
|
||||
ListRouteHealth(context.Context, string, time.Time) ([]RouteHealth, error)
|
||||
RecordRouteHealth(context.Context, RouteHealth) error
|
||||
RecordRouteObservation(context.Context, RouteObservation) error
|
||||
}
|
||||
|
||||
type CapacityProvider interface {
|
||||
ListCapacity(context.Context, time.Time) ([]CapacitySnapshot, error)
|
||||
PublishDesiredCapacity(context.Context, DesiredCapacity) error
|
||||
}
|
||||
@@ -0,0 +1,481 @@
|
||||
package httpapi
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
const (
|
||||
acceptanceRunHeader = "X-EasyAI-Acceptance-Run"
|
||||
acceptanceTokenHeader = "X-EasyAI-Acceptance-Token"
|
||||
acceptanceUpstreamHeader = "X-EasyAI-Acceptance-Upstream"
|
||||
taskTrafficAuthorizationCacheTTL = 5 * time.Second
|
||||
)
|
||||
|
||||
type taskTrafficAuthorizationCacheKey struct {
|
||||
runID string
|
||||
apiKeyID string
|
||||
userID string
|
||||
tokenHash [sha256.Size]byte
|
||||
}
|
||||
|
||||
type taskTrafficAuthorizationCacheEntry struct {
|
||||
runID string
|
||||
expiresAt time.Time
|
||||
}
|
||||
|
||||
type taskTrafficAuthorizationCall struct {
|
||||
done chan struct{}
|
||||
runID string
|
||||
err error
|
||||
}
|
||||
|
||||
type taskTrafficAdmission struct {
|
||||
RunMode string
|
||||
AcceptanceRunID string
|
||||
}
|
||||
|
||||
type taskTrafficError struct {
|
||||
Status int
|
||||
Code string
|
||||
Message string
|
||||
Err error
|
||||
}
|
||||
|
||||
func (s *Server) acceptanceStore() *store.Store {
|
||||
if s.coordinationStore != nil {
|
||||
return s.coordinationStore
|
||||
}
|
||||
return s.store
|
||||
}
|
||||
|
||||
func (e *taskTrafficError) Error() string {
|
||||
return e.Message
|
||||
}
|
||||
|
||||
func (e *taskTrafficError) ErrorCode() string {
|
||||
return e.Code
|
||||
}
|
||||
|
||||
func (s *Server) authorizeTaskTraffic(r *http.Request, user *auth.User) (taskTrafficAdmission, error) {
|
||||
runID, err := s.authorizeAcceptanceTask(r.Context(), r.Header.Get(acceptanceRunHeader), r.Header.Get(acceptanceTokenHeader), user)
|
||||
if err != nil {
|
||||
switch {
|
||||
case errors.Is(err, store.ErrProductionTrafficPaused):
|
||||
return taskTrafficAdmission{}, &taskTrafficError{
|
||||
Status: http.StatusServiceUnavailable, Code: "validation_in_progress",
|
||||
Message: "new production tasks are paused while validation is running", Err: err,
|
||||
}
|
||||
case errors.Is(err, store.ErrAcceptanceNotAuthorized):
|
||||
return taskTrafficAdmission{}, &taskTrafficError{
|
||||
Status: http.StatusForbidden, Code: "acceptance_not_authorized",
|
||||
Message: "acceptance credentials do not match the active run", Err: err,
|
||||
}
|
||||
case errors.Is(err, store.ErrAcceptanceRunNotActive):
|
||||
return taskTrafficAdmission{}, &taskTrafficError{
|
||||
Status: http.StatusConflict, Code: "acceptance_run_not_active",
|
||||
Message: "acceptance headers are not valid while live traffic is enabled", Err: err,
|
||||
}
|
||||
default:
|
||||
return taskTrafficAdmission{}, &taskTrafficError{
|
||||
Status: http.StatusServiceUnavailable, Code: "traffic_gate_unavailable",
|
||||
Message: "task traffic gate is unavailable", Err: err,
|
||||
}
|
||||
}
|
||||
}
|
||||
if runID != "" {
|
||||
runMode := "acceptance"
|
||||
if strings.EqualFold(strings.TrimSpace(r.Header.Get(acceptanceUpstreamHeader)), "real") {
|
||||
runMode = "acceptance_canary"
|
||||
}
|
||||
return taskTrafficAdmission{RunMode: runMode, AcceptanceRunID: runID}, nil
|
||||
}
|
||||
return taskTrafficAdmission{RunMode: "production"}, nil
|
||||
}
|
||||
|
||||
// authorizeAcceptanceTask collapses concurrent validation requests for the same
|
||||
// participant into one critical-pool lookup and briefly caches only successful
|
||||
// validation grants. Live requests and incomplete acceptance credentials always
|
||||
// consult PostgreSQL so a validation transition remains fail-closed.
|
||||
func (s *Server) authorizeAcceptanceTask(ctx context.Context, runID string, token string, user *auth.User) (string, error) {
|
||||
runID = strings.TrimSpace(runID)
|
||||
token = strings.TrimSpace(token)
|
||||
if runID == "" || token == "" || user == nil ||
|
||||
strings.TrimSpace(user.APIKeyID) == "" || strings.TrimSpace(user.ID) == "" {
|
||||
return s.acceptanceStore().AuthorizeAcceptanceTask(ctx, runID, token, user)
|
||||
}
|
||||
key := taskTrafficAuthorizationCacheKey{
|
||||
runID: runID, apiKeyID: strings.TrimSpace(user.APIKeyID), userID: strings.TrimSpace(user.ID),
|
||||
tokenHash: sha256.Sum256([]byte(token)),
|
||||
}
|
||||
return s.cachedTaskTrafficAuthorization(ctx, key, taskTrafficAuthorizationCacheTTL, func() (string, error) {
|
||||
return s.acceptanceStore().AuthorizeAcceptanceTask(ctx, runID, token, user)
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) cachedTaskTrafficAuthorization(
|
||||
ctx context.Context,
|
||||
key taskTrafficAuthorizationCacheKey,
|
||||
ttl time.Duration,
|
||||
load func() (string, error),
|
||||
) (string, error) {
|
||||
now := time.Now()
|
||||
s.taskTrafficCacheMu.Lock()
|
||||
if entry, ok := s.taskTrafficCache[key]; ok {
|
||||
if now.Before(entry.expiresAt) {
|
||||
s.taskTrafficCacheMu.Unlock()
|
||||
return entry.runID, nil
|
||||
}
|
||||
delete(s.taskTrafficCache, key)
|
||||
}
|
||||
if call, ok := s.taskTrafficInflight[key]; ok {
|
||||
s.taskTrafficCacheMu.Unlock()
|
||||
select {
|
||||
case <-call.done:
|
||||
return call.runID, call.err
|
||||
case <-ctx.Done():
|
||||
return "", ctx.Err()
|
||||
}
|
||||
}
|
||||
if s.taskTrafficInflight == nil {
|
||||
s.taskTrafficInflight = make(map[taskTrafficAuthorizationCacheKey]*taskTrafficAuthorizationCall)
|
||||
}
|
||||
call := &taskTrafficAuthorizationCall{done: make(chan struct{})}
|
||||
s.taskTrafficInflight[key] = call
|
||||
s.taskTrafficCacheMu.Unlock()
|
||||
|
||||
call.runID, call.err = load()
|
||||
|
||||
s.taskTrafficCacheMu.Lock()
|
||||
delete(s.taskTrafficInflight, key)
|
||||
if call.err == nil && call.runID != "" {
|
||||
if s.taskTrafficCache == nil {
|
||||
s.taskTrafficCache = make(map[taskTrafficAuthorizationCacheKey]taskTrafficAuthorizationCacheEntry)
|
||||
}
|
||||
callTTL := ttl
|
||||
if callTTL <= 0 {
|
||||
callTTL = taskTrafficAuthorizationCacheTTL
|
||||
}
|
||||
s.taskTrafficCache[key] = taskTrafficAuthorizationCacheEntry{
|
||||
runID: call.runID, expiresAt: time.Now().Add(callTTL),
|
||||
}
|
||||
}
|
||||
close(call.done)
|
||||
s.taskTrafficCacheMu.Unlock()
|
||||
return call.runID, call.err
|
||||
}
|
||||
|
||||
func (s *Server) admittedTaskRunMode(admission taskTrafficAdmission, body map[string]any) (string, error) {
|
||||
if admission.RunMode == "acceptance" || admission.RunMode == "acceptance_canary" {
|
||||
return admission.RunMode, nil
|
||||
}
|
||||
requested := runModeFromRequest(body)
|
||||
requestedMode := strings.ToLower(strings.TrimSpace(requested))
|
||||
if requestedMode == "acceptance" || requestedMode == "acceptance_canary" {
|
||||
return "", &taskTrafficError{
|
||||
Status: http.StatusForbidden, Code: "acceptance_not_authorized",
|
||||
Message: "acceptance mode is available only through an active acceptance run",
|
||||
}
|
||||
}
|
||||
if strings.EqualFold(strings.TrimSpace(s.cfg.AppEnv), "production") && requestedMode == "simulation" {
|
||||
return "", &taskTrafficError{
|
||||
Status: http.StatusForbidden, Code: "simulation_not_authorized",
|
||||
Message: "simulation mode is available only through an active acceptance run",
|
||||
}
|
||||
}
|
||||
return requested, nil
|
||||
}
|
||||
|
||||
func writeTaskTrafficError(w http.ResponseWriter, err error, protocol string) {
|
||||
var trafficErr *taskTrafficError
|
||||
if !errors.As(err, &trafficErr) {
|
||||
trafficErr = &taskTrafficError{
|
||||
Status: http.StatusServiceUnavailable, Code: "traffic_gate_unavailable",
|
||||
Message: "task traffic gate is unavailable", Err: err,
|
||||
}
|
||||
}
|
||||
if protocol != "" {
|
||||
writeProtocolError(w, protocol, trafficErr.Status, trafficErr.Message, nil, trafficErr.Code)
|
||||
return
|
||||
}
|
||||
writeErrorWithDetails(w, trafficErr.Status, trafficErr.Message, nil, trafficErr.Code)
|
||||
}
|
||||
|
||||
// getGatewayTrafficMode godoc
|
||||
// @Summary 获取 Gateway 流量模式
|
||||
// @Tags acceptance
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Success 200 {object} store.GatewayTrafficMode
|
||||
// @Router /api/admin/system/acceptance/traffic-mode [get]
|
||||
func (s *Server) getGatewayTrafficMode(w http.ResponseWriter, r *http.Request) {
|
||||
mode, err := s.acceptanceStore().GetGatewayTrafficMode(r.Context())
|
||||
if err != nil {
|
||||
writeError(w, http.StatusInternalServerError, "get gateway traffic mode failed")
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, mode)
|
||||
}
|
||||
|
||||
// pauseGatewayTraffic godoc
|
||||
// @Summary 上线后硬门禁异常时通过 CAS 暂停正式新任务
|
||||
// @Tags acceptance
|
||||
// @Accept json
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param body body store.PauseGatewayTrafficInput true "当前线上 release 与 revision CAS"
|
||||
// @Success 200 {object} store.GatewayTrafficMode
|
||||
// @Router /api/admin/system/acceptance/traffic-mode/pause [post]
|
||||
func (s *Server) pauseGatewayTraffic(w http.ResponseWriter, r *http.Request) {
|
||||
var input store.PauseGatewayTrafficInput
|
||||
if err := json.NewDecoder(r.Body).Decode(&input); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid json body")
|
||||
return
|
||||
}
|
||||
mode, err := s.acceptanceStore().PauseGatewayTraffic(r.Context(), input)
|
||||
if err != nil {
|
||||
writeAcceptanceMutationError(w, err)
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, mode)
|
||||
}
|
||||
|
||||
// listCapacityProfiles godoc
|
||||
// @Summary 获取经生产同构验收认证的容量配置
|
||||
// @Tags acceptance
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Success 200 {array} store.CapacityProfile
|
||||
// @Router /api/admin/system/acceptance/capacity-profiles [get]
|
||||
func (s *Server) listCapacityProfiles(w http.ResponseWriter, r *http.Request) {
|
||||
profiles, err := s.acceptanceStore().ListCapacityProfiles(r.Context())
|
||||
if err != nil {
|
||||
writeError(w, http.StatusInternalServerError, "list capacity profiles failed")
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, profiles)
|
||||
}
|
||||
|
||||
// getWorkerClusterRuntime godoc
|
||||
// @Summary 获取集群 Worker 实时负载与领取容量
|
||||
// @Tags acceptance
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Success 200 {object} store.WorkerClusterRuntime
|
||||
// @Router /api/admin/runtime/workers [get]
|
||||
func (s *Server) getWorkerClusterRuntime(w http.ResponseWriter, r *http.Request) {
|
||||
runtime, err := s.acceptanceStore().GetWorkerClusterRuntime(r.Context())
|
||||
if err != nil {
|
||||
writeError(w, http.StatusInternalServerError, "get worker cluster runtime failed")
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, runtime)
|
||||
}
|
||||
|
||||
// createAcceptanceRun godoc
|
||||
// @Summary 创建生产同构验收 Run
|
||||
// @Tags acceptance
|
||||
// @Accept json
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param body body store.CreateAcceptanceRunInput true "验收 Run"
|
||||
// @Success 201 {object} store.AcceptanceRun
|
||||
// @Router /api/admin/system/acceptance/runs [post]
|
||||
func (s *Server) createAcceptanceRun(w http.ResponseWriter, r *http.Request) {
|
||||
var input store.CreateAcceptanceRunInput
|
||||
if err := json.NewDecoder(r.Body).Decode(&input); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid json body")
|
||||
return
|
||||
}
|
||||
run, err := s.acceptanceStore().CreateAcceptanceRun(r.Context(), input)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusCreated, run)
|
||||
}
|
||||
|
||||
// getAcceptanceRun godoc
|
||||
// @Summary 获取生产同构验收 Run
|
||||
// @Tags acceptance
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param runID path string true "验收 Run ID"
|
||||
// @Success 200 {object} store.AcceptanceRun
|
||||
// @Router /api/admin/system/acceptance/runs/{runID} [get]
|
||||
func (s *Server) getAcceptanceRun(w http.ResponseWriter, r *http.Request) {
|
||||
run, err := s.acceptanceStore().GetAcceptanceRun(r.Context(), r.PathValue("runID"))
|
||||
if err != nil {
|
||||
if store.IsNotFound(err) {
|
||||
writeError(w, http.StatusNotFound, "acceptance run not found")
|
||||
return
|
||||
}
|
||||
writeError(w, http.StatusInternalServerError, "get acceptance run failed")
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, run)
|
||||
}
|
||||
|
||||
// activateAcceptanceRun godoc
|
||||
// @Summary 切换到 validation 并激活 Run
|
||||
// @Tags acceptance
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param runID path string true "验收 Run ID"
|
||||
// @Success 200 {object} store.GatewayTrafficMode
|
||||
// @Router /api/admin/system/acceptance/runs/{runID}/activate [post]
|
||||
func (s *Server) activateAcceptanceRun(w http.ResponseWriter, r *http.Request) {
|
||||
mode, err := s.acceptanceStore().ActivateAcceptanceRun(r.Context(), r.PathValue("runID"))
|
||||
if err != nil {
|
||||
writeAcceptanceMutationError(w, err)
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, mode)
|
||||
}
|
||||
|
||||
// stageAcceptanceCapacityProfiles godoc
|
||||
// @Summary 为当前 validation Run 暂存 80% 容量门禁
|
||||
// @Tags acceptance
|
||||
// @Accept json
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param runID path string true "验收 Run ID"
|
||||
// @Param body body store.StageAcceptanceCapacityProfilesInput true "待验证容量配置"
|
||||
// @Success 200 {object} store.AcceptanceRun
|
||||
// @Router /api/admin/system/acceptance/runs/{runID}/capacity-profiles [post]
|
||||
func (s *Server) stageAcceptanceCapacityProfiles(w http.ResponseWriter, r *http.Request) {
|
||||
var input store.StageAcceptanceCapacityProfilesInput
|
||||
if err := json.NewDecoder(r.Body).Decode(&input); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid json body")
|
||||
return
|
||||
}
|
||||
input.RunID = r.PathValue("runID")
|
||||
run, err := s.acceptanceStore().StageAcceptanceCapacityProfiles(r.Context(), input)
|
||||
if err != nil {
|
||||
switch {
|
||||
case store.IsNotFound(err):
|
||||
writeError(w, http.StatusConflict, "acceptance run is not running")
|
||||
default:
|
||||
writeError(w, http.StatusBadRequest, err.Error())
|
||||
}
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, run)
|
||||
}
|
||||
|
||||
// finishAcceptanceRun godoc
|
||||
// @Summary 记录验收门禁结果
|
||||
// @Tags acceptance
|
||||
// @Accept json
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param runID path string true "验收 Run ID"
|
||||
// @Param body body store.FinishAcceptanceRunInput true "门禁结果"
|
||||
// @Success 200 {object} store.AcceptanceRun
|
||||
// @Router /api/admin/system/acceptance/runs/{runID}/finish [post]
|
||||
func (s *Server) finishAcceptanceRun(w http.ResponseWriter, r *http.Request) {
|
||||
var input store.FinishAcceptanceRunInput
|
||||
if err := json.NewDecoder(r.Body).Decode(&input); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid json body")
|
||||
return
|
||||
}
|
||||
input.RunID = r.PathValue("runID")
|
||||
run, err := s.acceptanceStore().FinishAcceptanceRun(r.Context(), input)
|
||||
if err != nil {
|
||||
if store.IsNotFound(err) {
|
||||
writeError(w, http.StatusConflict, "acceptance run is not running")
|
||||
return
|
||||
}
|
||||
writeError(w, http.StatusInternalServerError, "finish acceptance run failed")
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, run)
|
||||
}
|
||||
|
||||
// retryAcceptanceRun godoc
|
||||
// @Summary 在 validation 关闭正式流量的状态下重试失败 Run
|
||||
// @Tags acceptance
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param runID path string true "验收 Run ID"
|
||||
// @Success 200 {object} store.AcceptanceRun
|
||||
// @Router /api/admin/system/acceptance/runs/{runID}/retry [post]
|
||||
func (s *Server) retryAcceptanceRun(w http.ResponseWriter, r *http.Request) {
|
||||
run, err := s.acceptanceStore().RetryAcceptanceRun(r.Context(), r.PathValue("runID"))
|
||||
if err != nil {
|
||||
writeAcceptanceMutationError(w, err)
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, run)
|
||||
}
|
||||
|
||||
// promoteAcceptanceRun godoc
|
||||
// @Summary 通过 CAS 切回 live
|
||||
// @Tags acceptance
|
||||
// @Accept json
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param runID path string true "验收 Run ID"
|
||||
// @Param body body store.PromoteAcceptanceRunInput true "线上 release 与 digest CAS"
|
||||
// @Success 200 {object} store.GatewayTrafficMode
|
||||
// @Router /api/admin/system/acceptance/runs/{runID}/promote [post]
|
||||
func (s *Server) promoteAcceptanceRun(w http.ResponseWriter, r *http.Request) {
|
||||
var input store.PromoteAcceptanceRunInput
|
||||
if err := json.NewDecoder(r.Body).Decode(&input); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid json body")
|
||||
return
|
||||
}
|
||||
input.RunID = r.PathValue("runID")
|
||||
mode, err := s.acceptanceStore().PromoteAcceptanceRun(r.Context(), input)
|
||||
if err != nil {
|
||||
writeAcceptanceMutationError(w, err)
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, mode)
|
||||
}
|
||||
|
||||
// abortAcceptanceRun godoc
|
||||
// @Summary 人工 CAS 中止验收并切回 live
|
||||
// @Tags acceptance
|
||||
// @Accept json
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param runID path string true "验收 Run ID"
|
||||
// @Param body body store.PromoteAcceptanceRunInput true "线上 release 与 digest CAS"
|
||||
// @Success 200 {object} store.GatewayTrafficMode
|
||||
// @Router /api/admin/system/acceptance/runs/{runID}/abort [post]
|
||||
func (s *Server) abortAcceptanceRun(w http.ResponseWriter, r *http.Request) {
|
||||
var input store.PromoteAcceptanceRunInput
|
||||
if err := json.NewDecoder(r.Body).Decode(&input); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid json body")
|
||||
return
|
||||
}
|
||||
input.RunID = r.PathValue("runID")
|
||||
mode, err := s.acceptanceStore().AbortAcceptanceRun(r.Context(), input)
|
||||
if err != nil {
|
||||
writeAcceptanceMutationError(w, err)
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, mode)
|
||||
}
|
||||
|
||||
func writeAcceptanceMutationError(w http.ResponseWriter, err error) {
|
||||
switch {
|
||||
case errors.Is(err, store.ErrAcceptanceStateConflict):
|
||||
writeError(w, http.StatusConflict, err.Error())
|
||||
case errors.Is(err, store.ErrAcceptancePromotionGates):
|
||||
writeError(w, http.StatusPreconditionFailed, err.Error())
|
||||
case store.IsNotFound(err):
|
||||
writeError(w, http.StatusNotFound, "acceptance run not found")
|
||||
default:
|
||||
writeError(w, http.StatusInternalServerError, "acceptance state update failed")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,103 @@
|
||||
package httpapi
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
|
||||
)
|
||||
|
||||
func TestAdmittedTaskRunModeReservesAcceptanceAndProductionSimulation(t *testing.T) {
|
||||
server := &Server{cfg: config.Config{AppEnv: "production"}}
|
||||
for _, runMode := range []string{"simulation", "SIMULATION", "acceptance", "acceptance_canary"} {
|
||||
_, err := server.admittedTaskRunMode(taskTrafficAdmission{}, map[string]any{"runMode": runMode})
|
||||
var trafficErr *taskTrafficError
|
||||
if !errors.As(err, &trafficErr) || trafficErr.Status != 403 {
|
||||
t.Fatalf("run mode %q error=%v", runMode, err)
|
||||
}
|
||||
}
|
||||
for _, runMode := range []string{"acceptance", "acceptance_canary"} {
|
||||
got, err := server.admittedTaskRunMode(taskTrafficAdmission{RunMode: runMode}, map[string]any{
|
||||
"runMode": "production",
|
||||
})
|
||||
if err != nil || got != runMode {
|
||||
t.Fatalf("admitted mode=%q got=%q err=%v", runMode, got, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestCachedTaskTrafficAuthorizationCollapsesConcurrentLoads(t *testing.T) {
|
||||
server := &Server{}
|
||||
key := taskTrafficAuthorizationCacheKey{runID: "run", apiKeyID: "key", userID: "user"}
|
||||
var loads atomic.Int32
|
||||
start := make(chan struct{})
|
||||
const callers = 64
|
||||
var wg sync.WaitGroup
|
||||
errCh := make(chan error, callers)
|
||||
for range callers {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
<-start
|
||||
runID, err := server.cachedTaskTrafficAuthorization(
|
||||
context.Background(), key, time.Minute,
|
||||
func() (string, error) {
|
||||
loads.Add(1)
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
return "run", nil
|
||||
},
|
||||
)
|
||||
if err != nil || runID != "run" {
|
||||
errCh <- errors.New("unexpected cached authorization result")
|
||||
}
|
||||
}()
|
||||
}
|
||||
close(start)
|
||||
wg.Wait()
|
||||
close(errCh)
|
||||
for err := range errCh {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := loads.Load(); got != 1 {
|
||||
t.Fatalf("authorization loads=%d, want 1", got)
|
||||
}
|
||||
|
||||
if _, err := server.cachedTaskTrafficAuthorization(
|
||||
context.Background(), key, time.Minute,
|
||||
func() (string, error) {
|
||||
loads.Add(1)
|
||||
return "run", nil
|
||||
},
|
||||
); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := loads.Load(); got != 1 {
|
||||
t.Fatalf("cached authorization loads=%d, want 1", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCachedTaskTrafficAuthorizationDoesNotCacheFailures(t *testing.T) {
|
||||
server := &Server{}
|
||||
key := taskTrafficAuthorizationCacheKey{runID: "run", apiKeyID: "key", userID: "user"}
|
||||
wantErr := errors.New("database unavailable")
|
||||
loads := 0
|
||||
for range 2 {
|
||||
_, err := server.cachedTaskTrafficAuthorization(
|
||||
context.Background(), key, time.Minute,
|
||||
func() (string, error) {
|
||||
loads++
|
||||
return "", wantErr
|
||||
},
|
||||
)
|
||||
if !errors.Is(err, wantErr) {
|
||||
t.Fatalf("authorization error=%v, want %v", err, wantErr)
|
||||
}
|
||||
}
|
||||
if loads != 2 {
|
||||
t.Fatalf("failed authorization loads=%d, want 2", loads)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,165 @@
|
||||
package httpapi
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
type acceptanceStatusSnapshot struct {
|
||||
ID string
|
||||
Status string
|
||||
}
|
||||
|
||||
type acceptanceUserGroupSnapshot struct {
|
||||
ID string
|
||||
RateLimitPolicy []byte
|
||||
}
|
||||
|
||||
// isolateHTTPAcceptanceState lets a stress test disable unrelated runtime
|
||||
// resources without leaking those mutations into later integration tests that
|
||||
// intentionally share the same dedicated test database.
|
||||
func isolateHTTPAcceptanceState(
|
||||
t *testing.T,
|
||||
ctx context.Context,
|
||||
pool *pgxpool.Pool,
|
||||
includeStorage bool,
|
||||
) func() {
|
||||
t.Helper()
|
||||
startedAt := time.Now()
|
||||
platforms := readAcceptanceStatusSnapshots(t, ctx, pool, `
|
||||
SELECT id::text, status
|
||||
FROM integration_platforms
|
||||
WHERE deleted_at IS NULL`)
|
||||
userGroups := make([]acceptanceUserGroupSnapshot, 0)
|
||||
rows, err := pool.Query(ctx, `
|
||||
SELECT id::text, rate_limit_policy
|
||||
FROM gateway_user_groups`)
|
||||
if err != nil {
|
||||
t.Fatalf("snapshot acceptance user groups: %v", err)
|
||||
}
|
||||
for rows.Next() {
|
||||
var item acceptanceUserGroupSnapshot
|
||||
if err := rows.Scan(&item.ID, &item.RateLimitPolicy); err != nil {
|
||||
rows.Close()
|
||||
t.Fatalf("scan acceptance user group snapshot: %v", err)
|
||||
}
|
||||
userGroups = append(userGroups, item)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
rows.Close()
|
||||
t.Fatalf("read acceptance user group snapshots: %v", err)
|
||||
}
|
||||
rows.Close()
|
||||
|
||||
storageChannels := []acceptanceStatusSnapshot(nil)
|
||||
if includeStorage {
|
||||
storageChannels = readAcceptanceStatusSnapshots(t, ctx, pool, `
|
||||
SELECT id::text, status
|
||||
FROM file_storage_channels
|
||||
WHERE deleted_at IS NULL`)
|
||||
}
|
||||
if _, err := pool.Exec(ctx, `
|
||||
UPDATE integration_platforms
|
||||
SET status = 'disabled'
|
||||
WHERE deleted_at IS NULL`); err != nil {
|
||||
t.Fatalf("isolate acceptance platforms: %v", err)
|
||||
}
|
||||
if includeStorage {
|
||||
if _, err := pool.Exec(ctx, `
|
||||
UPDATE file_storage_channels
|
||||
SET status = 'disabled'
|
||||
WHERE deleted_at IS NULL`); err != nil {
|
||||
t.Fatalf("isolate acceptance storage channels: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
return func() {
|
||||
restoreCtx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
|
||||
defer cancel()
|
||||
tx, err := pool.Begin(restoreCtx)
|
||||
if err != nil {
|
||||
t.Errorf("begin acceptance state restore: %v", err)
|
||||
return
|
||||
}
|
||||
defer func() {
|
||||
_ = tx.Rollback(restoreCtx)
|
||||
}()
|
||||
for _, item := range platforms {
|
||||
if _, err := tx.Exec(restoreCtx, `
|
||||
UPDATE integration_platforms
|
||||
SET status = $2
|
||||
WHERE id = $1::uuid`, item.ID, item.Status); err != nil {
|
||||
t.Errorf("restore acceptance platform %s: %v", item.ID, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
if _, err := tx.Exec(restoreCtx, `
|
||||
UPDATE integration_platforms
|
||||
SET status = 'disabled'
|
||||
WHERE created_at >= $1
|
||||
AND deleted_at IS NULL`, startedAt); err != nil {
|
||||
t.Errorf("disable acceptance-created platforms: %v", err)
|
||||
return
|
||||
}
|
||||
for _, item := range userGroups {
|
||||
if _, err := tx.Exec(restoreCtx, `
|
||||
UPDATE gateway_user_groups
|
||||
SET rate_limit_policy = $2::jsonb
|
||||
WHERE id = $1::uuid`, item.ID, string(item.RateLimitPolicy)); err != nil {
|
||||
t.Errorf("restore acceptance user group %s: %v", item.ID, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
if includeStorage {
|
||||
for _, item := range storageChannels {
|
||||
if _, err := tx.Exec(restoreCtx, `
|
||||
UPDATE file_storage_channels
|
||||
SET status = $2
|
||||
WHERE id = $1::uuid`, item.ID, item.Status); err != nil {
|
||||
t.Errorf("restore acceptance storage channel %s: %v", item.ID, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
if _, err := tx.Exec(restoreCtx, `
|
||||
UPDATE file_storage_channels
|
||||
SET status = 'disabled'
|
||||
WHERE created_at >= $1
|
||||
AND deleted_at IS NULL`, startedAt); err != nil {
|
||||
t.Errorf("disable acceptance-created storage channels: %v", err)
|
||||
return
|
||||
}
|
||||
}
|
||||
if err := tx.Commit(restoreCtx); err != nil {
|
||||
t.Errorf("commit acceptance state restore: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func readAcceptanceStatusSnapshots(
|
||||
t *testing.T,
|
||||
ctx context.Context,
|
||||
pool *pgxpool.Pool,
|
||||
query string,
|
||||
) []acceptanceStatusSnapshot {
|
||||
t.Helper()
|
||||
rows, err := pool.Query(ctx, query)
|
||||
if err != nil {
|
||||
t.Fatalf("snapshot acceptance statuses: %v", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
items := make([]acceptanceStatusSnapshot, 0)
|
||||
for rows.Next() {
|
||||
var item acceptanceStatusSnapshot
|
||||
if err := rows.Scan(&item.ID, &item.Status); err != nil {
|
||||
t.Fatalf("scan acceptance status snapshot: %v", err)
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
t.Fatalf("read acceptance status snapshots: %v", err)
|
||||
}
|
||||
return items
|
||||
}
|
||||
@@ -12,7 +12,7 @@ import (
|
||||
|
||||
// listAccessRules godoc
|
||||
// @Summary 列出访问规则
|
||||
// @Description 管理端返回用户组、租户、用户或 API Key 到平台、平台模型、基础模型的访问规则。
|
||||
// @Description 管理端返回用户组、租户、用户或 API Key 到平台、平台模型、基础模型的分层访问规则。主体当前层无 allow 时继承上级,存在 allow 时仅允许白名单,deny 始终优先。
|
||||
// @Tags access-rules
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
@@ -33,7 +33,7 @@ func (s *Server) listAccessRules(w http.ResponseWriter, r *http.Request) {
|
||||
|
||||
// listAPIKeyAccessRules godoc
|
||||
// @Summary 列出 API Key 访问规则
|
||||
// @Description 返回当前本地用户可管理的 API Key 访问规则。
|
||||
// @Description 返回当前本地用户拥有的 API Key 访问规则;不会混入其他用户或其他 API Key 的规则。
|
||||
// @Tags api-keys
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
@@ -60,16 +60,18 @@ func (s *Server) listAPIKeyAccessRules(w http.ResponseWriter, r *http.Request) {
|
||||
|
||||
// listAPIKeyAssignableModels godoc
|
||||
// @Summary 列出 API Key 可分配模型
|
||||
// @Description 按当前用户自身的用户、租户和用户组权限返回可分配给 API Key 的启用模型,不受任何 API Key 权限规则影响。
|
||||
// @Description 按当前用户自身的租户、用户组和用户分层白名单返回可分配给 API Key 的启用模型,不应用任何 API Key 层规则。
|
||||
// @Tags api-keys
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Deprecated
|
||||
// @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) {
|
||||
w.Header().Set("Deprecation", "true")
|
||||
user, _ := auth.UserFromContext(r.Context())
|
||||
models, err := s.store.ListAPIKeyAssignablePlatformModels(r.Context(), user)
|
||||
if err != nil {
|
||||
@@ -84,9 +86,44 @@ func (s *Server) listAPIKeyAssignableModels(w http.ResponseWriter, r *http.Reque
|
||||
writeJSON(w, http.StatusOK, map[string]any{"items": s.platformModelResponses(r.Context(), models)})
|
||||
}
|
||||
|
||||
// listAPIKeyAssignableModelsForKey godoc
|
||||
// @Summary 列出指定 API Key 可分配模型
|
||||
// @Description 返回全局启用、命中指定 API Key 的租户/用户组/用户基线且符合 Key scope 的可分配平台来源;当前 Key 的 allow/deny 不缩减候选,仅作为已有规则有效性诊断返回。
|
||||
// @Tags api-keys
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param apiKeyID path string true "API Key ID"
|
||||
// @Success 200 {object} APIKeyAssignableModelsResponse
|
||||
// @Failure 401 {object} ErrorEnvelope
|
||||
// @Failure 404 {object} ErrorEnvelope
|
||||
// @Failure 503 {object} ErrorEnvelope
|
||||
// @Failure 500 {object} ErrorEnvelope
|
||||
// @Router /api/v1/api-keys/{apiKeyID}/assignable-models [get]
|
||||
func (s *Server) listAPIKeyAssignableModelsForKey(w http.ResponseWriter, r *http.Request) {
|
||||
user, _ := auth.UserFromContext(r.Context())
|
||||
models, diagnostics, err := s.store.ListAPIKeyAssignablePlatformModelsForKey(r.Context(), user, r.PathValue("apiKeyID"))
|
||||
if err != nil {
|
||||
if errors.Is(err, store.ErrLocalUserRequired) {
|
||||
writeLocalUserRequired(w)
|
||||
return
|
||||
}
|
||||
if store.IsNotFound(err) {
|
||||
writeError(w, http.StatusNotFound, "api key not found")
|
||||
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, APIKeyAssignableModelsResponse{
|
||||
Items: s.platformModelResponses(r.Context(), models),
|
||||
RuleDiagnostics: diagnostics,
|
||||
})
|
||||
}
|
||||
|
||||
// createAccessRule godoc
|
||||
// @Summary 创建访问规则
|
||||
// @Description 管理端创建一条访问控制规则。
|
||||
// @Description 管理端创建一条访问控制规则;同一主体层存在任意有效 allow 后该层启用白名单,deny 始终优先。
|
||||
// @Tags access-rules
|
||||
// @Accept json
|
||||
// @Produce json
|
||||
@@ -124,7 +161,7 @@ func (s *Server) createAccessRule(w http.ResponseWriter, r *http.Request) {
|
||||
|
||||
// batchAccessRules godoc
|
||||
// @Summary 批量写入访问规则
|
||||
// @Description 管理端为同一主体批量新增、更新或删除资源访问规则。
|
||||
// @Description 管理端为同一主体批量新增、更新或删除资源访问规则。清空该主体全部 allow 会恢复上级继承。
|
||||
// @Tags access-rules
|
||||
// @Accept json
|
||||
// @Produce json
|
||||
@@ -157,7 +194,7 @@ func (s *Server) batchAccessRules(w http.ResponseWriter, r *http.Request) {
|
||||
|
||||
// batchAPIKeyAccessRules godoc
|
||||
// @Summary 批量写入 API Key 访问规则
|
||||
// @Description 当前本地用户为自己的 API Key 批量新增、更新或删除可访问资源。
|
||||
// @Description 当前本地用户为自己的 API Key 批量新增、更新或删除白名单/拒绝资源;Key 无 allow 时继承父级范围,存在 allow 后仅允许命中项。
|
||||
// @Tags api-keys
|
||||
// @Accept json
|
||||
// @Produce json
|
||||
@@ -192,7 +229,7 @@ func (s *Server) batchAPIKeyAccessRules(w http.ResponseWriter, r *http.Request)
|
||||
return
|
||||
}
|
||||
if errors.Is(err, store.ErrAccessRuleResourceDenied) {
|
||||
writeError(w, http.StatusForbidden, "resource is not available for current user group")
|
||||
writeError(w, http.StatusForbidden, "resource is not available for current user group or API key scope")
|
||||
return
|
||||
}
|
||||
s.logger.Error("batch api key access rules failed", "error", err)
|
||||
|
||||
@@ -0,0 +1,143 @@
|
||||
package httpapi
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/executionpool"
|
||||
)
|
||||
|
||||
type AdminExecutionPoolWorker struct {
|
||||
WorkerID string `json:"workerId"`
|
||||
InstanceID string `json:"instanceId"`
|
||||
Revision string `json:"revision,omitempty"`
|
||||
ProtocolVersion string `json:"protocolVersion"`
|
||||
Capabilities map[string]any `json:"capabilities"`
|
||||
Allocated int `json:"allocated"`
|
||||
SafeCapacity int `json:"safeCapacity"`
|
||||
ActiveTasks int `json:"activeTasks"`
|
||||
PressureState string `json:"pressureState"`
|
||||
HeartbeatAt time.Time `json:"heartbeatAt"`
|
||||
}
|
||||
|
||||
type AdminExecutionPool struct {
|
||||
PoolID string `json:"poolId"`
|
||||
Labels map[string]string `json:"labels"`
|
||||
Capabilities map[string]any `json:"capabilities"`
|
||||
State string `json:"state"`
|
||||
Capacity *AdminExecutionPoolCapacity `json:"capacity,omitempty"`
|
||||
Workers []AdminExecutionPoolWorker `json:"workers"`
|
||||
}
|
||||
|
||||
type AdminExecutionPoolCapacity struct {
|
||||
WorkerCount int `json:"workerCount"`
|
||||
Allocated int `json:"allocated"`
|
||||
SafeCapacity int `json:"safeCapacity"`
|
||||
ActiveTasks int `json:"activeTasks"`
|
||||
HeavyCapacity int `json:"heavyCapacity"`
|
||||
HeavyTasks int `json:"heavyTasks"`
|
||||
QueueDepth int `json:"queueDepth"`
|
||||
EstimatedQueueWaitMilli int64 `json:"estimatedQueueWaitMs"`
|
||||
ResourceHeadroom float64 `json:"resourceHeadroom"`
|
||||
Stability float64 `json:"stability"`
|
||||
Critical bool `json:"critical"`
|
||||
SampledAt time.Time `json:"sampledAt"`
|
||||
}
|
||||
|
||||
type AdminRouteHealth struct {
|
||||
PoolID string `json:"poolId"`
|
||||
RouteProfileKey string `json:"routeProfileKey"`
|
||||
State string `json:"state"`
|
||||
SuccessRate float64 `json:"successRate"`
|
||||
ConnectTLSP95Milli int64 `json:"connectTlsP95Ms"`
|
||||
FirstByteP95Milli int64 `json:"firstByteP95Ms"`
|
||||
UploadBytesPerSecond float64 `json:"uploadBytesPerSecond"`
|
||||
JitterP95Milli int64 `json:"jitterP95Ms"`
|
||||
ConsecutiveFailures int `json:"consecutiveFailures"`
|
||||
ConsecutiveSuccesses int `json:"consecutiveSuccesses"`
|
||||
SampleCount int `json:"sampleCount"`
|
||||
SampledAt time.Time `json:"sampledAt"`
|
||||
ExpiresAt time.Time `json:"expiresAt"`
|
||||
}
|
||||
|
||||
type AdminExecutionPoolResponse struct {
|
||||
Pools []AdminExecutionPool `json:"pools"`
|
||||
RouteHealth []AdminRouteHealth `json:"routeHealth"`
|
||||
CapturedAt time.Time `json:"capturedAt"`
|
||||
}
|
||||
|
||||
// listExecutionPools godoc
|
||||
// @Summary 查询逻辑执行池、Worker 容量和上游路由健康状态
|
||||
// @Tags Admin Runtime
|
||||
// @Produce json
|
||||
// @Success 200 {object} AdminExecutionPoolResponse
|
||||
// @Failure 500 {object} ErrorEnvelope
|
||||
// @Security BearerAuth
|
||||
// @Router /api/admin/runtime/execution-pools [get]
|
||||
func (s *Server) listExecutionPools(w http.ResponseWriter, r *http.Request) {
|
||||
now := time.Now()
|
||||
pools, err := s.coordinationStore.ListExecutionPools(r.Context())
|
||||
if err != nil {
|
||||
writeError(w, http.StatusInternalServerError, err.Error())
|
||||
return
|
||||
}
|
||||
workers, err := s.coordinationStore.ListWorkers(r.Context(), now)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusInternalServerError, err.Error())
|
||||
return
|
||||
}
|
||||
capacities, err := s.coordinationStore.ListCapacity(r.Context(), now)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusInternalServerError, err.Error())
|
||||
return
|
||||
}
|
||||
health, err := s.coordinationStore.ListAllRouteHealth(r.Context())
|
||||
if err != nil {
|
||||
writeError(w, http.StatusInternalServerError, err.Error())
|
||||
return
|
||||
}
|
||||
capacityByPool := make(map[string]executionpool.CapacitySnapshot, len(capacities))
|
||||
for _, capacity := range capacities {
|
||||
capacityByPool[capacity.PoolID] = capacity
|
||||
}
|
||||
workersByPool := make(map[string][]AdminExecutionPoolWorker)
|
||||
for _, worker := range workers {
|
||||
workersByPool[worker.PoolID] = append(workersByPool[worker.PoolID], AdminExecutionPoolWorker{
|
||||
WorkerID: worker.WorkerID, InstanceID: worker.InstanceID, Revision: worker.Revision,
|
||||
ProtocolVersion: worker.ProtocolVersion, Capabilities: worker.Capabilities,
|
||||
Allocated: worker.Allocated, SafeCapacity: worker.SafeCapacity,
|
||||
ActiveTasks: worker.ActiveTasks, PressureState: worker.PressureState,
|
||||
HeartbeatAt: worker.HeartbeatAt,
|
||||
})
|
||||
}
|
||||
response := AdminExecutionPoolResponse{
|
||||
Pools: make([]AdminExecutionPool, 0, len(pools)), RouteHealth: make([]AdminRouteHealth, 0, len(health)), CapturedAt: now,
|
||||
}
|
||||
for _, item := range health {
|
||||
response.RouteHealth = append(response.RouteHealth, AdminRouteHealth{
|
||||
PoolID: item.PoolID, RouteProfileKey: item.RouteProfileKey, State: string(item.State),
|
||||
SuccessRate: item.SuccessRate, ConnectTLSP95Milli: item.ConnectTLSP95.Milliseconds(),
|
||||
FirstByteP95Milli: item.FirstByteP95.Milliseconds(), UploadBytesPerSecond: item.UploadBytesPerSecond,
|
||||
JitterP95Milli: item.JitterP95.Milliseconds(), ConsecutiveFailures: item.ConsecutiveFailures,
|
||||
ConsecutiveSuccesses: item.ConsecutiveSuccesses, SampleCount: item.SampleCount,
|
||||
SampledAt: item.SampledAt, ExpiresAt: item.ExpiresAt,
|
||||
})
|
||||
}
|
||||
for _, pool := range pools {
|
||||
item := AdminExecutionPool{
|
||||
PoolID: pool.ID, Labels: pool.Labels, Capabilities: pool.Capabilities,
|
||||
State: string(pool.State), Workers: workersByPool[pool.ID],
|
||||
}
|
||||
if capacity, ok := capacityByPool[pool.ID]; ok {
|
||||
item.Capacity = &AdminExecutionPoolCapacity{
|
||||
WorkerCount: capacity.WorkerCount, Allocated: capacity.Allocated, SafeCapacity: capacity.SafeCapacity,
|
||||
ActiveTasks: capacity.ActiveTasks, HeavyCapacity: capacity.HeavyCapacity, HeavyTasks: capacity.HeavyTasks,
|
||||
QueueDepth: capacity.QueueDepth, EstimatedQueueWaitMilli: capacity.EstimatedQueueWait.Milliseconds(),
|
||||
ResourceHeadroom: capacity.ResourceHeadroom, Stability: capacity.Stability,
|
||||
Critical: capacity.Critical, SampledAt: capacity.SampledAt,
|
||||
}
|
||||
}
|
||||
response.Pools = append(response.Pools, item)
|
||||
}
|
||||
writeJSON(w, http.StatusOK, response)
|
||||
}
|
||||
@@ -0,0 +1,213 @@
|
||||
package httpapi
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// listAdminTasks godoc
|
||||
// @Summary 管理员列出任务
|
||||
// @Description 跨租户分页查询任务;请求、结果和执行快照中的常见敏感字段默认脱敏。
|
||||
// @Tags admin-tasks
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param q query string false "关键词,匹配任务、用户、租户、模型、平台和 API Key"
|
||||
// @Param tenantId query string false "网关租户 ID"
|
||||
// @Param userId query string false "网关用户 ID"
|
||||
// @Param userGroupId query string false "用户组 ID"
|
||||
// @Param status query string false "任务状态"
|
||||
// @Param platformId query string false "执行平台 ID"
|
||||
// @Param model query string false "调用或实际模型名称"
|
||||
// @Param modelType query string false "模型类型"
|
||||
// @Param runMode query string false "运行模式"
|
||||
// @Param billingStatus query string false "计费状态"
|
||||
// @Param apiKey query string false "API Key ID、名称或前缀"
|
||||
// @Param createdFrom query string false "创建时间起点,支持 RFC3339 或日期格式"
|
||||
// @Param createdTo query string false "创建时间终点,支持 RFC3339 或日期格式"
|
||||
// @Param page query int false "页码" default(1)
|
||||
// @Param pageSize query int false "每页数量" default(50)
|
||||
// @Success 200 {object} AdminTaskListResponse
|
||||
// @Failure 400 {object} ErrorEnvelope
|
||||
// @Failure 401 {object} ErrorEnvelope
|
||||
// @Failure 403 {object} ErrorEnvelope
|
||||
// @Failure 500 {object} ErrorEnvelope
|
||||
// @Router /api/admin/tasks [get]
|
||||
func (s *Server) listAdminTasks(w http.ResponseWriter, r *http.Request) {
|
||||
query := r.URL.Query()
|
||||
page, err := positiveQueryInt(query.Get("page"), 1)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid page")
|
||||
return
|
||||
}
|
||||
pageSize, err := positiveQueryInt(query.Get("pageSize"), 50)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid pageSize")
|
||||
return
|
||||
}
|
||||
createdFrom, err := parseTaskListTime(query.Get("createdFrom"), query.Get("from"))
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid createdFrom")
|
||||
return
|
||||
}
|
||||
createdTo, err := parseTaskListTime(query.Get("createdTo"), query.Get("to"))
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid createdTo")
|
||||
return
|
||||
}
|
||||
for name, value := range map[string]string{
|
||||
"tenantId": query.Get("tenantId"),
|
||||
"userId": query.Get("userId"),
|
||||
"userGroupId": query.Get("userGroupId"),
|
||||
"platformId": query.Get("platformId"),
|
||||
} {
|
||||
if value != "" && !validUUID(value) {
|
||||
writeError(w, http.StatusBadRequest, "invalid "+name)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
result, err := s.store.ListAdminTasks(r.Context(), store.AdminTaskListFilter{
|
||||
Query: firstNonEmpty(query.Get("q"), query.Get("query")),
|
||||
GatewayTenant: query.Get("tenantId"),
|
||||
GatewayUser: query.Get("userId"),
|
||||
UserGroup: query.Get("userGroupId"),
|
||||
Status: query.Get("status"),
|
||||
Platform: query.Get("platformId"),
|
||||
Model: query.Get("model"),
|
||||
ModelType: query.Get("modelType"),
|
||||
RunMode: query.Get("runMode"),
|
||||
BillingStatus: query.Get("billingStatus"),
|
||||
APIKey: query.Get("apiKey"),
|
||||
CreatedFrom: createdFrom,
|
||||
CreatedTo: createdTo,
|
||||
Page: page,
|
||||
PageSize: pageSize,
|
||||
})
|
||||
if err != nil {
|
||||
s.logger.Error("list admin tasks failed", "error", err)
|
||||
writeError(w, http.StatusInternalServerError, "list admin tasks failed")
|
||||
return
|
||||
}
|
||||
for index := range result.Items {
|
||||
result.Items[index] = store.MaskAdminGatewayTask(result.Items[index])
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]any{
|
||||
"items": result.Items,
|
||||
"total": result.Total,
|
||||
"page": result.Page,
|
||||
"pageSize": result.PageSize,
|
||||
})
|
||||
}
|
||||
|
||||
// getAdminTask godoc
|
||||
// @Summary 管理员获取任务详情
|
||||
// @Description 返回任意租户的任务详情;默认脱敏,sensitive=full 时返回完整存储内容。
|
||||
// @Tags admin-tasks
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param taskID path string true "任务 ID"
|
||||
// @Param sensitive query string false "敏感字段模式" Enums(masked,full) default(masked)
|
||||
// @Success 200 {object} store.AdminGatewayTask
|
||||
// @Failure 400 {object} ErrorEnvelope
|
||||
// @Failure 401 {object} ErrorEnvelope
|
||||
// @Failure 403 {object} ErrorEnvelope
|
||||
// @Failure 404 {object} ErrorEnvelope
|
||||
// @Failure 500 {object} ErrorEnvelope
|
||||
// @Router /api/admin/tasks/{taskID} [get]
|
||||
func (s *Server) getAdminTask(w http.ResponseWriter, r *http.Request) {
|
||||
taskID := r.PathValue("taskID")
|
||||
if !validUUID(taskID) {
|
||||
writeError(w, http.StatusBadRequest, "invalid taskID")
|
||||
return
|
||||
}
|
||||
full, ok := parseAdminTaskSensitiveMode(r.URL.Query().Get("sensitive"))
|
||||
if !ok {
|
||||
writeError(w, http.StatusBadRequest, "invalid sensitive mode")
|
||||
return
|
||||
}
|
||||
task, err := s.store.GetAdminTask(r.Context(), taskID)
|
||||
if err != nil {
|
||||
if store.IsNotFound(err) {
|
||||
writeError(w, http.StatusNotFound, "task not found")
|
||||
return
|
||||
}
|
||||
s.logger.Error("get admin task failed", "taskID", taskID, "error", err)
|
||||
writeError(w, http.StatusInternalServerError, "get admin task failed")
|
||||
return
|
||||
}
|
||||
task.GatewayTask, err = s.hydrateTaskResult(r.Context(), task.GatewayTask)
|
||||
if err != nil {
|
||||
writeStoredBinaryResultError(w, err)
|
||||
return
|
||||
}
|
||||
if !full {
|
||||
task = store.MaskAdminGatewayTask(task)
|
||||
}
|
||||
writeJSON(w, http.StatusOK, task)
|
||||
}
|
||||
|
||||
// adminTaskParamPreprocessing godoc
|
||||
// @Summary 管理员获取任务参数预处理日志
|
||||
// @Description 返回任意租户任务的参数改写、校验或模板处理日志;默认脱敏。
|
||||
// @Tags admin-tasks
|
||||
// @Produce json
|
||||
// @Security BearerAuth
|
||||
// @Param taskID path string true "任务 ID"
|
||||
// @Param sensitive query string false "敏感字段模式" Enums(masked,full) default(masked)
|
||||
// @Success 200 {object} TaskParamPreprocessingLogListResponse
|
||||
// @Failure 400 {object} ErrorEnvelope
|
||||
// @Failure 401 {object} ErrorEnvelope
|
||||
// @Failure 403 {object} ErrorEnvelope
|
||||
// @Failure 404 {object} ErrorEnvelope
|
||||
// @Failure 500 {object} ErrorEnvelope
|
||||
// @Router /api/admin/tasks/{taskID}/param-preprocessing [get]
|
||||
func (s *Server) adminTaskParamPreprocessing(w http.ResponseWriter, r *http.Request) {
|
||||
taskID := r.PathValue("taskID")
|
||||
if !validUUID(taskID) {
|
||||
writeError(w, http.StatusBadRequest, "invalid taskID")
|
||||
return
|
||||
}
|
||||
full, ok := parseAdminTaskSensitiveMode(r.URL.Query().Get("sensitive"))
|
||||
if !ok {
|
||||
writeError(w, http.StatusBadRequest, "invalid sensitive mode")
|
||||
return
|
||||
}
|
||||
if _, err := s.store.GetAdminTask(r.Context(), taskID); err != nil {
|
||||
if store.IsNotFound(err) {
|
||||
writeError(w, http.StatusNotFound, "task not found")
|
||||
return
|
||||
}
|
||||
s.logger.Error("get admin task failed", "taskID", taskID, "error", err)
|
||||
writeError(w, http.StatusInternalServerError, "get admin task failed")
|
||||
return
|
||||
}
|
||||
items, err := s.store.ListTaskParamPreprocessingLogs(r.Context(), taskID)
|
||||
if err != nil {
|
||||
s.logger.Error("list admin task parameter preprocessing logs failed", "taskID", taskID, "error", err)
|
||||
writeError(w, http.StatusInternalServerError, "list admin task parameter preprocessing logs failed")
|
||||
return
|
||||
}
|
||||
if !full {
|
||||
items = store.MaskTaskParamPreprocessingLogs(items)
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]any{"items": items})
|
||||
}
|
||||
|
||||
func parseAdminTaskSensitiveMode(value string) (bool, bool) {
|
||||
switch strings.ToLower(strings.TrimSpace(value)) {
|
||||
case "", "masked":
|
||||
return false, true
|
||||
case "full":
|
||||
return true, true
|
||||
default:
|
||||
return false, false
|
||||
}
|
||||
}
|
||||
|
||||
func validUUID(value string) bool {
|
||||
_, err := uuid.Parse(strings.TrimSpace(value))
|
||||
return err == nil
|
||||
}
|
||||
@@ -0,0 +1,22 @@
|
||||
package httpapi
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestParseAdminTaskSensitiveMode(t *testing.T) {
|
||||
tests := []struct {
|
||||
value string
|
||||
full bool
|
||||
accepted bool
|
||||
}{
|
||||
{value: "", full: false, accepted: true},
|
||||
{value: "masked", full: false, accepted: true},
|
||||
{value: "FULL", full: true, accepted: true},
|
||||
{value: "invalid", full: false, accepted: false},
|
||||
}
|
||||
for _, test := range tests {
|
||||
full, accepted := parseAdminTaskSensitiveMode(test.value)
|
||||
if full != test.full || accepted != test.accepted {
|
||||
t.Fatalf("parseAdminTaskSensitiveMode(%q)=(%v,%v), want (%v,%v)", test.value, full, accepted, test.full, test.accepted)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,321 @@
|
||||
package httpapi
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func TestAdminTasksCrossTenantQueryAndRedaction(t *testing.T) {
|
||||
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
|
||||
if databaseURL == "" {
|
||||
t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run the PostgreSQL integration flow")
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
applyMigration(t, ctx, databaseURL)
|
||||
db, err := store.Connect(ctx, databaseURL)
|
||||
if err != nil {
|
||||
t.Fatalf("connect store: %v", err)
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
suffix := strconv.FormatInt(time.Now().UnixNano(), 10)
|
||||
createTenant := func(key, name string) store.GatewayTenant {
|
||||
t.Helper()
|
||||
tenant, err := db.CreateTenant(ctx, store.GatewayTenantInput{
|
||||
TenantKey: key + "-" + suffix,
|
||||
Name: name,
|
||||
Source: "gateway",
|
||||
Status: "active",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("create tenant %s: %v", name, err)
|
||||
}
|
||||
return tenant
|
||||
}
|
||||
createUser := func(tenant store.GatewayTenant, prefix string, roles []string) store.GatewayUser {
|
||||
t.Helper()
|
||||
username := prefix + "_" + suffix
|
||||
user, err := db.CreateGatewayUser(ctx, store.GatewayUserInput{
|
||||
UserKey: "gateway:" + username,
|
||||
Username: username,
|
||||
DisplayName: prefix + " display",
|
||||
Email: username + "@example.com",
|
||||
Source: "gateway",
|
||||
GatewayTenantID: tenant.ID,
|
||||
TenantKey: tenant.TenantKey,
|
||||
Roles: roles,
|
||||
Status: "active",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("create gateway user %s: %v", prefix, err)
|
||||
}
|
||||
return user
|
||||
}
|
||||
authUser := func(user store.GatewayUser, apiKeyName string) *auth.User {
|
||||
return &auth.User{
|
||||
ID: user.ID,
|
||||
Username: user.Username,
|
||||
Roles: user.Roles,
|
||||
Source: user.Source,
|
||||
GatewayUserID: user.ID,
|
||||
GatewayTenantID: user.GatewayTenantID,
|
||||
TenantKey: user.TenantKey,
|
||||
APIKeyID: "key-" + apiKeyName + "-" + suffix,
|
||||
APIKeyName: apiKeyName,
|
||||
APIKeyPrefix: "sk-" + apiKeyName,
|
||||
}
|
||||
}
|
||||
createTask := func(user *auth.User, model, secret string) store.GatewayTask {
|
||||
t.Helper()
|
||||
task, err := db.CreateTask(ctx, store.CreateTaskInput{
|
||||
Kind: "chat.completions",
|
||||
Model: model,
|
||||
RunMode: "simulation",
|
||||
Request: map[string]any{
|
||||
"model": model,
|
||||
"diagnostic": map[string]any{
|
||||
"apiKey": secret,
|
||||
},
|
||||
},
|
||||
}, user)
|
||||
if err != nil {
|
||||
t.Fatalf("create task for %s: %v", user.Username, err)
|
||||
}
|
||||
return task
|
||||
}
|
||||
|
||||
tenantA := createTenant("admin-task-a", "Admin Task Tenant A")
|
||||
tenantB := createTenant("admin-task-b", "Admin Task Tenant B")
|
||||
operator := createUser(tenantA, "admin_task_operator", []string{"operator"})
|
||||
userA := createUser(tenantA, "admin_task_user_a", []string{"user"})
|
||||
userB := createUser(tenantB, "admin_task_user_b", []string{"user"})
|
||||
userAAuth := authUser(userA, "alpha-key")
|
||||
userBAuth := authUser(userB, "bravo-key")
|
||||
|
||||
modelA := "admin-task-model-a-" + suffix
|
||||
modelB := "admin-task-model-b-" + suffix
|
||||
taskA := createTask(userAAuth, modelA, "alpha-secret-"+suffix)
|
||||
taskBSecret := "bravo-secret-" + suffix
|
||||
taskB := createTask(userBAuth, modelB, taskBSecret)
|
||||
|
||||
platform, err := db.CreatePlatform(ctx, store.CreatePlatformInput{
|
||||
Provider: "openai",
|
||||
PlatformKey: "admin-task-platform-" + suffix,
|
||||
Name: "Admin Task Platform " + suffix,
|
||||
BaseURL: "https://example.invalid/v1",
|
||||
AuthType: "bearer",
|
||||
Credentials: map[string]any{"apiKey": "platform-secret-" + suffix},
|
||||
Status: "enabled",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("create platform: %v", err)
|
||||
}
|
||||
attemptID, err := db.CreateTaskAttempt(ctx, store.CreateTaskAttemptInput{
|
||||
TaskID: taskB.ID,
|
||||
AttemptNo: 1,
|
||||
PlatformID: platform.ID,
|
||||
QueueKey: "admin-task-integration-" + suffix,
|
||||
Status: "succeeded",
|
||||
Simulated: true,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("create task attempt: %v", err)
|
||||
}
|
||||
if _, err := db.Pool().Exec(ctx, `
|
||||
UPDATE gateway_tasks
|
||||
SET status = 'succeeded',
|
||||
model_type = 'text_generate',
|
||||
requested_model = $2,
|
||||
resolved_model = $2,
|
||||
request_id = $3,
|
||||
billing_status = 'settled',
|
||||
result = $4::jsonb,
|
||||
finished_at = now(),
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid`,
|
||||
taskB.ID,
|
||||
modelB,
|
||||
"request-admin-task-"+suffix,
|
||||
`{"authorization":"Bearer result-secret","output":"ok"}`,
|
||||
); err != nil {
|
||||
t.Fatalf("enrich task: %v", err)
|
||||
}
|
||||
if _, err := db.Pool().Exec(ctx, `
|
||||
UPDATE gateway_task_attempts
|
||||
SET request_snapshot = $2::jsonb,
|
||||
response_snapshot = $3::jsonb,
|
||||
finished_at = now()
|
||||
WHERE id = $1::uuid`,
|
||||
attemptID,
|
||||
`{"headers":{"authorization":"Bearer attempt-secret"}}`,
|
||||
`{"cookie":"attempt-cookie","status":"ok"}`,
|
||||
); err != nil {
|
||||
t.Fatalf("enrich task attempt: %v", err)
|
||||
}
|
||||
if _, err := db.Pool().Exec(ctx, `
|
||||
INSERT INTO gateway_task_param_preprocessing_logs (
|
||||
task_id, attempt_id, attempt_no, model_type, platform_id, client_id,
|
||||
changed, change_count, actual_input, converted_output, changes, model_snapshot
|
||||
)
|
||||
VALUES (
|
||||
$1::uuid, $2::uuid, 1, 'text_generate', $3::uuid, $4,
|
||||
true, 1, $5::jsonb, $6::jsonb, $7::jsonb, '{}'::jsonb
|
||||
)`,
|
||||
taskB.ID,
|
||||
attemptID,
|
||||
platform.ID,
|
||||
"admin-task-client-"+suffix,
|
||||
`{"password":"input-secret","prompt":"hello"}`,
|
||||
`{"token":"converted-secret","prompt":"hello"}`,
|
||||
`[{"path":"credentials.secret","before":"old-secret","after":"new-secret"}]`,
|
||||
); err != nil {
|
||||
t.Fatalf("create parameter preprocessing log: %v", err)
|
||||
}
|
||||
|
||||
const jwtSecret = "admin-task-integration-jwt-secret"
|
||||
handlerCtx, cancelHandler := context.WithCancel(ctx)
|
||||
defer cancelHandler()
|
||||
server := httptest.NewServer(NewServerWithContext(handlerCtx, config.Config{
|
||||
AppEnv: "test",
|
||||
HTTPAddr: ":0",
|
||||
DatabaseURL: databaseURL,
|
||||
IdentityMode: "hybrid",
|
||||
JWTSecret: jwtSecret,
|
||||
CORSAllowedOrigin: "*",
|
||||
}, db, slog.New(slog.NewTextHandler(io.Discard, nil))))
|
||||
defer server.Close()
|
||||
|
||||
authenticator := auth.New(jwtSecret, "", "")
|
||||
signToken := func(user *auth.User) string {
|
||||
t.Helper()
|
||||
token, err := authenticator.SignJWT(user, time.Hour)
|
||||
if err != nil {
|
||||
t.Fatalf("sign JWT for %s: %v", user.Username, err)
|
||||
}
|
||||
return token
|
||||
}
|
||||
operatorToken := signToken(&auth.User{
|
||||
ID: operator.ID,
|
||||
Username: operator.Username,
|
||||
Roles: operator.Roles,
|
||||
Source: operator.Source,
|
||||
GatewayUserID: operator.ID,
|
||||
GatewayTenantID: operator.GatewayTenantID,
|
||||
TenantKey: operator.TenantKey,
|
||||
})
|
||||
userAToken := signToken(userAAuth)
|
||||
|
||||
doJSON(t, server.URL, http.MethodGet, "/api/admin/tasks", userAToken, nil, http.StatusForbidden, nil)
|
||||
|
||||
query := url.Values{
|
||||
"q": {taskB.ID},
|
||||
"tenantId": {tenantB.ID},
|
||||
"userId": {userB.ID},
|
||||
"status": {"succeeded"},
|
||||
"platformId": {platform.ID},
|
||||
"model": {modelB},
|
||||
"modelType": {"text_generate"},
|
||||
"runMode": {"simulation"},
|
||||
"billingStatus": {"settled"},
|
||||
"apiKey": {"bravo-key"},
|
||||
"page": {"1"},
|
||||
"pageSize": {"10"},
|
||||
}
|
||||
var listResponse struct {
|
||||
Items []store.AdminGatewayTask `json:"items"`
|
||||
Total int `json:"total"`
|
||||
Page int `json:"page"`
|
||||
PageSize int `json:"pageSize"`
|
||||
}
|
||||
doJSON(t, server.URL, http.MethodGet, "/api/admin/tasks?"+query.Encode(), operatorToken, nil, http.StatusOK, &listResponse)
|
||||
if len(listResponse.Items) != 1 || listResponse.Total != 1 || listResponse.Page != 1 || listResponse.PageSize != 10 {
|
||||
t.Fatalf("unexpected filtered admin task list: %+v", listResponse)
|
||||
}
|
||||
listed := listResponse.Items[0]
|
||||
if listed.ID != taskB.ID || listed.AdminContext.User == nil || listed.AdminContext.User.ID != userB.ID ||
|
||||
listed.AdminContext.Tenant == nil || listed.AdminContext.Tenant.ID != tenantB.ID ||
|
||||
listed.AdminContext.LatestPlatform == nil || listed.AdminContext.LatestPlatform.ID != platform.ID {
|
||||
t.Fatalf("admin task list should expose cross-tenant identity and platform summaries: %+v", listed)
|
||||
}
|
||||
if nestedString(listed.Request, "diagnostic", "apiKey") != "***" ||
|
||||
nestedString(listed.Result, "authorization") != "***" ||
|
||||
nestedString(listed.Attempts[0].RequestSnapshot, "headers", "authorization") != "***" {
|
||||
t.Fatalf("admin task list should mask task and attempt secrets: %+v", listed)
|
||||
}
|
||||
|
||||
var maskedDetail store.AdminGatewayTask
|
||||
doJSON(t, server.URL, http.MethodGet, "/api/admin/tasks/"+taskB.ID, operatorToken, nil, http.StatusOK, &maskedDetail)
|
||||
if nestedString(maskedDetail.Request, "diagnostic", "apiKey") != "***" {
|
||||
t.Fatalf("admin task detail should be masked by default: %+v", maskedDetail.Request)
|
||||
}
|
||||
var fullDetail store.AdminGatewayTask
|
||||
doJSON(t, server.URL, http.MethodGet, "/api/admin/tasks/"+taskB.ID+"?sensitive=full", operatorToken, nil, http.StatusOK, &fullDetail)
|
||||
if nestedString(fullDetail.Request, "diagnostic", "apiKey") != taskBSecret {
|
||||
t.Fatalf("explicit full task detail should expose stored content: %+v", fullDetail.Request)
|
||||
}
|
||||
|
||||
var maskedLogs struct {
|
||||
Items []store.TaskParamPreprocessingLog `json:"items"`
|
||||
}
|
||||
doJSON(t, server.URL, http.MethodGet, "/api/admin/tasks/"+taskB.ID+"/param-preprocessing", operatorToken, nil, http.StatusOK, &maskedLogs)
|
||||
if len(maskedLogs.Items) != 1 ||
|
||||
nestedString(maskedLogs.Items[0].ActualInput, "password") != "***" ||
|
||||
nestedString(maskedLogs.Items[0].ConvertedOutput, "token") != "***" {
|
||||
t.Fatalf("admin task preprocessing logs should be masked by default: %+v", maskedLogs.Items)
|
||||
}
|
||||
var fullLogs struct {
|
||||
Items []store.TaskParamPreprocessingLog `json:"items"`
|
||||
}
|
||||
doJSON(t, server.URL, http.MethodGet, "/api/admin/tasks/"+taskB.ID+"/param-preprocessing?sensitive=full", operatorToken, nil, http.StatusOK, &fullLogs)
|
||||
if len(fullLogs.Items) != 1 ||
|
||||
nestedString(fullLogs.Items[0].ActualInput, "password") != "input-secret" ||
|
||||
nestedString(fullLogs.Items[0].ConvertedOutput, "token") != "converted-secret" {
|
||||
t.Fatalf("explicit full preprocessing logs should expose stored content: %+v", fullLogs.Items)
|
||||
}
|
||||
|
||||
var ownTasks struct {
|
||||
Items []store.GatewayTask `json:"items"`
|
||||
}
|
||||
doJSON(t, server.URL, http.MethodGet, "/api/workspace/tasks?pageSize=100", userAToken, nil, http.StatusOK, &ownTasks)
|
||||
if !gatewayTaskListHasID(ownTasks.Items, taskA.ID) || gatewayTaskListHasID(ownTasks.Items, taskB.ID) {
|
||||
t.Fatalf("workspace task list must remain scoped to the current user: %+v", ownTasks.Items)
|
||||
}
|
||||
doJSON(t, server.URL, http.MethodGet, "/api/workspace/tasks/"+taskB.ID, userAToken, nil, http.StatusNotFound, nil)
|
||||
doJSON(t, server.URL, http.MethodGet, "/api/workspace/tasks/"+taskB.ID+"/param-preprocessing", userAToken, nil, http.StatusNotFound, nil)
|
||||
}
|
||||
|
||||
func nestedString(value map[string]any, path ...string) string {
|
||||
var current any = value
|
||||
for _, key := range path {
|
||||
object, ok := current.(map[string]any)
|
||||
if !ok {
|
||||
return ""
|
||||
}
|
||||
current = object[key]
|
||||
}
|
||||
result, _ := current.(string)
|
||||
return result
|
||||
}
|
||||
|
||||
func gatewayTaskListHasID(items []store.GatewayTask, taskID string) bool {
|
||||
for _, item := range items {
|
||||
if item.ID == taskID {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
package httpapi
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
||||
)
|
||||
|
||||
func TestAdvancedMediaScopeAliases(t *testing.T) {
|
||||
tests := []struct {
|
||||
kind string
|
||||
scope string
|
||||
}{
|
||||
{kind: "images.vectorize", scope: "image_vectorize"},
|
||||
{kind: "images.vectorize", scope: "image"},
|
||||
{kind: "images.vectorize", scope: "vectorize"},
|
||||
{kind: "videos.upscales", scope: "video_enhance"},
|
||||
{kind: "videos.upscales", scope: "video"},
|
||||
{kind: "videos.upscales", scope: "video_upscale"},
|
||||
}
|
||||
for _, test := range tests {
|
||||
user := &auth.User{APIKeyID: "key", APIKeyScopes: []string{test.scope}}
|
||||
if !apiKeyScopeAllowed(user, test.kind) {
|
||||
t.Fatalf("scope %q should allow %q", test.scope, test.kind)
|
||||
}
|
||||
}
|
||||
if apiKeyScopeAllowed(&auth.User{APIKeyID: "key", APIKeyScopes: []string{"chat"}}, "videos.upscales") {
|
||||
t.Fatal("chat scope must not allow video upscale")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFillDailyTokenUsageDaysKeepsGapsAndStreaks(t *testing.T) {
|
||||
location, err := time.LoadLocation("Asia/Shanghai")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
from := time.Date(2026, 7, 20, 0, 0, 0, 0, location)
|
||||
to := time.Date(2026, 7, 24, 0, 0, 0, 0, location)
|
||||
items, summary := fillDailyTokenUsageDays(map[string]store.DailyTokenUsage{
|
||||
"2026-07-20": {Date: "2026-07-20", TotalTokens: 10, TaskCount: 1},
|
||||
"2026-07-22": {Date: "2026-07-22", TotalTokens: 20, TaskCount: 1},
|
||||
"2026-07-23": {Date: "2026-07-23", TotalTokens: 30, TaskCount: 2},
|
||||
"2026-07-24": {Date: "2026-07-24", TotalTokens: 40, TaskCount: 1},
|
||||
}, from, to)
|
||||
if len(items) != 5 || items[1].Date != "2026-07-21" || items[1].TaskCount != 0 {
|
||||
t.Fatalf("daily usage should contain continuous zero days: %+v", items)
|
||||
}
|
||||
if summary.CumulativeTokens != 100 || summary.PeakDailyTokens != 40 || summary.CurrentStreakDays != 3 || summary.LongestStreakDays != 3 {
|
||||
t.Fatalf("unexpected daily usage summary: %+v", summary)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdvancedMediaDefaultModels(t *testing.T) {
|
||||
if got := canonicalTaskModelName("images.vectorize", ""); got != "easy-image-vectorizer-1" {
|
||||
t.Fatalf("vectorizer default model=%q", got)
|
||||
}
|
||||
if got := canonicalTaskModelName("videos.upscales", ""); got != "easy-proteus-standard-4" {
|
||||
t.Fatalf("Topaz default model=%q", got)
|
||||
}
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user