mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-05 00:47:48 +08:00
Compare commits
304
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
4e96b9d870 | ||
|
|
d3c0b1aa7f | ||
|
|
66bfd3e592 | ||
|
|
392c557831 | ||
|
|
c34565b02b | ||
|
|
f57fe6e13e | ||
|
|
7c678b715f | ||
|
|
bd4e5f3a5d | ||
|
|
165d9eab8f | ||
|
|
dfb95f09e1 | ||
|
|
4d5c591654 | ||
|
|
46737d32f8 | ||
|
|
25d38ae632 | ||
|
|
aa83b4a7a7 | ||
|
|
913ce2dbcb | ||
|
|
cae5e520ac | ||
|
|
772f2ea601 | ||
|
|
28fa03451c | ||
|
|
6984984c22 | ||
|
|
1209c835c7 | ||
|
|
e4ebd5cca1 | ||
|
|
f573110725 | ||
|
|
a8620e133a | ||
|
|
ddd6adbcf7 | ||
|
|
b570aaac48 | ||
|
|
086efe6efe | ||
|
|
56f3c95763 | ||
|
|
8a6a961900 | ||
|
|
6e55968487 | ||
|
|
8d8cddcef6 | ||
|
|
b90d5095f1 | ||
|
|
a4505b1281 | ||
|
|
1d72a8f9c1 | ||
|
|
3d5b6141a5 | ||
|
|
203cd5a9d5 | ||
|
|
696ec65175 | ||
|
|
cbb66a5667 | ||
|
|
53ef35ec80 | ||
|
|
1af3067303 | ||
|
|
d026398bab | ||
|
|
684689a82b | ||
|
|
eeb5f41bad | ||
|
|
7180eaea88 | ||
|
|
7cb204f18a | ||
|
|
0342f609d0 | ||
|
|
59840fa419 | ||
|
|
d390d46ee8 | ||
|
|
37eada9682 | ||
|
|
73a5325a38 | ||
|
|
460eb5434d | ||
|
|
8a21cb9a55 | ||
|
|
d0df52ce35 | ||
|
|
c1ed42fd3a | ||
|
|
c4bb6b8161 | ||
|
|
d480aa11f3 | ||
|
|
d63d5eff85 | ||
|
|
5dae2a4792 | ||
|
|
dcbd7dc219 | ||
|
|
2dcf8b8414 | ||
|
|
438f16094f | ||
|
|
40e0b82fa0 | ||
|
|
b9b0a75fe4 | ||
|
|
b6cc0bc3a7 | ||
|
|
d2f1431269 | ||
|
|
4ecaefbade | ||
|
|
c97c9332eb | ||
|
|
c070e5a9f6 | ||
|
|
8cd3a69803 | ||
|
|
791c9c98dc | ||
|
|
025e979935 | ||
|
|
131471a13f | ||
|
|
7ff63077c3 | ||
|
|
faba0cbd07 | ||
|
|
75f17935f9 | ||
|
|
60842fbbb5 | ||
|
|
d58c27d22d | ||
|
|
65550159bb | ||
|
|
3c7ad81d62 | ||
|
|
6b23c9b3ce | ||
|
|
900e54d740 | ||
|
|
8cc70934da | ||
|
|
f92b0943b5 | ||
|
|
920a383136 | ||
|
|
693e37d2df | ||
|
|
9c6036a103 | ||
|
|
751a4d9111 | ||
|
|
aafd332198 | ||
|
|
337cd0c505 | ||
|
|
b15ce9977a | ||
|
|
e0286aebe3 | ||
|
|
45985f1c04 | ||
|
|
776dd2f8ea | ||
|
|
bdfe4adc98 | ||
|
|
00a0371997 | ||
|
|
ded6b5b081 | ||
|
|
e9678ea899 | ||
|
|
8d69f72e2a | ||
|
|
127b4e11de | ||
|
|
22093bed4d | ||
|
|
ebd53ad679 | ||
|
|
280c604327 | ||
|
|
ad31cdbf85 | ||
|
|
0112ab752b | ||
|
|
66fec80e79 | ||
|
|
fddfaecf5e | ||
|
|
71ae1a2307 | ||
|
|
4d338ebd3d | ||
|
|
dc440f1507 | ||
|
|
5c732f844a | ||
|
|
c4044ba0b1 | ||
|
|
a8159b7bda | ||
|
|
0ab20be667 | ||
|
|
3f048d373f | ||
|
|
8b49a3d264 | ||
|
|
6e51a3f45d | ||
|
|
0d770d1c4d | ||
|
|
b45f021bba | ||
|
|
1a1bce3e8c | ||
|
|
6aeb5d40ab | ||
|
|
31ef2d134e | ||
|
|
85b50e67e1 | ||
|
|
9353f89af0 | ||
|
|
380d69e096 | ||
|
|
02e2f4f500 | ||
|
|
0dbfefa834 | ||
|
|
04ab4bd9f2 | ||
|
|
4955166b85 | ||
|
|
6235c772ac | ||
|
|
e2ec3f7942 | ||
|
|
a816235efb | ||
|
|
1e39ab3c2e | ||
|
|
85aa66c76d | ||
|
|
9a4817faf8 | ||
|
|
7b0c80a0c4 | ||
|
|
6ec8df97e8 | ||
|
|
f82964217e | ||
|
|
cfa5535f6e | ||
|
|
2d846b2c58 | ||
|
|
f40e8037dd | ||
|
|
2b21a75982 | ||
|
|
9ee27308db | ||
|
|
1518223de6 | ||
|
|
3063938a82 | ||
|
|
86449cae52 | ||
|
|
7c580e843f | ||
|
|
c9f0685b40 | ||
|
|
afd0dcf2ff | ||
|
|
68d4df71d8 | ||
|
|
596227659a | ||
|
|
3b0dbadb1e | ||
|
|
8a8bc999d2 | ||
|
|
57c7cca556 | ||
|
|
9e6578a71b | ||
|
|
bf818e3b61 | ||
|
|
9e7f291aaf | ||
|
|
46bab1b97f | ||
|
|
0258d01ee6 | ||
|
|
4999a1a0a8 | ||
|
|
8cb8666456 | ||
|
|
48f3f481db | ||
|
|
e60462e068 | ||
|
|
c4877e3b6a | ||
|
|
afbb1b9a5d | ||
|
|
9516619b92 | ||
|
|
e106a65c1d | ||
|
|
d84c9d4b71 | ||
|
|
0046123e22 | ||
|
|
40736a8334 | ||
|
|
0a256adc94 | ||
|
|
0bddc7965b | ||
|
|
258be3b640 | ||
|
|
4dbfeb87b8 | ||
|
|
fa69287449 | ||
|
|
654ce89541 | ||
|
|
fd32597015 | ||
|
|
84cf07b7a2 | ||
|
|
e4476d0bc6 | ||
|
|
0379f01ce8 | ||
|
|
95e72594ea | ||
|
|
f9ffb1cae5 | ||
|
|
91b6e0a382 | ||
|
|
f5f7a23bb0 | ||
|
|
8be9601963 | ||
|
|
eaad1579e6 | ||
|
|
1d8bf56efd | ||
|
|
73db997e92 | ||
|
|
2c5654d694 | ||
|
|
f490c3a5cd | ||
|
|
2d9158b321 | ||
|
|
48d13762d9 | ||
|
|
bd3f73c2fc | ||
|
|
25c33846be | ||
|
|
124c4ca403 | ||
|
|
ef0f8dd4d0 | ||
|
|
d0eca509d4 | ||
|
|
90663793a2 | ||
|
|
783f654953 | ||
|
|
9cdcce1b5f | ||
|
|
06b483f79d | ||
|
|
a00e137ffc | ||
|
|
239238fe47 | ||
|
|
4cd6e0d10f | ||
|
|
9b29a65c68 | ||
|
|
df4a49a9fb | ||
|
|
bb268310e2 | ||
|
|
7b0908cd87 | ||
|
|
95bc742057 | ||
|
|
e40c890a8e | ||
|
|
357c4fd61f | ||
|
|
d28fea80df | ||
|
|
fb1aeb789a | ||
|
|
1f3693d3a2 | ||
|
|
90760da499 | ||
|
|
7950ba7dc5 | ||
|
|
a7088ee538 | ||
|
|
050cba9563 | ||
|
|
2269617a9f | ||
|
|
bdccfa6e78 | ||
|
|
d97ec3fde2 | ||
|
|
d17472f09e | ||
|
|
f8b7cd2925 | ||
|
|
e9c3ac94c6 | ||
|
|
fa71cddb60 | ||
|
|
228cbc8f87 | ||
|
|
da915208a8 | ||
|
|
694167f78f | ||
|
|
a7697032a4 | ||
|
|
e86c8edd4b | ||
|
|
b1be413dc0 | ||
|
|
fdb50a065b | ||
|
|
1ac59d4894 | ||
|
|
32ccf61baa | ||
|
|
1d04c41ae7 | ||
|
|
b2dcf82ca8 | ||
|
|
57b86034cf | ||
|
|
095e312ab3 | ||
|
|
82a9fb3c39 | ||
|
|
e181329a81 | ||
|
|
cdce817928 | ||
|
|
97b0146ce9 | ||
|
|
0a60492146 | ||
|
|
4ea187cfac | ||
|
|
dcba7c62a2 | ||
|
|
dba99455a7 | ||
|
|
7ebce161e8 | ||
|
|
022aec5720 | ||
|
|
11997c024e | ||
|
|
f787b1b02a | ||
|
|
a03368a3fe | ||
|
|
e26ed8481f | ||
|
|
8e98eed5c8 | ||
|
|
0bff15f964 | ||
|
|
3384c6d666 | ||
|
|
5f1c74aca0 | ||
|
|
310355bcc1 | ||
|
|
ab07d83aaf | ||
|
|
11cdf52f7b | ||
|
|
8a2c3596ce | ||
|
|
a61d16d120 | ||
|
|
01df063cc1 | ||
|
|
68bae686da | ||
|
|
f978888759 | ||
|
|
f3b9f42202 | ||
|
|
0564893c4f | ||
|
|
c9dbe7936b | ||
|
|
3f7a3d7600 | ||
|
|
5b9d9452c9 | ||
|
|
b7ef900181 | ||
|
|
3eef673885 | ||
|
|
039a18c243 | ||
|
|
97d42703da | ||
|
|
4bf3a453e7 | ||
|
|
a137601728 | ||
|
|
d4840df447 | ||
|
|
2c1e51a490 | ||
|
|
2dcccc9820 | ||
|
|
8fa97ab8da | ||
|
|
a66fa9792d | ||
|
|
dc2bc83c17 | ||
|
|
8b0e92e408 | ||
|
|
b61fc4eb6b | ||
|
|
fea3d183bf | ||
|
|
20ed9cd123 | ||
|
|
d7a8a89aa7 | ||
|
|
005cc3e388 | ||
|
|
fbcb54a8a5 | ||
|
|
85a126f48a | ||
|
|
44cd35c10e | ||
|
|
a2d1cff3b0 | ||
|
|
3ff67fec2f | ||
|
|
6a8b5e6c8e | ||
|
|
92e9caf57e | ||
|
|
10ccbf109c | ||
|
|
f2dc51434d | ||
|
|
9d751e3525 | ||
|
|
27047e5880 | ||
|
|
4fdebfc78e | ||
|
|
11832edf49 | ||
|
|
87d50052cc | ||
|
|
08b89b7ef8 | ||
|
|
4b02078b60 | ||
|
|
1d644de500 | ||
|
|
daecdb7676 | ||
|
|
423eb95f7a |
@@ -11,6 +11,7 @@ ENV/
|
||||
.uv/
|
||||
*.egg-info/
|
||||
dist/
|
||||
!aether-hub/dist/aether-hub
|
||||
build/
|
||||
*.egg
|
||||
|
||||
|
||||
+53
-3
@@ -21,6 +21,10 @@ JWT_SECRET_KEY=change-this-to-a-secure-random-string
|
||||
# 注意:更换此密钥后需要在管理面板重新配置所有 Provider API Key
|
||||
ENCRYPTION_KEY=change-this-to-another-secure-random-string
|
||||
|
||||
# 支付回调共享密钥(公开 /api/payment/callback/* 入口必须携带 x-payment-callback-token)
|
||||
# 建议使用 32+ 位随机字符串
|
||||
PAYMENT_CALLBACK_SECRET=change-this-to-a-secure-callback-secret
|
||||
|
||||
# 管理员账号(仅首次初始化时使用, 创建完成后可在系统内修改密码)
|
||||
ADMIN_EMAIL=[email protected]
|
||||
ADMIN_USERNAME=admin
|
||||
@@ -32,15 +36,53 @@ ADMIN_PASSWORD=admin123456
|
||||
# 应用端口(默认 8084)
|
||||
# APP_PORT=8084
|
||||
|
||||
# Gunicorn Worker 数量(默认 4)
|
||||
# 建议最小设置为 2
|
||||
# GUNICORN_WORKERS=4
|
||||
# 生产部署镜像(deploy.sh 会读取)
|
||||
# APP_IMAGE=ghcr.io/fawney19/aether:latest
|
||||
|
||||
# Gunicorn Worker 数量(默认 2)
|
||||
# Tunnel 请求统一经 Hub 转发,可安全使用多 worker。
|
||||
# 非 Docker 运行时若使用 ProxyNode tunnel,请确保 aether-hub 可达(默认 ws://127.0.0.1:8085)。
|
||||
# GUNICORN_WORKERS=2
|
||||
|
||||
# Gunicorn Max Requests(默认 4000)
|
||||
# Worker 处理指定数量请求后自动重启,防止内存泄漏
|
||||
# max-requests-jitter 会自动设置为 MAX_REQUESTS/20 (5%)
|
||||
# MAX_REQUESTS=4000
|
||||
|
||||
# glibc malloc arena 上限(默认 2)
|
||||
# 降低 malloc 内存碎片,减少 gunicorn worker RSS
|
||||
# MALLOC_ARENA_MAX=2
|
||||
|
||||
# HTTP 连接池上限(默认总预算约 200,按 worker 平分)
|
||||
# 如果容器内存偏高,可继续下调;例如 2 worker 时设为 80-100
|
||||
# HTTP_MAX_CONNECTIONS=100
|
||||
|
||||
# HTTP 保活连接数(默认约为 max_connections 的 30%)
|
||||
# HTTP_KEEPALIVE_CONNECTIONS=30
|
||||
|
||||
# HTTP 代理/Tunnel 客户端空闲清理(默认每 5 分钟扫描,空闲 600 秒即关闭)
|
||||
# HTTP_CLIENT_IDLE_CLEANUP_INTERVAL_MINUTES=5
|
||||
# HTTP_CLIENT_IDLE_CLEANUP_MAX_SECONDS=600
|
||||
|
||||
# curl_cffi session 池上限(默认 20,按 impersonate + proxy 组合缓存)
|
||||
# CURL_CFFI_MAX_SESSIONS=20
|
||||
|
||||
# 流式响应块缓存上限(单位 MB,默认 2)
|
||||
# 说明:
|
||||
# - 这是单个流式请求可保留的“解析后响应块”内存上限,不是全局上限
|
||||
# - 粗略峰值内存 ≈ 并发流数量 × RESPONSE_CHUNKS_MAX_SIZE_MB
|
||||
# 例如:100 并发、2MB 上限,理论峰值约 200MB
|
||||
# - 建议:
|
||||
# - 内存敏感环境:1
|
||||
# - 通用生产环境:2(默认)
|
||||
# - 需要更多调试上下文:4
|
||||
# RESPONSE_CHUNKS_MAX_SIZE_MB=2
|
||||
|
||||
# 流式空闲超时(单位秒,默认 30)
|
||||
# 当流已经开始但连续一段时间没有任何新 chunk 时,提前中断并返回 504,
|
||||
# 避免一直等到 worker 超时(如 300s)
|
||||
# STREAM_IDLE_TIMEOUT_SECONDS=30
|
||||
|
||||
# API Key 前缀(默认 sk)
|
||||
# API_KEY_PREFIX=sk
|
||||
|
||||
@@ -52,6 +94,14 @@ ADMIN_PASSWORD=admin123456
|
||||
# 默认: * (允许所有源)
|
||||
# CORS_ORIGINS=*
|
||||
|
||||
# 启动预热配置(默认启用,降低首请求冷启动延迟)
|
||||
# 是否启用启动期预热任务(默认 true)
|
||||
# STARTUP_WARMUP_ENABLED=true
|
||||
# /readyz 是否等待预热完成(默认 true)
|
||||
# STARTUP_WARMUP_GATE_READINESS=true
|
||||
# 预热时优先 bootstrap 的 provider_type 列表(逗号分隔;留空表示自动探测)
|
||||
# STARTUP_WARMUP_PROVIDER_TYPES=codex,kiro
|
||||
|
||||
# ==================== 计费系统(可选) ====================
|
||||
# Video/Image/Audio 缺失 billing_rule 时是否拒绝请求(默认 false:允许请求但 cost=0 并告警)
|
||||
# BILLING_REQUIRE_RULE=false
|
||||
|
||||
@@ -0,0 +1,91 @@
|
||||
name: Build aether-hub
|
||||
|
||||
on:
|
||||
push:
|
||||
tags: ['hub-v*']
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
contents: write
|
||||
|
||||
jobs:
|
||||
build:
|
||||
name: ${{ matrix.name }}
|
||||
runs-on: ubuntu-latest
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
include:
|
||||
- name: linux-amd64
|
||||
target: x86_64-unknown-linux-gnu
|
||||
use_cross: true
|
||||
- name: linux-arm64
|
||||
target: aarch64-unknown-linux-gnu
|
||||
use_cross: true
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v5
|
||||
|
||||
- name: Install Rust toolchain
|
||||
uses: dtolnay/rust-toolchain@stable
|
||||
with:
|
||||
targets: ${{ matrix.target }}
|
||||
|
||||
- name: Rust cache
|
||||
uses: Swatinem/rust-cache@v2
|
||||
with:
|
||||
workspaces: aether-hub -> target
|
||||
key: ${{ matrix.target }}
|
||||
|
||||
- name: Install cross
|
||||
if: matrix.use_cross
|
||||
uses: taiki-e/install-action@cross
|
||||
|
||||
- name: Build
|
||||
working-directory: aether-hub
|
||||
shell: bash
|
||||
run: |
|
||||
if [ "${{ matrix.use_cross }}" = "true" ]; then
|
||||
cross build --release --target ${{ matrix.target }}
|
||||
else
|
||||
cargo build --release --target ${{ matrix.target }}
|
||||
fi
|
||||
|
||||
- name: Package
|
||||
shell: bash
|
||||
run: |
|
||||
cd aether-hub/target/${{ matrix.target }}/release
|
||||
chmod +x aether-hub
|
||||
tar czf ../../../../aether-hub-${{ matrix.name }}.tar.gz aether-hub
|
||||
|
||||
- name: Upload artifact
|
||||
uses: actions/upload-artifact@v5
|
||||
with:
|
||||
name: aether-hub-${{ matrix.name }}
|
||||
path: aether-hub-*.tar.gz
|
||||
if-no-files-found: error
|
||||
|
||||
release:
|
||||
needs: build
|
||||
runs-on: ubuntu-latest
|
||||
if: startsWith(github.ref, 'refs/tags/')
|
||||
steps:
|
||||
- name: Download all artifacts
|
||||
uses: actions/download-artifact@v5
|
||||
with:
|
||||
merge-multiple: true
|
||||
path: artifacts
|
||||
|
||||
- name: Generate checksums
|
||||
working-directory: artifacts
|
||||
run: sha256sum aether-hub-* > SHA256SUMS.txt
|
||||
|
||||
- name: Create GitHub Release
|
||||
uses: softprops/action-gh-release@v2
|
||||
with:
|
||||
name: "${{ github.ref_name }}"
|
||||
generate_release_notes: true
|
||||
files: |
|
||||
artifacts/aether-hub-*
|
||||
artifacts/SHA256SUMS.txt
|
||||
fail_on_unmatched_files: true
|
||||
@@ -44,7 +44,7 @@ jobs:
|
||||
use_cross: false
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v5
|
||||
|
||||
- name: Install Rust toolchain
|
||||
uses: dtolnay/rust-toolchain@stable
|
||||
@@ -87,7 +87,7 @@ jobs:
|
||||
7z a ../../../../aether-proxy-${{ matrix.name }}.zip aether-proxy.exe
|
||||
|
||||
- name: Upload artifact
|
||||
uses: actions/upload-artifact@v4
|
||||
uses: actions/upload-artifact@v5
|
||||
with:
|
||||
name: aether-proxy-${{ matrix.name }}
|
||||
path: |
|
||||
@@ -101,7 +101,7 @@ jobs:
|
||||
if: startsWith(github.ref, 'refs/tags/')
|
||||
steps:
|
||||
- name: Download all artifacts
|
||||
uses: actions/download-artifact@v4
|
||||
uses: actions/download-artifact@v5
|
||||
with:
|
||||
merge-multiple: true
|
||||
path: artifacts
|
||||
@@ -113,7 +113,7 @@ jobs:
|
||||
- name: Create GitHub Release
|
||||
uses: softprops/action-gh-release@v2
|
||||
with:
|
||||
name: "aether-proxy ${{ github.ref_name }}"
|
||||
name: "${{ github.ref_name }}"
|
||||
generate_release_notes: true
|
||||
files: |
|
||||
artifacts/aether-proxy-*
|
||||
@@ -125,10 +125,10 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
if: startsWith(github.ref, 'refs/tags/')
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v5
|
||||
|
||||
- name: Download Linux artifacts
|
||||
uses: actions/download-artifact@v4
|
||||
uses: actions/download-artifact@v5
|
||||
with:
|
||||
pattern: aether-proxy-linux-*
|
||||
merge-multiple: true
|
||||
@@ -174,7 +174,7 @@ jobs:
|
||||
latest=auto
|
||||
|
||||
- name: Build and push
|
||||
uses: docker/build-push-action@v5
|
||||
uses: docker/build-push-action@v6
|
||||
with:
|
||||
context: ./aether-proxy
|
||||
push: true
|
||||
@@ -187,7 +187,7 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
if: startsWith(github.ref, 'refs/tags/')
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v5
|
||||
with:
|
||||
ref: master
|
||||
|
||||
|
||||
@@ -18,12 +18,12 @@ jobs:
|
||||
build:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v5
|
||||
|
||||
- name: Setup Node.js
|
||||
uses: actions/setup-node@v4
|
||||
uses: actions/setup-node@v5
|
||||
with:
|
||||
node-version: '20'
|
||||
node-version: '22'
|
||||
cache: 'npm'
|
||||
cache-dependency-path: frontend/package-lock.json
|
||||
|
||||
@@ -41,7 +41,7 @@ jobs:
|
||||
run: cp frontend/dist/index.html frontend/dist/404.html
|
||||
|
||||
- name: Setup Pages
|
||||
uses: actions/configure-pages@v4
|
||||
uses: actions/configure-pages@v5
|
||||
|
||||
- name: Upload artifact
|
||||
uses: actions/upload-pages-artifact@v3
|
||||
|
||||
@@ -15,6 +15,7 @@ env:
|
||||
REGISTRY: ghcr.io
|
||||
BASE_IMAGE_NAME: fawney19/aether-base
|
||||
APP_IMAGE_NAME: fawney19/aether
|
||||
GITHUB_REPO: fawney19/Aether
|
||||
# Base image hash inputs:
|
||||
# - Dockerfile.base
|
||||
# - pyproject.toml (dependency fingerprint only; ignores tool/optional deps)
|
||||
@@ -29,7 +30,7 @@ jobs:
|
||||
outputs:
|
||||
base_changed: ${{ steps.check.outputs.base_changed }}
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v5
|
||||
|
||||
- name: Log in to Container Registry
|
||||
uses: docker/login-action@v3
|
||||
@@ -107,7 +108,7 @@ jobs:
|
||||
contents: read
|
||||
packages: write
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v5
|
||||
|
||||
- name: Set up Docker Buildx
|
||||
uses: docker/setup-buildx-action@v3
|
||||
@@ -163,7 +164,7 @@ jobs:
|
||||
org.opencontainers.image.base.hash=${{ steps.hash.outputs.hash }}
|
||||
|
||||
- name: Build and push base image
|
||||
uses: docker/build-push-action@v5
|
||||
uses: docker/build-push-action@v6
|
||||
with:
|
||||
context: .
|
||||
file: ./Dockerfile.base
|
||||
@@ -174,15 +175,36 @@ jobs:
|
||||
cache-to: type=gha,mode=max,scope=base
|
||||
platforms: linux/amd64,linux/arm64
|
||||
|
||||
download-hub:
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
contents: read
|
||||
outputs:
|
||||
hub_tag: ${{ steps.hub-tag.outputs.tag }}
|
||||
steps:
|
||||
- name: Get latest hub release tag
|
||||
id: hub-tag
|
||||
env:
|
||||
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
run: |
|
||||
TAG=$(gh release list --repo "${{ env.GITHUB_REPO }}" --limit 50 --json tagName,isDraft,isPrerelease \
|
||||
--jq '[.[] | select(.tagName | startswith("hub-v")) | select(.isDraft == false and .isPrerelease == false)] | .[0].tagName')
|
||||
if [ -z "$TAG" ] || [ "$TAG" = "null" ]; then
|
||||
echo "No hub release found"
|
||||
exit 1
|
||||
fi
|
||||
echo "tag=$TAG" >> $GITHUB_OUTPUT
|
||||
echo "Hub release tag: $TAG"
|
||||
|
||||
build-app:
|
||||
needs: [check-base-changes, build-base]
|
||||
if: always() && (needs.build-base.result == 'success' || needs.build-base.result == 'skipped')
|
||||
needs: [check-base-changes, build-base, download-hub]
|
||||
if: always() && (needs.build-base.result == 'success' || needs.build-base.result == 'skipped') && needs.download-hub.result == 'success'
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
contents: read
|
||||
packages: write
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v5
|
||||
|
||||
- name: Set up Docker Buildx
|
||||
uses: docker/setup-buildx-action@v3
|
||||
@@ -243,15 +265,63 @@ jobs:
|
||||
version_tuple = __version_tuple__
|
||||
EOF
|
||||
|
||||
- name: Build and push app image
|
||||
uses: docker/build-push-action@v5
|
||||
- name: Resolve hub release for build args
|
||||
run: |
|
||||
echo "Hub release tag: ${{ needs.download-hub.outputs.hub_tag }}"
|
||||
|
||||
- name: Build and push app image (amd64)
|
||||
id: build-amd64
|
||||
uses: docker/build-push-action@v6
|
||||
with:
|
||||
context: .
|
||||
file: ./Dockerfile.app
|
||||
push: true
|
||||
tags: ${{ steps.meta.outputs.tags }}
|
||||
labels: ${{ steps.meta.outputs.labels }}
|
||||
no-cache-filters: builder
|
||||
cache-from: type=gha,scope=app
|
||||
cache-to: type=gha,mode=min,scope=app
|
||||
platforms: linux/amd64,linux/arm64
|
||||
cache-from: type=gha,scope=app-amd64
|
||||
cache-to: type=gha,mode=min,scope=app-amd64
|
||||
build-args: |
|
||||
HUB_RELEASE_REPO=${{ env.GITHUB_REPO }}
|
||||
HUB_TAG=${{ needs.download-hub.outputs.hub_tag }}
|
||||
platforms: linux/amd64
|
||||
outputs: type=image,"name=${{ env.REGISTRY }}/${{ env.APP_IMAGE_NAME }},docker.io/fawney19/aether",push-by-digest=true,name-canonical=true,push=true
|
||||
|
||||
- name: Build and push app image (arm64)
|
||||
id: build-arm64
|
||||
uses: docker/build-push-action@v6
|
||||
with:
|
||||
context: .
|
||||
file: ./Dockerfile.app
|
||||
labels: ${{ steps.meta.outputs.labels }}
|
||||
no-cache-filters: builder
|
||||
cache-from: type=gha,scope=app-arm64
|
||||
cache-to: type=gha,mode=min,scope=app-arm64
|
||||
build-args: |
|
||||
HUB_RELEASE_REPO=${{ env.GITHUB_REPO }}
|
||||
HUB_TAG=${{ needs.download-hub.outputs.hub_tag }}
|
||||
platforms: linux/arm64
|
||||
outputs: type=image,"name=${{ env.REGISTRY }}/${{ env.APP_IMAGE_NAME }},docker.io/fawney19/aether",push-by-digest=true,name-canonical=true,push=true
|
||||
|
||||
- name: Create multi-arch manifest and push
|
||||
run: |
|
||||
# Extract digests
|
||||
AMD64_DIGEST="${{ steps.build-amd64.outputs.digest }}"
|
||||
ARM64_DIGEST="${{ steps.build-arm64.outputs.digest }}"
|
||||
echo "amd64 digest: $AMD64_DIGEST"
|
||||
echo "arm64 digest: $ARM64_DIGEST"
|
||||
|
||||
# For each tag, create multi-arch manifest on each registry
|
||||
TAGS=$(echo "${{ steps.meta.outputs.tags }}" | tr '\n' ' ')
|
||||
for FULL_TAG in $TAGS; do
|
||||
# Determine which registry this tag belongs to
|
||||
if [[ "$FULL_TAG" == ghcr.io/* ]]; then
|
||||
REPO="${{ env.REGISTRY }}/${{ env.APP_IMAGE_NAME }}"
|
||||
elif [[ "$FULL_TAG" == docker.io/* ]]; then
|
||||
REPO="docker.io/fawney19/aether"
|
||||
else
|
||||
continue
|
||||
fi
|
||||
echo "Creating manifest for $FULL_TAG"
|
||||
docker buildx imagetools create -t "$FULL_TAG" \
|
||||
"$REPO@$AMD64_DIGEST" \
|
||||
"$REPO@$ARM64_DIGEST"
|
||||
done
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
# Edit at https://www.toptal.com/developers/gitignore?templates=python
|
||||
|
||||
# AI Assistant Configuration
|
||||
.codex/
|
||||
.claude/
|
||||
.serena/
|
||||
.gemini*/
|
||||
@@ -203,6 +204,7 @@ logs/
|
||||
|
||||
# Git backup
|
||||
.git.backup/
|
||||
.worktrees/
|
||||
|
||||
# Database backups
|
||||
backups/
|
||||
@@ -226,6 +228,10 @@ test.py
|
||||
.deps-hash
|
||||
.code-hash
|
||||
.migration-hash
|
||||
.hub-hash
|
||||
|
||||
# Hub prebuilt binaries
|
||||
aether-hub/dist/
|
||||
|
||||
# Version file (auto-generated by hatch-vcs)
|
||||
src/_version.py
|
||||
|
||||
+124
-6
@@ -2,14 +2,22 @@
|
||||
# 运行镜像:从 base 提取产物到精简运行时
|
||||
# 构建命令: docker build -f Dockerfile.app -t aether-app:latest .
|
||||
# 用于 GitHub Actions CI(官方源)
|
||||
|
||||
FROM aether-base:latest AS builder
|
||||
WORKDIR /app
|
||||
# 复制前端源码并构建(CI 通过 no-cache-filters=builder 确保每次重建)
|
||||
COPY frontend/ ./frontend/
|
||||
RUN cd frontend && npm run build
|
||||
|
||||
# ==================== 运行时镜像 ====================
|
||||
FROM python:3.13-slim
|
||||
WORKDIR /app
|
||||
|
||||
ARG HUB_RELEASE_REPO=fawney19/Aether
|
||||
ARG HUB_TAG
|
||||
ARG TARGETARCH
|
||||
ARG GITHUB_TOKEN
|
||||
|
||||
# 运行时依赖(无 gcc/nodejs/npm,使用 BuildKit 缓存加速)
|
||||
RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
|
||||
--mount=type=cache,target=/var/lib/apt,sharing=locked \
|
||||
@@ -17,13 +25,49 @@ RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
|
||||
nginx \
|
||||
supervisor \
|
||||
libpq5 \
|
||||
curl
|
||||
curl \
|
||||
libjemalloc2
|
||||
RUN set -eux; \
|
||||
jemalloc_path="$(find /usr/lib -type f -name 'libjemalloc.so.2' | head -n1)"; \
|
||||
[ -n "$jemalloc_path" ]; \
|
||||
ln -sf "$jemalloc_path" /usr/local/lib/libjemalloc.so.2
|
||||
# 从 base 镜像复制 Python 包
|
||||
COPY --from=builder /usr/local/lib/python3.13/site-packages /usr/local/lib/python3.13/site-packages
|
||||
# 只复制需要的 Python 可执行文件
|
||||
COPY --from=builder /usr/local/bin/gunicorn /usr/local/bin/
|
||||
COPY --from=builder /usr/local/bin/uvicorn /usr/local/bin/
|
||||
COPY --from=builder /usr/local/bin/alembic /usr/local/bin/
|
||||
# Hub 预编译二进制(构建时从 GitHub Release 下载)
|
||||
# GITHUB_TOKEN 可选:未认证 API 限流 60 次/小时,认证后 5000 次/小时
|
||||
RUN set -eux; \
|
||||
auth_header=""; \
|
||||
if [ -n "${GITHUB_TOKEN:-}" ]; then \
|
||||
auth_header="Authorization: token ${GITHUB_TOKEN}"; \
|
||||
fi; \
|
||||
tag="${HUB_TAG:-}"; \
|
||||
if [ -z "$tag" ]; then \
|
||||
tag="$(curl -sL ${auth_header:+-H "$auth_header"} "https://api.github.com/repos/${HUB_RELEASE_REPO}/releases" | python3 -c "import json,sys;print(next((r['tag_name'] for r in json.load(sys.stdin) if r.get('tag_name','').startswith('hub-v') and not r.get('draft') and not r.get('prerelease')),''))")"; \
|
||||
fi; \
|
||||
if [ -z "$tag" ]; then \
|
||||
echo "Failed to resolve hub release tag"; \
|
||||
exit 1; \
|
||||
fi; \
|
||||
arch="${TARGETARCH:-}"; \
|
||||
if [ -z "$arch" ]; then \
|
||||
arch="$(dpkg --print-architecture)"; \
|
||||
fi; \
|
||||
case "$arch" in \
|
||||
amd64|arm64) ;; \
|
||||
x86_64) arch="amd64" ;; \
|
||||
aarch64) arch="arm64" ;; \
|
||||
*) echo "Unsupported architecture: $arch"; exit 1 ;; \
|
||||
esac; \
|
||||
echo "Using Hub release tag: $tag"; \
|
||||
url="https://github.com/${HUB_RELEASE_REPO}/releases/download/${tag}/aether-hub-linux-${arch}.tar.gz"; \
|
||||
curl -L --fail -o /tmp/aether-hub.tar.gz "$url"; \
|
||||
tar xzf /tmp/aether-hub.tar.gz -C /usr/local/bin; \
|
||||
chmod +x /usr/local/bin/aether-hub; \
|
||||
rm -f /tmp/aether-hub.tar.gz
|
||||
# 从 builder 阶段复制前端构建产物
|
||||
COPY --from=builder /app/frontend/dist /usr/share/nginx/html
|
||||
RUN chmod -R 755 /usr/share/nginx/html
|
||||
@@ -46,6 +90,11 @@ RUN printf '%s\n' \
|
||||
' "" $remote_addr;' \
|
||||
'}' \
|
||||
'' \
|
||||
'map $http_upgrade $connection_upgrade {' \
|
||||
' default upgrade;' \
|
||||
' "" "";' \
|
||||
'}' \
|
||||
'' \
|
||||
'server {' \
|
||||
' listen 80;' \
|
||||
' server_name _;' \
|
||||
@@ -77,7 +126,7 @@ RUN printf '%s\n' \
|
||||
'' \
|
||||
' # WebSocket 隧道端点(aether-proxy tunnel 模式)' \
|
||||
' location = /api/internal/proxy-tunnel {' \
|
||||
' proxy_pass http://127.0.0.1:PORT_PLACEHOLDER;' \
|
||||
' proxy_pass http://127.0.0.1:8085/proxy;' \
|
||||
' proxy_http_version 1.1;' \
|
||||
' proxy_set_header Host $host;' \
|
||||
' proxy_set_header X-Real-IP $real_ip;' \
|
||||
@@ -85,6 +134,25 @@ RUN printf '%s\n' \
|
||||
' proxy_set_header X-Forwarded-Proto $scheme;' \
|
||||
' proxy_set_header Upgrade $http_upgrade;' \
|
||||
' proxy_set_header Connection "upgrade";' \
|
||||
' # 剥离 CF 头,防止泄露给上游或返回给客户端' \
|
||||
' proxy_hide_header CF-Connecting-IP;' \
|
||||
' proxy_hide_header CF-IPCountry;' \
|
||||
' proxy_hide_header CF-Ray;' \
|
||||
' proxy_hide_header CF-Visitor;' \
|
||||
' proxy_hide_header CDN-Loop;' \
|
||||
' proxy_hide_header True-Client-IP;' \
|
||||
' proxy_hide_header CF-Worker;' \
|
||||
' proxy_hide_header CF-EW-Via;' \
|
||||
' proxy_hide_header CF-Warp-Tag-ID;' \
|
||||
' proxy_set_header CF-Connecting-IP "";' \
|
||||
' proxy_set_header CF-IPCountry "";' \
|
||||
' proxy_set_header CF-Ray "";' \
|
||||
' proxy_set_header CF-Visitor "";' \
|
||||
' proxy_set_header CDN-Loop "";' \
|
||||
' proxy_set_header True-Client-IP "";' \
|
||||
' proxy_set_header CF-Worker "";' \
|
||||
' proxy_set_header CF-EW-Via "";' \
|
||||
' proxy_set_header CF-Warp-Tag-ID "";' \
|
||||
' proxy_read_timeout 86400s;' \
|
||||
' proxy_send_timeout 86400s;' \
|
||||
' }' \
|
||||
@@ -97,11 +165,31 @@ RUN printf '%s\n' \
|
||||
' proxy_set_header X-Real-IP $real_ip;' \
|
||||
' proxy_set_header X-Forwarded-For $forwarded_for;' \
|
||||
' proxy_set_header X-Forwarded-Proto $scheme;' \
|
||||
' proxy_set_header Connection "";' \
|
||||
' proxy_set_header Upgrade $http_upgrade;' \
|
||||
' proxy_set_header Connection $connection_upgrade;' \
|
||||
' proxy_set_header Accept $http_accept;' \
|
||||
' proxy_set_header Content-Type $content_type;' \
|
||||
' proxy_set_header Authorization $http_authorization;' \
|
||||
' proxy_set_header X-Api-Key $http_x_api_key;' \
|
||||
' # 剥离 CF 头,防止泄露给上游或返回给客户端' \
|
||||
' proxy_hide_header CF-Connecting-IP;' \
|
||||
' proxy_hide_header CF-IPCountry;' \
|
||||
' proxy_hide_header CF-Ray;' \
|
||||
' proxy_hide_header CF-Visitor;' \
|
||||
' proxy_hide_header CDN-Loop;' \
|
||||
' proxy_hide_header True-Client-IP;' \
|
||||
' proxy_hide_header CF-Worker;' \
|
||||
' proxy_hide_header CF-EW-Via;' \
|
||||
' proxy_hide_header CF-Warp-Tag-ID;' \
|
||||
' proxy_set_header CF-Connecting-IP "";' \
|
||||
' proxy_set_header CF-IPCountry "";' \
|
||||
' proxy_set_header CF-Ray "";' \
|
||||
' proxy_set_header CF-Visitor "";' \
|
||||
' proxy_set_header CDN-Loop "";' \
|
||||
' proxy_set_header True-Client-IP "";' \
|
||||
' proxy_set_header CF-Worker "";' \
|
||||
' proxy_set_header CF-EW-Via "";' \
|
||||
' proxy_set_header CF-Warp-Tag-ID "";' \
|
||||
' proxy_buffering off;' \
|
||||
' proxy_cache off;' \
|
||||
' proxy_request_buffering off;' \
|
||||
@@ -121,6 +209,25 @@ RUN printf '%s\n' \
|
||||
' proxy_set_header X-Real-IP $real_ip;' \
|
||||
' proxy_set_header X-Forwarded-For $forwarded_for;' \
|
||||
' proxy_set_header X-Forwarded-Proto $scheme;' \
|
||||
' # 剥离 CF 头,防止泄露给上游或返回给客户端' \
|
||||
' proxy_hide_header CF-Connecting-IP;' \
|
||||
' proxy_hide_header CF-IPCountry;' \
|
||||
' proxy_hide_header CF-Ray;' \
|
||||
' proxy_hide_header CF-Visitor;' \
|
||||
' proxy_hide_header CDN-Loop;' \
|
||||
' proxy_hide_header True-Client-IP;' \
|
||||
' proxy_hide_header CF-Worker;' \
|
||||
' proxy_hide_header CF-EW-Via;' \
|
||||
' proxy_hide_header CF-Warp-Tag-ID;' \
|
||||
' proxy_set_header CF-Connecting-IP "";' \
|
||||
' proxy_set_header CF-IPCountry "";' \
|
||||
' proxy_set_header CF-Ray "";' \
|
||||
' proxy_set_header CF-Visitor "";' \
|
||||
' proxy_set_header CDN-Loop "";' \
|
||||
' proxy_set_header True-Client-IP "";' \
|
||||
' proxy_set_header CF-Worker "";' \
|
||||
' proxy_set_header CF-EW-Via "";' \
|
||||
' proxy_set_header CF-Warp-Tag-ID "";' \
|
||||
' }' \
|
||||
'' \
|
||||
' # 所有其他路由 → 前端 SPA(先尝试静态文件,再回退到 index.html)' \
|
||||
@@ -143,7 +250,7 @@ RUN printf '%s\n' \
|
||||
'stderr_logfile=/var/log/nginx/error.log' \
|
||||
'' \
|
||||
'[program:app]' \
|
||||
'command=/bin/bash -c "MAX_REQUESTS_JITTER=$((${MAX_REQUESTS:-4000}/20)); exec gunicorn src.main:app -c gunicorn_conf.py --preload -w %(ENV_GUNICORN_WORKERS)s -k uvicorn.workers.UvicornWorker --bind 127.0.0.1:8084 --timeout 120 --max-requests ${MAX_REQUESTS:-4000} --max-requests-jitter $MAX_REQUESTS_JITTER --access-logfile - --error-logfile - --log-level info"' \
|
||||
'command=/bin/bash -c "MAX_REQUESTS_JITTER=$((${MAX_REQUESTS:-50000}/20)); exec gunicorn src.main:app -c gunicorn_conf.py --preload -w %(ENV_GUNICORN_WORKERS)s -k uvicorn.workers.UvicornWorker --bind 127.0.0.1:8084 --max-requests ${MAX_REQUESTS:-50000} --max-requests-jitter $MAX_REQUESTS_JITTER --access-logfile - --error-logfile - --log-level info"' \
|
||||
'directory=/app' \
|
||||
'autostart=true' \
|
||||
'autorestart=true' \
|
||||
@@ -151,7 +258,16 @@ RUN printf '%s\n' \
|
||||
'stdout_logfile_maxbytes=0' \
|
||||
'stderr_logfile=/dev/stderr' \
|
||||
'stderr_logfile_maxbytes=0' \
|
||||
'environment=PYTHONUNBUFFERED=1,PYTHONIOENCODING=utf-8,LANG=C.UTF-8,LC_ALL=C.UTF-8,DOCKER_CONTAINER=true' > /etc/supervisor/conf.d/supervisord.conf
|
||||
'environment=PYTHONUNBUFFERED=1,PYTHONIOENCODING=utf-8,LANG=C.UTF-8,LC_ALL=C.UTF-8,DOCKER_CONTAINER=true,LD_PRELOAD=/usr/local/lib/libjemalloc.so.2,MALLOC_CONF="background_thread:true,dirty_decay_ms:5000,muzzy_decay_ms:5000"' \
|
||||
'' \
|
||||
'[program:tunnel-hub]' \
|
||||
'command=/usr/local/bin/aether-hub --bind 0.0.0.0:8085' \
|
||||
'autostart=true' \
|
||||
'autorestart=true' \
|
||||
'stdout_logfile=/dev/stdout' \
|
||||
'stdout_logfile_maxbytes=0' \
|
||||
'stderr_logfile=/dev/stderr' \
|
||||
'stderr_logfile_maxbytes=0' > /etc/supervisor/conf.d/supervisord.conf
|
||||
# 创建目录
|
||||
RUN mkdir -p /var/log/supervisor /app/logs /app/data
|
||||
# 入口脚本(启动前执行迁移)
|
||||
@@ -163,8 +279,10 @@ ENV PYTHONUNBUFFERED=1 \
|
||||
PYTHONIOENCODING=utf-8 \
|
||||
LANG=C.UTF-8 \
|
||||
LC_ALL=C.UTF-8 \
|
||||
LD_PRELOAD=/usr/local/lib/libjemalloc.so.2 \
|
||||
MALLOC_CONF=background_thread:true,dirty_decay_ms:5000,muzzy_decay_ms:5000 \
|
||||
PORT=8084 \
|
||||
GUNICORN_WORKERS=4 \
|
||||
GUNICORN_WORKERS=2 \
|
||||
MAX_REQUESTS=4000
|
||||
EXPOSE 80
|
||||
HEALTHCHECK --interval=30s --timeout=10s --start-period=5s --retries=3 \
|
||||
|
||||
+135
-6
@@ -2,6 +2,7 @@
|
||||
# 运行镜像:从 base 提取产物到精简运行时(国内镜像源版本)
|
||||
# 构建命令: docker build -f Dockerfile.app.local -t aether-app:latest .
|
||||
# 用于本地/国内服务器部署
|
||||
|
||||
FROM aether-base:latest AS builder
|
||||
|
||||
WORKDIR /app
|
||||
@@ -15,6 +16,16 @@ FROM python:3.13-slim
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
ARG HUB_RELEASE_REPO=fawney19/Aether
|
||||
ARG HUB_TAG
|
||||
ARG TARGETARCH
|
||||
ARG GITHUB_TOKEN
|
||||
# GitHub 下载镜像前缀,国内构建时传入可用的镜像加速地址
|
||||
# 用法: --build-arg GITHUB_MIRROR=https://ghfast.top
|
||||
# 或: --build-arg GITHUB_MIRROR=https://gh-proxy.com
|
||||
# 或: --build-arg GITHUB_MIRROR=https://mirror.ghproxy.com
|
||||
ARG GITHUB_MIRROR
|
||||
|
||||
# 运行时依赖(使用清华镜像源 + BuildKit 缓存加速)
|
||||
RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
|
||||
--mount=type=cache,target=/var/lib/apt,sharing=locked \
|
||||
@@ -23,7 +34,12 @@ RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
|
||||
nginx \
|
||||
supervisor \
|
||||
libpq5 \
|
||||
curl
|
||||
curl \
|
||||
libjemalloc2
|
||||
RUN set -eux; \
|
||||
jemalloc_path="$(find /usr/lib -type f -name 'libjemalloc.so.2' | head -n1)"; \
|
||||
[ -n "$jemalloc_path" ]; \
|
||||
ln -sf "$jemalloc_path" /usr/local/lib/libjemalloc.so.2
|
||||
|
||||
# 从 base 镜像复制 Python 包
|
||||
COPY --from=builder /usr/local/lib/python3.13/site-packages /usr/local/lib/python3.13/site-packages
|
||||
@@ -33,6 +49,45 @@ COPY --from=builder /usr/local/bin/gunicorn /usr/local/bin/
|
||||
COPY --from=builder /usr/local/bin/uvicorn /usr/local/bin/
|
||||
COPY --from=builder /usr/local/bin/alembic /usr/local/bin/
|
||||
|
||||
# Hub 预编译二进制
|
||||
# 国内构建: --build-arg GITHUB_MIRROR=https://ghfast.top 即可走镜像下载
|
||||
# GITHUB_TOKEN 可选:未认证 API 限流 60 次/小时,认证后 5000 次/小时
|
||||
RUN set -eux; \
|
||||
arch="${TARGETARCH:-}"; \
|
||||
if [ -z "$arch" ]; then \
|
||||
arch="$(dpkg --print-architecture)"; \
|
||||
fi; \
|
||||
case "$arch" in \
|
||||
amd64|arm64) ;; \
|
||||
x86_64) arch="amd64" ;; \
|
||||
aarch64) arch="arm64" ;; \
|
||||
*) echo "Unsupported architecture: $arch"; exit 1 ;; \
|
||||
esac; \
|
||||
auth_header=""; \
|
||||
if [ -n "${GITHUB_TOKEN:-}" ]; then \
|
||||
auth_header="Authorization: token ${GITHUB_TOKEN}"; \
|
||||
fi; \
|
||||
tag="${HUB_TAG:-}"; \
|
||||
if [ -z "$tag" ]; then \
|
||||
tag="$(curl -sL ${auth_header:+-H "$auth_header"} "https://api.github.com/repos/${HUB_RELEASE_REPO}/releases" | python3 -c "import json,sys;print(next((r['tag_name'] for r in json.load(sys.stdin) if r.get('tag_name','').startswith('hub-v') and not r.get('draft') and not r.get('prerelease')),''))")"; \
|
||||
fi; \
|
||||
if [ -z "$tag" ]; then \
|
||||
echo "Failed to resolve hub release tag"; \
|
||||
exit 1; \
|
||||
fi; \
|
||||
echo "Using Hub release tag: $tag"; \
|
||||
origin_url="https://github.com/${HUB_RELEASE_REPO}/releases/download/${tag}/aether-hub-linux-${arch}.tar.gz"; \
|
||||
if [ -n "${GITHUB_MIRROR:-}" ]; then \
|
||||
url="${GITHUB_MIRROR}/https://github.com/${HUB_RELEASE_REPO}/releases/download/${tag}/aether-hub-linux-${arch}.tar.gz"; \
|
||||
echo "Using mirror: ${GITHUB_MIRROR}"; \
|
||||
else \
|
||||
url="$origin_url"; \
|
||||
fi; \
|
||||
curl -L --fail -o /tmp/aether-hub.tar.gz "$url"; \
|
||||
tar xzf /tmp/aether-hub.tar.gz -C /usr/local/bin; \
|
||||
chmod +x /usr/local/bin/aether-hub; \
|
||||
rm -f /tmp/aether-hub.tar.gz
|
||||
|
||||
# 从 builder 阶段复制前端构建产物
|
||||
COPY --from=builder /app/frontend/dist /usr/share/nginx/html
|
||||
RUN chmod -R 755 /usr/share/nginx/html
|
||||
@@ -57,6 +112,11 @@ RUN printf '%s\n' \
|
||||
' "" $remote_addr;' \
|
||||
'}' \
|
||||
'' \
|
||||
'map $http_upgrade $connection_upgrade {' \
|
||||
' default upgrade;' \
|
||||
' "" "";' \
|
||||
'}' \
|
||||
'' \
|
||||
'server {' \
|
||||
' listen 80;' \
|
||||
' server_name _;' \
|
||||
@@ -88,7 +148,7 @@ RUN printf '%s\n' \
|
||||
'' \
|
||||
' # WebSocket 隧道端点(aether-proxy tunnel 模式)' \
|
||||
' location = /api/internal/proxy-tunnel {' \
|
||||
' proxy_pass http://127.0.0.1:PORT_PLACEHOLDER;' \
|
||||
' proxy_pass http://127.0.0.1:8085/proxy;' \
|
||||
' proxy_http_version 1.1;' \
|
||||
' proxy_set_header Host $host;' \
|
||||
' proxy_set_header X-Real-IP $real_ip;' \
|
||||
@@ -96,6 +156,25 @@ RUN printf '%s\n' \
|
||||
' proxy_set_header X-Forwarded-Proto $scheme;' \
|
||||
' proxy_set_header Upgrade $http_upgrade;' \
|
||||
' proxy_set_header Connection "upgrade";' \
|
||||
' # 剥离 CF 头,防止泄露给上游或返回给客户端' \
|
||||
' proxy_hide_header CF-Connecting-IP;' \
|
||||
' proxy_hide_header CF-IPCountry;' \
|
||||
' proxy_hide_header CF-Ray;' \
|
||||
' proxy_hide_header CF-Visitor;' \
|
||||
' proxy_hide_header CDN-Loop;' \
|
||||
' proxy_hide_header True-Client-IP;' \
|
||||
' proxy_hide_header CF-Worker;' \
|
||||
' proxy_hide_header CF-EW-Via;' \
|
||||
' proxy_hide_header CF-Warp-Tag-ID;' \
|
||||
' proxy_set_header CF-Connecting-IP "";' \
|
||||
' proxy_set_header CF-IPCountry "";' \
|
||||
' proxy_set_header CF-Ray "";' \
|
||||
' proxy_set_header CF-Visitor "";' \
|
||||
' proxy_set_header CDN-Loop "";' \
|
||||
' proxy_set_header True-Client-IP "";' \
|
||||
' proxy_set_header CF-Worker "";' \
|
||||
' proxy_set_header CF-EW-Via "";' \
|
||||
' proxy_set_header CF-Warp-Tag-ID "";' \
|
||||
' proxy_read_timeout 86400s;' \
|
||||
' proxy_send_timeout 86400s;' \
|
||||
' }' \
|
||||
@@ -108,11 +187,31 @@ RUN printf '%s\n' \
|
||||
' proxy_set_header X-Real-IP $real_ip;' \
|
||||
' proxy_set_header X-Forwarded-For $forwarded_for;' \
|
||||
' proxy_set_header X-Forwarded-Proto $scheme;' \
|
||||
' proxy_set_header Connection "";' \
|
||||
' proxy_set_header Upgrade $http_upgrade;' \
|
||||
' proxy_set_header Connection $connection_upgrade;' \
|
||||
' proxy_set_header Accept $http_accept;' \
|
||||
' proxy_set_header Content-Type $content_type;' \
|
||||
' proxy_set_header Authorization $http_authorization;' \
|
||||
' proxy_set_header X-Api-Key $http_x_api_key;' \
|
||||
' # 剥离 CF 头,防止泄露给上游或返回给客户端' \
|
||||
' proxy_hide_header CF-Connecting-IP;' \
|
||||
' proxy_hide_header CF-IPCountry;' \
|
||||
' proxy_hide_header CF-Ray;' \
|
||||
' proxy_hide_header CF-Visitor;' \
|
||||
' proxy_hide_header CDN-Loop;' \
|
||||
' proxy_hide_header True-Client-IP;' \
|
||||
' proxy_hide_header CF-Worker;' \
|
||||
' proxy_hide_header CF-EW-Via;' \
|
||||
' proxy_hide_header CF-Warp-Tag-ID;' \
|
||||
' proxy_set_header CF-Connecting-IP "";' \
|
||||
' proxy_set_header CF-IPCountry "";' \
|
||||
' proxy_set_header CF-Ray "";' \
|
||||
' proxy_set_header CF-Visitor "";' \
|
||||
' proxy_set_header CDN-Loop "";' \
|
||||
' proxy_set_header True-Client-IP "";' \
|
||||
' proxy_set_header CF-Worker "";' \
|
||||
' proxy_set_header CF-EW-Via "";' \
|
||||
' proxy_set_header CF-Warp-Tag-ID "";' \
|
||||
' proxy_buffering off;' \
|
||||
' proxy_cache off;' \
|
||||
' proxy_request_buffering off;' \
|
||||
@@ -132,6 +231,25 @@ RUN printf '%s\n' \
|
||||
' proxy_set_header X-Real-IP $real_ip;' \
|
||||
' proxy_set_header X-Forwarded-For $forwarded_for;' \
|
||||
' proxy_set_header X-Forwarded-Proto $scheme;' \
|
||||
' # 剥离 CF 头,防止泄露给上游或返回给客户端' \
|
||||
' proxy_hide_header CF-Connecting-IP;' \
|
||||
' proxy_hide_header CF-IPCountry;' \
|
||||
' proxy_hide_header CF-Ray;' \
|
||||
' proxy_hide_header CF-Visitor;' \
|
||||
' proxy_hide_header CDN-Loop;' \
|
||||
' proxy_hide_header True-Client-IP;' \
|
||||
' proxy_hide_header CF-Worker;' \
|
||||
' proxy_hide_header CF-EW-Via;' \
|
||||
' proxy_hide_header CF-Warp-Tag-ID;' \
|
||||
' proxy_set_header CF-Connecting-IP "";' \
|
||||
' proxy_set_header CF-IPCountry "";' \
|
||||
' proxy_set_header CF-Ray "";' \
|
||||
' proxy_set_header CF-Visitor "";' \
|
||||
' proxy_set_header CDN-Loop "";' \
|
||||
' proxy_set_header True-Client-IP "";' \
|
||||
' proxy_set_header CF-Worker "";' \
|
||||
' proxy_set_header CF-EW-Via "";' \
|
||||
' proxy_set_header CF-Warp-Tag-ID "";' \
|
||||
' }' \
|
||||
'' \
|
||||
' # 所有其他路由 → 前端 SPA(先尝试静态文件,再回退到 index.html)' \
|
||||
@@ -155,7 +273,7 @@ RUN printf '%s\n' \
|
||||
'stderr_logfile=/var/log/nginx/error.log' \
|
||||
'' \
|
||||
'[program:app]' \
|
||||
'command=/bin/bash -c "MAX_REQUESTS_JITTER=$((${MAX_REQUESTS:-4000}/20)); exec gunicorn src.main:app -c gunicorn_conf.py --preload -w %(ENV_GUNICORN_WORKERS)s -k uvicorn.workers.UvicornWorker --bind 0.0.0.0:%(ENV_PORT)s --timeout 120 --max-requests ${MAX_REQUESTS:-4000} --max-requests-jitter $MAX_REQUESTS_JITTER --access-logfile - --error-logfile - --log-level info"' \
|
||||
'command=/bin/bash -c "MAX_REQUESTS_JITTER=$((${MAX_REQUESTS:-50000}/20)); exec gunicorn src.main:app -c gunicorn_conf.py --preload -w %(ENV_GUNICORN_WORKERS)s -k uvicorn.workers.UvicornWorker --bind 0.0.0.0:%(ENV_PORT)s --max-requests ${MAX_REQUESTS:-50000} --max-requests-jitter $MAX_REQUESTS_JITTER --access-logfile - --error-logfile - --log-level info"' \
|
||||
'directory=/app' \
|
||||
'autostart=true' \
|
||||
'autorestart=true' \
|
||||
@@ -163,7 +281,16 @@ RUN printf '%s\n' \
|
||||
'stdout_logfile_maxbytes=0' \
|
||||
'stderr_logfile=/dev/stderr' \
|
||||
'stderr_logfile_maxbytes=0' \
|
||||
'environment=PYTHONUNBUFFERED=1,PYTHONIOENCODING=utf-8,LANG=C.UTF-8,LC_ALL=C.UTF-8,DOCKER_CONTAINER=true' > /etc/supervisor/conf.d/supervisord.conf
|
||||
'environment=PYTHONUNBUFFERED=1,PYTHONIOENCODING=utf-8,LANG=C.UTF-8,LC_ALL=C.UTF-8,DOCKER_CONTAINER=true,LD_PRELOAD=/usr/local/lib/libjemalloc.so.2,MALLOC_CONF="background_thread:true,dirty_decay_ms:5000,muzzy_decay_ms:5000"' \
|
||||
'' \
|
||||
'[program:tunnel-hub]' \
|
||||
'command=/usr/local/bin/aether-hub --bind 0.0.0.0:8085' \
|
||||
'autostart=true' \
|
||||
'autorestart=true' \
|
||||
'stdout_logfile=/dev/stdout' \
|
||||
'stdout_logfile_maxbytes=0' \
|
||||
'stderr_logfile=/dev/stderr' \
|
||||
'stderr_logfile_maxbytes=0' > /etc/supervisor/conf.d/supervisord.conf
|
||||
|
||||
# 创建目录
|
||||
RUN mkdir -p /var/log/supervisor /app/logs /app/data
|
||||
@@ -178,8 +305,10 @@ ENV PYTHONUNBUFFERED=1 \
|
||||
PYTHONIOENCODING=utf-8 \
|
||||
LANG=C.UTF-8 \
|
||||
LC_ALL=C.UTF-8 \
|
||||
LD_PRELOAD=/usr/local/lib/libjemalloc.so.2 \
|
||||
MALLOC_CONF=background_thread:true,dirty_decay_ms:5000,muzzy_decay_ms:5000 \
|
||||
PORT=8084 \
|
||||
GUNICORN_WORKERS=4 \
|
||||
GUNICORN_WORKERS=2 \
|
||||
MAX_REQUESTS=4000
|
||||
|
||||
EXPOSE 80
|
||||
|
||||
@@ -48,7 +48,7 @@ python generate_keys.py # 生成密钥, 并将生成的密钥填入 .env
|
||||
# 3. 部署 / 更新(自动执行数据库迁移)
|
||||
docker compose pull && docker compose up -d
|
||||
|
||||
# 4. 升级前备份
|
||||
# 4. 升级前备份 (可选)
|
||||
docker compose exec postgres pg_dump -U postgres aether | gzip > backup_$(date +%Y%m%d_%H%M%S).sql.gz
|
||||
```
|
||||
|
||||
@@ -64,6 +64,7 @@ cp .env.example .env
|
||||
python generate_keys.py # 生成密钥, 并将生成的密钥填入 .env
|
||||
|
||||
# 3. 部署 / 更新(自动构建、启动、迁移)
|
||||
git pull
|
||||
./deploy.sh
|
||||
```
|
||||
|
||||
@@ -81,7 +82,7 @@ uv sync
|
||||
cd frontend && npm install && npm run dev
|
||||
```
|
||||
|
||||
## Aether Proxy
|
||||
## Aether Proxy (可选)
|
||||
|
||||
Aether Proxy 是配套的正向代理节点,部署在海外 VPS 上,为墙内的 Aether 实例中转 API 流量。或者部署在其他服务器为指定的提供商、账号、Key使用不同的节点访问。支持 TUI 向导一键配置、systemd 服务管理、TLS 加密、DNS 缓存及连接池调优。
|
||||
|
||||
@@ -110,7 +111,7 @@ Aether Proxy 是配套的正向代理节点,部署在海外 VPS 上,为墙
|
||||
| `APP_PORT` | 8084 | 应用端口 |
|
||||
| `API_KEY_PREFIX` | sk | API Key 前缀 |
|
||||
| `LOG_LEVEL` | INFO | 日志级别 (DEBUG/INFO/WARNING/ERROR) |
|
||||
| `GUNICORN_WORKERS` | 4 | Gunicorn 工作进程数 |
|
||||
| `GUNICORN_WORKERS` | 2 | Gunicorn 工作进程数 |
|
||||
| `DB_PORT` | 5432 | PostgreSQL 端口 |
|
||||
| `REDIS_PORT` | 6379 | Redis 端口 |
|
||||
|
||||
@@ -180,4 +181,3 @@ docker compose up -d app
|
||||
## Star History
|
||||
|
||||
[](https://star-history.com/#fawney19/Aether&Date)
|
||||
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
target/
|
||||
.git/
|
||||
.DS_Store
|
||||
Generated
+2012
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,27 @@
|
||||
[package]
|
||||
name = "aether-hub"
|
||||
version = "0.2.0"
|
||||
edition = "2021"
|
||||
description = "Tunnel Hub for Aether - frame router between workers and proxies"
|
||||
|
||||
[dependencies]
|
||||
tokio = { version = "1", features = ["full"] }
|
||||
axum = { version = "0.8", features = ["ws"] }
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
tracing = "0.1"
|
||||
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
|
||||
clap = { version = "4", features = ["derive", "env"] }
|
||||
dashmap = "6"
|
||||
parking_lot = "0.12"
|
||||
flate2 = "1"
|
||||
futures-util = "0.3"
|
||||
bytes = "1"
|
||||
async-stream = "0.3"
|
||||
http-body-util = "0.1"
|
||||
reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls"] }
|
||||
|
||||
[profile.release]
|
||||
lto = true
|
||||
strip = true
|
||||
codegen-units = 1
|
||||
@@ -0,0 +1,34 @@
|
||||
# syntax=docker/dockerfile:1
|
||||
|
||||
FROM rust:1.85-slim AS builder
|
||||
WORKDIR /build/aether-hub
|
||||
|
||||
# 可选:配置国内 Cargo 镜像源(本地构建时传 --build-arg CARGO_MIRROR=1)
|
||||
ARG CARGO_MIRROR
|
||||
RUN if [ -n "$CARGO_MIRROR" ]; then \
|
||||
printf '[source.crates-io]\nreplace-with = "tuna"\n\n[source.tuna]\nregistry = "sparse+https://mirrors.tuna.tsinghua.edu.cn/crates.io-index/"\n' \
|
||||
> /usr/local/cargo/config.toml; \
|
||||
fi
|
||||
|
||||
# 先构建依赖层,最大化后续代码变更时的缓存命中
|
||||
COPY Cargo.toml Cargo.lock ./
|
||||
RUN mkdir src && printf 'fn main() {}\n' > src/main.rs
|
||||
RUN --mount=type=cache,target=/usr/local/cargo/registry,sharing=locked \
|
||||
--mount=type=cache,target=/build/aether-hub/target,sharing=locked \
|
||||
cargo build --release --locked
|
||||
RUN rm -rf src
|
||||
|
||||
COPY src ./src
|
||||
RUN --mount=type=cache,target=/usr/local/cargo/registry,sharing=locked \
|
||||
--mount=type=cache,target=/build/aether-hub/target,sharing=locked \
|
||||
cargo build --release --locked && \
|
||||
cp target/release/aether-hub /tmp/aether-hub
|
||||
|
||||
FROM debian:bookworm-slim
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends ca-certificates && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
COPY --from=builder /tmp/aether-hub /usr/local/bin/aether-hub
|
||||
|
||||
EXPOSE 8085
|
||||
ENTRYPOINT ["/usr/local/bin/aether-hub"]
|
||||
CMD ["--bind", "0.0.0.0:8085"]
|
||||
@@ -0,0 +1,38 @@
|
||||
# aether-hub
|
||||
|
||||
`aether-hub` 是 Tunnel Hub 服务,负责在 proxy 与 worker 之间路由帧。
|
||||
|
||||
已集成在Docker镜像中, 无需单独部署。
|
||||
|
||||
## 部署端指定 Hub 版本并构建
|
||||
|
||||
```bash
|
||||
cd /path/to/Aether
|
||||
./deploy.sh --hub-tag hub-v0.1.0
|
||||
```
|
||||
|
||||
不指定 `--hub-tag` 时,`./deploy.sh` 会自动解析最新 `hub-v*` release,并在构建 app 镜像时从 GitHub Release 下载对应架构的 Hub 二进制。
|
||||
|
||||
## build.sh 模式说明
|
||||
|
||||
- 默认是 `binary` 模式(`cross` 构建二进制)。
|
||||
- `--upload <hub-vX.Y.Z>` 会把构建产物上传到 GitHub Release。
|
||||
- 加 `--image` 后进入镜像模式(`docker buildx`,可选)。
|
||||
|
||||
常用参数:
|
||||
|
||||
- `--tag <tag>`: 镜像 tag
|
||||
- `--image-name <name>`: 镜像名(默认 `ghcr.io/fawney19/aether-hub`)
|
||||
- `--platforms <list>`: 例如 `linux/amd64,linux/arm64`
|
||||
- `--push`: 推送镜像
|
||||
- `--load`: 加载到本地 Docker(单平台)
|
||||
- `--latest`: 额外打 `latest` tag
|
||||
|
||||
## 运行时参数
|
||||
|
||||
- `TUNNEL_HUB_WORKER_IDLE_TIMEOUT`:worker 心跳空闲超时,默认 `60` 秒
|
||||
- `TUNNEL_HUB_OUTBOUND_QUEUE_CAPACITY`:单连接出站队列容量,默认 `128`;队列打满时会把连接视为拥塞并主动关闭,避免 Hub 内存无限增长
|
||||
|
||||
## 与部署脚本关系
|
||||
|
||||
- `./deploy.sh`: 本地构建部署(会本地构建 app/base,并在构建 app 时从 GitHub Release 下载 Hub,可用 `--hub-tag` 固定版本)。
|
||||
Executable
+272
@@ -0,0 +1,272 @@
|
||||
#!/bin/bash
|
||||
# aether-hub 构建脚本
|
||||
#
|
||||
# 支持两种模式:
|
||||
# 1) binary 模式(默认): 构建多架构二进制并可上传 GitHub Release
|
||||
# 2) image 模式: 构建并推送/加载 Docker 镜像(推荐生产发布用)
|
||||
#
|
||||
# 示例:
|
||||
# # binary 模式(兼容旧行为)
|
||||
# ./build.sh
|
||||
# ./build.sh amd64
|
||||
# ./build.sh --upload hub-v0.1.0
|
||||
#
|
||||
# # image 模式(多架构推送)
|
||||
# ./build.sh --image --tag v0.2.5 --push --latest
|
||||
# ./build.sh --image --tag sha-abc123 --image-name ghcr.io/fawney19/aether-hub --push
|
||||
# ./build.sh --image --tag local-test --platforms linux/amd64 --load
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
|
||||
PROJECT_DIR="$(cd "$SCRIPT_DIR/.." && pwd)"
|
||||
DIST_DIR="$SCRIPT_DIR/dist"
|
||||
|
||||
# -------------------------------
|
||||
# Defaults
|
||||
# -------------------------------
|
||||
MODE="binary" # binary | image
|
||||
|
||||
# binary mode options
|
||||
UPLOAD=false
|
||||
UPLOAD_TAG=""
|
||||
BINARY_TARGETS=""
|
||||
|
||||
# image mode options
|
||||
IMAGE_NAME="${IMAGE_NAME:-ghcr.io/fawney19/aether-hub}"
|
||||
IMAGE_TAG=""
|
||||
IMAGE_PLATFORMS="linux/amd64,linux/arm64"
|
||||
IMAGE_PUSH=false
|
||||
IMAGE_LOAD=false
|
||||
IMAGE_LATEST=false
|
||||
|
||||
usage() {
|
||||
cat <<'EOF'
|
||||
用法:
|
||||
./build.sh [binary-args]
|
||||
./build.sh --image [image-args]
|
||||
|
||||
binary 模式(默认):
|
||||
amd64|arm64 仅构建指定架构(可重复)
|
||||
--upload <hub-vX.Y.Z> 上传到 GitHub Release(需要 gh CLI)
|
||||
|
||||
image 模式:
|
||||
--image 启用镜像模式
|
||||
--tag <tag> 镜像 tag(默认自动从 git describe 推导)
|
||||
--image-name <name> 镜像名(默认 ghcr.io/fawney19/aether-hub)
|
||||
--platforms <list> 平台列表,逗号分隔(默认 linux/amd64,linux/arm64)
|
||||
--push 推送镜像到仓库
|
||||
--load 加载到本地 Docker(仅单平台)
|
||||
--latest 额外打 latest tag
|
||||
|
||||
通用:
|
||||
-h, --help 显示帮助
|
||||
EOF
|
||||
}
|
||||
|
||||
while [ $# -gt 0 ]; do
|
||||
case "$1" in
|
||||
--image)
|
||||
MODE="image"
|
||||
shift
|
||||
;;
|
||||
--tag)
|
||||
IMAGE_TAG="${2:-}"
|
||||
shift 2
|
||||
;;
|
||||
--image-name)
|
||||
IMAGE_NAME="${2:-}"
|
||||
shift 2
|
||||
;;
|
||||
--platforms)
|
||||
IMAGE_PLATFORMS="${2:-}"
|
||||
shift 2
|
||||
;;
|
||||
--push)
|
||||
IMAGE_PUSH=true
|
||||
shift
|
||||
;;
|
||||
--load)
|
||||
IMAGE_LOAD=true
|
||||
shift
|
||||
;;
|
||||
--latest)
|
||||
IMAGE_LATEST=true
|
||||
shift
|
||||
;;
|
||||
--upload)
|
||||
UPLOAD=true
|
||||
UPLOAD_TAG="${2:-}"
|
||||
shift 2
|
||||
;;
|
||||
amd64|arm64)
|
||||
BINARY_TARGETS="$BINARY_TARGETS $1"
|
||||
shift
|
||||
;;
|
||||
-h|--help)
|
||||
usage
|
||||
exit 0
|
||||
;;
|
||||
*)
|
||||
echo "❌ 未知参数: $1"
|
||||
usage
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
build_binary() {
|
||||
if [ -z "$BINARY_TARGETS" ]; then
|
||||
BINARY_TARGETS="amd64 arm64"
|
||||
fi
|
||||
|
||||
if ! command -v cross >/dev/null 2>&1; then
|
||||
echo "❌ 需要安装 cross: cargo install cross --git https://github.com/cross-rs/cross"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
mkdir -p "$DIST_DIR"
|
||||
|
||||
echo "🔨 开始构建 aether-hub 二进制..."
|
||||
echo " 目标平台: $BINARY_TARGETS"
|
||||
echo ""
|
||||
|
||||
ARTIFACTS=""
|
||||
for arch in $BINARY_TARGETS; do
|
||||
case "$arch" in
|
||||
amd64) target="x86_64-unknown-linux-gnu" ;;
|
||||
arm64) target="aarch64-unknown-linux-gnu" ;;
|
||||
*) echo "❌ 未知架构: $arch"; exit 1 ;;
|
||||
esac
|
||||
|
||||
echo ">>> 构建 $arch ($target)..."
|
||||
cd "$SCRIPT_DIR"
|
||||
cross build --release --target "$target" --locked
|
||||
|
||||
BIN="target/$target/release/aether-hub"
|
||||
if [ ! -f "$BIN" ]; then
|
||||
echo "❌ 未找到二进制文件: $BIN"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
ARCHIVE="$DIST_DIR/aether-hub-linux-$arch.tar.gz"
|
||||
tar czf "$ARCHIVE" -C "target/$target/release" aether-hub
|
||||
ARTIFACTS="$ARTIFACTS $ARCHIVE"
|
||||
|
||||
SIZE=$(du -h "$ARCHIVE" | cut -f1)
|
||||
echo "✅ $arch 构建完成: $ARCHIVE ($SIZE)"
|
||||
echo ""
|
||||
done
|
||||
|
||||
cd "$DIST_DIR"
|
||||
shasum -a 256 aether-hub-*.tar.gz > SHA256SUMS.txt
|
||||
echo "📋 SHA256 校验和:"
|
||||
cat SHA256SUMS.txt
|
||||
echo ""
|
||||
|
||||
if [ "$UPLOAD" = true ]; then
|
||||
if [ -z "$UPLOAD_TAG" ]; then
|
||||
echo "❌ --upload 需要指定 tag,例如: ./build.sh --upload hub-v0.1.0"
|
||||
exit 1
|
||||
fi
|
||||
if ! command -v gh >/dev/null 2>&1; then
|
||||
echo "❌ 需要安装 GitHub CLI: brew install gh"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "📦 上传到 GitHub Release: $UPLOAD_TAG"
|
||||
cd "$PROJECT_DIR"
|
||||
|
||||
if ! git rev-parse "$UPLOAD_TAG" >/dev/null 2>&1; then
|
||||
git tag "$UPLOAD_TAG"
|
||||
git push origin "$UPLOAD_TAG"
|
||||
fi
|
||||
|
||||
gh release create "$UPLOAD_TAG" \
|
||||
--title "aether-hub ${UPLOAD_TAG#hub-}" \
|
||||
--generate-notes \
|
||||
$ARTIFACTS \
|
||||
"$DIST_DIR/SHA256SUMS.txt"
|
||||
|
||||
echo "✅ 上传完成!"
|
||||
fi
|
||||
|
||||
echo "🎉 binary 模式完成!"
|
||||
}
|
||||
|
||||
build_image() {
|
||||
if ! command -v docker >/dev/null 2>&1; then
|
||||
echo "❌ 未找到 docker,请先安装 Docker"
|
||||
exit 1
|
||||
fi
|
||||
if ! docker buildx version >/dev/null 2>&1; then
|
||||
echo "❌ 未找到 docker buildx,请先启用 buildx"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if [ "$IMAGE_PUSH" = true ] && [ "$IMAGE_LOAD" = true ]; then
|
||||
echo "❌ --push 与 --load 不能同时使用"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if [ "$IMAGE_PUSH" = false ] && [ "$IMAGE_LOAD" = false ]; then
|
||||
# image 模式默认走 push,符合发布场景
|
||||
IMAGE_PUSH=true
|
||||
fi
|
||||
|
||||
if [ -z "$IMAGE_TAG" ]; then
|
||||
IMAGE_TAG=$(git -C "$PROJECT_DIR" describe --tags --always 2>/dev/null | sed 's/^v//')
|
||||
if [ -z "$IMAGE_TAG" ]; then
|
||||
IMAGE_TAG=$(date +%Y%m%d%H%M%S)
|
||||
fi
|
||||
fi
|
||||
|
||||
if [ "$IMAGE_LOAD" = true ] && [[ "$IMAGE_PLATFORMS" == *,* ]]; then
|
||||
echo "❌ --load 仅支持单平台,请用 --platforms linux/amd64(或 arm64)"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
local ref="${IMAGE_NAME}:${IMAGE_TAG}"
|
||||
local cmd=(docker buildx build
|
||||
--platform "$IMAGE_PLATFORMS"
|
||||
-f "$SCRIPT_DIR/Dockerfile"
|
||||
-t "$ref"
|
||||
)
|
||||
|
||||
if [ "$IMAGE_LATEST" = true ]; then
|
||||
cmd+=(-t "${IMAGE_NAME}:latest")
|
||||
fi
|
||||
|
||||
if [ "$IMAGE_PUSH" = true ]; then
|
||||
cmd+=(--push)
|
||||
else
|
||||
cmd+=(--load)
|
||||
fi
|
||||
|
||||
cmd+=("$SCRIPT_DIR")
|
||||
|
||||
echo "🔨 开始构建 aether-hub 镜像..."
|
||||
echo " image: $ref"
|
||||
echo " platforms: $IMAGE_PLATFORMS"
|
||||
echo " mode: $([ "$IMAGE_PUSH" = true ] && echo push || echo load)"
|
||||
echo ""
|
||||
|
||||
"${cmd[@]}"
|
||||
|
||||
if [ "$IMAGE_PUSH" = true ]; then
|
||||
echo "✅ 镜像已推送: $ref"
|
||||
if [ "$IMAGE_LATEST" = true ]; then
|
||||
echo "✅ 镜像已推送: ${IMAGE_NAME}:latest"
|
||||
fi
|
||||
else
|
||||
echo "✅ 镜像已加载到本地: $ref"
|
||||
fi
|
||||
|
||||
echo "🎉 image 模式完成!"
|
||||
}
|
||||
|
||||
if [ "$MODE" = "image" ]; then
|
||||
build_image
|
||||
else
|
||||
build_binary
|
||||
fi
|
||||
@@ -0,0 +1,85 @@
|
||||
use reqwest::Client;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct ControlPlaneClient {
|
||||
client: Option<Client>,
|
||||
base_url: String,
|
||||
}
|
||||
|
||||
impl ControlPlaneClient {
|
||||
pub fn new(base_url: String) -> Self {
|
||||
let client = Client::builder()
|
||||
.timeout(std::time::Duration::from_secs(10))
|
||||
.build()
|
||||
.ok();
|
||||
Self { client, base_url }
|
||||
}
|
||||
|
||||
pub fn disabled() -> Self {
|
||||
Self {
|
||||
client: None,
|
||||
base_url: String::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn heartbeat_ack(&self, payload: &[u8]) -> Result<Vec<u8>, String> {
|
||||
let Some(client) = &self.client else {
|
||||
return Ok(b"{}".to_vec());
|
||||
};
|
||||
let url = format!(
|
||||
"{}/api/internal/hub/heartbeat",
|
||||
self.base_url.trim_end_matches('/')
|
||||
);
|
||||
let response = client
|
||||
.post(&url)
|
||||
.header("content-type", "application/json")
|
||||
.body(payload.to_vec())
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("heartbeat callback request failed: {e}"))?;
|
||||
if !response.status().is_success() {
|
||||
return Err(format!(
|
||||
"heartbeat callback failed with status {}",
|
||||
response.status()
|
||||
));
|
||||
}
|
||||
response
|
||||
.bytes()
|
||||
.await
|
||||
.map(|bytes| bytes.to_vec())
|
||||
.map_err(|e| format!("heartbeat callback body read failed: {e}"))
|
||||
}
|
||||
|
||||
pub async fn push_node_status(
|
||||
&self,
|
||||
node_id: &str,
|
||||
connected: bool,
|
||||
conn_count: usize,
|
||||
) -> Result<(), String> {
|
||||
let Some(client) = &self.client else {
|
||||
return Ok(());
|
||||
};
|
||||
let url = format!(
|
||||
"{}/api/internal/hub/node-status",
|
||||
self.base_url.trim_end_matches('/')
|
||||
);
|
||||
let response = client
|
||||
.post(&url)
|
||||
.json(&serde_json::json!({
|
||||
"node_id": node_id,
|
||||
"connected": connected,
|
||||
"conn_count": conn_count,
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("node-status callback request failed: {e}"))?;
|
||||
if response.status().is_success() {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(format!(
|
||||
"node-status callback failed with status {}",
|
||||
response.status()
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,878 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, AtomicUsize, Ordering};
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use axum::extract::ws::Message;
|
||||
use bytes::Bytes;
|
||||
use dashmap::DashMap;
|
||||
use parking_lot::{Mutex, RwLock};
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::sync::mpsc::error::TrySendError;
|
||||
use tokio::sync::{watch, Notify};
|
||||
use tracing::{debug, info, warn};
|
||||
|
||||
use crate::control_plane::ControlPlaneClient;
|
||||
use crate::protocol;
|
||||
|
||||
const MAX_REQUEST_BODY_FRAME_SIZE: usize = 32 * 1024;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum SendStatus {
|
||||
Queued,
|
||||
Closed,
|
||||
Congested,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct ConnConfig {
|
||||
pub ping_interval: Duration,
|
||||
pub idle_timeout: Duration,
|
||||
pub outbound_queue_capacity: usize,
|
||||
}
|
||||
|
||||
pub struct BoundedOutbound {
|
||||
tx: mpsc::Sender<Message>,
|
||||
close_tx: watch::Sender<bool>,
|
||||
closing: AtomicBool,
|
||||
}
|
||||
|
||||
impl BoundedOutbound {
|
||||
pub fn new(tx: mpsc::Sender<Message>, close_tx: watch::Sender<bool>) -> Self {
|
||||
Self {
|
||||
tx,
|
||||
close_tx,
|
||||
closing: AtomicBool::new(false),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn send(&self, msg: Message) -> SendStatus {
|
||||
if self.is_closing() {
|
||||
return SendStatus::Closed;
|
||||
}
|
||||
|
||||
match self.tx.try_send(msg) {
|
||||
Ok(()) => SendStatus::Queued,
|
||||
Err(TrySendError::Closed(_)) => {
|
||||
self.mark_closing();
|
||||
SendStatus::Closed
|
||||
}
|
||||
Err(TrySendError::Full(_)) => {
|
||||
self.mark_closing();
|
||||
SendStatus::Congested
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_closing(&self) -> bool {
|
||||
self.closing.load(Ordering::Acquire)
|
||||
}
|
||||
|
||||
pub fn mark_closing(&self) -> bool {
|
||||
if self.closing.swap(true, Ordering::AcqRel) {
|
||||
return false;
|
||||
}
|
||||
let _ = self.close_tx.send(true);
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
pub struct ProxyConn {
|
||||
pub id: u64,
|
||||
pub node_id: String,
|
||||
pub node_name: String,
|
||||
pub outbound: BoundedOutbound,
|
||||
next_stream_id: AtomicU32,
|
||||
pub stream_count: AtomicUsize,
|
||||
pub max_streams: usize,
|
||||
}
|
||||
|
||||
impl ProxyConn {
|
||||
pub fn new(
|
||||
id: u64,
|
||||
node_id: String,
|
||||
node_name: String,
|
||||
tx: mpsc::Sender<Message>,
|
||||
close_tx: watch::Sender<bool>,
|
||||
max_streams: usize,
|
||||
) -> Self {
|
||||
Self {
|
||||
id,
|
||||
node_id,
|
||||
node_name,
|
||||
outbound: BoundedOutbound::new(tx, close_tx),
|
||||
next_stream_id: AtomicU32::new(2),
|
||||
stream_count: AtomicUsize::new(0),
|
||||
max_streams,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn alloc_stream_id(&self) -> Option<u32> {
|
||||
let mut current = self.stream_count.load(Ordering::Relaxed);
|
||||
loop {
|
||||
if current >= self.max_streams || !self.is_available() {
|
||||
return None;
|
||||
}
|
||||
match self.stream_count.compare_exchange_weak(
|
||||
current,
|
||||
current + 1,
|
||||
Ordering::AcqRel,
|
||||
Ordering::Relaxed,
|
||||
) {
|
||||
Ok(_) => break,
|
||||
Err(observed) => current = observed,
|
||||
}
|
||||
}
|
||||
|
||||
let sid = loop {
|
||||
let current_sid = self.next_stream_id.load(Ordering::Relaxed);
|
||||
let next_sid = if current_sid >= 0xFFFF_FFFE {
|
||||
2
|
||||
} else {
|
||||
current_sid + 2
|
||||
};
|
||||
if self
|
||||
.next_stream_id
|
||||
.compare_exchange_weak(current_sid, next_sid, Ordering::AcqRel, Ordering::Relaxed)
|
||||
.is_ok()
|
||||
{
|
||||
break current_sid;
|
||||
}
|
||||
};
|
||||
|
||||
Some(sid)
|
||||
}
|
||||
|
||||
pub fn release_stream(&self) {
|
||||
let mut current = self.stream_count.load(Ordering::Relaxed);
|
||||
while current > 0 {
|
||||
match self.stream_count.compare_exchange_weak(
|
||||
current,
|
||||
current - 1,
|
||||
Ordering::AcqRel,
|
||||
Ordering::Relaxed,
|
||||
) {
|
||||
Ok(_) => return,
|
||||
Err(observed) => current = observed,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_available(&self) -> bool {
|
||||
!self.outbound.is_closing()
|
||||
}
|
||||
|
||||
pub fn request_close(&self) {
|
||||
self.outbound.mark_closing();
|
||||
}
|
||||
|
||||
pub fn send(&self, msg: Message) -> SendStatus {
|
||||
let was_closing = self.outbound.is_closing();
|
||||
let status = self.outbound.send(msg);
|
||||
if status == SendStatus::Congested && !was_closing {
|
||||
warn!(
|
||||
conn_id = self.id,
|
||||
node_id = %self.node_id,
|
||||
node_name = %self.node_name,
|
||||
queued_streams = self.stream_count.load(Ordering::Relaxed),
|
||||
"proxy outbound queue full, closing congested connection"
|
||||
);
|
||||
}
|
||||
status
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct LocalResponseHead {
|
||||
pub status: u16,
|
||||
pub headers: Vec<(String, String)>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum LocalBodyEvent {
|
||||
Chunk(Bytes),
|
||||
End,
|
||||
Error(String),
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct LocalWaitState {
|
||||
response: Option<LocalResponseHead>,
|
||||
error: Option<String>,
|
||||
}
|
||||
|
||||
pub struct LocalStream {
|
||||
pub id: u64,
|
||||
proxy_conn_id: u64,
|
||||
proxy_stream_id: u32,
|
||||
wait_state: Mutex<LocalWaitState>,
|
||||
headers_notify: Notify,
|
||||
body_tx: mpsc::Sender<LocalBodyEvent>,
|
||||
body_rx: Mutex<Option<mpsc::Receiver<LocalBodyEvent>>>,
|
||||
terminal: AtomicBool,
|
||||
}
|
||||
|
||||
impl LocalStream {
|
||||
fn new(id: u64, proxy_conn_id: u64, proxy_stream_id: u32) -> Self {
|
||||
let (body_tx, body_rx) = mpsc::channel(128);
|
||||
Self {
|
||||
id,
|
||||
proxy_conn_id,
|
||||
proxy_stream_id,
|
||||
wait_state: Mutex::new(LocalWaitState::default()),
|
||||
headers_notify: Notify::new(),
|
||||
body_tx,
|
||||
body_rx: Mutex::new(Some(body_rx)),
|
||||
terminal: AtomicBool::new(false),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn wait_headers(&self, timeout: Duration) -> Result<LocalResponseHead, String> {
|
||||
tokio::time::timeout(timeout, async {
|
||||
loop {
|
||||
let outcome = {
|
||||
let state = self.wait_state.lock();
|
||||
if let Some(response) = &state.response {
|
||||
return Ok(response.clone());
|
||||
}
|
||||
state.error.clone()
|
||||
};
|
||||
if let Some(error) = outcome {
|
||||
return Err(error);
|
||||
}
|
||||
self.headers_notify.notified().await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.map_err(|_| "timed out waiting for response headers".to_string())?
|
||||
}
|
||||
|
||||
pub fn take_body_receiver(&self) -> Option<mpsc::Receiver<LocalBodyEvent>> {
|
||||
self.body_rx.lock().take()
|
||||
}
|
||||
|
||||
fn set_response_headers(&self, meta: protocol::ResponseMeta) {
|
||||
let mut notify = false;
|
||||
{
|
||||
let mut state = self.wait_state.lock();
|
||||
if state.response.is_none() && state.error.is_none() {
|
||||
state.response = Some(LocalResponseHead {
|
||||
status: meta.status,
|
||||
headers: meta.headers,
|
||||
});
|
||||
notify = true;
|
||||
}
|
||||
}
|
||||
if notify {
|
||||
self.headers_notify.notify_waiters();
|
||||
}
|
||||
}
|
||||
|
||||
fn push_body_chunk(&self, payload: Bytes) -> bool {
|
||||
if self.terminal.load(Ordering::Acquire) {
|
||||
return false;
|
||||
}
|
||||
self.body_tx
|
||||
.try_send(LocalBodyEvent::Chunk(payload))
|
||||
.is_ok()
|
||||
}
|
||||
|
||||
fn finish(&self) {
|
||||
if self.terminal.swap(true, Ordering::AcqRel) {
|
||||
return;
|
||||
}
|
||||
let mut notify = false;
|
||||
{
|
||||
let mut state = self.wait_state.lock();
|
||||
if state.response.is_none() && state.error.is_none() {
|
||||
state.error = Some("stream ended before response headers".to_string());
|
||||
notify = true;
|
||||
}
|
||||
}
|
||||
if notify {
|
||||
self.headers_notify.notify_waiters();
|
||||
}
|
||||
let _ = self.body_tx.try_send(LocalBodyEvent::End);
|
||||
}
|
||||
|
||||
fn fail(&self, error: impl Into<String>) {
|
||||
if self.terminal.swap(true, Ordering::AcqRel) {
|
||||
return;
|
||||
}
|
||||
|
||||
let error = error.into();
|
||||
let mut notify = false;
|
||||
{
|
||||
let mut state = self.wait_state.lock();
|
||||
if state.response.is_none() && state.error.is_none() {
|
||||
state.error = Some(error.clone());
|
||||
notify = true;
|
||||
}
|
||||
}
|
||||
if notify {
|
||||
self.headers_notify.notify_waiters();
|
||||
}
|
||||
let _ = self.body_tx.try_send(LocalBodyEvent::Error(error));
|
||||
}
|
||||
}
|
||||
|
||||
pub struct HubRouter {
|
||||
proxy_conns: RwLock<HashMap<String, Vec<Arc<ProxyConn>>>>,
|
||||
proxy_conns_by_id: DashMap<u64, Arc<ProxyConn>>,
|
||||
local_streams: DashMap<u64, Arc<LocalStream>>,
|
||||
proxy_to_local: DashMap<(u64, u32), u64>,
|
||||
next_conn_id: AtomicU64,
|
||||
next_local_stream_id: AtomicU64,
|
||||
control_plane: ControlPlaneClient,
|
||||
}
|
||||
|
||||
impl HubRouter {
|
||||
pub fn new(control_plane: ControlPlaneClient) -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
proxy_conns: RwLock::new(HashMap::new()),
|
||||
proxy_conns_by_id: DashMap::new(),
|
||||
local_streams: DashMap::new(),
|
||||
proxy_to_local: DashMap::new(),
|
||||
next_conn_id: AtomicU64::new(1),
|
||||
next_local_stream_id: AtomicU64::new(1),
|
||||
control_plane,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn alloc_conn_id(&self) -> u64 {
|
||||
self.next_conn_id.fetch_add(1, Ordering::Relaxed)
|
||||
}
|
||||
|
||||
pub fn register_proxy(&self, conn: Arc<ProxyConn>) {
|
||||
let node_id = conn.node_id.clone();
|
||||
let node_name = conn.node_name.clone();
|
||||
let conn_id = conn.id;
|
||||
self.proxy_conns_by_id.insert(conn_id, conn.clone());
|
||||
|
||||
let pool_size = {
|
||||
let mut map = self.proxy_conns.write();
|
||||
map.entry(node_id.clone()).or_default().push(conn);
|
||||
map.get(&node_id).map(|v| v.len()).unwrap_or(0)
|
||||
};
|
||||
|
||||
info!(
|
||||
node_id = %node_id,
|
||||
node_name = %node_name,
|
||||
conn_id = conn_id,
|
||||
pool_size = pool_size,
|
||||
"proxy connected"
|
||||
);
|
||||
|
||||
self.notify_node_status(node_id, true, pool_size);
|
||||
}
|
||||
|
||||
pub fn unregister_proxy(&self, conn_id: u64, node_id: &str) {
|
||||
self.proxy_conns_by_id.remove(&conn_id);
|
||||
|
||||
let pool_size = {
|
||||
let mut map = self.proxy_conns.write();
|
||||
if let Some(conns) = map.get_mut(node_id) {
|
||||
conns.retain(|c| c.id != conn_id);
|
||||
if conns.is_empty() {
|
||||
map.remove(node_id);
|
||||
}
|
||||
}
|
||||
map.get(node_id).map(|v| v.len()).unwrap_or(0)
|
||||
};
|
||||
|
||||
info!(
|
||||
node_id = %node_id,
|
||||
conn_id = conn_id,
|
||||
remaining = pool_size,
|
||||
"proxy disconnected"
|
||||
);
|
||||
|
||||
self.cancel_streams_for_proxy(conn_id);
|
||||
self.notify_node_status(node_id.to_string(), pool_size > 0, pool_size);
|
||||
}
|
||||
|
||||
fn notify_node_status(&self, node_id: String, connected: bool, conn_count: usize) {
|
||||
let control_plane = self.control_plane.clone();
|
||||
tokio::spawn(async move {
|
||||
if let Err(error) = control_plane
|
||||
.push_node_status(&node_id, connected, conn_count)
|
||||
.await
|
||||
{
|
||||
warn!(
|
||||
node_id = %node_id,
|
||||
connected = connected,
|
||||
conn_count = conn_count,
|
||||
error = %error,
|
||||
"failed to push node status to app control plane"
|
||||
);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
fn get_proxy_conn(&self, node_id: &str) -> Option<Arc<ProxyConn>> {
|
||||
let map = self.proxy_conns.read();
|
||||
let conns = map.get(node_id)?;
|
||||
conns
|
||||
.iter()
|
||||
.filter(|c| c.is_available())
|
||||
.min_by_key(|c| c.stream_count.load(Ordering::Relaxed))
|
||||
.cloned()
|
||||
}
|
||||
|
||||
pub fn open_local_stream(
|
||||
&self,
|
||||
node_id: &str,
|
||||
meta: &protocol::RequestMeta,
|
||||
) -> Result<Arc<LocalStream>, String> {
|
||||
let proxy_conn = self
|
||||
.get_proxy_conn(node_id)
|
||||
.ok_or_else(|| format!("no proxy connection for node {node_id}"))?;
|
||||
let proxy_stream_id = proxy_conn
|
||||
.alloc_stream_id()
|
||||
.ok_or_else(|| format!("stream limit reached for node {node_id}"))?;
|
||||
|
||||
// Encode frames before registering the stream so that encoding failures
|
||||
// (practically impossible but theoretically possible) don't leak a stream
|
||||
// slot or orphan map entries.
|
||||
let meta_json = match serde_json::to_vec(meta) {
|
||||
Ok(json) => json,
|
||||
Err(e) => {
|
||||
proxy_conn.release_stream();
|
||||
return Err(format!("failed to encode request metadata: {e}"));
|
||||
}
|
||||
};
|
||||
let (meta_payload, meta_flags) = match protocol::compress_payload(&meta_json) {
|
||||
Ok(result) => result,
|
||||
Err(e) => {
|
||||
proxy_conn.release_stream();
|
||||
return Err(format!("failed to compress request metadata: {e}"));
|
||||
}
|
||||
};
|
||||
let header_frame = protocol::encode_frame(
|
||||
proxy_stream_id,
|
||||
protocol::REQUEST_HEADERS,
|
||||
meta_flags,
|
||||
&meta_payload,
|
||||
);
|
||||
|
||||
// Frames encoded successfully -- now register the stream.
|
||||
let local_stream_id = self.next_local_stream_id.fetch_add(1, Ordering::Relaxed);
|
||||
let local_stream = Arc::new(LocalStream::new(
|
||||
local_stream_id,
|
||||
proxy_conn.id,
|
||||
proxy_stream_id,
|
||||
));
|
||||
self.local_streams
|
||||
.insert(local_stream_id, local_stream.clone());
|
||||
self.proxy_to_local
|
||||
.insert((proxy_conn.id, proxy_stream_id), local_stream_id);
|
||||
|
||||
match proxy_conn.send(Message::Binary(header_frame.into())) {
|
||||
SendStatus::Queued => Ok(local_stream),
|
||||
SendStatus::Closed | SendStatus::Congested => {
|
||||
self.cleanup_local_stream(local_stream_id);
|
||||
proxy_conn.release_stream();
|
||||
Err("proxy connection congested".to_string())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn push_local_request_body(
|
||||
&self,
|
||||
local_stream_id: u64,
|
||||
payload: Bytes,
|
||||
end_stream: bool,
|
||||
) -> Result<(), String> {
|
||||
let stream = self
|
||||
.local_streams
|
||||
.get(&local_stream_id)
|
||||
.map(|entry| entry.value().clone())
|
||||
.ok_or_else(|| "local stream not found".to_string())?;
|
||||
let proxy_conn = self
|
||||
.proxy_conns_by_id
|
||||
.get(&stream.proxy_conn_id)
|
||||
.map(|entry| entry.value().clone())
|
||||
.ok_or_else(|| "proxy connection unavailable".to_string())?;
|
||||
|
||||
let total_chunks = payload.len().div_ceil(MAX_REQUEST_BODY_FRAME_SIZE);
|
||||
if total_chunks == 0 {
|
||||
if end_stream {
|
||||
self.send_request_body_frame(&proxy_conn, stream.proxy_stream_id, &[], true)?;
|
||||
}
|
||||
} else {
|
||||
for (index, chunk) in payload.chunks(MAX_REQUEST_BODY_FRAME_SIZE).enumerate() {
|
||||
let is_last_chunk = index + 1 == total_chunks;
|
||||
self.send_request_body_frame(
|
||||
&proxy_conn,
|
||||
stream.proxy_stream_id,
|
||||
chunk,
|
||||
end_stream && is_last_chunk,
|
||||
)?;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn send_request_body_frame(
|
||||
&self,
|
||||
proxy_conn: &Arc<ProxyConn>,
|
||||
proxy_stream_id: u32,
|
||||
payload: &[u8],
|
||||
end_stream: bool,
|
||||
) -> Result<(), String> {
|
||||
let (body_payload, body_flags) = protocol::compress_payload(payload)
|
||||
.map_err(|e| format!("failed to compress request body: {e}"))?;
|
||||
let body_frame = protocol::encode_frame(
|
||||
proxy_stream_id,
|
||||
protocol::REQUEST_BODY,
|
||||
body_flags
|
||||
| if end_stream {
|
||||
protocol::FLAG_END_STREAM
|
||||
} else {
|
||||
0
|
||||
},
|
||||
&body_payload,
|
||||
);
|
||||
match proxy_conn.send(Message::Binary(body_frame.into())) {
|
||||
SendStatus::Queued => Ok(()),
|
||||
SendStatus::Closed | SendStatus::Congested => {
|
||||
Err("proxy connection congested".to_string())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn cancel_local_stream(&self, local_stream_id: u64, reason: &str) {
|
||||
let Some((_, stream)) = self.local_streams.remove(&local_stream_id) else {
|
||||
return;
|
||||
};
|
||||
|
||||
self.proxy_to_local
|
||||
.remove(&(stream.proxy_conn_id, stream.proxy_stream_id));
|
||||
if let Some(pc) = self.proxy_conns_by_id.get(&stream.proxy_conn_id) {
|
||||
pc.release_stream();
|
||||
let frame = protocol::encode_stream_error(stream.proxy_stream_id, reason);
|
||||
let _ = pc.send(Message::Binary(frame.into()));
|
||||
}
|
||||
stream.fail(reason.to_string());
|
||||
}
|
||||
|
||||
fn cleanup_local_stream(&self, local_stream_id: u64) {
|
||||
let Some((_, stream)) = self.local_streams.remove(&local_stream_id) else {
|
||||
return;
|
||||
};
|
||||
self.proxy_to_local
|
||||
.remove(&(stream.proxy_conn_id, stream.proxy_stream_id));
|
||||
}
|
||||
|
||||
pub async fn handle_proxy_frame(&self, proxy_conn_id: u64, data: &mut [u8]) {
|
||||
let header = match protocol::FrameHeader::parse(data) {
|
||||
Some(h) => h,
|
||||
None => return,
|
||||
};
|
||||
let expected_len = protocol::HEADER_SIZE + header.payload_len as usize;
|
||||
if data.len() < expected_len {
|
||||
return;
|
||||
}
|
||||
|
||||
match header.msg_type {
|
||||
protocol::RESPONSE_HEADERS => {
|
||||
self.route_response_headers(proxy_conn_id, header, data);
|
||||
}
|
||||
protocol::RESPONSE_BODY => {
|
||||
self.route_response_body(proxy_conn_id, header, data);
|
||||
}
|
||||
protocol::STREAM_END => {
|
||||
self.finish_proxy_stream(proxy_conn_id, header.stream_id);
|
||||
}
|
||||
protocol::STREAM_ERROR => {
|
||||
let message = protocol::decode_payload(data, &header)
|
||||
.ok()
|
||||
.and_then(|payload| String::from_utf8(payload).ok())
|
||||
.unwrap_or_else(|| "stream error".to_string());
|
||||
self.fail_proxy_stream(proxy_conn_id, header.stream_id, message);
|
||||
}
|
||||
protocol::HEARTBEAT_DATA => {
|
||||
self.handle_heartbeat(proxy_conn_id, header.stream_id, data, &header)
|
||||
.await;
|
||||
}
|
||||
protocol::PING => {
|
||||
let payload = protocol::frame_payload_by_header(data, &header).unwrap_or(&[]);
|
||||
let pong = protocol::encode_pong(payload);
|
||||
if let Some(pc) = self.proxy_conns_by_id.get(&proxy_conn_id) {
|
||||
let _ = pc.send(Message::Binary(pong.into()));
|
||||
}
|
||||
}
|
||||
protocol::PONG => {}
|
||||
protocol::GOAWAY => {
|
||||
warn!(
|
||||
proxy_conn_id = proxy_conn_id,
|
||||
"received GOAWAY from proxy connection"
|
||||
);
|
||||
}
|
||||
_ => {
|
||||
debug!(
|
||||
msg_type = header.msg_type,
|
||||
proxy_conn_id = proxy_conn_id,
|
||||
"unexpected frame type from proxy"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn route_response_headers(
|
||||
&self,
|
||||
proxy_conn_id: u64,
|
||||
header: protocol::FrameHeader,
|
||||
data: &[u8],
|
||||
) {
|
||||
let Some(local_id) = self.lookup_local_stream(proxy_conn_id, header.stream_id) else {
|
||||
return;
|
||||
};
|
||||
let Ok(payload) = protocol::decode_payload(data, &header) else {
|
||||
self.fail_proxy_stream(
|
||||
proxy_conn_id,
|
||||
header.stream_id,
|
||||
"failed to decode response headers",
|
||||
);
|
||||
return;
|
||||
};
|
||||
let Ok(meta) = serde_json::from_slice::<protocol::ResponseMeta>(&payload) else {
|
||||
self.fail_proxy_stream(
|
||||
proxy_conn_id,
|
||||
header.stream_id,
|
||||
"invalid response headers payload",
|
||||
);
|
||||
return;
|
||||
};
|
||||
if let Some(entry) = self.local_streams.get(&local_id) {
|
||||
entry.value().set_response_headers(meta);
|
||||
}
|
||||
}
|
||||
|
||||
fn route_response_body(&self, proxy_conn_id: u64, header: protocol::FrameHeader, data: &[u8]) {
|
||||
let Some(local_id) = self.lookup_local_stream(proxy_conn_id, header.stream_id) else {
|
||||
return;
|
||||
};
|
||||
let Ok(payload) = protocol::decode_payload(data, &header) else {
|
||||
self.fail_proxy_stream(
|
||||
proxy_conn_id,
|
||||
header.stream_id,
|
||||
"failed to decode response body",
|
||||
);
|
||||
return;
|
||||
};
|
||||
|
||||
let stream = match self.local_streams.get(&local_id) {
|
||||
Some(entry) => entry.value().clone(),
|
||||
None => return,
|
||||
};
|
||||
|
||||
if !stream.push_body_chunk(Bytes::from(payload)) {
|
||||
self.cancel_local_stream(local_id, "local relay response congested");
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_stream_cleanup(
|
||||
&self,
|
||||
proxy_conn_id: u64,
|
||||
proxy_stream_id: u32,
|
||||
) -> Option<Arc<LocalStream>> {
|
||||
let local_id = self
|
||||
.proxy_to_local
|
||||
.remove(&(proxy_conn_id, proxy_stream_id))
|
||||
.map(|(_, local_id)| local_id)?;
|
||||
|
||||
let stream = self
|
||||
.local_streams
|
||||
.remove(&local_id)
|
||||
.map(|(_, stream)| stream)?;
|
||||
if let Some(pc) = self.proxy_conns_by_id.get(&proxy_conn_id) {
|
||||
pc.release_stream();
|
||||
}
|
||||
Some(stream)
|
||||
}
|
||||
|
||||
fn finish_proxy_stream(&self, proxy_conn_id: u64, proxy_stream_id: u32) {
|
||||
if let Some(stream) = self.handle_stream_cleanup(proxy_conn_id, proxy_stream_id) {
|
||||
stream.finish();
|
||||
}
|
||||
}
|
||||
|
||||
fn fail_proxy_stream(
|
||||
&self,
|
||||
proxy_conn_id: u64,
|
||||
proxy_stream_id: u32,
|
||||
error: impl Into<String>,
|
||||
) {
|
||||
if let Some(stream) = self.handle_stream_cleanup(proxy_conn_id, proxy_stream_id) {
|
||||
stream.fail(error.into());
|
||||
}
|
||||
}
|
||||
|
||||
fn lookup_local_stream(&self, proxy_conn_id: u64, proxy_stream_id: u32) -> Option<u64> {
|
||||
self.proxy_to_local
|
||||
.get(&(proxy_conn_id, proxy_stream_id))
|
||||
.map(|entry| *entry.value())
|
||||
}
|
||||
|
||||
async fn handle_heartbeat(
|
||||
&self,
|
||||
proxy_conn_id: u64,
|
||||
stream_id: u32,
|
||||
data: &[u8],
|
||||
header: &protocol::FrameHeader,
|
||||
) {
|
||||
let payload = match protocol::decode_payload(data, header) {
|
||||
Ok(payload) => payload,
|
||||
Err(error) => {
|
||||
warn!(proxy_conn_id = proxy_conn_id, error = %error, "failed to decode heartbeat payload");
|
||||
return;
|
||||
}
|
||||
};
|
||||
let ack_payload = match self.control_plane.heartbeat_ack(&payload).await {
|
||||
Ok(payload) => payload,
|
||||
Err(error) => {
|
||||
warn!(proxy_conn_id = proxy_conn_id, error = %error, "control-plane heartbeat callback failed");
|
||||
b"{}".to_vec()
|
||||
}
|
||||
};
|
||||
if let Some(pc) = self.proxy_conns_by_id.get(&proxy_conn_id) {
|
||||
let frame = protocol::encode_frame(stream_id, protocol::HEARTBEAT_ACK, 0, &ack_payload);
|
||||
let _ = pc.send(Message::Binary(frame.into()));
|
||||
}
|
||||
}
|
||||
|
||||
fn cancel_streams_for_proxy(&self, proxy_conn_id: u64) {
|
||||
let mut cancelled = 0usize;
|
||||
self.proxy_to_local.retain(|key, local_id| {
|
||||
if key.0 != proxy_conn_id {
|
||||
return true;
|
||||
}
|
||||
if let Some((_, stream)) = self.local_streams.remove(local_id) {
|
||||
stream.fail("proxy disconnected".to_string());
|
||||
}
|
||||
cancelled += 1;
|
||||
false
|
||||
});
|
||||
|
||||
if cancelled > 0 {
|
||||
warn!(
|
||||
proxy_conn_id = proxy_conn_id,
|
||||
streams_cancelled = cancelled,
|
||||
"cancelled in-flight streams due to proxy disconnect"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn stats(&self) -> HubStats {
|
||||
let proxy_conns = self.proxy_conns.read();
|
||||
let total_proxy = proxy_conns.values().map(|v| v.len()).sum();
|
||||
let nodes = proxy_conns.len();
|
||||
drop(proxy_conns);
|
||||
|
||||
HubStats {
|
||||
proxy_connections: total_proxy,
|
||||
nodes,
|
||||
active_streams: self.local_streams.len(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(serde::Serialize)]
|
||||
pub struct HubStats {
|
||||
pub proxy_connections: usize,
|
||||
pub nodes: usize,
|
||||
pub active_streams: usize,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn build_meta() -> protocol::RequestMeta {
|
||||
protocol::RequestMeta {
|
||||
method: "GET".to_string(),
|
||||
url: "https://example.com".to_string(),
|
||||
headers: HashMap::new(),
|
||||
timeout: 30,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cancel_local_stream_notifies_proxy() {
|
||||
let hub = HubRouter::new(ControlPlaneClient::disabled());
|
||||
|
||||
let (proxy_tx, mut proxy_rx) = mpsc::channel(8);
|
||||
let (proxy_close_tx, _) = watch::channel(false);
|
||||
let proxy = Arc::new(ProxyConn::new(
|
||||
100,
|
||||
"node-1".to_string(),
|
||||
"Node 1".to_string(),
|
||||
proxy_tx,
|
||||
proxy_close_tx,
|
||||
16,
|
||||
));
|
||||
hub.register_proxy(proxy);
|
||||
|
||||
let stream = hub
|
||||
.open_local_stream("node-1", &build_meta())
|
||||
.expect("open local stream");
|
||||
let _ = proxy_rx.try_recv().expect("headers frame");
|
||||
hub.push_local_request_body(stream.id, Bytes::new(), true)
|
||||
.expect("finish empty body");
|
||||
let _ = proxy_rx.try_recv().expect("body frame");
|
||||
|
||||
hub.cancel_local_stream(stream.id, "client dropped");
|
||||
|
||||
let cancelled = proxy_rx.try_recv().expect("cancel frame");
|
||||
let cancelled_data = match cancelled {
|
||||
Message::Binary(data) => data.to_vec(),
|
||||
other => panic!("unexpected message: {other:?}"),
|
||||
};
|
||||
let header = protocol::FrameHeader::parse(&cancelled_data).expect("cancel frame header");
|
||||
assert_eq!(header.msg_type, protocol::STREAM_ERROR);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn push_local_request_body_splits_large_payload_and_marks_end() {
|
||||
let hub = HubRouter::new(ControlPlaneClient::disabled());
|
||||
|
||||
let (proxy_tx, mut proxy_rx) = mpsc::channel(8);
|
||||
let (proxy_close_tx, _) = watch::channel(false);
|
||||
let proxy = Arc::new(ProxyConn::new(
|
||||
200,
|
||||
"node-2".to_string(),
|
||||
"Node 2".to_string(),
|
||||
proxy_tx,
|
||||
proxy_close_tx,
|
||||
16,
|
||||
));
|
||||
hub.register_proxy(proxy);
|
||||
|
||||
let stream = hub
|
||||
.open_local_stream("node-2", &build_meta())
|
||||
.expect("open local stream");
|
||||
let _ = proxy_rx.try_recv().expect("headers frame");
|
||||
|
||||
let payload = Bytes::from(vec![b'x'; MAX_REQUEST_BODY_FRAME_SIZE + 17]);
|
||||
hub.push_local_request_body(stream.id, payload, true)
|
||||
.expect("push request body");
|
||||
|
||||
let first = match proxy_rx.try_recv().expect("first body frame") {
|
||||
Message::Binary(data) => data.to_vec(),
|
||||
other => panic!("unexpected message: {other:?}"),
|
||||
};
|
||||
let first_header = protocol::FrameHeader::parse(&first).expect("first body header");
|
||||
assert_eq!(first_header.msg_type, protocol::REQUEST_BODY);
|
||||
assert_eq!(first_header.flags & protocol::FLAG_END_STREAM, 0);
|
||||
|
||||
let second = match proxy_rx.try_recv().expect("second body frame") {
|
||||
Message::Binary(data) => data.to_vec(),
|
||||
other => panic!("unexpected message: {other:?}"),
|
||||
};
|
||||
let second_header = protocol::FrameHeader::parse(&second).expect("second body header");
|
||||
assert_eq!(second_header.msg_type, protocol::REQUEST_BODY);
|
||||
assert_ne!(second_header.flags & protocol::FLAG_END_STREAM, 0);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,262 @@
|
||||
use std::io;
|
||||
use std::net::SocketAddr;
|
||||
use std::time::Duration;
|
||||
|
||||
use async_stream::stream;
|
||||
use axum::body::{Body, Bytes};
|
||||
use axum::extract::{ConnectInfo, Path, Request, State};
|
||||
use axum::http::{HeaderMap, HeaderName, HeaderValue, Response, StatusCode};
|
||||
use axum::response::IntoResponse;
|
||||
use bytes::BytesMut;
|
||||
use futures_util::StreamExt;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::hub::{LocalBodyEvent, LocalStream};
|
||||
use crate::protocol;
|
||||
use crate::AppState;
|
||||
|
||||
pub const TUNNEL_ERROR_HEADER: &str = "x-aether-tunnel-error";
|
||||
const MAX_RELAY_META_LEN: usize = 256 * 1024;
|
||||
|
||||
struct StreamGuard {
|
||||
hub: std::sync::Arc<crate::hub::HubRouter>,
|
||||
stream_id: u64,
|
||||
finished: bool,
|
||||
}
|
||||
|
||||
impl Drop for StreamGuard {
|
||||
fn drop(&mut self) {
|
||||
if !self.finished {
|
||||
self.hub
|
||||
.cancel_local_stream(self.stream_id, "local relay client dropped");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn relay_request(
|
||||
Path(node_id): Path<String>,
|
||||
State(state): State<AppState>,
|
||||
ConnectInfo(addr): ConnectInfo<SocketAddr>,
|
||||
request: Request,
|
||||
) -> impl IntoResponse {
|
||||
if !addr.ip().is_loopback() {
|
||||
return tunnel_error_response(
|
||||
StatusCode::FORBIDDEN,
|
||||
"forbidden",
|
||||
"local relay only accepts loopback requests",
|
||||
);
|
||||
}
|
||||
|
||||
let mut body_stream = request.into_body().into_data_stream();
|
||||
let mut envelope_buf = BytesMut::new();
|
||||
let mut meta: Option<protocol::RequestMeta> = None;
|
||||
let mut stream: Option<std::sync::Arc<LocalStream>> = None;
|
||||
|
||||
while let Some(chunk_result) = body_stream.next().await {
|
||||
let chunk = match chunk_result {
|
||||
Ok(chunk) => chunk,
|
||||
Err(error) => {
|
||||
if let Some(active_stream) = &stream {
|
||||
state
|
||||
.hub
|
||||
.cancel_local_stream(active_stream.id, "failed to read relay request body");
|
||||
}
|
||||
warn!(error = %error, "failed to read local relay request body");
|
||||
return tunnel_error_response(
|
||||
StatusCode::BAD_GATEWAY,
|
||||
"relay",
|
||||
"failed to read relay request body",
|
||||
);
|
||||
}
|
||||
};
|
||||
|
||||
if stream.is_none() {
|
||||
envelope_buf.extend_from_slice(&chunk);
|
||||
let Some((parsed_meta, body_offset)) = (match try_decode_envelope_meta(&envelope_buf) {
|
||||
Ok(result) => result,
|
||||
Err(error) => {
|
||||
return tunnel_error_response(StatusCode::BAD_REQUEST, "bad_request", &error);
|
||||
}
|
||||
}) else {
|
||||
continue;
|
||||
};
|
||||
|
||||
let opened_stream = match state.hub.open_local_stream(&node_id, &parsed_meta) {
|
||||
Ok(stream) => stream,
|
||||
Err(error) => {
|
||||
return tunnel_error_response(
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
"connect",
|
||||
&error,
|
||||
);
|
||||
}
|
||||
};
|
||||
|
||||
if envelope_buf.len() > body_offset {
|
||||
let first_body_chunk = Bytes::copy_from_slice(&envelope_buf[body_offset..]);
|
||||
if let Err(error) =
|
||||
state
|
||||
.hub
|
||||
.push_local_request_body(opened_stream.id, first_body_chunk, false)
|
||||
{
|
||||
state.hub.cancel_local_stream(opened_stream.id, &error);
|
||||
return tunnel_error_response(
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
"connect",
|
||||
&error,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
envelope_buf.clear();
|
||||
meta = Some(parsed_meta);
|
||||
stream = Some(opened_stream);
|
||||
continue;
|
||||
}
|
||||
|
||||
let Some(active_stream) = &stream else {
|
||||
continue;
|
||||
};
|
||||
if let Err(error) = state
|
||||
.hub
|
||||
.push_local_request_body(active_stream.id, chunk, false)
|
||||
{
|
||||
state.hub.cancel_local_stream(active_stream.id, &error);
|
||||
return tunnel_error_response(StatusCode::SERVICE_UNAVAILABLE, "connect", &error);
|
||||
}
|
||||
}
|
||||
|
||||
let (meta, stream) = match (meta, stream) {
|
||||
(Some(meta), Some(stream)) => (meta, stream),
|
||||
_ => {
|
||||
return tunnel_error_response(
|
||||
StatusCode::BAD_REQUEST,
|
||||
"bad_request",
|
||||
"relay envelope metadata truncated",
|
||||
);
|
||||
}
|
||||
};
|
||||
|
||||
if let Err(error) = state
|
||||
.hub
|
||||
.push_local_request_body(stream.id, Bytes::new(), true)
|
||||
{
|
||||
state.hub.cancel_local_stream(stream.id, &error);
|
||||
return tunnel_error_response(StatusCode::SERVICE_UNAVAILABLE, "connect", &error);
|
||||
}
|
||||
|
||||
let request_guard = StreamGuard {
|
||||
hub: state.hub.clone(),
|
||||
stream_id: stream.id,
|
||||
finished: false,
|
||||
};
|
||||
|
||||
let wait_timeout = Duration::from_secs(meta.timeout.clamp(5, 300));
|
||||
let response_head = match stream.wait_headers(wait_timeout).await {
|
||||
Ok(response) => response,
|
||||
Err(error) => {
|
||||
state.hub.cancel_local_stream(stream.id, &error);
|
||||
return tunnel_error_response(StatusCode::GATEWAY_TIMEOUT, "timeout", &error);
|
||||
}
|
||||
};
|
||||
|
||||
let Some(mut body_rx) = stream.take_body_receiver() else {
|
||||
state
|
||||
.hub
|
||||
.cancel_local_stream(stream.id, "missing relay response body receiver");
|
||||
return tunnel_error_response(
|
||||
StatusCode::BAD_GATEWAY,
|
||||
"relay",
|
||||
"missing relay response body receiver",
|
||||
);
|
||||
};
|
||||
|
||||
let hub = state.hub.clone();
|
||||
let stream_id = stream.id;
|
||||
let body_stream = stream! {
|
||||
let mut guard = request_guard;
|
||||
guard.hub = hub;
|
||||
guard.stream_id = stream_id;
|
||||
while let Some(event) = body_rx.recv().await {
|
||||
match event {
|
||||
LocalBodyEvent::Chunk(chunk) => yield Ok::<Bytes, io::Error>(chunk),
|
||||
LocalBodyEvent::End => {
|
||||
guard.finished = true;
|
||||
break;
|
||||
}
|
||||
LocalBodyEvent::Error(error) => {
|
||||
guard.finished = true;
|
||||
yield Err(io::Error::other(error));
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
guard.finished = true;
|
||||
};
|
||||
|
||||
let mut builder = Response::builder().status(response_head.status);
|
||||
if let Some(headers) = builder.headers_mut() {
|
||||
append_headers(headers, &response_head.headers);
|
||||
}
|
||||
match builder.body(Body::from_stream(body_stream)) {
|
||||
Ok(response) => response,
|
||||
Err(error) => {
|
||||
warn!(error = %error, "failed to build relay response");
|
||||
tunnel_error_response(
|
||||
StatusCode::BAD_GATEWAY,
|
||||
"relay",
|
||||
"failed to build relay response",
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn try_decode_envelope_meta(
|
||||
buffer: &BytesMut,
|
||||
) -> Result<Option<(protocol::RequestMeta, usize)>, String> {
|
||||
if buffer.len() < 4 {
|
||||
return Ok(None);
|
||||
}
|
||||
let meta_len = u32::from_be_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]) as usize;
|
||||
if meta_len > MAX_RELAY_META_LEN {
|
||||
return Err("relay metadata too large".to_string());
|
||||
}
|
||||
let meta_end = 4usize
|
||||
.checked_add(meta_len)
|
||||
.ok_or_else(|| "relay envelope length overflow".to_string())?;
|
||||
if buffer.len() < meta_end {
|
||||
return Ok(None);
|
||||
}
|
||||
let meta = serde_json::from_slice::<protocol::RequestMeta>(&buffer[4..meta_end])
|
||||
.map_err(|e| format!("invalid relay metadata: {e}"))?;
|
||||
Ok(Some((meta, meta_end)))
|
||||
}
|
||||
|
||||
fn append_headers(target: &mut HeaderMap, headers: &[(String, String)]) {
|
||||
for (name, value) in headers {
|
||||
let Ok(name) = HeaderName::from_bytes(name.as_bytes()) else {
|
||||
continue;
|
||||
};
|
||||
let Ok(value) = HeaderValue::from_str(value) else {
|
||||
continue;
|
||||
};
|
||||
target.append(name, value);
|
||||
}
|
||||
}
|
||||
|
||||
fn tunnel_error_response(status: StatusCode, kind: &str, message: &str) -> Response<Body> {
|
||||
let mut builder = Response::builder().status(status);
|
||||
if let Some(headers) = builder.headers_mut() {
|
||||
headers.insert(
|
||||
HeaderName::from_static(TUNNEL_ERROR_HEADER),
|
||||
HeaderValue::from_str(kind).unwrap_or_else(|_| HeaderValue::from_static("relay")),
|
||||
);
|
||||
headers.insert(
|
||||
axum::http::header::CONTENT_TYPE,
|
||||
HeaderValue::from_static("text/plain; charset=utf-8"),
|
||||
);
|
||||
}
|
||||
builder
|
||||
.body(Body::from(message.to_string()))
|
||||
.unwrap_or_else(|_| Response::new(Body::from("relay error")))
|
||||
}
|
||||
@@ -0,0 +1,167 @@
|
||||
mod control_plane;
|
||||
mod hub;
|
||||
mod local_relay;
|
||||
mod protocol;
|
||||
mod proxy_conn;
|
||||
|
||||
use std::net::SocketAddr;
|
||||
use std::time::Duration;
|
||||
|
||||
use axum::extract::ws::WebSocketUpgrade;
|
||||
use axum::extract::State;
|
||||
use axum::response::{IntoResponse, Json};
|
||||
use axum::routing::{get, post};
|
||||
use axum::Router;
|
||||
use clap::Parser;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use crate::control_plane::ControlPlaneClient;
|
||||
use crate::hub::{ConnConfig, HubRouter};
|
||||
use crate::local_relay::relay_request;
|
||||
|
||||
#[derive(Parser, Debug)]
|
||||
#[command(name = "aether-hub", about = "Tunnel Hub for Aether")]
|
||||
struct Args {
|
||||
/// Bind address
|
||||
#[arg(long, default_value = "0.0.0.0:8085", env = "TUNNEL_HUB_BIND")]
|
||||
bind: String,
|
||||
|
||||
/// Proxy-side idle timeout in seconds (0 to disable)
|
||||
#[arg(long, default_value_t = 0, env = "TUNNEL_HUB_PROXY_IDLE_TIMEOUT")]
|
||||
proxy_idle_timeout: u64,
|
||||
|
||||
/// Ping interval in seconds (for both sides)
|
||||
#[arg(long, default_value_t = 15, env = "TUNNEL_HUB_PING_INTERVAL")]
|
||||
ping_interval: u64,
|
||||
|
||||
/// Max concurrent streams per proxy connection
|
||||
#[arg(long, default_value_t = 2048, env = "TUNNEL_HUB_MAX_STREAMS")]
|
||||
max_streams: usize,
|
||||
|
||||
/// Per-connection outbound queue capacity before treating the socket as congested
|
||||
#[arg(
|
||||
long,
|
||||
default_value_t = 128,
|
||||
env = "TUNNEL_HUB_OUTBOUND_QUEUE_CAPACITY"
|
||||
)]
|
||||
outbound_queue_capacity: usize,
|
||||
|
||||
/// Local Aether app base URL for control-plane callbacks
|
||||
#[arg(
|
||||
long,
|
||||
default_value = "http://127.0.0.1:8084",
|
||||
env = "TUNNEL_HUB_APP_BASE_URL"
|
||||
)]
|
||||
app_base_url: String,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct AppState {
|
||||
pub hub: std::sync::Arc<HubRouter>,
|
||||
pub proxy_conn_cfg: ConnConfig,
|
||||
pub max_streams: usize,
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
// Initialize tracing
|
||||
tracing_subscriber::fmt()
|
||||
.with_env_filter(
|
||||
tracing_subscriber::EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| "aether_hub=info".into()),
|
||||
)
|
||||
.init();
|
||||
|
||||
let args = Args::parse();
|
||||
|
||||
let hub = HubRouter::new(ControlPlaneClient::new(args.app_base_url));
|
||||
let outbound_queue_capacity = args.outbound_queue_capacity.clamp(8, 4096);
|
||||
let ping_interval = Duration::from_secs(args.ping_interval);
|
||||
let state = AppState {
|
||||
hub,
|
||||
proxy_conn_cfg: ConnConfig {
|
||||
ping_interval,
|
||||
idle_timeout: Duration::from_secs(args.proxy_idle_timeout),
|
||||
outbound_queue_capacity,
|
||||
},
|
||||
max_streams: args.max_streams,
|
||||
};
|
||||
|
||||
let app = Router::new()
|
||||
.route("/health", get(health))
|
||||
.route("/stats", get(stats))
|
||||
.route("/proxy", get(ws_proxy))
|
||||
.route("/local/relay/{node_id}", post(relay_request))
|
||||
.with_state(state);
|
||||
|
||||
let listener = tokio::net::TcpListener::bind(&args.bind).await?;
|
||||
info!(bind = %args.bind, "aether-hub started");
|
||||
|
||||
axum::serve(
|
||||
listener,
|
||||
app.into_make_service_with_connect_info::<SocketAddr>(),
|
||||
)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// HTTP endpoints
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
async fn health() -> impl IntoResponse {
|
||||
Json(serde_json::json!({"status": "ok"}))
|
||||
}
|
||||
|
||||
async fn stats(State(state): State<AppState>) -> impl IntoResponse {
|
||||
Json(state.hub.stats())
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// WebSocket endpoints
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
async fn ws_proxy(
|
||||
ws: WebSocketUpgrade,
|
||||
State(state): State<AppState>,
|
||||
headers: axum::http::HeaderMap,
|
||||
) -> impl IntoResponse {
|
||||
let node_id = headers
|
||||
.get("x-node-id")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.unwrap_or("")
|
||||
.trim()
|
||||
.to_string();
|
||||
|
||||
let node_name = headers
|
||||
.get("x-node-name")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.unwrap_or(&node_id)
|
||||
.trim()
|
||||
.to_string();
|
||||
|
||||
let max_streams: usize = headers
|
||||
.get("x-tunnel-max-streams")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.and_then(|v| v.parse().ok())
|
||||
.unwrap_or(state.max_streams)
|
||||
.clamp(64, 2048);
|
||||
|
||||
if node_id.is_empty() {
|
||||
warn!("proxy connection rejected: missing X-Node-ID header");
|
||||
return axum::http::StatusCode::BAD_REQUEST.into_response();
|
||||
}
|
||||
|
||||
ws.max_frame_size(64 * 1024 * 1024)
|
||||
.on_upgrade(move |socket| {
|
||||
proxy_conn::handle_proxy_connection(
|
||||
socket,
|
||||
state.hub,
|
||||
node_id,
|
||||
node_name,
|
||||
max_streams,
|
||||
state.proxy_conn_cfg,
|
||||
)
|
||||
})
|
||||
.into_response()
|
||||
}
|
||||
@@ -0,0 +1,176 @@
|
||||
/// Tunnel binary frame protocol
|
||||
///
|
||||
/// Frame format (10-byte header + payload):
|
||||
/// | stream_id (4B) | msg_type (1B) | flags (1B) | payload_len (4B) | payload (NB) |
|
||||
use std::io::Read;
|
||||
|
||||
use flate2::read::GzDecoder;
|
||||
use flate2::write::GzEncoder;
|
||||
use flate2::Compression;
|
||||
|
||||
pub const HEADER_SIZE: usize = 10;
|
||||
|
||||
// Message types
|
||||
pub const REQUEST_HEADERS: u8 = 0x01;
|
||||
pub const REQUEST_BODY: u8 = 0x02;
|
||||
pub const RESPONSE_HEADERS: u8 = 0x03;
|
||||
pub const RESPONSE_BODY: u8 = 0x04;
|
||||
pub const STREAM_END: u8 = 0x05;
|
||||
pub const STREAM_ERROR: u8 = 0x06;
|
||||
pub const PING: u8 = 0x10;
|
||||
pub const PONG: u8 = 0x11;
|
||||
pub const GOAWAY: u8 = 0x12;
|
||||
pub const HEARTBEAT_DATA: u8 = 0x13;
|
||||
pub const HEARTBEAT_ACK: u8 = 0x14;
|
||||
// Flags
|
||||
pub const FLAG_END_STREAM: u8 = 0x01;
|
||||
pub const FLAG_GZIP_COMPRESSED: u8 = 0x02;
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct FrameHeader {
|
||||
pub stream_id: u32,
|
||||
pub msg_type: u8,
|
||||
pub flags: u8,
|
||||
pub payload_len: u32,
|
||||
}
|
||||
|
||||
impl FrameHeader {
|
||||
/// Parse frame header from raw bytes (must be >= HEADER_SIZE)
|
||||
#[inline]
|
||||
pub fn parse(data: &[u8]) -> Option<Self> {
|
||||
if data.len() < HEADER_SIZE {
|
||||
return None;
|
||||
}
|
||||
Some(Self {
|
||||
stream_id: u32::from_be_bytes([data[0], data[1], data[2], data[3]]),
|
||||
msg_type: data[4],
|
||||
flags: data[5],
|
||||
payload_len: u32::from_be_bytes([data[6], data[7], data[8], data[9]]),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
pub struct RequestMeta {
|
||||
pub method: String,
|
||||
pub url: String,
|
||||
pub headers: std::collections::HashMap<String, String>,
|
||||
#[serde(default = "default_timeout", deserialize_with = "deserialize_timeout")]
|
||||
pub timeout: u64,
|
||||
}
|
||||
|
||||
fn default_timeout() -> u64 {
|
||||
60
|
||||
}
|
||||
|
||||
fn deserialize_timeout<'de, D>(deserializer: D) -> Result<u64, D::Error>
|
||||
where
|
||||
D: serde::Deserializer<'de>,
|
||||
{
|
||||
#[derive(serde::Deserialize)]
|
||||
#[serde(untagged)]
|
||||
enum TimeoutValue {
|
||||
Int(u64),
|
||||
Float(f64),
|
||||
}
|
||||
|
||||
match <TimeoutValue as serde::Deserialize>::deserialize(deserializer)? {
|
||||
TimeoutValue::Int(v) => Ok(v),
|
||||
TimeoutValue::Float(v) => {
|
||||
if !v.is_finite() || v < 0.0 {
|
||||
return Err(serde::de::Error::custom(
|
||||
"timeout must be a non-negative finite number",
|
||||
));
|
||||
}
|
||||
if v.fract() != 0.0 {
|
||||
return Err(serde::de::Error::custom("timeout must be integer seconds"));
|
||||
}
|
||||
if v > (u64::MAX as f64) {
|
||||
return Err(serde::de::Error::custom("timeout is too large"));
|
||||
}
|
||||
Ok(v as u64)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
pub struct ResponseMeta {
|
||||
pub status: u16,
|
||||
pub headers: Vec<(String, String)>,
|
||||
}
|
||||
|
||||
pub fn encode_frame(stream_id: u32, msg_type: u8, flags: u8, payload: &[u8]) -> Vec<u8> {
|
||||
let mut buf = Vec::with_capacity(HEADER_SIZE + payload.len());
|
||||
buf.extend_from_slice(&stream_id.to_be_bytes());
|
||||
buf.push(msg_type);
|
||||
buf.push(flags);
|
||||
buf.extend_from_slice(&(payload.len() as u32).to_be_bytes());
|
||||
buf.extend_from_slice(payload);
|
||||
buf
|
||||
}
|
||||
|
||||
/// Encode a STREAM_ERROR frame for a given stream_id with an error message
|
||||
pub fn encode_stream_error(stream_id: u32, msg: &str) -> Vec<u8> {
|
||||
encode_frame(stream_id, STREAM_ERROR, 0, msg.as_bytes())
|
||||
}
|
||||
|
||||
/// Encode a PING frame (stream_id=0)
|
||||
pub fn encode_ping() -> Vec<u8> {
|
||||
encode_frame(0, PING, 0, &[])
|
||||
}
|
||||
|
||||
/// Encode a PONG frame (stream_id=0, echo payload)
|
||||
pub fn encode_pong(payload: &[u8]) -> Vec<u8> {
|
||||
encode_frame(0, PONG, 0, payload)
|
||||
}
|
||||
|
||||
/// Encode a GOAWAY frame (stream_id=0)
|
||||
pub fn encode_goaway() -> Vec<u8> {
|
||||
encode_frame(0, GOAWAY, 0, &[])
|
||||
}
|
||||
|
||||
#[inline]
|
||||
pub fn frame_payload_by_header<'a>(data: &'a [u8], header: &FrameHeader) -> Option<&'a [u8]> {
|
||||
let payload_len = header.payload_len as usize;
|
||||
let end = HEADER_SIZE.checked_add(payload_len)?;
|
||||
if data.len() < end {
|
||||
return None;
|
||||
}
|
||||
Some(&data[HEADER_SIZE..end])
|
||||
}
|
||||
|
||||
pub fn decode_payload(data: &[u8], header: &FrameHeader) -> Result<Vec<u8>, String> {
|
||||
let payload = frame_payload_by_header(data, header)
|
||||
.ok_or_else(|| "incomplete frame payload".to_string())?;
|
||||
if header.flags & FLAG_GZIP_COMPRESSED != 0 {
|
||||
let mut decoder = GzDecoder::new(payload);
|
||||
let mut decoded = Vec::new();
|
||||
decoder
|
||||
.read_to_end(&mut decoded)
|
||||
.map_err(|e| format!("failed to decompress payload: {e}"))?;
|
||||
Ok(decoded)
|
||||
} else {
|
||||
Ok(payload.to_vec())
|
||||
}
|
||||
}
|
||||
|
||||
pub fn compress_payload(payload: &[u8]) -> Result<(Vec<u8>, u8), std::io::Error> {
|
||||
maybe_recompress_payload(payload, true)
|
||||
}
|
||||
|
||||
fn maybe_recompress_payload(
|
||||
payload: &[u8],
|
||||
prefer_gzip: bool,
|
||||
) -> Result<(Vec<u8>, u8), std::io::Error> {
|
||||
if !prefer_gzip {
|
||||
return Ok((payload.to_vec(), 0));
|
||||
}
|
||||
let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
|
||||
std::io::Write::write_all(&mut encoder, payload)?;
|
||||
let compressed = encoder.finish()?;
|
||||
if compressed.len() < payload.len() {
|
||||
Ok((compressed, FLAG_GZIP_COMPRESSED))
|
||||
} else {
|
||||
Ok((payload.to_vec(), 0))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,156 @@
|
||||
/// Proxy-side WebSocket connection handler
|
||||
///
|
||||
/// Handles the lifecycle of a single aether-proxy connection:
|
||||
/// accept -> authenticate (headers) -> read loop -> cleanup
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use axum::extract::ws::{Message, WebSocket};
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use tokio::sync::{mpsc, watch};
|
||||
use tracing::{debug, info, warn};
|
||||
|
||||
use crate::hub::{ConnConfig, HubRouter, ProxyConn, SendStatus};
|
||||
use crate::protocol;
|
||||
|
||||
/// Maximum single frame size: 64 MB
|
||||
const MAX_FRAME_SIZE: usize = 64 * 1024 * 1024;
|
||||
|
||||
pub async fn handle_proxy_connection(
|
||||
ws: WebSocket,
|
||||
hub: Arc<HubRouter>,
|
||||
node_id: String,
|
||||
node_name: String,
|
||||
max_streams: usize,
|
||||
cfg: ConnConfig,
|
||||
) {
|
||||
let conn_id = hub.alloc_conn_id();
|
||||
let (mut ws_tx, ws_rx) = ws.split();
|
||||
|
||||
let (tx, mut rx) = mpsc::channel::<Message>(cfg.outbound_queue_capacity);
|
||||
let (close_tx, mut close_rx) = watch::channel(false);
|
||||
|
||||
let conn = Arc::new(ProxyConn::new(
|
||||
conn_id,
|
||||
node_id.clone(),
|
||||
node_name.clone(),
|
||||
tx,
|
||||
close_tx,
|
||||
max_streams,
|
||||
));
|
||||
|
||||
hub.register_proxy(conn.clone());
|
||||
|
||||
let writer = tokio::spawn(async move {
|
||||
loop {
|
||||
tokio::select! {
|
||||
msg = rx.recv() => match msg {
|
||||
Some(msg) => {
|
||||
if ws_tx.send(msg).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
None => break,
|
||||
},
|
||||
changed = close_rx.changed() => {
|
||||
if changed.is_err() || *close_rx.borrow() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
let _ = ws_tx.close().await;
|
||||
});
|
||||
|
||||
let ping_conn = conn.clone();
|
||||
let ping_interval = cfg.ping_interval;
|
||||
let ping_task = tokio::spawn(async move {
|
||||
loop {
|
||||
tokio::time::sleep(ping_interval).await;
|
||||
let ping = protocol::encode_ping();
|
||||
if !matches!(
|
||||
ping_conn.send(Message::Binary(ping.into())),
|
||||
SendStatus::Queued
|
||||
) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let reader_hub = hub.clone();
|
||||
let reader_conn = conn.clone();
|
||||
let reader = tokio::spawn(async move {
|
||||
run_proxy_reader(ws_rx, reader_hub, reader_conn, cfg.idle_timeout).await;
|
||||
});
|
||||
|
||||
let _ = reader.await;
|
||||
ping_task.abort();
|
||||
conn.request_close();
|
||||
hub.unregister_proxy(conn_id, &node_id);
|
||||
drop(conn);
|
||||
tokio::time::sleep(Duration::from_millis(100)).await;
|
||||
writer.abort();
|
||||
let _ = writer.await;
|
||||
}
|
||||
|
||||
async fn run_proxy_reader(
|
||||
mut ws_rx: futures_util::stream::SplitStream<WebSocket>,
|
||||
hub: Arc<HubRouter>,
|
||||
conn: Arc<ProxyConn>,
|
||||
idle_timeout: Duration,
|
||||
) {
|
||||
let idle_enabled = !idle_timeout.is_zero();
|
||||
let mut oversized_count = 0u32;
|
||||
loop {
|
||||
let msg = if idle_enabled {
|
||||
tokio::select! {
|
||||
msg = ws_rx.next() => msg,
|
||||
_ = tokio::time::sleep(idle_timeout) => {
|
||||
warn!(conn_id = conn.id, node_id = %conn.node_id, "proxy idle timeout");
|
||||
let _ = conn.send(Message::Binary(protocol::encode_goaway().into()));
|
||||
conn.request_close();
|
||||
break;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
ws_rx.next().await
|
||||
};
|
||||
|
||||
match msg {
|
||||
Some(Ok(Message::Binary(data))) => {
|
||||
let mut data = data.to_vec();
|
||||
if data.len() > MAX_FRAME_SIZE {
|
||||
oversized_count += 1;
|
||||
warn!(
|
||||
conn_id = conn.id,
|
||||
size = data.len(),
|
||||
"oversized frame from proxy"
|
||||
);
|
||||
if oversized_count >= 5 {
|
||||
warn!(conn_id = conn.id, "too many oversized frames, closing");
|
||||
conn.request_close();
|
||||
break;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
oversized_count = 0;
|
||||
|
||||
if data.len() < protocol::HEADER_SIZE {
|
||||
debug!(conn_id = conn.id, "frame too small, skipping");
|
||||
continue;
|
||||
}
|
||||
|
||||
hub.handle_proxy_frame(conn.id, &mut data).await;
|
||||
}
|
||||
Some(Ok(Message::Close(_))) | None => {
|
||||
info!(conn_id = conn.id, node_id = %conn.node_id, "proxy WebSocket closed");
|
||||
break;
|
||||
}
|
||||
Some(Err(e)) => {
|
||||
warn!(conn_id = conn.id, error = %e, "proxy WebSocket error");
|
||||
break;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
Generated
+6
-1
@@ -10,7 +10,7 @@ checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa"
|
||||
|
||||
[[package]]
|
||||
name = "aether-proxy"
|
||||
version = "0.2.0"
|
||||
version = "0.2.4"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"arc-swap",
|
||||
@@ -21,6 +21,9 @@ dependencies = [
|
||||
"flate2",
|
||||
"futures-util",
|
||||
"hex",
|
||||
"http-body-util",
|
||||
"hyper",
|
||||
"hyper-util",
|
||||
"libc",
|
||||
"ratatui",
|
||||
"reqwest",
|
||||
@@ -33,8 +36,10 @@ dependencies = [
|
||||
"tar",
|
||||
"thiserror 2.0.18",
|
||||
"tokio",
|
||||
"tokio-rustls",
|
||||
"tokio-tungstenite",
|
||||
"toml",
|
||||
"tower-service",
|
||||
"tracing",
|
||||
"tracing-subscriber",
|
||||
"url",
|
||||
|
||||
@@ -1,13 +1,17 @@
|
||||
[package]
|
||||
name = "aether-proxy"
|
||||
version = "0.2.0"
|
||||
version = "0.2.5"
|
||||
edition = "2021"
|
||||
description = "Tunnel proxy for Aether"
|
||||
|
||||
[dependencies]
|
||||
tokio = { version = "1", features = ["full"] }
|
||||
reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls", "stream", "http2"] }
|
||||
hyper = { version = "1", features = ["client", "http1", "http2"] }
|
||||
hyper-util = { version = "0.1", features = ["client", "client-legacy", "http1", "http2", "tokio"] }
|
||||
http-body-util = "0.1"
|
||||
tokio-tungstenite = { version = "0.24", features = ["rustls-tls-webpki-roots"] }
|
||||
tokio-rustls = "0.26"
|
||||
futures-util = "0.3"
|
||||
base64 = "0.22"
|
||||
clap = { version = "4", features = ["derive", "env"] }
|
||||
@@ -31,6 +35,7 @@ libc = "0.2"
|
||||
flate2 = "1"
|
||||
tar = "0.4"
|
||||
socket2 = { version = "0.5", features = ["all"] }
|
||||
tower-service = "0.3"
|
||||
webpki-roots = "0.26"
|
||||
|
||||
[profile.release]
|
||||
|
||||
@@ -19,11 +19,11 @@ docker compose up -d
|
||||
<!-- DOWNLOAD_TABLE_START -->
|
||||
| Platform | Download |
|
||||
|----------|----------|
|
||||
| Linux x86_64 | [aether-proxy-linux-amd64.tar.gz](https://github.com/fawney19/Aether/releases/download/proxy-v0.2.0/aether-proxy-linux-amd64.tar.gz) |
|
||||
| Linux ARM64 | [aether-proxy-linux-arm64.tar.gz](https://github.com/fawney19/Aether/releases/download/proxy-v0.2.0/aether-proxy-linux-arm64.tar.gz) |
|
||||
| macOS x86_64 | [aether-proxy-macos-amd64.tar.gz](https://github.com/fawney19/Aether/releases/download/proxy-v0.2.0/aether-proxy-macos-amd64.tar.gz) |
|
||||
| macOS ARM64 | [aether-proxy-macos-arm64.tar.gz](https://github.com/fawney19/Aether/releases/download/proxy-v0.2.0/aether-proxy-macos-arm64.tar.gz) |
|
||||
| Windows x86_64 | [aether-proxy-windows-amd64.zip](https://github.com/fawney19/Aether/releases/download/proxy-v0.2.0/aether-proxy-windows-amd64.zip) |
|
||||
| Linux x86_64 | [aether-proxy-linux-amd64.tar.gz](https://github.com/fawney19/Aether/releases/download/proxy-v0.2.5/aether-proxy-linux-amd64.tar.gz) |
|
||||
| Linux ARM64 | [aether-proxy-linux-arm64.tar.gz](https://github.com/fawney19/Aether/releases/download/proxy-v0.2.5/aether-proxy-linux-arm64.tar.gz) |
|
||||
| macOS x86_64 | [aether-proxy-macos-amd64.tar.gz](https://github.com/fawney19/Aether/releases/download/proxy-v0.2.5/aether-proxy-macos-amd64.tar.gz) |
|
||||
| macOS ARM64 | [aether-proxy-macos-arm64.tar.gz](https://github.com/fawney19/Aether/releases/download/proxy-v0.2.5/aether-proxy-macos-arm64.tar.gz) |
|
||||
| Windows x86_64 | [aether-proxy-windows-amd64.zip](https://github.com/fawney19/Aether/releases/download/proxy-v0.2.5/aether-proxy-windows-amd64.zip) |
|
||||
<!-- DOWNLOAD_TABLE_END -->
|
||||
|
||||
## 快速开始
|
||||
|
||||
+6
-23
@@ -13,8 +13,8 @@ use crate::config::{Config, ServerEntry};
|
||||
use crate::net;
|
||||
use crate::registration::client::AetherClient;
|
||||
use crate::runtime::{self, DynamicConfig};
|
||||
use crate::safe_dns::SafeDnsResolver;
|
||||
use crate::state::{AppState, ProxyMetrics, ServerContext};
|
||||
use crate::upstream_client;
|
||||
use crate::{hardware, target_filter, tunnel};
|
||||
|
||||
/// Run the full application lifecycle after config has been parsed.
|
||||
@@ -67,27 +67,10 @@ pub async fn run(mut config: Config, servers: Vec<ServerEntry>) -> anyhow::Resul
|
||||
config.dns_cache_capacity,
|
||||
));
|
||||
|
||||
// Build reqwest client for tunnel upstream requests (shared).
|
||||
// Inject SafeDnsResolver so reqwest only connects to addresses that were
|
||||
// validated by validate_target() — this eliminates the DNS rebinding
|
||||
// TOCTTOU gap where a second DNS lookup could return a private IP.
|
||||
let safe_resolver = SafeDnsResolver::new(Arc::clone(&dns_cache));
|
||||
let mut reqwest_builder = reqwest::Client::builder()
|
||||
.dns_resolver(Arc::new(safe_resolver))
|
||||
.pool_max_idle_per_host(config.upstream_pool_max_idle_per_host)
|
||||
.pool_idle_timeout(Duration::from_secs(config.upstream_pool_idle_timeout_secs))
|
||||
.connect_timeout(Duration::from_secs(config.upstream_connect_timeout_secs))
|
||||
.tcp_nodelay(config.upstream_tcp_nodelay);
|
||||
|
||||
if config.upstream_tcp_keepalive_secs > 0 {
|
||||
reqwest_builder = reqwest_builder.tcp_keepalive(Some(Duration::from_secs(
|
||||
config.upstream_tcp_keepalive_secs,
|
||||
)));
|
||||
}
|
||||
|
||||
let reqwest_client = reqwest_builder
|
||||
.build()
|
||||
.expect("failed to build reqwest client");
|
||||
// Build Hyper client for tunnel upstream requests (shared).
|
||||
// DNS still flows through validated addresses from DnsCache, while the
|
||||
// custom connector exposes per-request connect/TLS timing when available.
|
||||
let upstream_client = upstream_client::build_upstream_client(&config, Arc::clone(&dns_cache));
|
||||
|
||||
// Register with each Aether server and build per-server contexts.
|
||||
// Wrapped in Arc<Mutex> so retry_failed_registrations can append later.
|
||||
@@ -160,7 +143,7 @@ pub async fn run(mut config: Config, servers: Vec<ServerEntry>) -> anyhow::Resul
|
||||
let state = Arc::new(AppState {
|
||||
config: Arc::new(config),
|
||||
dns_cache,
|
||||
reqwest_client,
|
||||
upstream_client,
|
||||
tunnel_tls_config,
|
||||
});
|
||||
|
||||
|
||||
@@ -202,21 +202,19 @@ pub struct Config {
|
||||
#[arg(long, env = "AETHER_PROXY_LOG_JSON", default_value_t = false)]
|
||||
pub log_json: bool,
|
||||
|
||||
/// Deprecated: reconnect now uses a fixed 1s delay. Kept for config compatibility.
|
||||
/// Tunnel reconnect base delay in milliseconds (used by exponential backoff)
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_PROXY_TUNNEL_RECONNECT_BASE_MS",
|
||||
default_value_t = 500,
|
||||
hide = true
|
||||
default_value_t = 500
|
||||
)]
|
||||
pub tunnel_reconnect_base_ms: u64,
|
||||
|
||||
/// Deprecated: reconnect now uses a fixed 1s delay. Kept for config compatibility.
|
||||
/// Tunnel reconnect max delay in milliseconds (cap for exponential backoff)
|
||||
#[arg(
|
||||
long,
|
||||
env = "AETHER_PROXY_TUNNEL_RECONNECT_MAX_MS",
|
||||
default_value_t = 30000,
|
||||
hide = true
|
||||
default_value_t = 30000
|
||||
)]
|
||||
pub tunnel_reconnect_max_ms: u64,
|
||||
|
||||
|
||||
@@ -4,11 +4,11 @@ mod hardware;
|
||||
mod net;
|
||||
mod registration;
|
||||
mod runtime;
|
||||
mod safe_dns;
|
||||
mod setup;
|
||||
mod state;
|
||||
mod target_filter;
|
||||
mod tunnel;
|
||||
mod upstream_client;
|
||||
|
||||
use std::path::PathBuf;
|
||||
|
||||
@@ -136,8 +136,8 @@ async fn run_proxy(config: Config) -> anyhow::Result<()> {
|
||||
// Skip this check when we ARE the systemd service (INVOCATION_ID is set by systemd).
|
||||
if std::env::var_os("INVOCATION_ID").is_none() && setup::service::is_service_active() {
|
||||
eprintln!("Warning: systemd service is already running.");
|
||||
eprintln!("Use `aether-proxy stop` to stop it first, or manage via subcommands:");
|
||||
eprintln!(" aether-proxy status / logs / restart / stop");
|
||||
eprintln!("Use `./aether-proxy stop` to stop it first, or manage via subcommands:");
|
||||
eprintln!(" ./aether-proxy status / logs / restart / stop");
|
||||
std::process::exit(1);
|
||||
}
|
||||
|
||||
|
||||
@@ -20,6 +20,8 @@ struct RegisterRequest {
|
||||
hardware_info: Option<serde_json::Value>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
estimated_max_concurrency: Option<u64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
proxy_metadata: Option<serde_json::Value>,
|
||||
tunnel_mode: bool,
|
||||
}
|
||||
|
||||
@@ -107,6 +109,9 @@ impl AetherClient {
|
||||
heartbeat_interval: config.heartbeat_interval,
|
||||
hardware_info: hw.and_then(|h| serde_json::to_value(h).ok()),
|
||||
estimated_max_concurrency: hw.map(|h| h.estimated_max_concurrency),
|
||||
proxy_metadata: Some(serde_json::json!({
|
||||
"version": env!("CARGO_PKG_VERSION"),
|
||||
})),
|
||||
tunnel_mode: true,
|
||||
};
|
||||
|
||||
|
||||
@@ -21,7 +21,7 @@ pub fn install_service(config_path: &Path) -> anyhow::Result<()> {
|
||||
anyhow::bail!("systemd not available");
|
||||
}
|
||||
if !is_root() {
|
||||
anyhow::bail!("root required, use: sudo aether-proxy setup");
|
||||
anyhow::bail!("root required, use: sudo ./aether-proxy setup");
|
||||
}
|
||||
|
||||
let exe_path = std::env::current_exe()?.canonicalize()?;
|
||||
@@ -94,11 +94,11 @@ pub fn install_service(config_path: &Path) -> anyhow::Result<()> {
|
||||
|
||||
eprintln!();
|
||||
eprintln!(" Commands:");
|
||||
eprintln!(" aether-proxy status # service status");
|
||||
eprintln!(" aether-proxy logs # tail logs");
|
||||
eprintln!(" sudo aether-proxy restart # restart");
|
||||
eprintln!(" sudo aether-proxy stop # stop");
|
||||
eprintln!(" sudo aether-proxy uninstall # remove service");
|
||||
eprintln!(" ./aether-proxy status # service status");
|
||||
eprintln!(" ./aether-proxy logs # tail logs");
|
||||
eprintln!(" sudo ./aether-proxy restart # restart");
|
||||
eprintln!(" sudo ./aether-proxy stop # stop");
|
||||
eprintln!(" sudo ./aether-proxy uninstall # remove service");
|
||||
eprintln!();
|
||||
|
||||
Ok(())
|
||||
@@ -166,7 +166,7 @@ pub fn is_service_active() -> bool {
|
||||
|
||||
fn ensure_service_installed() -> anyhow::Result<()> {
|
||||
if !std::path::Path::new(UNIT_PATH).exists() {
|
||||
anyhow::bail!("service not installed, run `sudo aether-proxy setup` first");
|
||||
anyhow::bail!("service not installed, run `sudo ./aether-proxy setup` first");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
@@ -174,7 +174,7 @@ fn ensure_service_installed() -> anyhow::Result<()> {
|
||||
fn ensure_root_and_service() -> anyhow::Result<()> {
|
||||
ensure_service_installed()?;
|
||||
if !is_root() {
|
||||
anyhow::bail!("root required, use: sudo aether-proxy <command>");
|
||||
anyhow::bail!("root required, use: sudo ./aether-proxy <command>");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -281,8 +281,17 @@ fn atomic_replace(new_binary: &Path) -> anyhow::Result<PathBuf> {
|
||||
|
||||
// ── Public entry point ───────────────────────────────────────────────────────
|
||||
|
||||
/// `aether-proxy upgrade [version]` -- self-upgrade from GitHub releases.
|
||||
pub async fn cmd_upgrade(version: Option<String>) -> anyhow::Result<()> {
|
||||
#[derive(Clone, Copy)]
|
||||
enum RestartMode {
|
||||
BestEffort,
|
||||
Required,
|
||||
}
|
||||
|
||||
async fn execute_upgrade(
|
||||
version: Option<&str>,
|
||||
require_root: bool,
|
||||
restart_mode: RestartMode,
|
||||
) -> anyhow::Result<()> {
|
||||
// Resolve exe path once; reuse throughout the function
|
||||
let current_exe = std::env::current_exe()?.canonicalize()?;
|
||||
let exe_dir = current_exe
|
||||
@@ -290,8 +299,12 @@ pub async fn cmd_upgrade(version: Option<String>) -> anyhow::Result<()> {
|
||||
.ok_or_else(|| anyhow::anyhow!("cannot determine binary directory"))?;
|
||||
let temp_path = exe_dir.join(".aether-proxy.upgrade.tmp");
|
||||
|
||||
// Check write permission to binary directory
|
||||
if require_root {
|
||||
if !super::service::is_root() {
|
||||
anyhow::bail!("automatic upgrade requires root privileges");
|
||||
}
|
||||
} else if !super::service::is_root() {
|
||||
// Check write permission to binary directory for manual upgrade mode.
|
||||
let test_path = exe_dir.join(".aether-proxy.write-test");
|
||||
match std::fs::File::create(&test_path) {
|
||||
Ok(_) => {
|
||||
@@ -311,7 +324,7 @@ pub async fn cmd_upgrade(version: Option<String>) -> anyhow::Result<()> {
|
||||
eprintln!(" Current version: {}", CURRENT_VERSION);
|
||||
|
||||
let client = build_github_client()?;
|
||||
let release = fetch_release(&client, version.as_deref()).await?;
|
||||
let release = fetch_release(&client, version).await?;
|
||||
let target_tag = &release.tag_name;
|
||||
let target_semver = target_tag.strip_prefix("proxy-v").unwrap_or(target_tag);
|
||||
|
||||
@@ -341,6 +354,8 @@ pub async fn cmd_upgrade(version: Option<String>) -> anyhow::Result<()> {
|
||||
}
|
||||
};
|
||||
|
||||
match restart_mode {
|
||||
RestartMode::BestEffort => {
|
||||
// Restart systemd service if running.
|
||||
// Use best-effort: binary is already replaced, so a restart failure should
|
||||
// not abort the whole upgrade -- the user can restart manually.
|
||||
@@ -362,6 +377,16 @@ pub async fn cmd_upgrade(version: Option<String>) -> anyhow::Result<()> {
|
||||
} else {
|
||||
eprintln!(" No active systemd service detected, skipping restart.");
|
||||
}
|
||||
}
|
||||
RestartMode::Required => {
|
||||
if !super::service::is_root() {
|
||||
anyhow::bail!("automatic upgrade requires root privileges");
|
||||
}
|
||||
eprintln!(" Restarting systemd service...");
|
||||
super::service::run_cmd("systemctl", &["restart", "aether-proxy"])?;
|
||||
eprintln!(" Service restarted.");
|
||||
}
|
||||
}
|
||||
|
||||
eprintln!();
|
||||
eprintln!(" Upgrade complete!");
|
||||
@@ -371,3 +396,16 @@ pub async fn cmd_upgrade(version: Option<String>) -> anyhow::Result<()> {
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// `aether-proxy upgrade [version]` -- self-upgrade from GitHub releases.
|
||||
pub async fn cmd_upgrade(version: Option<String>) -> anyhow::Result<()> {
|
||||
execute_upgrade(version.as_deref(), false, RestartMode::BestEffort).await
|
||||
}
|
||||
|
||||
/// Perform automatic upgrade to a specific version.
|
||||
///
|
||||
/// This path is designed for server-pushed upgrades in systemd/root scenarios:
|
||||
/// it requires root and requires a successful `systemctl restart aether-proxy`.
|
||||
pub async fn perform_upgrade(version: &str) -> anyhow::Result<()> {
|
||||
execute_upgrade(Some(version), true, RestartMode::Required).await
|
||||
}
|
||||
|
||||
@@ -8,14 +8,15 @@ use crate::config::Config;
|
||||
use crate::registration::client::AetherClient;
|
||||
use crate::runtime::SharedDynamicConfig;
|
||||
use crate::target_filter::DnsCache;
|
||||
use crate::upstream_client::UpstreamClient;
|
||||
|
||||
/// Central application state shared across all servers/tunnels.
|
||||
pub struct AppState {
|
||||
pub config: Arc<Config>,
|
||||
/// DNS cache for upstream target resolution (shared).
|
||||
pub dns_cache: Arc<DnsCache>,
|
||||
/// Reqwest client for tunnel upstream requests (shared).
|
||||
pub reqwest_client: reqwest::Client,
|
||||
/// Hyper client for tunnel upstream requests with validated DNS and connection timing.
|
||||
pub upstream_client: UpstreamClient,
|
||||
/// Shared TLS config for tunnel WebSocket connections (avoids re-parsing root CAs on each reconnect).
|
||||
pub tunnel_tls_config: Arc<rustls::ClientConfig>,
|
||||
}
|
||||
@@ -47,6 +48,8 @@ pub struct ServerContext {
|
||||
/// Aggregate metrics for reporting to Aether.
|
||||
pub struct ProxyMetrics {
|
||||
pub total_requests: AtomicU64,
|
||||
/// Cumulative connection-establishment latency in nanoseconds
|
||||
/// (DNS + TCP/TLS + TTFB, excludes response body streaming).
|
||||
pub total_latency_ns: AtomicU64,
|
||||
pub failed_requests: AtomicU64,
|
||||
pub dns_failures: AtomicU64,
|
||||
@@ -64,8 +67,10 @@ impl ProxyMetrics {
|
||||
}
|
||||
}
|
||||
|
||||
pub fn record_request(&self, elapsed: Duration) {
|
||||
let nanos = u64::try_from(elapsed.as_nanos()).unwrap_or(u64::MAX);
|
||||
/// Record a completed request with its connection-establishment latency
|
||||
/// (DNS + TCP/TLS + TTFB, excludes response body streaming).
|
||||
pub fn record_request(&self, connect_elapsed: Duration) {
|
||||
let nanos = u64::try_from(connect_elapsed.as_nanos()).unwrap_or(u64::MAX);
|
||||
self.total_requests.fetch_add(1, Ordering::Release);
|
||||
self.total_latency_ns.fetch_add(nanos, Ordering::Release);
|
||||
}
|
||||
|
||||
@@ -14,7 +14,7 @@ use tracing::{debug, error, info, warn};
|
||||
use crate::state::{AppState, ServerContext};
|
||||
|
||||
use super::heartbeat::HeartbeatHandle;
|
||||
use super::protocol::{Frame, MsgType, RequestMeta};
|
||||
use super::protocol::{decompress_if_gzip, Frame, MsgType, RequestMeta};
|
||||
use super::stream_handler;
|
||||
use super::writer::FrameSender;
|
||||
|
||||
@@ -92,8 +92,15 @@ where
|
||||
|
||||
match frame.msg_type {
|
||||
MsgType::RequestHeaders => {
|
||||
// Parse request metadata
|
||||
let meta: RequestMeta = match serde_json::from_slice(&frame.payload) {
|
||||
// Decompress if the frame is gzip-compressed, then parse metadata
|
||||
let payload = match decompress_if_gzip(&frame) {
|
||||
Ok(p) => p,
|
||||
Err(e) => {
|
||||
warn!(stream_id = frame.stream_id, error = %e, "frame decompress failed");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let meta: RequestMeta = match serde_json::from_slice(&payload) {
|
||||
Ok(m) => m,
|
||||
Err(e) => {
|
||||
warn!(stream_id = frame.stream_id, error = %e, "invalid request metadata");
|
||||
@@ -175,7 +182,9 @@ where
|
||||
|
||||
MsgType::StreamEnd | MsgType::StreamError => {
|
||||
// Client-side cancellation or end
|
||||
streams.remove(&frame.stream_id);
|
||||
if let Some(tx) = streams.remove(&frame.stream_id) {
|
||||
let _ = tx.send(frame).await;
|
||||
}
|
||||
}
|
||||
|
||||
MsgType::Ping => {
|
||||
|
||||
@@ -1,12 +1,14 @@
|
||||
//! Tunnel heartbeat: sends metrics over the tunnel, processes ACKs.
|
||||
|
||||
use std::sync::atomic::Ordering;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use std::time::SystemTime;
|
||||
use std::time::UNIX_EPOCH;
|
||||
|
||||
use bytes::Bytes;
|
||||
use tokio::sync::watch;
|
||||
use tracing::{debug, warn};
|
||||
use tracing::{debug, info, warn};
|
||||
|
||||
use crate::config::Config;
|
||||
use crate::registration::client::RemoteConfig;
|
||||
@@ -16,6 +18,18 @@ use crate::state::ServerContext;
|
||||
use super::protocol::{Frame, MsgType};
|
||||
use super::writer::FrameSender;
|
||||
|
||||
const CURRENT_VERSION: &str = env!("CARGO_PKG_VERSION");
|
||||
static UPGRADE_IN_PROGRESS: AtomicBool = AtomicBool::new(false);
|
||||
static NON_ROOT_UPGRADE_WARNED: AtomicBool = AtomicBool::new(false);
|
||||
|
||||
enum AckDecision {
|
||||
Accept {
|
||||
heartbeat_id: Option<u64>,
|
||||
upgrade_to: Option<String>,
|
||||
},
|
||||
Ignore,
|
||||
}
|
||||
|
||||
/// Handle for the dispatcher to forward HeartbeatAck frames.
|
||||
#[derive(Clone)]
|
||||
pub struct HeartbeatHandle {
|
||||
@@ -37,6 +51,15 @@ pub fn spawn_noop() -> HeartbeatHandle {
|
||||
HeartbeatHandle { ack_tx }
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default)]
|
||||
struct HeartbeatSnapshot {
|
||||
requests: u64,
|
||||
latency_ns: u64,
|
||||
failed: u64,
|
||||
dns_failures: u64,
|
||||
stream_errors: u64,
|
||||
}
|
||||
|
||||
/// Spawn the heartbeat task. Returns a handle for forwarding ACKs.
|
||||
pub fn spawn(
|
||||
_config: Arc<Config>,
|
||||
@@ -50,6 +73,19 @@ pub fn spawn(
|
||||
// Read initial interval from dynamic config (may be updated by remote config).
|
||||
let initial_interval = Duration::from_secs(server.dynamic.load().heartbeat_interval);
|
||||
let mut current_interval = initial_interval;
|
||||
// At most one in-flight heartbeat snapshot is tracked at a time.
|
||||
// Snapshot is only cleared after receiving an ACK, which avoids losing
|
||||
// interval counters when ACK/frame delivery is temporarily unstable.
|
||||
let mut pending: Option<(u64, HeartbeatSnapshot)> = None;
|
||||
let mut next_heartbeat_id: u64 = 1;
|
||||
let heartbeat_session_id = format!(
|
||||
"{}-{}",
|
||||
std::process::id(),
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_nanos()
|
||||
);
|
||||
|
||||
// Skip first immediate tick by sleeping first.
|
||||
tokio::time::sleep(current_interval).await;
|
||||
@@ -57,9 +93,30 @@ pub fn spawn(
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = tokio::time::sleep(current_interval) => {
|
||||
let payload = build_heartbeat_payload(&server);
|
||||
let (heartbeat_id, snapshot) = if let Some((id, snap)) = pending {
|
||||
(id, snap)
|
||||
} else {
|
||||
let snap = collect_snapshot(&server);
|
||||
let id = next_heartbeat_id;
|
||||
next_heartbeat_id = next_heartbeat_id.wrapping_add(1);
|
||||
if next_heartbeat_id == 0 {
|
||||
next_heartbeat_id = 1;
|
||||
}
|
||||
pending = Some((id, snap));
|
||||
(id, snap)
|
||||
};
|
||||
|
||||
let payload = build_heartbeat_payload(
|
||||
&server,
|
||||
&heartbeat_session_id,
|
||||
heartbeat_id,
|
||||
snapshot
|
||||
);
|
||||
let frame = Frame::control(MsgType::HeartbeatData, payload);
|
||||
if frame_tx.send(frame).await.is_err() {
|
||||
if let Some((_, snap)) = pending.take() {
|
||||
restore_snapshot(&server, snap);
|
||||
}
|
||||
break; // Writer closed
|
||||
}
|
||||
debug!("sent heartbeat data");
|
||||
@@ -79,10 +136,34 @@ pub fn spawn(
|
||||
}
|
||||
}
|
||||
Some(ack_payload) = ack_rx.recv() => {
|
||||
handle_ack(&server, &ack_payload);
|
||||
match handle_ack(&server, &ack_payload) {
|
||||
AckDecision::Accept {
|
||||
heartbeat_id: ack_id,
|
||||
upgrade_to,
|
||||
} => {
|
||||
if let Some((pending_id, _)) = pending {
|
||||
match ack_id {
|
||||
Some(id) if id == pending_id => {
|
||||
pending = None;
|
||||
}
|
||||
None => {
|
||||
// Backward-compatible with servers that don't echo
|
||||
// heartbeat_id in ACK payload yet.
|
||||
pending = None;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
maybe_trigger_upgrade(upgrade_to);
|
||||
}
|
||||
AckDecision::Ignore => {}
|
||||
}
|
||||
}
|
||||
_ = shutdown.changed() => {
|
||||
debug!("heartbeat task shutting down");
|
||||
if let Some((_, snap)) = pending.take() {
|
||||
restore_snapshot(&server, snap);
|
||||
}
|
||||
break;
|
||||
}
|
||||
}
|
||||
@@ -92,36 +173,87 @@ pub fn spawn(
|
||||
HeartbeatHandle { ack_tx }
|
||||
}
|
||||
|
||||
fn build_heartbeat_payload(server: &ServerContext) -> Bytes {
|
||||
fn collect_snapshot(server: &ServerContext) -> HeartbeatSnapshot {
|
||||
HeartbeatSnapshot {
|
||||
requests: server.metrics.total_requests.swap(0, Ordering::AcqRel),
|
||||
latency_ns: server.metrics.total_latency_ns.swap(0, Ordering::AcqRel),
|
||||
failed: server.metrics.failed_requests.swap(0, Ordering::AcqRel),
|
||||
dns_failures: server.metrics.dns_failures.swap(0, Ordering::AcqRel),
|
||||
stream_errors: server.metrics.stream_errors.swap(0, Ordering::AcqRel),
|
||||
}
|
||||
}
|
||||
|
||||
fn restore_snapshot(server: &ServerContext, snap: HeartbeatSnapshot) {
|
||||
if snap.requests > 0 {
|
||||
server
|
||||
.metrics
|
||||
.total_requests
|
||||
.fetch_add(snap.requests, Ordering::Release);
|
||||
}
|
||||
if snap.latency_ns > 0 {
|
||||
server
|
||||
.metrics
|
||||
.total_latency_ns
|
||||
.fetch_add(snap.latency_ns, Ordering::Release);
|
||||
}
|
||||
if snap.failed > 0 {
|
||||
server
|
||||
.metrics
|
||||
.failed_requests
|
||||
.fetch_add(snap.failed, Ordering::Release);
|
||||
}
|
||||
if snap.dns_failures > 0 {
|
||||
server
|
||||
.metrics
|
||||
.dns_failures
|
||||
.fetch_add(snap.dns_failures, Ordering::Release);
|
||||
}
|
||||
if snap.stream_errors > 0 {
|
||||
server
|
||||
.metrics
|
||||
.stream_errors
|
||||
.fetch_add(snap.stream_errors, Ordering::Release);
|
||||
}
|
||||
}
|
||||
|
||||
fn build_heartbeat_payload(
|
||||
server: &ServerContext,
|
||||
heartbeat_session_id: &str,
|
||||
heartbeat_id: u64,
|
||||
snapshot: HeartbeatSnapshot,
|
||||
) -> Bytes {
|
||||
let node_id = server.node_id.read().unwrap().clone();
|
||||
|
||||
let interval_requests = server.metrics.total_requests.swap(0, Ordering::AcqRel);
|
||||
let interval_latency_ns = server.metrics.total_latency_ns.swap(0, Ordering::AcqRel);
|
||||
let interval_failed = server.metrics.failed_requests.swap(0, Ordering::AcqRel);
|
||||
let interval_dns_failures = server.metrics.dns_failures.swap(0, Ordering::AcqRel);
|
||||
let interval_stream_errors = server.metrics.stream_errors.swap(0, Ordering::AcqRel);
|
||||
let avg_latency_ms = if interval_requests > 0 {
|
||||
Some(interval_latency_ns as f64 / interval_requests as f64 / 1_000_000.0)
|
||||
let avg_latency_ms = if snapshot.requests > 0 {
|
||||
Some(snapshot.latency_ns as f64 / snapshot.requests as f64 / 1_000_000.0)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let payload = serde_json::json!({
|
||||
"node_id": node_id,
|
||||
"heartbeat_session_id": heartbeat_session_id,
|
||||
"heartbeat_id": heartbeat_id,
|
||||
"active_connections": server.active_connections.load(Ordering::Acquire),
|
||||
"total_requests": interval_requests,
|
||||
"total_requests": snapshot.requests,
|
||||
"avg_latency_ms": avg_latency_ms,
|
||||
"failed_requests": interval_failed,
|
||||
"dns_failures": interval_dns_failures,
|
||||
"stream_errors": interval_stream_errors,
|
||||
"failed_requests": snapshot.failed,
|
||||
"dns_failures": snapshot.dns_failures,
|
||||
"stream_errors": snapshot.stream_errors,
|
||||
"proxy_metadata": {
|
||||
"version": CURRENT_VERSION,
|
||||
},
|
||||
});
|
||||
|
||||
Bytes::from(serde_json::to_vec(&payload).unwrap_or_default())
|
||||
}
|
||||
|
||||
fn handle_ack(server: &ServerContext, payload: &[u8]) {
|
||||
fn handle_ack(server: &ServerContext, payload: &[u8]) -> AckDecision {
|
||||
if payload.is_empty() {
|
||||
return;
|
||||
return AckDecision::Accept {
|
||||
heartbeat_id: None,
|
||||
upgrade_to: None,
|
||||
};
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
@@ -130,6 +262,10 @@ fn handle_ack(server: &ServerContext, payload: &[u8]) {
|
||||
remote_config: Option<RemoteConfig>,
|
||||
#[serde(default)]
|
||||
config_version: u64,
|
||||
#[serde(default)]
|
||||
heartbeat_id: Option<u64>,
|
||||
#[serde(default)]
|
||||
upgrade_to: Option<String>,
|
||||
}
|
||||
|
||||
match serde_json::from_slice::<AckPayload>(payload) {
|
||||
@@ -137,9 +273,68 @@ fn handle_ack(server: &ServerContext, payload: &[u8]) {
|
||||
if let Some(ref rc) = ack.remote_config {
|
||||
runtime::apply_remote_config(&server.dynamic, rc, ack.config_version);
|
||||
}
|
||||
AckDecision::Accept {
|
||||
heartbeat_id: ack.heartbeat_id,
|
||||
upgrade_to: ack.upgrade_to.and_then(normalize_upgrade_target),
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
warn!(error = %e, "failed to parse heartbeat ACK");
|
||||
AckDecision::Ignore
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn normalize_upgrade_target(raw: String) -> Option<String> {
|
||||
let trimmed = raw.trim();
|
||||
if trimmed.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let normalized = trimmed.strip_prefix("proxy-v").unwrap_or(trimmed);
|
||||
if normalized == CURRENT_VERSION {
|
||||
return None;
|
||||
}
|
||||
Some(normalized.to_string())
|
||||
}
|
||||
|
||||
fn maybe_trigger_upgrade(version: Option<String>) {
|
||||
let Some(target_version) = version else {
|
||||
return;
|
||||
};
|
||||
if !crate::setup::service::is_root() {
|
||||
if NON_ROOT_UPGRADE_WARNED
|
||||
.compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
|
||||
.is_ok()
|
||||
{
|
||||
warn!(
|
||||
target_version = %target_version,
|
||||
"remote upgrade skipped: root privileges are required"
|
||||
);
|
||||
}
|
||||
return;
|
||||
}
|
||||
if UPGRADE_IN_PROGRESS
|
||||
.compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
|
||||
.is_err()
|
||||
{
|
||||
debug!(target_version = %target_version, "upgrade already in progress, ignoring");
|
||||
return;
|
||||
}
|
||||
|
||||
tokio::spawn(async move {
|
||||
info!(target_version = %target_version, "received remote upgrade instruction");
|
||||
match crate::setup::upgrade::perform_upgrade(&target_version).await {
|
||||
Ok(()) => {
|
||||
info!(target_version = %target_version, "remote upgrade finished");
|
||||
}
|
||||
Err(e) => {
|
||||
warn!(
|
||||
target_version = %target_version,
|
||||
error = %e,
|
||||
"remote upgrade failed"
|
||||
);
|
||||
UPGRADE_IN_PROGRESS.store(false, Ordering::Release);
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
@@ -6,16 +6,26 @@ pub mod stream_handler;
|
||||
pub mod writer;
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use tokio::sync::watch;
|
||||
use tracing::{error, info};
|
||||
|
||||
use crate::state::{AppState, ServerContext};
|
||||
|
||||
/// Fixed reconnect delay -- short enough for fast recovery, long enough to
|
||||
/// avoid CPU spin when the network is completely down.
|
||||
const RECONNECT_DELAY: Duration = Duration::from_secs(1);
|
||||
/// If a tunnel stays connected at least this long, treat the next disconnect
|
||||
/// as a non-failure and reset reconnect backoff.
|
||||
const STABLE_SESSION_RESET_AFTER: Duration = Duration::from_secs(30);
|
||||
/// Startup staggering step per secondary connection, used to avoid
|
||||
/// simultaneous bursts when a pool of tunnels starts together.
|
||||
const STARTUP_STAGGER_STEP_MS: u64 = 150;
|
||||
/// Upper bound for startup staggering.
|
||||
const MAX_STARTUP_STAGGER_MS: u64 = 1_500;
|
||||
/// Keep a tiny floor for repeated reconnects; first retry is still immediate.
|
||||
const MIN_RECONNECT_DELAY_MS: u64 = 50;
|
||||
/// Even under sustained failures, keep probing frequently so recovery is fast
|
||||
/// once cross-border network quality improves.
|
||||
const RECONNECT_PROBE_MAX_DELAY_MS: u64 = 3_000;
|
||||
|
||||
/// Run the tunnel mode main loop (connect, dispatch, reconnect).
|
||||
///
|
||||
@@ -28,8 +38,29 @@ pub async fn run(
|
||||
mut shutdown: watch::Receiver<bool>,
|
||||
) {
|
||||
info!(server = %server.server_label, conn = conn_idx, "starting tunnel");
|
||||
let reconnect_salt = compute_connection_salt(server, conn_idx);
|
||||
|
||||
let startup_delay = compute_startup_stagger(conn_idx, reconnect_salt);
|
||||
if !startup_delay.is_zero() {
|
||||
info!(
|
||||
server = %server.server_label,
|
||||
conn = conn_idx,
|
||||
delay_ms = startup_delay.as_millis(),
|
||||
"startup stagger before first connect"
|
||||
);
|
||||
tokio::select! {
|
||||
_ = tokio::time::sleep(startup_delay) => {}
|
||||
_ = shutdown.changed() => {
|
||||
info!(server = %server.server_label, conn = conn_idx, "shutdown requested during startup stagger");
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mut consecutive_failures: u32 = 0;
|
||||
|
||||
loop {
|
||||
let started_at = Instant::now();
|
||||
match client::connect_and_run(state, server, conn_idx, &mut shutdown).await {
|
||||
Ok(client::TunnelOutcome::Shutdown) => {
|
||||
info!(server = %server.server_label, conn = conn_idx, "tunnel shut down gracefully");
|
||||
@@ -48,8 +79,31 @@ pub async fn run(
|
||||
return;
|
||||
}
|
||||
|
||||
// Reset backoff after a stable session to keep recovery snappy when
|
||||
// failures are only occasional.
|
||||
let connected_for = started_at.elapsed();
|
||||
if connected_for >= STABLE_SESSION_RESET_AFTER {
|
||||
consecutive_failures = 0;
|
||||
} else {
|
||||
consecutive_failures = consecutive_failures.saturating_add(1);
|
||||
}
|
||||
|
||||
let reconnect_delay = compute_reconnect_delay(
|
||||
state.config.tunnel_reconnect_base_ms,
|
||||
state.config.tunnel_reconnect_max_ms,
|
||||
consecutive_failures,
|
||||
reconnect_salt,
|
||||
);
|
||||
info!(
|
||||
server = %server.server_label,
|
||||
conn = conn_idx,
|
||||
failures = consecutive_failures,
|
||||
delay_ms = reconnect_delay.as_millis(),
|
||||
"waiting before reconnect"
|
||||
);
|
||||
|
||||
tokio::select! {
|
||||
_ = tokio::time::sleep(RECONNECT_DELAY) => {}
|
||||
_ = tokio::time::sleep(reconnect_delay) => {}
|
||||
_ = shutdown.changed() => {
|
||||
info!(server = %server.server_label, conn = conn_idx, "shutdown requested during reconnect wait");
|
||||
return;
|
||||
@@ -57,3 +111,126 @@ pub async fn run(
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn compute_connection_salt(server: &ServerContext, conn_idx: usize) -> u64 {
|
||||
// FNV-1a style hash over server label + connection index.
|
||||
let mut h: u64 = 0xcbf29ce484222325;
|
||||
for &b in server.server_label.as_bytes() {
|
||||
h ^= b as u64;
|
||||
h = h.wrapping_mul(0x100000001b3);
|
||||
}
|
||||
h ^= conn_idx as u64;
|
||||
mix_u64(h)
|
||||
}
|
||||
|
||||
fn compute_startup_stagger(conn_idx: usize, salt: u64) -> Duration {
|
||||
if conn_idx == 0 {
|
||||
return Duration::ZERO;
|
||||
}
|
||||
let base = (conn_idx as u64).saturating_mul(STARTUP_STAGGER_STEP_MS);
|
||||
let jitter = mix_u64(salt) % 301; // 0..=300ms
|
||||
Duration::from_millis((base + jitter).min(MAX_STARTUP_STAGGER_MS))
|
||||
}
|
||||
|
||||
fn compute_reconnect_delay(
|
||||
base_ms: u64,
|
||||
max_ms: u64,
|
||||
consecutive_failures: u32,
|
||||
salt: u64,
|
||||
) -> Duration {
|
||||
// First retry should be immediate to maximize recovery speed on transient
|
||||
// blips (the user's primary expectation in poor networks).
|
||||
if consecutive_failures <= 1 {
|
||||
return Duration::ZERO;
|
||||
}
|
||||
|
||||
// Keep a sane minimum for repeated failures.
|
||||
let base_ms = base_ms.max(MIN_RECONNECT_DELAY_MS);
|
||||
let max_ms = max_ms.max(base_ms);
|
||||
let cap_ms = compute_reconnect_cap_ms(base_ms, max_ms, consecutive_failures)
|
||||
.min(RECONNECT_PROBE_MAX_DELAY_MS.max(base_ms));
|
||||
|
||||
// Equal-jitter: randomize in [cap/2, cap], preventing synchronized reconnect
|
||||
// storms while keeping reconnect latency bounded.
|
||||
if cap_ms <= 1 {
|
||||
return Duration::from_millis(cap_ms);
|
||||
}
|
||||
|
||||
let half = cap_ms / 2;
|
||||
let span = cap_ms - half;
|
||||
let now_nanos = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.subsec_nanos() as u64)
|
||||
.unwrap_or(0);
|
||||
let mixed = mix_u64(now_nanos ^ salt);
|
||||
let jitter = if span == 0 { 0 } else { mixed % (span + 1) };
|
||||
Duration::from_millis(half + jitter)
|
||||
}
|
||||
|
||||
fn compute_reconnect_cap_ms(base_ms: u64, max_ms: u64, consecutive_failures: u32) -> u64 {
|
||||
if consecutive_failures <= 1 {
|
||||
return base_ms.min(max_ms);
|
||||
}
|
||||
|
||||
let shift = (consecutive_failures - 1).min(31);
|
||||
let factor = 1u64 << shift;
|
||||
base_ms.saturating_mul(factor).min(max_ms)
|
||||
}
|
||||
|
||||
fn mix_u64(mut x: u64) -> u64 {
|
||||
// SplitMix64 finalizer - cheap bit mixing for pseudo-random jitter.
|
||||
x ^= x >> 30;
|
||||
x = x.wrapping_mul(0xbf58476d1ce4e5b9);
|
||||
x ^= x >> 27;
|
||||
x = x.wrapping_mul(0x94d049bb133111eb);
|
||||
x ^ (x >> 31)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::time::Duration;
|
||||
|
||||
use super::{
|
||||
compute_reconnect_cap_ms, compute_reconnect_delay, compute_startup_stagger,
|
||||
MAX_STARTUP_STAGGER_MS, RECONNECT_PROBE_MAX_DELAY_MS, STARTUP_STAGGER_STEP_MS,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn reconnect_cap_grows_exponentially_and_caps() {
|
||||
let base = 500;
|
||||
let max = 30_000;
|
||||
assert_eq!(compute_reconnect_cap_ms(base, max, 0), 500);
|
||||
assert_eq!(compute_reconnect_cap_ms(base, max, 1), 500);
|
||||
assert_eq!(compute_reconnect_cap_ms(base, max, 2), 1_000);
|
||||
assert_eq!(compute_reconnect_cap_ms(base, max, 3), 2_000);
|
||||
assert_eq!(compute_reconnect_cap_ms(base, max, 4), 4_000);
|
||||
assert_eq!(compute_reconnect_cap_ms(base, max, 5), 8_000);
|
||||
assert_eq!(compute_reconnect_cap_ms(base, max, 6), 16_000);
|
||||
assert_eq!(compute_reconnect_cap_ms(base, max, 7), 30_000);
|
||||
assert_eq!(compute_reconnect_cap_ms(base, max, 20), 30_000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn startup_stagger_is_zero_for_primary_and_bounded_for_secondary() {
|
||||
assert_eq!(compute_startup_stagger(0, 42), Duration::ZERO);
|
||||
|
||||
let d1 = compute_startup_stagger(1, 42);
|
||||
let d2 = compute_startup_stagger(2, 42);
|
||||
|
||||
assert!(d1 >= Duration::from_millis(STARTUP_STAGGER_STEP_MS));
|
||||
assert!(d1 <= Duration::from_millis(MAX_STARTUP_STAGGER_MS));
|
||||
assert!(d2 >= Duration::from_millis(STARTUP_STAGGER_STEP_MS * 2));
|
||||
assert!(d2 <= Duration::from_millis(MAX_STARTUP_STAGGER_MS));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reconnect_delay_is_immediate_on_first_failure() {
|
||||
assert_eq!(compute_reconnect_delay(700, 45_000, 1, 123), Duration::ZERO);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reconnect_delay_stays_within_probe_ceiling_after_many_failures() {
|
||||
let d = compute_reconnect_delay(500, 45_000, 100, 12345);
|
||||
assert!(d <= Duration::from_millis(RECONNECT_PROBE_MAX_DELAY_MS));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -144,7 +144,7 @@ pub struct RequestMeta {
|
||||
pub method: String,
|
||||
pub url: String,
|
||||
pub headers: std::collections::HashMap<String, String>,
|
||||
#[serde(default = "default_timeout")]
|
||||
#[serde(default = "default_timeout", deserialize_with = "deserialize_timeout")]
|
||||
pub timeout: u64,
|
||||
}
|
||||
|
||||
@@ -152,6 +152,36 @@ fn default_timeout() -> u64 {
|
||||
60
|
||||
}
|
||||
|
||||
fn deserialize_timeout<'de, D>(deserializer: D) -> Result<u64, D::Error>
|
||||
where
|
||||
D: serde::Deserializer<'de>,
|
||||
{
|
||||
#[derive(serde::Deserialize)]
|
||||
#[serde(untagged)]
|
||||
enum TimeoutValue {
|
||||
Int(u64),
|
||||
Float(f64),
|
||||
}
|
||||
|
||||
match <TimeoutValue as serde::Deserialize>::deserialize(deserializer)? {
|
||||
TimeoutValue::Int(v) => Ok(v),
|
||||
TimeoutValue::Float(v) => {
|
||||
if !v.is_finite() || v < 0.0 {
|
||||
return Err(serde::de::Error::custom(
|
||||
"timeout must be a non-negative finite number",
|
||||
));
|
||||
}
|
||||
if v.fract() != 0.0 {
|
||||
return Err(serde::de::Error::custom("timeout must be integer seconds"));
|
||||
}
|
||||
if v > (u64::MAX as f64) {
|
||||
return Err(serde::de::Error::custom("timeout is too large"));
|
||||
}
|
||||
Ok(v as u64)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// JSON payload for RESPONSE_HEADERS frames.
|
||||
#[derive(Debug, serde::Serialize)]
|
||||
pub struct ResponseMeta {
|
||||
@@ -159,3 +189,72 @@ pub struct ResponseMeta {
|
||||
/// Header list preserving duplicates (e.g. multiple Set-Cookie).
|
||||
pub headers: Vec<(String, String)>,
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tunnel frame compression helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Minimum payload size to attempt gzip compression (bytes).
|
||||
const COMPRESS_MIN_SIZE: usize = 512;
|
||||
|
||||
/// If the frame has the GZIP_COMPRESSED flag, decompress the payload; otherwise
|
||||
/// return a clone of the raw payload bytes.
|
||||
pub fn decompress_if_gzip(frame: &Frame) -> Result<Bytes, std::io::Error> {
|
||||
if frame.is_gzip() {
|
||||
decompress_gzip(&frame.payload)
|
||||
} else {
|
||||
Ok(frame.payload.clone())
|
||||
}
|
||||
}
|
||||
|
||||
/// Gzip-compress `data` if it is large enough and compression actually shrinks
|
||||
/// the payload. Returns `(payload, extra_flags)` where `extra_flags` contains
|
||||
/// `GZIP_COMPRESSED` when compression was applied.
|
||||
pub fn compress_payload(data: Bytes) -> (Bytes, u8) {
|
||||
if data.len() >= COMPRESS_MIN_SIZE {
|
||||
if let Ok(compressed) = compress_gzip(&data) {
|
||||
if compressed.len() < data.len() {
|
||||
return (compressed, flags::GZIP_COMPRESSED);
|
||||
}
|
||||
}
|
||||
}
|
||||
(data, 0)
|
||||
}
|
||||
|
||||
fn decompress_gzip(data: &[u8]) -> Result<Bytes, std::io::Error> {
|
||||
use flate2::read::GzDecoder;
|
||||
use std::io::Read;
|
||||
let mut decoder = GzDecoder::new(data);
|
||||
let mut buf = Vec::new();
|
||||
decoder.read_to_end(&mut buf)?;
|
||||
Ok(Bytes::from(buf))
|
||||
}
|
||||
|
||||
fn compress_gzip(data: &[u8]) -> Result<Bytes, std::io::Error> {
|
||||
use flate2::write::GzEncoder;
|
||||
use flate2::Compression;
|
||||
use std::io::Write;
|
||||
let mut encoder = GzEncoder::new(Vec::new(), Compression::fast());
|
||||
encoder.write_all(data)?;
|
||||
let compressed = encoder.finish()?;
|
||||
Ok(Bytes::from(compressed))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::RequestMeta;
|
||||
|
||||
#[test]
|
||||
fn request_meta_accepts_integer_timeout() {
|
||||
let raw = br#"{"method":"GET","url":"https://example.com","headers":{},"timeout":15}"#;
|
||||
let meta: RequestMeta = serde_json::from_slice(raw).expect("parse request meta");
|
||||
assert_eq!(meta.timeout, 15);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_meta_accepts_integer_like_float_timeout() {
|
||||
let raw = br#"{"method":"GET","url":"https://example.com","headers":{},"timeout":15.0}"#;
|
||||
let meta: RequestMeta = serde_json::from_slice(raw).expect("parse request meta");
|
||||
assert_eq!(meta.timeout, 15);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3,19 +3,28 @@
|
||||
//! Receives request frames, executes the upstream HTTP request,
|
||||
//! and sends response frames back through the writer channel.
|
||||
|
||||
use std::io;
|
||||
use std::sync::atomic::AtomicUsize;
|
||||
use std::sync::atomic::Ordering;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use bytes::Bytes;
|
||||
use futures_util::stream;
|
||||
use futures_util::StreamExt;
|
||||
use http_body_util::BodyExt;
|
||||
use hyper::body::Frame as BodyFrame;
|
||||
use tokio::sync::mpsc;
|
||||
use tracing::{debug, warn};
|
||||
|
||||
use crate::state::{AppState, ServerContext};
|
||||
use crate::target_filter;
|
||||
use crate::upstream_client;
|
||||
|
||||
use super::protocol::{flags, Frame, MsgType, RequestMeta, ResponseMeta};
|
||||
use super::protocol::{
|
||||
compress_payload, decompress_if_gzip, flags, Frame as TunnelFrame, MsgType, RequestMeta,
|
||||
ResponseMeta,
|
||||
};
|
||||
use super::writer::FrameSender;
|
||||
|
||||
/// Maximum response body chunk size per frame (32 KB).
|
||||
@@ -32,8 +41,18 @@ const MIN_TIMEOUT_SECS: u64 = 5;
|
||||
const MAX_TIMEOUT_SECS: u64 = 300;
|
||||
|
||||
/// Headers that must not be forwarded to upstream (hop-by-hop or security-sensitive).
|
||||
///
|
||||
/// `host` and `content-length` are managed by the HTTP client (reqwest/hyper):
|
||||
/// - `host` → translated to `:authority` pseudo-header in HTTP/2; forwarding
|
||||
/// the original `host` alongside `:authority` triggers PROTOCOL_ERROR on
|
||||
/// strict H2 implementations (e.g. Google APIs).
|
||||
/// - `content-length` → recalculated by hyper from the actual body; a stale
|
||||
/// value from the tunnel (body may have been re-compressed) causes H2
|
||||
/// PROTOCOL_ERROR when it mismatches the real frame length.
|
||||
const BLOCKED_HEADERS: &[&str] = &[
|
||||
"connection",
|
||||
"content-length",
|
||||
"host",
|
||||
"keep-alive",
|
||||
"proxy-authenticate",
|
||||
"proxy-authorization",
|
||||
@@ -50,20 +69,22 @@ pub async fn handle_stream(
|
||||
server: Arc<ServerContext>,
|
||||
stream_id: u32,
|
||||
meta: RequestMeta,
|
||||
mut body_rx: mpsc::Receiver<Frame>,
|
||||
body_rx: mpsc::Receiver<TunnelFrame>,
|
||||
frame_tx: FrameSender,
|
||||
) {
|
||||
let start = Instant::now();
|
||||
server.active_connections.fetch_add(1, Ordering::Release);
|
||||
|
||||
handle_stream_inner(&state, &server, stream_id, meta, &mut body_rx, &frame_tx).await;
|
||||
let connect_elapsed =
|
||||
handle_stream_inner(&state, &server, stream_id, meta, body_rx, &frame_tx).await;
|
||||
|
||||
server.active_connections.fetch_sub(1, Ordering::Release);
|
||||
server.metrics.record_request(start.elapsed());
|
||||
if let Some(d) = connect_elapsed {
|
||||
server.metrics.record_request(d);
|
||||
}
|
||||
}
|
||||
|
||||
/// Send a frame to the writer with a timeout. Returns false if send failed.
|
||||
async fn send_frame(tx: &FrameSender, frame: Frame) -> bool {
|
||||
async fn send_frame(tx: &FrameSender, frame: TunnelFrame) -> bool {
|
||||
match tokio::time::timeout(FRAME_SEND_TIMEOUT, tx.send(frame)).await {
|
||||
Ok(Ok(())) => true,
|
||||
Ok(Err(_)) => {
|
||||
@@ -78,77 +99,23 @@ async fn send_frame(tx: &FrameSender, frame: Frame) -> bool {
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the connection-establishment duration (DNS + TCP/TLS + TTFB) if the
|
||||
/// upstream request succeeded, or `None` if the request never reached the
|
||||
/// response-headers stage.
|
||||
async fn handle_stream_inner(
|
||||
state: &AppState,
|
||||
server: &ServerContext,
|
||||
stream_id: u32,
|
||||
meta: RequestMeta,
|
||||
body_rx: &mut mpsc::Receiver<Frame>,
|
||||
body_rx: mpsc::Receiver<TunnelFrame>,
|
||||
frame_tx: &FrameSender,
|
||||
) {
|
||||
// Collect request body
|
||||
let mut body_parts: Vec<Bytes> = Vec::new();
|
||||
let mut body_done = false;
|
||||
|
||||
// Drain body frames
|
||||
while !body_done {
|
||||
match body_rx.recv().await {
|
||||
Some(frame) => {
|
||||
if frame.msg_type == MsgType::RequestBody {
|
||||
let payload = if frame.is_gzip() {
|
||||
match decompress_gzip(&frame.payload) {
|
||||
Ok(d) => d,
|
||||
Err(e) => {
|
||||
send_error(
|
||||
frame_tx,
|
||||
stream_id,
|
||||
&format!("gzip decompress failed: {e}"),
|
||||
)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
frame.payload.clone()
|
||||
};
|
||||
if !payload.is_empty() {
|
||||
body_parts.push(payload);
|
||||
}
|
||||
if frame.is_end_stream() {
|
||||
body_done = true;
|
||||
}
|
||||
} else if frame.msg_type == MsgType::StreamEnd
|
||||
|| frame.msg_type == MsgType::StreamError
|
||||
{
|
||||
body_done = true;
|
||||
if frame.msg_type == MsgType::StreamError {
|
||||
return; // Client cancelled
|
||||
}
|
||||
}
|
||||
}
|
||||
None => return, // Channel closed
|
||||
}
|
||||
}
|
||||
|
||||
let body: Bytes = if body_parts.is_empty() {
|
||||
Bytes::new()
|
||||
} else if body_parts.len() == 1 {
|
||||
body_parts.into_iter().next().unwrap()
|
||||
} else {
|
||||
let total: usize = body_parts.iter().map(|b| b.len()).sum();
|
||||
let mut combined = Vec::with_capacity(total);
|
||||
for part in &body_parts {
|
||||
combined.extend_from_slice(part);
|
||||
}
|
||||
Bytes::from(combined)
|
||||
};
|
||||
|
||||
) -> Option<Duration> {
|
||||
// Validate target
|
||||
let target_url = match url::Url::parse(&meta.url) {
|
||||
Ok(u) => u,
|
||||
Err(e) => {
|
||||
send_error(frame_tx, stream_id, &format!("invalid URL: {e}")).await;
|
||||
return;
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
@@ -162,7 +129,7 @@ async fn handle_stream_inner(
|
||||
&format!("unsupported URL scheme: {other}"),
|
||||
)
|
||||
.await;
|
||||
return;
|
||||
return None;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -170,13 +137,13 @@ async fn handle_stream_inner(
|
||||
Some(h) => h.to_string(),
|
||||
None => {
|
||||
send_error(frame_tx, stream_id, "missing host in URL").await;
|
||||
return;
|
||||
return None;
|
||||
}
|
||||
};
|
||||
let port = target_url.port_or_known_default().unwrap_or(443);
|
||||
|
||||
// DNS + target validation (populates dns_cache for SafeDnsResolver)
|
||||
let dns_start = Instant::now();
|
||||
let connect_start = Instant::now();
|
||||
{
|
||||
let allowed_ports = Arc::clone(&server.dynamic.load().allowed_ports);
|
||||
if let Err(e) =
|
||||
@@ -184,73 +151,122 @@ async fn handle_stream_inner(
|
||||
{
|
||||
server.metrics.dns_failures.fetch_add(1, Ordering::Release);
|
||||
send_error(frame_tx, stream_id, &format!("target blocked: {e}")).await;
|
||||
return;
|
||||
return None;
|
||||
}
|
||||
}
|
||||
let dns_ms = dns_start.elapsed().as_millis() as u64;
|
||||
let dns_ms = connect_start.elapsed().as_millis() as u64;
|
||||
|
||||
// Execute upstream request
|
||||
let client = &state.reqwest_client;
|
||||
let client = &state.upstream_client;
|
||||
let timeout = Duration::from_secs(meta.timeout.clamp(MIN_TIMEOUT_SECS, MAX_TIMEOUT_SECS));
|
||||
let request_body_size = Arc::new(AtomicUsize::new(0));
|
||||
let request_body = build_streaming_request_body(body_rx, Arc::clone(&request_body_size));
|
||||
|
||||
let method: reqwest::Method = meta.method.parse().unwrap_or(reqwest::Method::GET);
|
||||
let mut req = client.request(method, &meta.url);
|
||||
let method: hyper::Method = meta.method.parse().unwrap_or(hyper::Method::GET);
|
||||
let mut request = match hyper::Request::builder()
|
||||
.method(method)
|
||||
.uri(meta.url.as_str())
|
||||
.body(request_body)
|
||||
{
|
||||
Ok(request) => request,
|
||||
Err(e) => {
|
||||
send_error(
|
||||
frame_tx,
|
||||
stream_id,
|
||||
&format!("invalid upstream request: {e}"),
|
||||
)
|
||||
.await;
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
let headers = request.headers_mut();
|
||||
for (k, v) in &meta.headers {
|
||||
let k_lower = k.to_ascii_lowercase();
|
||||
// Skip hop-by-hop and security-sensitive headers
|
||||
if BLOCKED_HEADERS.contains(&k_lower.as_str()) {
|
||||
continue;
|
||||
}
|
||||
// Validate header name/value are valid HTTP
|
||||
if let (Ok(name), Ok(value)) = (
|
||||
reqwest::header::HeaderName::from_bytes(k.as_bytes()),
|
||||
reqwest::header::HeaderValue::from_str(v),
|
||||
hyper::header::HeaderName::from_bytes(k.as_bytes()),
|
||||
hyper::header::HeaderValue::from_str(v),
|
||||
) {
|
||||
req = req.header(name, value);
|
||||
headers.insert(name, value);
|
||||
}
|
||||
}
|
||||
let body_size = body.len();
|
||||
if !body.is_empty() {
|
||||
req = req.body(body);
|
||||
}
|
||||
req = req.timeout(timeout);
|
||||
|
||||
let mut captured_connection = upstream_client::capture_connection(&mut request);
|
||||
let connection_start = Instant::now();
|
||||
let connection_capture = tokio::spawn(async move {
|
||||
let connected = captured_connection.wait_for_connection_metadata().await;
|
||||
connected
|
||||
.as_ref()
|
||||
.map(|_| connection_start.elapsed().as_millis() as u64)
|
||||
});
|
||||
|
||||
let upstream_start = Instant::now();
|
||||
let response = match req.send().await {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
let response = match tokio::time::timeout(timeout, client.request(request)).await {
|
||||
Ok(Ok(response)) => response,
|
||||
Ok(Err(e)) => {
|
||||
connection_capture.abort();
|
||||
server
|
||||
.metrics
|
||||
.failed_requests
|
||||
.fetch_add(1, Ordering::Release);
|
||||
let msg = if e.is_timeout() {
|
||||
"upstream timeout".to_string()
|
||||
} else if e.is_connect() {
|
||||
let msg = if e.is_connect() {
|
||||
format!("upstream connect error: {e}")
|
||||
} else {
|
||||
format!("upstream error: {e}")
|
||||
};
|
||||
send_error(frame_tx, stream_id, &msg).await;
|
||||
return;
|
||||
return None;
|
||||
}
|
||||
Err(_) => {
|
||||
connection_capture.abort();
|
||||
server
|
||||
.metrics
|
||||
.failed_requests
|
||||
.fetch_add(1, Ordering::Release);
|
||||
send_error(frame_tx, stream_id, "upstream timeout").await;
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
// Capture connection-establishment duration (DNS + TCP/TLS + TTFB)
|
||||
// before proceeding to stream the response body.
|
||||
let connect_elapsed = connect_start.elapsed();
|
||||
|
||||
// Send RESPONSE_HEADERS
|
||||
let status = response.status().as_u16();
|
||||
let ttfb_ms = upstream_start.elapsed().as_millis() as u64;
|
||||
// Short timeout: on connection reuse hyper may never fire the connect
|
||||
// callback, so avoid blocking indefinitely.
|
||||
let connection_acquire_ms =
|
||||
match tokio::time::timeout(Duration::from_millis(100), connection_capture).await {
|
||||
Ok(Ok(ms)) => ms,
|
||||
Ok(Err(_)) => None, // JoinError (task panicked / cancelled)
|
||||
Err(_) => None, // timeout -- task is detached but lightweight
|
||||
};
|
||||
let request_timing =
|
||||
upstream_client::resolve_request_timing(&response, connection_acquire_ms, ttfb_ms);
|
||||
let mut resp_headers: Vec<(String, String)> = Vec::with_capacity(response.headers().len() + 1);
|
||||
for (k, v) in response.headers() {
|
||||
if let Ok(vs) = v.to_str() {
|
||||
resp_headers.push((k.as_str().to_string(), vs.to_string()));
|
||||
}
|
||||
}
|
||||
// Inject proxy timing (same format as delegate mode)
|
||||
let timing = serde_json::json!({
|
||||
"dns_ms": dns_ms,
|
||||
"connection_acquire_ms": request_timing.connection_acquire_ms,
|
||||
"connection_reused": request_timing.connection_reused,
|
||||
"connect_ms": request_timing.connect_ms,
|
||||
"tls_ms": request_timing.tls_ms,
|
||||
"ttfb_ms": ttfb_ms,
|
||||
"upstream_ms": ttfb_ms,
|
||||
"upstream_processing_ms": ttfb_ms.saturating_sub(dns_ms),
|
||||
"body_size": body_size,
|
||||
"response_wait_ms": request_timing.response_wait_ms,
|
||||
"upstream_processing_ms": request_timing.response_wait_ms,
|
||||
"timing_source": "instrumented_connector",
|
||||
"total_ms": connect_elapsed.as_millis() as u64,
|
||||
"body_size": request_body_size.load(Ordering::Relaxed),
|
||||
"mode": "tunnel",
|
||||
});
|
||||
resp_headers.push(("x-proxy-timing".to_string(), timing.to_string()));
|
||||
@@ -258,43 +274,60 @@ async fn handle_stream_inner(
|
||||
status,
|
||||
headers: resp_headers,
|
||||
};
|
||||
let meta_json = serde_json::to_vec(&resp_meta).unwrap_or_default();
|
||||
let meta_json: Bytes = serde_json::to_vec(&resp_meta).unwrap_or_default().into();
|
||||
let (meta_payload, meta_flags) = compress_payload(meta_json);
|
||||
if !send_frame(
|
||||
frame_tx,
|
||||
Frame::new(stream_id, MsgType::ResponseHeaders, 0, meta_json),
|
||||
TunnelFrame::new(
|
||||
stream_id,
|
||||
MsgType::ResponseHeaders,
|
||||
meta_flags,
|
||||
meta_payload,
|
||||
),
|
||||
)
|
||||
.await
|
||||
{
|
||||
return;
|
||||
return Some(connect_elapsed);
|
||||
}
|
||||
|
||||
// Stream response body
|
||||
let mut stream = response.bytes_stream();
|
||||
// Stream response body — relay upstream bytes through the tunnel.
|
||||
// Apply tunnel-level frame compression for chunks that benefit from it
|
||||
// (e.g. uncompressed SSE text). Already-compressed data (gzip/br from
|
||||
// upstream Content-Encoding) won't shrink further and will be sent as-is
|
||||
// thanks to the size check in compress_payload().
|
||||
let mut stream = response.into_body().into_data_stream();
|
||||
while let Some(chunk_result) = stream.next().await {
|
||||
match chunk_result {
|
||||
Ok(chunk) => {
|
||||
if chunk.len() <= MAX_CHUNK_SIZE {
|
||||
let (payload, extra_flags) = compress_payload(chunk);
|
||||
if !send_frame(
|
||||
frame_tx,
|
||||
Frame::new(stream_id, MsgType::ResponseBody, 0, chunk),
|
||||
TunnelFrame::new(stream_id, MsgType::ResponseBody, extra_flags, payload),
|
||||
)
|
||||
.await
|
||||
{
|
||||
return;
|
||||
return Some(connect_elapsed);
|
||||
}
|
||||
} else {
|
||||
// Split oversized chunks
|
||||
// Split oversized chunks, compress each slice
|
||||
let mut offset = 0;
|
||||
while offset < chunk.len() {
|
||||
let end = (offset + MAX_CHUNK_SIZE).min(chunk.len());
|
||||
let slice = chunk.slice(offset..end);
|
||||
let (payload, extra_flags) = compress_payload(slice);
|
||||
if !send_frame(
|
||||
frame_tx,
|
||||
Frame::new(stream_id, MsgType::ResponseBody, 0, slice),
|
||||
TunnelFrame::new(
|
||||
stream_id,
|
||||
MsgType::ResponseBody,
|
||||
extra_flags,
|
||||
payload,
|
||||
),
|
||||
)
|
||||
.await
|
||||
{
|
||||
return;
|
||||
return Some(connect_elapsed);
|
||||
}
|
||||
offset = end;
|
||||
}
|
||||
@@ -304,7 +337,7 @@ async fn handle_stream_inner(
|
||||
server.metrics.stream_errors.fetch_add(1, Ordering::Release);
|
||||
warn!(stream_id, error = %e, "upstream body read error");
|
||||
send_error(frame_tx, stream_id, &format!("body read error: {e}")).await;
|
||||
return;
|
||||
return Some(connect_elapsed);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -312,7 +345,7 @@ async fn handle_stream_inner(
|
||||
// Send STREAM_END
|
||||
let _ = send_frame(
|
||||
frame_tx,
|
||||
Frame::new(
|
||||
TunnelFrame::new(
|
||||
stream_id,
|
||||
MsgType::StreamEnd,
|
||||
flags::END_STREAM,
|
||||
@@ -322,13 +355,14 @@ async fn handle_stream_inner(
|
||||
.await;
|
||||
|
||||
debug!(stream_id, status, "stream completed");
|
||||
Some(connect_elapsed)
|
||||
}
|
||||
|
||||
async fn send_error(tx: &FrameSender, stream_id: u32, msg: &str) {
|
||||
// Error frames use best-effort delivery — don't block if writer is congested
|
||||
let _ = send_frame(
|
||||
tx,
|
||||
Frame::new(
|
||||
TunnelFrame::new(
|
||||
stream_id,
|
||||
MsgType::StreamError,
|
||||
0,
|
||||
@@ -338,11 +372,135 @@ async fn send_error(tx: &FrameSender, stream_id: u32, msg: &str) {
|
||||
.await;
|
||||
}
|
||||
|
||||
fn decompress_gzip(data: &[u8]) -> Result<Bytes, std::io::Error> {
|
||||
use flate2::read::GzDecoder;
|
||||
use std::io::Read;
|
||||
let mut decoder = GzDecoder::new(data);
|
||||
let mut buf = Vec::new();
|
||||
decoder.read_to_end(&mut buf)?;
|
||||
Ok(Bytes::from(buf))
|
||||
fn build_streaming_request_body(
|
||||
body_rx: mpsc::Receiver<TunnelFrame>,
|
||||
body_size: Arc<AtomicUsize>,
|
||||
) -> upstream_client::UpstreamRequestBody {
|
||||
let body_stream = stream::unfold(
|
||||
(body_rx, body_size, false),
|
||||
|(mut body_rx, body_size, finished)| async move {
|
||||
if finished {
|
||||
return None;
|
||||
}
|
||||
|
||||
loop {
|
||||
let frame = match body_rx.recv().await {
|
||||
Some(frame) => frame,
|
||||
None => return None,
|
||||
};
|
||||
|
||||
match frame.msg_type {
|
||||
MsgType::RequestBody => {
|
||||
let end_stream = frame.is_end_stream();
|
||||
let payload = match decompress_if_gzip(&frame) {
|
||||
Ok(payload) => payload,
|
||||
Err(error) => {
|
||||
let err =
|
||||
io::Error::other(format!("gzip decompress failed: {error}"));
|
||||
return Some((Err(err), (body_rx, body_size, true)));
|
||||
}
|
||||
};
|
||||
|
||||
if payload.is_empty() {
|
||||
if end_stream {
|
||||
return None;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
body_size.fetch_add(payload.len(), Ordering::Relaxed);
|
||||
return Some((
|
||||
Ok(BodyFrame::data(payload)),
|
||||
(body_rx, body_size, end_stream),
|
||||
));
|
||||
}
|
||||
MsgType::StreamError => {
|
||||
let message = String::from_utf8(frame.payload.to_vec())
|
||||
.unwrap_or_else(|_| "client cancelled request body".to_string());
|
||||
return Some((Err(io::Error::other(message)), (body_rx, body_size, true)));
|
||||
}
|
||||
MsgType::StreamEnd => return None,
|
||||
_ => continue,
|
||||
}
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
upstream_client::stream_request_body(body_stream)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn streaming_request_body_yields_chunks_and_tracks_size() {
|
||||
let (tx, rx) = mpsc::channel(4);
|
||||
let body_size = Arc::new(AtomicUsize::new(0));
|
||||
let mut body = build_streaming_request_body(rx, Arc::clone(&body_size));
|
||||
|
||||
tx.send(TunnelFrame::new(
|
||||
1,
|
||||
MsgType::RequestBody,
|
||||
0,
|
||||
Bytes::from_static(b"abc"),
|
||||
))
|
||||
.await
|
||||
.expect("send first chunk");
|
||||
tx.send(TunnelFrame::new(
|
||||
1,
|
||||
MsgType::RequestBody,
|
||||
flags::END_STREAM,
|
||||
Bytes::from_static(b"def"),
|
||||
))
|
||||
.await
|
||||
.expect("send final chunk");
|
||||
drop(tx);
|
||||
|
||||
let first = body
|
||||
.frame()
|
||||
.await
|
||||
.expect("first frame")
|
||||
.expect("first frame ok")
|
||||
.into_data()
|
||||
.expect("first data frame");
|
||||
let second = body
|
||||
.frame()
|
||||
.await
|
||||
.expect("second frame")
|
||||
.expect("second frame ok")
|
||||
.into_data()
|
||||
.expect("second data frame");
|
||||
|
||||
assert_eq!(first, Bytes::from_static(b"abc"));
|
||||
assert_eq!(second, Bytes::from_static(b"def"));
|
||||
assert!(body.frame().await.is_none());
|
||||
assert_eq!(body_size.load(Ordering::Relaxed), 6);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn streaming_request_body_surfaces_client_cancel_as_error() {
|
||||
let (tx, rx) = mpsc::channel(4);
|
||||
let body_size = Arc::new(AtomicUsize::new(0));
|
||||
let mut body = build_streaming_request_body(rx, Arc::clone(&body_size));
|
||||
|
||||
tx.send(TunnelFrame::new(
|
||||
1,
|
||||
MsgType::StreamError,
|
||||
0,
|
||||
Bytes::from_static(b"client cancelled"),
|
||||
))
|
||||
.await
|
||||
.expect("send cancel frame");
|
||||
drop(tx);
|
||||
|
||||
let err = body
|
||||
.frame()
|
||||
.await
|
||||
.expect("error frame present")
|
||||
.expect_err("body should surface cancellation error");
|
||||
assert!(err.to_string().contains("client cancelled"));
|
||||
assert!(body.frame().await.is_none());
|
||||
assert_eq!(body_size.load(Ordering::Relaxed), 0);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,444 @@
|
||||
use std::future::Future;
|
||||
use std::io;
|
||||
use std::net::IpAddr;
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
use std::task::{Context, Poll};
|
||||
use std::time::Duration;
|
||||
|
||||
use bytes::Bytes;
|
||||
use futures_util::Stream;
|
||||
use http_body_util::combinators::UnsyncBoxBody;
|
||||
use http_body_util::{BodyExt, StreamBody};
|
||||
use hyper::body::Frame;
|
||||
use hyper::rt;
|
||||
use hyper::Response;
|
||||
use hyper::Uri;
|
||||
pub use hyper_util::client::legacy::connect::capture_connection;
|
||||
use hyper_util::client::legacy::connect::dns::Name;
|
||||
use hyper_util::client::legacy::connect::{Connected, Connection, HttpConnector};
|
||||
use hyper_util::client::legacy::Client;
|
||||
use hyper_util::rt::{TokioExecutor, TokioIo, TokioTimer};
|
||||
use rustls::pki_types::ServerName;
|
||||
use rustls::ClientConfig;
|
||||
use tokio::net::TcpStream;
|
||||
use tokio_rustls::TlsConnector;
|
||||
use tower_service::Service;
|
||||
|
||||
use crate::config::Config;
|
||||
use crate::target_filter::{self, DnsCache};
|
||||
|
||||
type BoxError = Box<dyn std::error::Error + Send + Sync>;
|
||||
|
||||
type PlainStream = TokioIo<TcpStream>;
|
||||
type TlsStream = TokioIo<tokio_rustls::client::TlsStream<TcpStream>>;
|
||||
|
||||
pub type UpstreamRequestBody = UnsyncBoxBody<Bytes, io::Error>;
|
||||
pub type UpstreamClient = Client<InstrumentedConnector, UpstreamRequestBody>;
|
||||
|
||||
pub fn stream_request_body<S>(stream: S) -> UpstreamRequestBody
|
||||
where
|
||||
S: Stream<Item = Result<Frame<Bytes>, io::Error>> + Send + 'static,
|
||||
{
|
||||
StreamBody::new(stream).boxed_unsync()
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default)]
|
||||
pub struct ConnectTiming {
|
||||
pub connect_ms: u64,
|
||||
pub tls_ms: u64,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default)]
|
||||
pub struct RequestTiming {
|
||||
pub connection_acquire_ms: u64,
|
||||
pub connect_ms: u64,
|
||||
pub tls_ms: u64,
|
||||
pub response_wait_ms: u64,
|
||||
pub connection_reused: bool,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct ValidatedResolver {
|
||||
dns_cache: Arc<DnsCache>,
|
||||
}
|
||||
|
||||
impl ValidatedResolver {
|
||||
pub fn new(dns_cache: Arc<DnsCache>) -> Self {
|
||||
Self { dns_cache }
|
||||
}
|
||||
}
|
||||
|
||||
pub struct ValidatedAddrs {
|
||||
inner: std::vec::IntoIter<std::net::SocketAddr>,
|
||||
}
|
||||
|
||||
impl Iterator for ValidatedAddrs {
|
||||
type Item = std::net::SocketAddr;
|
||||
|
||||
fn next(&mut self) -> Option<Self::Item> {
|
||||
self.inner.next()
|
||||
}
|
||||
}
|
||||
|
||||
impl Service<Name> for ValidatedResolver {
|
||||
type Response = ValidatedAddrs;
|
||||
type Error = io::Error;
|
||||
type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
|
||||
|
||||
fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
||||
Poll::Ready(Ok(()))
|
||||
}
|
||||
|
||||
fn call(&mut self, name: Name) -> Self::Future {
|
||||
let dns_cache = Arc::clone(&self.dns_cache);
|
||||
let host = name.as_str().to_string();
|
||||
Box::pin(async move {
|
||||
if let Some(addrs) = dns_cache.get_by_host(&host).await {
|
||||
return Ok(ValidatedAddrs {
|
||||
inner: (*addrs).clone().into_iter(),
|
||||
});
|
||||
}
|
||||
|
||||
let resolved = target_filter::resolve_public_addrs(&host, 0, dns_cache.as_ref())
|
||||
.await
|
||||
.map_err(|err| io::Error::other(err.to_string()))?;
|
||||
Ok(ValidatedAddrs {
|
||||
inner: resolved.into_iter(),
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct InstrumentedConnector {
|
||||
http: HttpConnector<ValidatedResolver>,
|
||||
tls_config: Arc<ClientConfig>,
|
||||
}
|
||||
|
||||
impl Service<Uri> for InstrumentedConnector {
|
||||
type Response = TimedConn;
|
||||
type Error = BoxError;
|
||||
type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
|
||||
|
||||
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
||||
self.http.poll_ready(cx).map_err(Into::into)
|
||||
}
|
||||
|
||||
fn call(&mut self, dst: Uri) -> Self::Future {
|
||||
let scheme = dst.scheme_str().map(|value| value.to_ascii_lowercase());
|
||||
let tls_config = Arc::clone(&self.tls_config);
|
||||
let connecting = self.http.call(dst.clone());
|
||||
let connect_start = std::time::Instant::now();
|
||||
|
||||
Box::pin(async move {
|
||||
match scheme.as_deref() {
|
||||
Some("http") => {
|
||||
let tcp = connecting.await.map_err(|err| Box::new(err) as BoxError)?;
|
||||
let connect_ms = connect_start.elapsed().as_millis() as u64;
|
||||
Ok(TimedConn::new(
|
||||
MaybeHttpsStream::Http(tcp),
|
||||
ConnectTiming {
|
||||
connect_ms,
|
||||
tls_ms: 0,
|
||||
},
|
||||
))
|
||||
}
|
||||
Some("https") => {
|
||||
let server_name = resolve_server_name(&dst)?;
|
||||
let tcp = connecting.await.map_err(|err| Box::new(err) as BoxError)?;
|
||||
let connect_ms = connect_start.elapsed().as_millis() as u64;
|
||||
|
||||
let tls_start = std::time::Instant::now();
|
||||
let tls_stream = TlsConnector::from(tls_config)
|
||||
.connect(server_name, tcp.into_inner())
|
||||
.await
|
||||
.map_err(io::Error::other)?;
|
||||
let tls_ms = tls_start.elapsed().as_millis() as u64;
|
||||
|
||||
Ok(TimedConn::new(
|
||||
MaybeHttpsStream::Https(TokioIo::new(tls_stream)),
|
||||
ConnectTiming { connect_ms, tls_ms },
|
||||
))
|
||||
}
|
||||
Some(other) => Err(io::Error::other(format!("unsupported scheme {other}")).into()),
|
||||
None => Err(io::Error::other("missing scheme").into()),
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub fn build_upstream_client(config: &Config, dns_cache: Arc<DnsCache>) -> UpstreamClient {
|
||||
let mut http = HttpConnector::new_with_resolver(ValidatedResolver::new(dns_cache));
|
||||
http.enforce_http(false);
|
||||
http.set_connect_timeout(Some(Duration::from_secs(
|
||||
config.upstream_connect_timeout_secs,
|
||||
)));
|
||||
http.set_nodelay(config.upstream_tcp_nodelay);
|
||||
if config.upstream_tcp_keepalive_secs > 0 {
|
||||
http.set_keepalive(Some(Duration::from_secs(
|
||||
config.upstream_tcp_keepalive_secs,
|
||||
)));
|
||||
} else {
|
||||
http.set_keepalive(None);
|
||||
}
|
||||
|
||||
let connector = InstrumentedConnector {
|
||||
http,
|
||||
tls_config: build_tls_config(),
|
||||
};
|
||||
|
||||
let mut builder = Client::builder(TokioExecutor::new());
|
||||
builder.pool_max_idle_per_host(config.upstream_pool_max_idle_per_host);
|
||||
builder.pool_idle_timeout(Duration::from_secs(config.upstream_pool_idle_timeout_secs));
|
||||
builder.pool_timer(TokioTimer::new());
|
||||
builder.build(connector)
|
||||
}
|
||||
|
||||
pub fn resolve_request_timing<B>(
|
||||
response: &Response<B>,
|
||||
connection_acquire_ms: Option<u64>,
|
||||
ttfb_ms: u64,
|
||||
) -> RequestTiming {
|
||||
let raw = response
|
||||
.extensions()
|
||||
.get::<ConnectTiming>()
|
||||
.copied()
|
||||
.unwrap_or_default();
|
||||
|
||||
let raw_connection_ms = raw.connect_ms.saturating_add(raw.tls_ms);
|
||||
let measured_acquire_ms = connection_acquire_ms.unwrap_or(raw_connection_ms.min(ttfb_ms));
|
||||
let likely_reused = measured_acquire_ms <= 5 && raw_connection_ms > 0;
|
||||
let connector_matches_request = raw_connection_ms <= measured_acquire_ms.saturating_add(25);
|
||||
|
||||
let (connect_ms, tls_ms) = if likely_reused || !connector_matches_request {
|
||||
(0, 0)
|
||||
} else {
|
||||
(raw.connect_ms, raw.tls_ms)
|
||||
};
|
||||
|
||||
RequestTiming {
|
||||
connection_acquire_ms: measured_acquire_ms,
|
||||
connect_ms,
|
||||
tls_ms,
|
||||
response_wait_ms: ttfb_ms.saturating_sub(measured_acquire_ms),
|
||||
connection_reused: likely_reused,
|
||||
}
|
||||
}
|
||||
|
||||
fn build_tls_config() -> Arc<ClientConfig> {
|
||||
let root_store =
|
||||
rustls::RootCertStore::from_iter(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
|
||||
let mut config = ClientConfig::builder()
|
||||
.with_root_certificates(root_store)
|
||||
.with_no_client_auth();
|
||||
config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()];
|
||||
Arc::new(config)
|
||||
}
|
||||
|
||||
fn resolve_server_name(uri: &Uri) -> Result<ServerName<'static>, BoxError> {
|
||||
let host = uri.host().ok_or_else(|| io::Error::other("missing host"))?;
|
||||
let host = host.trim_start_matches('[').trim_end_matches(']');
|
||||
|
||||
if let Ok(ip) = host.parse::<IpAddr>() {
|
||||
return Ok(ServerName::from(ip));
|
||||
}
|
||||
|
||||
Ok(ServerName::try_from(host.to_string())?)
|
||||
}
|
||||
|
||||
pub struct TimedConn {
|
||||
inner: MaybeHttpsStream,
|
||||
timing: ConnectTiming,
|
||||
}
|
||||
|
||||
impl TimedConn {
|
||||
fn new(inner: MaybeHttpsStream, timing: ConnectTiming) -> Self {
|
||||
Self { inner, timing }
|
||||
}
|
||||
}
|
||||
|
||||
impl Connection for TimedConn {
|
||||
fn connected(&self) -> Connected {
|
||||
self.inner.connected().extra(self.timing)
|
||||
}
|
||||
}
|
||||
|
||||
impl rt::Read for TimedConn {
|
||||
fn poll_read(
|
||||
mut self: Pin<&mut Self>,
|
||||
cx: &mut Context<'_>,
|
||||
buf: rt::ReadBufCursor<'_>,
|
||||
) -> Poll<Result<(), io::Error>> {
|
||||
Pin::new(&mut self.inner).poll_read(cx, buf)
|
||||
}
|
||||
}
|
||||
|
||||
impl rt::Write for TimedConn {
|
||||
fn poll_write(
|
||||
mut self: Pin<&mut Self>,
|
||||
cx: &mut Context<'_>,
|
||||
buf: &[u8],
|
||||
) -> Poll<Result<usize, io::Error>> {
|
||||
Pin::new(&mut self.inner).poll_write(cx, buf)
|
||||
}
|
||||
|
||||
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
|
||||
Pin::new(&mut self.inner).poll_flush(cx)
|
||||
}
|
||||
|
||||
fn poll_shutdown(
|
||||
mut self: Pin<&mut Self>,
|
||||
cx: &mut Context<'_>,
|
||||
) -> Poll<Result<(), io::Error>> {
|
||||
Pin::new(&mut self.inner).poll_shutdown(cx)
|
||||
}
|
||||
|
||||
fn is_write_vectored(&self) -> bool {
|
||||
self.inner.is_write_vectored()
|
||||
}
|
||||
|
||||
fn poll_write_vectored(
|
||||
mut self: Pin<&mut Self>,
|
||||
cx: &mut Context<'_>,
|
||||
bufs: &[std::io::IoSlice<'_>],
|
||||
) -> Poll<Result<usize, io::Error>> {
|
||||
Pin::new(&mut self.inner).poll_write_vectored(cx, bufs)
|
||||
}
|
||||
}
|
||||
|
||||
pub enum MaybeHttpsStream {
|
||||
Http(PlainStream),
|
||||
Https(TlsStream),
|
||||
}
|
||||
|
||||
impl Connection for MaybeHttpsStream {
|
||||
fn connected(&self) -> Connected {
|
||||
match self {
|
||||
Self::Http(stream) => stream.connected(),
|
||||
Self::Https(stream) => {
|
||||
let (tcp, tls) = stream.inner().get_ref();
|
||||
if tls.alpn_protocol() == Some(b"h2") {
|
||||
tcp.connected().negotiated_h2()
|
||||
} else {
|
||||
tcp.connected()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl rt::Read for MaybeHttpsStream {
|
||||
fn poll_read(
|
||||
self: Pin<&mut Self>,
|
||||
cx: &mut Context<'_>,
|
||||
buf: rt::ReadBufCursor<'_>,
|
||||
) -> Poll<Result<(), io::Error>> {
|
||||
match Pin::get_mut(self) {
|
||||
Self::Http(stream) => Pin::new(stream).poll_read(cx, buf),
|
||||
Self::Https(stream) => Pin::new(stream).poll_read(cx, buf),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl rt::Write for MaybeHttpsStream {
|
||||
fn poll_write(
|
||||
self: Pin<&mut Self>,
|
||||
cx: &mut Context<'_>,
|
||||
buf: &[u8],
|
||||
) -> Poll<Result<usize, io::Error>> {
|
||||
match Pin::get_mut(self) {
|
||||
Self::Http(stream) => Pin::new(stream).poll_write(cx, buf),
|
||||
Self::Https(stream) => Pin::new(stream).poll_write(cx, buf),
|
||||
}
|
||||
}
|
||||
|
||||
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
|
||||
match Pin::get_mut(self) {
|
||||
Self::Http(stream) => Pin::new(stream).poll_flush(cx),
|
||||
Self::Https(stream) => Pin::new(stream).poll_flush(cx),
|
||||
}
|
||||
}
|
||||
|
||||
fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
|
||||
match Pin::get_mut(self) {
|
||||
Self::Http(stream) => Pin::new(stream).poll_shutdown(cx),
|
||||
Self::Https(stream) => Pin::new(stream).poll_shutdown(cx),
|
||||
}
|
||||
}
|
||||
|
||||
fn is_write_vectored(&self) -> bool {
|
||||
match self {
|
||||
Self::Http(stream) => stream.is_write_vectored(),
|
||||
Self::Https(stream) => stream.is_write_vectored(),
|
||||
}
|
||||
}
|
||||
|
||||
fn poll_write_vectored(
|
||||
self: Pin<&mut Self>,
|
||||
cx: &mut Context<'_>,
|
||||
bufs: &[std::io::IoSlice<'_>],
|
||||
) -> Poll<Result<usize, io::Error>> {
|
||||
match Pin::get_mut(self) {
|
||||
Self::Http(stream) => Pin::new(stream).poll_write_vectored(cx, bufs),
|
||||
Self::Https(stream) => Pin::new(stream).poll_write_vectored(cx, bufs),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use hyper::Response;
|
||||
|
||||
#[test]
|
||||
fn fresh_connection_uses_connector_breakdown() {
|
||||
let mut response = Response::new(());
|
||||
response.extensions_mut().insert(ConnectTiming {
|
||||
connect_ms: 80,
|
||||
tls_ms: 40,
|
||||
});
|
||||
|
||||
let timing = resolve_request_timing(&response, Some(125), 600);
|
||||
|
||||
assert_eq!(timing.connection_acquire_ms, 125);
|
||||
assert_eq!(timing.connect_ms, 80);
|
||||
assert_eq!(timing.tls_ms, 40);
|
||||
assert_eq!(timing.response_wait_ms, 475);
|
||||
assert!(!timing.connection_reused);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reused_connection_zeroes_stale_connect_timings() {
|
||||
let mut response = Response::new(());
|
||||
response.extensions_mut().insert(ConnectTiming {
|
||||
connect_ms: 70,
|
||||
tls_ms: 30,
|
||||
});
|
||||
|
||||
let timing = resolve_request_timing(&response, Some(0), 310);
|
||||
|
||||
assert_eq!(timing.connection_acquire_ms, 0);
|
||||
assert_eq!(timing.connect_ms, 0);
|
||||
assert_eq!(timing.tls_ms, 0);
|
||||
assert_eq!(timing.response_wait_ms, 310);
|
||||
assert!(timing.connection_reused);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn falls_back_to_connector_timings_when_capture_missing() {
|
||||
let mut response = Response::new(());
|
||||
response.extensions_mut().insert(ConnectTiming {
|
||||
connect_ms: 55,
|
||||
tls_ms: 25,
|
||||
});
|
||||
|
||||
let timing = resolve_request_timing(&response, None, 400);
|
||||
|
||||
assert_eq!(timing.connection_acquire_ms, 80);
|
||||
assert_eq!(timing.connect_ms, 55);
|
||||
assert_eq!(timing.tls_ms, 25);
|
||||
assert_eq!(timing.response_wait_ms, 320);
|
||||
assert!(!timing.connection_reused);
|
||||
}
|
||||
}
|
||||
+25
-5
@@ -3,13 +3,15 @@ Alembic 环境配置
|
||||
用于数据库迁移的运行时环境设置
|
||||
"""
|
||||
|
||||
from logging.config import fileConfig
|
||||
from sqlalchemy import engine_from_config, pool
|
||||
from alembic import context
|
||||
import os
|
||||
import sys
|
||||
from logging.config import fileConfig
|
||||
from pathlib import Path
|
||||
|
||||
from sqlalchemy import engine_from_config, pool, text
|
||||
|
||||
from alembic import context
|
||||
|
||||
# 添加项目根目录到 Python 路径
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
|
||||
|
||||
@@ -48,6 +50,11 @@ if config.config_file_name is not None:
|
||||
# 目标元数据(包含所有表定义)
|
||||
target_metadata = Base.metadata
|
||||
|
||||
# PostgreSQL 全局迁移锁,避免多进程并发执行 Alembic 导致竞态(重复加列/索引等)
|
||||
# 使用会话级 advisory lock(pg_advisory_lock),在迁移完成后手动释放。
|
||||
# ID 由 crc32("aether-alembic-migration") 拼接生成,仅需全局唯一即可。
|
||||
MIGRATION_ADVISORY_LOCK_ID = 582694137405821
|
||||
|
||||
|
||||
def run_migrations_offline() -> None:
|
||||
"""
|
||||
@@ -83,15 +90,28 @@ def run_migrations_online() -> None:
|
||||
)
|
||||
|
||||
with connectable.connect() as connection:
|
||||
try:
|
||||
# 使用会话级 advisory lock(非事务级),避免干扰 Alembic 的事务管理。
|
||||
# pg_advisory_lock 在会话结束时自动释放,不受 COMMIT/ROLLBACK 影响。
|
||||
if connection.dialect.name == "postgresql":
|
||||
connection.execute(
|
||||
text("SELECT pg_advisory_lock(:lock_id)"),
|
||||
{"lock_id": MIGRATION_ADVISORY_LOCK_ID},
|
||||
)
|
||||
connection.commit()
|
||||
|
||||
context.configure(
|
||||
connection=connection,
|
||||
target_metadata=target_metadata,
|
||||
compare_type=True, # 比较列类型变更
|
||||
compare_server_default=True, # 比较默认值变更
|
||||
compare_type=True,
|
||||
compare_server_default=True,
|
||||
transaction_per_migration=True, # 每个迁移文件独立事务,完成即提交
|
||||
)
|
||||
|
||||
with context.begin_transaction():
|
||||
context.run_migrations()
|
||||
except Exception:
|
||||
raise
|
||||
|
||||
|
||||
# 根据模式选择运行方式
|
||||
|
||||
+10
-20
@@ -10,8 +10,7 @@ from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy import inspect
|
||||
from sqlalchemy import text
|
||||
|
||||
from alembic import op
|
||||
|
||||
@@ -24,33 +23,24 @@ depends_on: str | Sequence[str] | None = None
|
||||
|
||||
def upgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
inspector = inspect(conn)
|
||||
existing_columns = {col["name"] for col in inspector.get_columns("usage")}
|
||||
|
||||
if "provider_request_body" not in existing_columns:
|
||||
op.add_column("usage", sa.Column("provider_request_body", sa.JSON(), nullable=True))
|
||||
if "provider_request_body_compressed" not in existing_columns:
|
||||
op.add_column(
|
||||
"usage", sa.Column("provider_request_body_compressed", sa.LargeBinary(), nullable=True)
|
||||
# Use PostgreSQL native IF NOT EXISTS to avoid duplicate-column races
|
||||
# when migrations are triggered concurrently (e.g. startup + manual run).
|
||||
conn.execute(text("ALTER TABLE usage ADD COLUMN IF NOT EXISTS provider_request_body JSON"))
|
||||
conn.execute(
|
||||
text("ALTER TABLE usage ADD COLUMN IF NOT EXISTS provider_request_body_compressed BYTEA")
|
||||
)
|
||||
if "client_response_body" not in existing_columns:
|
||||
op.add_column("usage", sa.Column("client_response_body", sa.JSON(), nullable=True))
|
||||
if "client_response_body_compressed" not in existing_columns:
|
||||
op.add_column(
|
||||
"usage", sa.Column("client_response_body_compressed", sa.LargeBinary(), nullable=True)
|
||||
conn.execute(text("ALTER TABLE usage ADD COLUMN IF NOT EXISTS client_response_body JSON"))
|
||||
conn.execute(
|
||||
text("ALTER TABLE usage ADD COLUMN IF NOT EXISTS client_response_body_compressed BYTEA")
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
inspector = inspect(conn)
|
||||
existing_columns = {col["name"] for col in inspector.get_columns("usage")}
|
||||
|
||||
for col in (
|
||||
"client_response_body_compressed",
|
||||
"client_response_body",
|
||||
"provider_request_body_compressed",
|
||||
"provider_request_body",
|
||||
):
|
||||
if col in existing_columns:
|
||||
op.drop_column("usage", col)
|
||||
conn.execute(text(f"ALTER TABLE usage DROP COLUMN IF EXISTS {col}"))
|
||||
|
||||
@@ -33,11 +33,23 @@ depends_on: str | Sequence[str] | None = None
|
||||
|
||||
def column_exists(table_name: str, column_name: str) -> bool:
|
||||
bind = op.get_bind()
|
||||
inspector = inspect(bind)
|
||||
columns = [c["name"] for c in inspector.get_columns(table_name)]
|
||||
insp = inspect(bind)
|
||||
columns = [c["name"] for c in insp.get_columns(table_name)]
|
||||
return column_name in columns
|
||||
|
||||
|
||||
def table_exists(table_name: str) -> bool:
|
||||
bind = op.get_bind()
|
||||
insp = inspect(bind)
|
||||
return table_name in insp.get_table_names()
|
||||
|
||||
|
||||
def index_exists(table_name: str, index_name: str) -> bool:
|
||||
bind = op.get_bind()
|
||||
insp = inspect(bind)
|
||||
return any(idx["name"] == index_name for idx in insp.get_indexes(table_name))
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# --- 1. Add cache_creation columns ---
|
||||
if not column_exists("usage", "cache_creation_input_tokens_5m"):
|
||||
@@ -126,6 +138,7 @@ def upgrade() -> None:
|
||||
)
|
||||
|
||||
# --- 3. Create user_model_usage_counts table ---
|
||||
if not table_exists("user_model_usage_counts"):
|
||||
op.create_table(
|
||||
"user_model_usage_counts",
|
||||
sa.Column("id", sa.String(36), primary_key=True),
|
||||
@@ -151,10 +164,13 @@ def upgrade() -> None:
|
||||
),
|
||||
sa.UniqueConstraint("user_id", "model", name="uq_user_model_usage_count"),
|
||||
)
|
||||
if not index_exists("user_model_usage_counts", "idx_user_model_usage_user"):
|
||||
op.create_index("idx_user_model_usage_user", "user_model_usage_counts", ["user_id"])
|
||||
if not index_exists("user_model_usage_counts", "idx_user_model_usage_model"):
|
||||
op.create_index("idx_user_model_usage_model", "user_model_usage_counts", ["model"])
|
||||
|
||||
# Backfill from existing usage records
|
||||
# Backfill from existing usage records (truncate first for idempotency)
|
||||
conn.execute(sa.text("DELETE FROM user_model_usage_counts"))
|
||||
rows = conn.execute(
|
||||
sa.text(
|
||||
"SELECT user_id, model, COUNT(*) AS cnt FROM usage"
|
||||
@@ -179,16 +195,24 @@ def upgrade() -> None:
|
||||
)
|
||||
|
||||
# --- 4. Enforce models.global_model_id NOT NULL ---
|
||||
conn = op.get_bind()
|
||||
insp = inspect(conn)
|
||||
model_cols = {c["name"]: c for c in insp.get_columns("models")}
|
||||
if model_cols.get("global_model_id", {}).get("nullable", True):
|
||||
op.execute("DELETE FROM models WHERE global_model_id IS NULL")
|
||||
op.alter_column("models", "global_model_id", existing_type=sa.String(36), nullable=False)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
# Revert models.global_model_id to nullable
|
||||
if column_exists("models", "global_model_id"):
|
||||
op.alter_column("models", "global_model_id", existing_type=sa.String(36), nullable=True)
|
||||
|
||||
# Drop user_model_usage_counts
|
||||
if table_exists("user_model_usage_counts"):
|
||||
if index_exists("user_model_usage_counts", "idx_user_model_usage_model"):
|
||||
op.drop_index("idx_user_model_usage_model", table_name="user_model_usage_counts")
|
||||
if index_exists("user_model_usage_counts", "idx_user_model_usage_user"):
|
||||
op.drop_index("idx_user_model_usage_user", table_name="user_model_usage_counts")
|
||||
op.drop_table("user_model_usage_counts")
|
||||
|
||||
|
||||
@@ -31,7 +31,34 @@ def _table_exists(table_name: str) -> bool:
|
||||
return table_name in insp.get_table_names()
|
||||
|
||||
|
||||
def _enum_has_value(enum_name: str, value: str) -> bool:
|
||||
"""检查 PostgreSQL 枚举类型是否包含指定值"""
|
||||
bind = op.get_bind()
|
||||
result = bind.execute(
|
||||
sa.text(
|
||||
"SELECT 1 FROM pg_enum e JOIN pg_type t ON e.enumtypid = t.oid"
|
||||
" WHERE t.typname = :enum_name AND e.enumlabel = :value"
|
||||
),
|
||||
{"enum_name": enum_name, "value": value},
|
||||
)
|
||||
return result.fetchone() is not None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# proxy_nodes: 将已废弃的 unhealthy 状态迁移为 offline,然后从枚举中移除
|
||||
if _enum_has_value("proxynodestatus", "unhealthy"):
|
||||
op.execute("UPDATE proxy_nodes SET status = 'offline' WHERE status = 'unhealthy'")
|
||||
op.execute("ALTER TYPE proxynodestatus RENAME TO proxynodestatus_old")
|
||||
op.execute("CREATE TYPE proxynodestatus AS ENUM ('online', 'offline')")
|
||||
# 必须先移除旧枚举类型的 DEFAULT,否则 ALTER TYPE 会因无法转换默认值而报错
|
||||
op.execute("ALTER TABLE proxy_nodes ALTER COLUMN status DROP DEFAULT")
|
||||
op.execute(
|
||||
"ALTER TABLE proxy_nodes ALTER COLUMN status TYPE proxynodestatus"
|
||||
" USING status::text::proxynodestatus"
|
||||
)
|
||||
op.execute("ALTER TABLE proxy_nodes ALTER COLUMN status SET DEFAULT 'online'::proxynodestatus")
|
||||
op.execute("DROP TYPE proxynodestatus_old")
|
||||
|
||||
# proxy_nodes: 新增错误指标字段
|
||||
if not _column_exists("proxy_nodes", "failed_requests"):
|
||||
op.add_column(
|
||||
@@ -102,6 +129,18 @@ def upgrade() -> None:
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
# 恢复 proxynodestatus 枚举,加回 unhealthy
|
||||
if not _enum_has_value("proxynodestatus", "unhealthy"):
|
||||
op.execute("ALTER TYPE proxynodestatus RENAME TO proxynodestatus_old")
|
||||
op.execute("CREATE TYPE proxynodestatus AS ENUM ('online', 'unhealthy', 'offline')")
|
||||
op.execute("ALTER TABLE proxy_nodes ALTER COLUMN status DROP DEFAULT")
|
||||
op.execute(
|
||||
"ALTER TABLE proxy_nodes ALTER COLUMN status TYPE proxynodestatus"
|
||||
" USING status::text::proxynodestatus"
|
||||
)
|
||||
op.execute("ALTER TABLE proxy_nodes ALTER COLUMN status SET DEFAULT 'online'::proxynodestatus")
|
||||
op.execute("DROP TYPE proxynodestatus_old")
|
||||
|
||||
if _table_exists("proxy_node_events"):
|
||||
op.drop_index(op.f("ix_proxy_node_events_node_id"), table_name="proxy_node_events")
|
||||
op.drop_index("idx_proxy_node_events_node_created", table_name="proxy_node_events")
|
||||
|
||||
+49
@@ -0,0 +1,49 @@
|
||||
"""add_request_candidates_composite_indexes
|
||||
|
||||
Revision ID: 00b9161b8729
|
||||
Revises: 48afe197cc15
|
||||
Create Date: 2026-02-28 14:48:00.000000+00:00
|
||||
|
||||
"""
|
||||
|
||||
from sqlalchemy import inspect
|
||||
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "00b9161b8729"
|
||||
down_revision = "48afe197cc15"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def _index_exists(index_name: str) -> bool:
|
||||
bind = op.get_bind()
|
||||
insp = inspect(bind)
|
||||
indexes = insp.get_indexes("request_candidates")
|
||||
return any(idx["name"] == index_name for idx in indexes)
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# (request_id, status) - fallback/retry 查询优化
|
||||
if not _index_exists("idx_rc_request_id_status"):
|
||||
op.create_index(
|
||||
"idx_rc_request_id_status",
|
||||
"request_candidates",
|
||||
["request_id", "status"],
|
||||
)
|
||||
|
||||
# (provider_id, status, created_at) - provider 聚合统计优化
|
||||
if not _index_exists("idx_rc_provider_status_created"):
|
||||
op.create_index(
|
||||
"idx_rc_provider_status_created",
|
||||
"request_candidates",
|
||||
["provider_id", "status", "created_at"],
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
if _index_exists("idx_rc_provider_status_created"):
|
||||
op.drop_index("idx_rc_provider_status_created", table_name="request_candidates")
|
||||
if _index_exists("idx_rc_request_id_status"):
|
||||
op.drop_index("idx_rc_request_id_status", table_name="request_candidates")
|
||||
@@ -0,0 +1,263 @@
|
||||
"""vertex_ai_provider_type
|
||||
|
||||
Migrate legacy Vertex auth_type/provider_type into the new model:
|
||||
- provider_type=vertex_ai
|
||||
- auth_type=service_account (legacy vertex_ai renamed)
|
||||
- fixed Vertex endpoints: gemini:chat + claude:chat
|
||||
|
||||
Revision ID: 2a624af8dd3a
|
||||
Revises: 00b9161b8729
|
||||
Create Date: 2026-02-28 15:00:00.000000+00:00
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
|
||||
import sqlalchemy as sa
|
||||
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "2a624af8dd3a"
|
||||
down_revision = "00b9161b8729"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
_VERTEX_BASE_URL = "https://aiplatform.googleapis.com"
|
||||
_VERTEX_ENDPOINTS: tuple[tuple[str, str, str], ...] = (
|
||||
("gemini:chat", "gemini", "chat"),
|
||||
("claude:chat", "claude", "chat"),
|
||||
)
|
||||
_VERTEX_KEY_FORMATS_SA = '["gemini:chat","claude:chat"]'
|
||||
_VERTEX_KEY_FORMATS_API_KEY = '["gemini:chat"]'
|
||||
|
||||
|
||||
def _select_vertex_provider_ids(conn: sa.Connection) -> list[str]:
|
||||
"""Collect providers that should be treated as Vertex after migration."""
|
||||
rows = conn.execute(sa.text("""
|
||||
SELECT DISTINCT p.id
|
||||
FROM providers p
|
||||
LEFT JOIN provider_api_keys pak ON pak.provider_id = p.id
|
||||
WHERE lower(COALESCE(p.provider_type, '')) = 'vertex_ai'
|
||||
OR pak.auth_type = 'vertex_ai'
|
||||
"""))
|
||||
return [str(row[0]) for row in rows if row[0]]
|
||||
|
||||
|
||||
def _ensure_fixed_vertex_endpoints(conn: sa.Connection, provider_ids: list[str]) -> None:
|
||||
"""Ensure every Vertex provider has fixed gemini:chat + claude:chat endpoints."""
|
||||
for provider_id in provider_ids:
|
||||
provider_max_retries = (
|
||||
conn.execute(
|
||||
sa.text("""
|
||||
SELECT COALESCE(max_retries, 2)
|
||||
FROM providers
|
||||
WHERE id = :provider_id
|
||||
"""),
|
||||
{"provider_id": provider_id},
|
||||
).scalar()
|
||||
or 2
|
||||
)
|
||||
|
||||
for api_format, api_family, endpoint_kind in _VERTEX_ENDPOINTS:
|
||||
# Normalize existing fixed endpoint fields.
|
||||
conn.execute(
|
||||
sa.text("""
|
||||
UPDATE provider_endpoints
|
||||
SET
|
||||
api_family = :api_family,
|
||||
endpoint_kind = :endpoint_kind,
|
||||
base_url = :base_url,
|
||||
custom_path = NULL,
|
||||
is_active = TRUE,
|
||||
updated_at = CURRENT_TIMESTAMP
|
||||
WHERE provider_id = :provider_id
|
||||
AND api_format = :api_format
|
||||
"""),
|
||||
{
|
||||
"provider_id": provider_id,
|
||||
"api_format": api_format,
|
||||
"api_family": api_family,
|
||||
"endpoint_kind": endpoint_kind,
|
||||
"base_url": _VERTEX_BASE_URL,
|
||||
},
|
||||
)
|
||||
|
||||
exists = conn.execute(
|
||||
sa.text("""
|
||||
SELECT 1
|
||||
FROM provider_endpoints
|
||||
WHERE provider_id = :provider_id
|
||||
AND api_format = :api_format
|
||||
LIMIT 1
|
||||
"""),
|
||||
{"provider_id": provider_id, "api_format": api_format},
|
||||
).first()
|
||||
|
||||
if not exists:
|
||||
conn.execute(
|
||||
sa.text("""
|
||||
INSERT INTO provider_endpoints (
|
||||
id,
|
||||
provider_id,
|
||||
api_format,
|
||||
api_family,
|
||||
endpoint_kind,
|
||||
base_url,
|
||||
custom_path,
|
||||
header_rules,
|
||||
body_rules,
|
||||
max_retries,
|
||||
is_active,
|
||||
config,
|
||||
format_acceptance_config,
|
||||
proxy,
|
||||
created_at,
|
||||
updated_at
|
||||
)
|
||||
VALUES (
|
||||
:id,
|
||||
:provider_id,
|
||||
:api_format,
|
||||
:api_family,
|
||||
:endpoint_kind,
|
||||
:base_url,
|
||||
NULL,
|
||||
NULL,
|
||||
NULL,
|
||||
:max_retries,
|
||||
TRUE,
|
||||
NULL,
|
||||
NULL,
|
||||
NULL,
|
||||
CURRENT_TIMESTAMP,
|
||||
CURRENT_TIMESTAMP
|
||||
)
|
||||
"""),
|
||||
{
|
||||
"id": str(uuid.uuid4()),
|
||||
"provider_id": provider_id,
|
||||
"api_format": api_format,
|
||||
"api_family": api_family,
|
||||
"endpoint_kind": endpoint_kind,
|
||||
"base_url": _VERTEX_BASE_URL,
|
||||
"max_retries": int(provider_max_retries),
|
||||
},
|
||||
)
|
||||
|
||||
# Vertex fixed-provider model: disable non-fixed endpoints.
|
||||
conn.execute(
|
||||
sa.text("""
|
||||
UPDATE provider_endpoints
|
||||
SET
|
||||
is_active = FALSE,
|
||||
updated_at = CURRENT_TIMESTAMP
|
||||
WHERE provider_id = :provider_id
|
||||
AND api_format NOT IN ('gemini:chat', 'claude:chat')
|
||||
"""),
|
||||
{"provider_id": provider_id},
|
||||
)
|
||||
|
||||
|
||||
def _normalize_vertex_key_formats(conn: sa.Connection, provider_ids: list[str]) -> None:
|
||||
"""Normalize key.api_formats for Vertex keys by auth type."""
|
||||
for provider_id in provider_ids:
|
||||
# Service Account (and legacy vertex_ai) keys: allow Gemini + Claude models.
|
||||
conn.execute(
|
||||
sa.text("""
|
||||
UPDATE provider_api_keys
|
||||
SET
|
||||
api_formats = CAST(:api_formats AS json),
|
||||
updated_at = CURRENT_TIMESTAMP
|
||||
WHERE provider_id = :provider_id
|
||||
AND auth_type IN ('service_account', 'vertex_ai')
|
||||
"""),
|
||||
{
|
||||
"provider_id": provider_id,
|
||||
"api_formats": _VERTEX_KEY_FORMATS_SA,
|
||||
},
|
||||
)
|
||||
|
||||
# API Key mode on Vertex 仅支持 Gemini(Google publisher)。
|
||||
conn.execute(
|
||||
sa.text("""
|
||||
UPDATE provider_api_keys
|
||||
SET
|
||||
api_formats = CAST(:api_formats AS json),
|
||||
updated_at = CURRENT_TIMESTAMP
|
||||
WHERE provider_id = :provider_id
|
||||
AND auth_type = 'api_key'
|
||||
"""),
|
||||
{
|
||||
"provider_id": provider_id,
|
||||
"api_formats": _VERTEX_KEY_FORMATS_API_KEY,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
|
||||
# 1) 收集目标 Provider(兼容重复执行,先识别 legacy/new 两种来源)。
|
||||
provider_ids = _select_vertex_provider_ids(conn)
|
||||
|
||||
# 2) 先重命名 auth_type(legacy vertex_ai -> service_account)。
|
||||
conn.execute(sa.text("""
|
||||
UPDATE provider_api_keys
|
||||
SET auth_type = 'service_account'
|
||||
WHERE auth_type = 'vertex_ai'
|
||||
"""))
|
||||
|
||||
if not provider_ids:
|
||||
return
|
||||
|
||||
# 3) 归一 provider_type,并启用格式转换(Vertex 同时承载 Gemini/Claude)。
|
||||
for provider_id in provider_ids:
|
||||
conn.execute(
|
||||
sa.text("""
|
||||
UPDATE providers
|
||||
SET
|
||||
provider_type = 'vertex_ai',
|
||||
enable_format_conversion = TRUE
|
||||
WHERE id = :provider_id
|
||||
"""),
|
||||
{"provider_id": provider_id},
|
||||
)
|
||||
|
||||
# 4) 固定端点落地:gemini:chat + claude:chat。
|
||||
_ensure_fixed_vertex_endpoints(conn, provider_ids)
|
||||
|
||||
# 5) 归一 key 的 api_formats,避免调度命中旧格式。
|
||||
_normalize_vertex_key_formats(conn, provider_ids)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
|
||||
provider_rows = conn.execute(sa.text("""
|
||||
SELECT id
|
||||
FROM providers
|
||||
WHERE lower(COALESCE(provider_type, '')) = 'vertex_ai'
|
||||
"""))
|
||||
provider_ids = [str(row[0]) for row in provider_rows if row[0]]
|
||||
|
||||
if provider_ids:
|
||||
for provider_id in provider_ids:
|
||||
conn.execute(
|
||||
sa.text("""
|
||||
UPDATE provider_api_keys
|
||||
SET auth_type = 'vertex_ai'
|
||||
WHERE provider_id = :provider_id
|
||||
AND auth_type = 'service_account'
|
||||
"""),
|
||||
{"provider_id": provider_id},
|
||||
)
|
||||
conn.execute(
|
||||
sa.text("""
|
||||
UPDATE providers
|
||||
SET provider_type = 'custom'
|
||||
WHERE id = :provider_id
|
||||
"""),
|
||||
{"provider_id": provider_id},
|
||||
)
|
||||
@@ -0,0 +1,199 @@
|
||||
"""backfill_codex_compact_endpoint
|
||||
|
||||
Backfill Codex reverse-proxy endpoints:
|
||||
- ensure `openai:cli` endpoint is pinned to force_stream
|
||||
- ensure `openai:compact` endpoint exists
|
||||
|
||||
Revision ID: f0c3a7b9d1e2
|
||||
Revises: 2a624af8dd3a
|
||||
Create Date: 2026-03-01 17:00:00.000000+00:00
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from typing import Any
|
||||
|
||||
import sqlalchemy as sa
|
||||
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "f0c3a7b9d1e2"
|
||||
down_revision = "2a624af8dd3a"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
_CODEX_BASE_URL = "https://chatgpt.com/backend-api/codex"
|
||||
_COMPACT_FORMAT = "openai:compact"
|
||||
_CLI_FORMAT = "openai:cli"
|
||||
_FORCE_STREAM = "force_stream"
|
||||
|
||||
|
||||
def _find_codex_provider_ids(conn: sa.Connection) -> list[str]:
|
||||
"""Find Codex providers (by provider_type or legacy base_url pattern)."""
|
||||
rows = conn.execute(sa.text("""
|
||||
SELECT DISTINCT p.id
|
||||
FROM providers p
|
||||
LEFT JOIN provider_endpoints pe ON pe.provider_id = p.id
|
||||
WHERE lower(COALESCE(p.provider_type, '')) = 'codex'
|
||||
OR (
|
||||
lower(COALESCE(pe.api_format, '')) = 'openai:cli'
|
||||
AND lower(COALESCE(pe.base_url, '')) LIKE '%/backend-api/codex%'
|
||||
)
|
||||
"""))
|
||||
return [str(r[0]) for r in rows if r[0]]
|
||||
|
||||
|
||||
def _get_cli_endpoint(conn: sa.Connection, provider_id: str) -> dict[str, Any] | None:
|
||||
"""Load existing openai:cli endpoint for the provider."""
|
||||
row = (
|
||||
conn.execute(
|
||||
sa.text("""
|
||||
SELECT base_url, header_rules, body_rules, max_retries, proxy, config
|
||||
FROM provider_endpoints
|
||||
WHERE provider_id = :pid AND api_format = :fmt
|
||||
LIMIT 1
|
||||
"""),
|
||||
{"pid": provider_id, "fmt": _CLI_FORMAT},
|
||||
)
|
||||
.mappings()
|
||||
.first()
|
||||
)
|
||||
return dict(row) if row else None
|
||||
|
||||
|
||||
def _pin_cli_force_stream(conn: sa.Connection, provider_id: str, cli: dict[str, Any]) -> None:
|
||||
"""Set upstream_stream_policy=force_stream on existing cli endpoint."""
|
||||
cfg = dict(cli.get("config") or {}) if isinstance(cli.get("config"), dict) else {}
|
||||
cfg.pop("upstreamStreamPolicy", None)
|
||||
cfg.pop("upstream_stream", None)
|
||||
cfg["upstream_stream_policy"] = _FORCE_STREAM
|
||||
conn.execute(
|
||||
sa.text("""
|
||||
UPDATE provider_endpoints
|
||||
SET api_family = 'openai',
|
||||
endpoint_kind = 'cli',
|
||||
config = CAST(:config AS json),
|
||||
updated_at = CURRENT_TIMESTAMP
|
||||
WHERE provider_id = :pid AND api_format = :fmt
|
||||
"""),
|
||||
{
|
||||
"pid": provider_id,
|
||||
"fmt": _CLI_FORMAT,
|
||||
"config": json.dumps(cfg, ensure_ascii=False),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _ensure_compact_endpoint(conn: sa.Connection, provider_id: str, cli: dict[str, Any]) -> None:
|
||||
"""Create openai:compact endpoint if missing (clone from cli)."""
|
||||
exists = conn.execute(
|
||||
sa.text(
|
||||
"SELECT 1 FROM provider_endpoints WHERE provider_id = :pid AND api_format = :fmt LIMIT 1"
|
||||
),
|
||||
{"pid": provider_id, "fmt": _COMPACT_FORMAT},
|
||||
).first()
|
||||
if exists:
|
||||
# Already exists, just ensure api_family/endpoint_kind are set.
|
||||
conn.execute(
|
||||
sa.text("""
|
||||
UPDATE provider_endpoints
|
||||
SET api_family = 'openai', endpoint_kind = 'compact',
|
||||
updated_at = CURRENT_TIMESTAMP
|
||||
WHERE provider_id = :pid AND api_format = :fmt
|
||||
"""),
|
||||
{"pid": provider_id, "fmt": _COMPACT_FORMAT},
|
||||
)
|
||||
return
|
||||
|
||||
# Clone from cli endpoint, strip stream policy.
|
||||
cfg = dict(cli.get("config") or {}) if isinstance(cli.get("config"), dict) else {}
|
||||
for k in ("upstream_stream_policy", "upstreamStreamPolicy", "upstream_stream"):
|
||||
cfg.pop(k, None)
|
||||
|
||||
def _json(val: Any) -> str | None:
|
||||
return json.dumps(val, ensure_ascii=False) if val is not None else None
|
||||
|
||||
conn.execute(
|
||||
sa.text("""
|
||||
INSERT INTO provider_endpoints (
|
||||
id, provider_id, api_format, api_family, endpoint_kind,
|
||||
base_url, custom_path, header_rules, body_rules,
|
||||
max_retries, is_active, config, format_acceptance_config,
|
||||
proxy, created_at, updated_at
|
||||
) VALUES (
|
||||
:id, :pid, :fmt, 'openai', 'compact',
|
||||
:base_url, NULL, CAST(:header_rules AS json), CAST(:body_rules AS json),
|
||||
:max_retries, TRUE, CAST(:config AS json), NULL,
|
||||
CAST(:proxy AS jsonb), CURRENT_TIMESTAMP, CURRENT_TIMESTAMP
|
||||
)
|
||||
"""),
|
||||
{
|
||||
"id": str(uuid.uuid4()),
|
||||
"pid": provider_id,
|
||||
"fmt": _COMPACT_FORMAT,
|
||||
"base_url": cli.get("base_url") or _CODEX_BASE_URL,
|
||||
"header_rules": _json(cli.get("header_rules")),
|
||||
"body_rules": _json(cli.get("body_rules")),
|
||||
"max_retries": cli.get("max_retries") or 2,
|
||||
"config": _json(cfg or None),
|
||||
"proxy": _json(cli.get("proxy")),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _add_compact_to_key_formats(conn: sa.Connection, provider_id: str) -> None:
|
||||
"""Ensure provider keys include openai:compact in api_formats."""
|
||||
rows = (
|
||||
conn.execute(
|
||||
sa.text("SELECT id, api_formats FROM provider_api_keys WHERE provider_id = :pid"),
|
||||
{"pid": provider_id},
|
||||
)
|
||||
.mappings()
|
||||
.all()
|
||||
)
|
||||
for row in rows:
|
||||
raw = row["api_formats"]
|
||||
formats: list[str] = []
|
||||
if isinstance(raw, list):
|
||||
for item in raw:
|
||||
v = str(item or "").strip().lower()
|
||||
if v and v not in formats:
|
||||
formats.append(v)
|
||||
|
||||
if _COMPACT_FORMAT in formats:
|
||||
continue
|
||||
|
||||
# Insert compact right after cli, or at end.
|
||||
if _CLI_FORMAT in formats:
|
||||
idx = formats.index(_CLI_FORMAT) + 1
|
||||
formats.insert(idx, _COMPACT_FORMAT)
|
||||
else:
|
||||
formats.append(_COMPACT_FORMAT)
|
||||
|
||||
conn.execute(
|
||||
sa.text("""
|
||||
UPDATE provider_api_keys
|
||||
SET api_formats = CAST(:fmts AS json), updated_at = CURRENT_TIMESTAMP
|
||||
WHERE id = :id
|
||||
"""),
|
||||
{"id": row["id"], "fmts": json.dumps(formats, ensure_ascii=False)},
|
||||
)
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
for provider_id in _find_codex_provider_ids(conn):
|
||||
cli = _get_cli_endpoint(conn, provider_id)
|
||||
if not cli:
|
||||
continue # No cli endpoint to clone from; skip.
|
||||
_pin_cli_force_stream(conn, provider_id, cli)
|
||||
_ensure_compact_endpoint(conn, provider_id, cli)
|
||||
_add_compact_to_key_formats(conn, provider_id)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
# Data backfill: no-op to avoid deleting user-managed data.
|
||||
return
|
||||
@@ -0,0 +1,29 @@
|
||||
"""add_proxy_metadata_to_proxy_nodes
|
||||
|
||||
Revision ID: 1d2e3f4a5b6c
|
||||
Revises: f0c3a7b9d1e2
|
||||
Create Date: 2026-03-02 13:00:00.000000+00:00
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "1d2e3f4a5b6c"
|
||||
down_revision: str | None = "f0c3a7b9d1e2"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.execute("ALTER TABLE public.proxy_nodes ADD COLUMN IF NOT EXISTS proxy_metadata json")
|
||||
op.execute(
|
||||
"COMMENT ON COLUMN public.proxy_nodes.proxy_metadata IS 'aether-proxy 上报元数据(版本等)'"
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.execute("ALTER TABLE public.proxy_nodes DROP COLUMN IF EXISTS proxy_metadata")
|
||||
@@ -0,0 +1,72 @@
|
||||
"""backfill_codex_default_body_rules
|
||||
|
||||
Backfill default body_rules for codex providers with openai:cli endpoints
|
||||
that currently have body_rules IS NULL.
|
||||
|
||||
Rules:
|
||||
- drop max_output_tokens
|
||||
- drop temperature
|
||||
- drop top_p
|
||||
- set store = false
|
||||
- set instructions = "You are GPT-5." (when instructions not exists)
|
||||
|
||||
Revision ID: dd0278c0a28c
|
||||
Revises: 1d2e3f4a5b6c
|
||||
Create Date: 2026-03-02 15:00:00.000000+00:00
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
import sqlalchemy as sa
|
||||
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "dd0278c0a28c"
|
||||
down_revision = "1d2e3f4a5b6c"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
_TARGET_FORMATS = ("openai:cli",)
|
||||
|
||||
_DEFAULT_BODY_RULES = [
|
||||
{"action": "drop", "path": "max_output_tokens"},
|
||||
{"action": "drop", "path": "temperature"},
|
||||
{"action": "drop", "path": "top_p"},
|
||||
{"action": "set", "path": "store", "value": False},
|
||||
{
|
||||
"action": "set",
|
||||
"path": "instructions",
|
||||
"value": "You are GPT-5.",
|
||||
"condition": {"path": "instructions", "op": "not_exists"},
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
|
||||
# 幂等性: 仅回填 codex 提供商中 body_rules 为空(SQL NULL 或 JSON null)的记录
|
||||
rules_json = json.dumps(_DEFAULT_BODY_RULES, ensure_ascii=False)
|
||||
result = conn.execute(
|
||||
sa.text("""
|
||||
UPDATE provider_endpoints pe
|
||||
SET body_rules = CAST(:rules AS json),
|
||||
updated_at = CURRENT_TIMESTAMP
|
||||
FROM providers p
|
||||
WHERE pe.provider_id = p.id
|
||||
AND p.provider_type = :ptype
|
||||
AND pe.api_format = :fmt
|
||||
AND (pe.body_rules IS NULL OR pe.body_rules::text = 'null')
|
||||
"""),
|
||||
{"rules": rules_json, "ptype": "codex", "fmt": _TARGET_FORMATS[0]},
|
||||
)
|
||||
if result.rowcount:
|
||||
print(f" backfilled body_rules for {result.rowcount} endpoint(s)")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
# Data backfill: no-op to avoid removing user-customized rules.
|
||||
return
|
||||
@@ -0,0 +1,45 @@
|
||||
"""add_idx_usage_provider_key
|
||||
|
||||
Add composite index on usage(provider_id, provider_api_key_id) to support
|
||||
the pool management page's per-key usage stats aggregation query.
|
||||
|
||||
Revision ID: 0ba031f328de
|
||||
Revises: dd0278c0a28c
|
||||
Create Date: 2026-03-03 10:00:00.000000+00:00
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
revision = "0ba031f328de"
|
||||
down_revision = "dd0278c0a28c"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
INDEX_NAME = "idx_usage_provider_key"
|
||||
TABLE = "usage"
|
||||
COLUMNS = ["provider_id", "provider_api_key_id"]
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
bind = op.get_bind()
|
||||
result = bind.execute(
|
||||
sa.text("SELECT 1 FROM pg_indexes WHERE indexname = :name"),
|
||||
{"name": INDEX_NAME},
|
||||
).fetchone()
|
||||
if result:
|
||||
return
|
||||
op.create_index(INDEX_NAME, TABLE, COLUMNS)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
bind = op.get_bind()
|
||||
result = bind.execute(
|
||||
sa.text("SELECT 1 FROM pg_indexes WHERE indexname = :name"),
|
||||
{"name": INDEX_NAME},
|
||||
).fetchone()
|
||||
if not result:
|
||||
return
|
||||
op.drop_index(INDEX_NAME, table_name=TABLE)
|
||||
@@ -0,0 +1,45 @@
|
||||
"""add_idx_usage_status_user_created
|
||||
|
||||
Add composite index on usage(status, user_id, created_at) to speed up
|
||||
interval timeline and active usage analytics queries.
|
||||
|
||||
Revision ID: 5f1d2e3c4b5a
|
||||
Revises: 0ba031f328de
|
||||
Create Date: 2026-03-03 17:30:00.000000+00:00
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
revision = "5f1d2e3c4b5a"
|
||||
down_revision = "0ba031f328de"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
INDEX_NAME = "idx_usage_status_user_created"
|
||||
TABLE = "usage"
|
||||
COLUMNS = ["status", "user_id", "created_at"]
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
bind = op.get_bind()
|
||||
result = bind.execute(
|
||||
sa.text("SELECT 1 FROM pg_indexes WHERE indexname = :name"),
|
||||
{"name": INDEX_NAME},
|
||||
).fetchone()
|
||||
if result:
|
||||
return
|
||||
op.create_index(INDEX_NAME, TABLE, COLUMNS)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
bind = op.get_bind()
|
||||
result = bind.execute(
|
||||
sa.text("SELECT 1 FROM pg_indexes WHERE indexname = :name"),
|
||||
{"name": INDEX_NAME},
|
||||
).fetchone()
|
||||
if not result:
|
||||
return
|
||||
op.drop_index(INDEX_NAME, table_name=TABLE)
|
||||
@@ -0,0 +1,41 @@
|
||||
"""add fingerprint column to provider_api_keys
|
||||
|
||||
Revision ID: 6a9b8c7d5e4f
|
||||
Revises: 5f1d2e3c4b5a
|
||||
Create Date: 2026-03-04 23:50:00.000000
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy import inspect
|
||||
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "6a9b8c7d5e4f"
|
||||
down_revision: str | None = "5f1d2e3c4b5a"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def column_exists(table_name: str, column_name: str) -> bool:
|
||||
bind = op.get_bind()
|
||||
inspector = inspect(bind)
|
||||
columns = [c["name"] for c in inspector.get_columns(table_name)]
|
||||
return column_name in columns
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
if not column_exists("provider_api_keys", "fingerprint"):
|
||||
op.add_column(
|
||||
"provider_api_keys",
|
||||
sa.Column("fingerprint", sa.JSON(), nullable=True),
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
if column_exists("provider_api_keys", "fingerprint"):
|
||||
op.drop_column("provider_api_keys", "fingerprint")
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,65 @@
|
||||
"""remove standalone api key locking
|
||||
|
||||
Revision ID: 7c91d2e4f8a1
|
||||
Revises: 6f7a8b9c0d1e
|
||||
Create Date: 2026-03-05 17:00:00.000000+00:00
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy import inspect
|
||||
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "7c91d2e4f8a1"
|
||||
down_revision: str | None = "6f7a8b9c0d1e"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
_CONSTRAINT_NAME = "ck_api_keys_standalone_not_locked"
|
||||
|
||||
|
||||
def _column_exists(table_name: str, column_name: str) -> bool:
|
||||
bind = op.get_bind()
|
||||
insp = inspect(bind)
|
||||
insp.clear_cache()
|
||||
return column_name in [c["name"] for c in insp.get_columns(table_name)]
|
||||
|
||||
|
||||
def _check_constraint_exists(table_name: str, constraint_name: str) -> bool:
|
||||
bind = op.get_bind()
|
||||
insp = inspect(bind)
|
||||
insp.clear_cache()
|
||||
return any(c.get("name") == constraint_name for c in insp.get_check_constraints(table_name))
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
if not (
|
||||
_column_exists("api_keys", "is_standalone")
|
||||
and _column_exists("api_keys", "is_locked")
|
||||
and _column_exists("api_keys", "is_active")
|
||||
):
|
||||
return
|
||||
|
||||
op.execute(sa.text("""
|
||||
UPDATE api_keys
|
||||
SET is_active = FALSE,
|
||||
is_locked = FALSE
|
||||
WHERE is_standalone IS TRUE AND is_locked IS TRUE
|
||||
"""))
|
||||
|
||||
if not _check_constraint_exists("api_keys", _CONSTRAINT_NAME):
|
||||
op.create_check_constraint(
|
||||
_CONSTRAINT_NAME,
|
||||
"api_keys",
|
||||
"(NOT is_standalone) OR (NOT is_locked)",
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
if _check_constraint_exists("api_keys", _CONSTRAINT_NAME):
|
||||
op.drop_constraint(_CONSTRAINT_NAME, "api_keys", type_="check")
|
||||
+220
@@ -0,0 +1,220 @@
|
||||
"""tighten wallet transaction snapshots and remove wallet version
|
||||
|
||||
Revision ID: 8e71f2a4c9b0
|
||||
Revises: 7c91d2e4f8a1
|
||||
Create Date: 2026-03-07 13:00:00.000000+00:00
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy import inspect
|
||||
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "8e71f2a4c9b0"
|
||||
down_revision: str | None = "7c91d2e4f8a1"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
_WALLET_TX_BEFORE_CHECK = "ck_wallet_tx_balance_before_consistent"
|
||||
_WALLET_TX_AFTER_CHECK = "ck_wallet_tx_balance_after_consistent"
|
||||
_WALLET_LIMIT_MODE_INDEX = "idx_wallets_limit_mode"
|
||||
|
||||
|
||||
def _table_exists(table_name: str) -> bool:
|
||||
bind = op.get_bind()
|
||||
insp = inspect(bind)
|
||||
insp.clear_cache()
|
||||
return table_name in insp.get_table_names()
|
||||
|
||||
|
||||
def _column_exists(table_name: str, column_name: str) -> bool:
|
||||
bind = op.get_bind()
|
||||
insp = inspect(bind)
|
||||
insp.clear_cache()
|
||||
return column_name in [c["name"] for c in insp.get_columns(table_name)]
|
||||
|
||||
|
||||
def _index_exists(table_name: str, index_name: str) -> bool:
|
||||
bind = op.get_bind()
|
||||
insp = inspect(bind)
|
||||
insp.clear_cache()
|
||||
return any(index.get("name") == index_name for index in insp.get_indexes(table_name))
|
||||
|
||||
|
||||
def _check_constraint_exists(table_name: str, constraint_name: str) -> bool:
|
||||
bind = op.get_bind()
|
||||
insp = inspect(bind)
|
||||
insp.clear_cache()
|
||||
return any(c.get("name") == constraint_name for c in insp.get_check_constraints(table_name))
|
||||
|
||||
|
||||
def _tighten_wallet_transaction_snapshots() -> None:
|
||||
if not _table_exists("wallet_transactions"):
|
||||
return
|
||||
|
||||
required_columns = {
|
||||
"balance_before",
|
||||
"balance_after",
|
||||
"recharge_balance_before",
|
||||
"recharge_balance_after",
|
||||
"gift_balance_before",
|
||||
"gift_balance_after",
|
||||
}
|
||||
existing_columns = {
|
||||
column["name"] for column in inspect(op.get_bind()).get_columns("wallet_transactions")
|
||||
}
|
||||
if not required_columns.issubset(existing_columns):
|
||||
return
|
||||
|
||||
op.execute(
|
||||
sa.text(
|
||||
"""
|
||||
UPDATE wallet_transactions
|
||||
SET recharge_balance_before = balance_before
|
||||
WHERE recharge_balance_before IS NULL
|
||||
"""
|
||||
)
|
||||
)
|
||||
op.execute(
|
||||
sa.text(
|
||||
"""
|
||||
UPDATE wallet_transactions
|
||||
SET recharge_balance_after = balance_after
|
||||
WHERE recharge_balance_after IS NULL
|
||||
"""
|
||||
)
|
||||
)
|
||||
op.execute(
|
||||
sa.text(
|
||||
"""
|
||||
UPDATE wallet_transactions
|
||||
SET gift_balance_before = 0
|
||||
WHERE gift_balance_before IS NULL
|
||||
"""
|
||||
)
|
||||
)
|
||||
op.execute(
|
||||
sa.text(
|
||||
"""
|
||||
UPDATE wallet_transactions
|
||||
SET gift_balance_after = 0
|
||||
WHERE gift_balance_after IS NULL
|
||||
"""
|
||||
)
|
||||
)
|
||||
op.execute(
|
||||
sa.text(
|
||||
"""
|
||||
UPDATE wallet_transactions
|
||||
SET balance_before = recharge_balance_before + gift_balance_before,
|
||||
balance_after = recharge_balance_after + gift_balance_after
|
||||
"""
|
||||
)
|
||||
)
|
||||
|
||||
if not _check_constraint_exists("wallet_transactions", _WALLET_TX_BEFORE_CHECK):
|
||||
op.create_check_constraint(
|
||||
_WALLET_TX_BEFORE_CHECK,
|
||||
"wallet_transactions",
|
||||
"balance_before = recharge_balance_before + gift_balance_before",
|
||||
)
|
||||
if not _check_constraint_exists("wallet_transactions", _WALLET_TX_AFTER_CHECK):
|
||||
op.create_check_constraint(
|
||||
_WALLET_TX_AFTER_CHECK,
|
||||
"wallet_transactions",
|
||||
"balance_after = recharge_balance_after + gift_balance_after",
|
||||
)
|
||||
|
||||
op.alter_column(
|
||||
"wallet_transactions",
|
||||
"recharge_balance_before",
|
||||
existing_type=sa.Numeric(20, 8),
|
||||
nullable=False,
|
||||
)
|
||||
op.alter_column(
|
||||
"wallet_transactions",
|
||||
"recharge_balance_after",
|
||||
existing_type=sa.Numeric(20, 8),
|
||||
nullable=False,
|
||||
)
|
||||
op.alter_column(
|
||||
"wallet_transactions",
|
||||
"gift_balance_before",
|
||||
existing_type=sa.Numeric(20, 8),
|
||||
nullable=False,
|
||||
)
|
||||
op.alter_column(
|
||||
"wallet_transactions",
|
||||
"gift_balance_after",
|
||||
existing_type=sa.Numeric(20, 8),
|
||||
nullable=False,
|
||||
)
|
||||
|
||||
|
||||
def _drop_wallet_cleanup_artifacts() -> None:
|
||||
if not _table_exists("wallets"):
|
||||
return
|
||||
|
||||
if _index_exists("wallets", _WALLET_LIMIT_MODE_INDEX):
|
||||
op.drop_index(_WALLET_LIMIT_MODE_INDEX, table_name="wallets")
|
||||
|
||||
if _column_exists("wallets", "version"):
|
||||
op.drop_column("wallets", "version")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
_tighten_wallet_transaction_snapshots()
|
||||
_drop_wallet_cleanup_artifacts()
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
if _table_exists("wallets"):
|
||||
if not _column_exists("wallets", "version"):
|
||||
op.add_column(
|
||||
"wallets",
|
||||
sa.Column("version", sa.Integer(), nullable=False, server_default="0"),
|
||||
)
|
||||
if not _index_exists("wallets", _WALLET_LIMIT_MODE_INDEX):
|
||||
op.create_index(_WALLET_LIMIT_MODE_INDEX, "wallets", ["limit_mode"])
|
||||
|
||||
if not _table_exists("wallet_transactions"):
|
||||
return
|
||||
|
||||
if _column_exists("wallet_transactions", "recharge_balance_before"):
|
||||
op.alter_column(
|
||||
"wallet_transactions",
|
||||
"recharge_balance_before",
|
||||
existing_type=sa.Numeric(20, 8),
|
||||
nullable=True,
|
||||
)
|
||||
if _column_exists("wallet_transactions", "recharge_balance_after"):
|
||||
op.alter_column(
|
||||
"wallet_transactions",
|
||||
"recharge_balance_after",
|
||||
existing_type=sa.Numeric(20, 8),
|
||||
nullable=True,
|
||||
)
|
||||
if _column_exists("wallet_transactions", "gift_balance_before"):
|
||||
op.alter_column(
|
||||
"wallet_transactions",
|
||||
"gift_balance_before",
|
||||
existing_type=sa.Numeric(20, 8),
|
||||
nullable=True,
|
||||
)
|
||||
if _column_exists("wallet_transactions", "gift_balance_after"):
|
||||
op.alter_column(
|
||||
"wallet_transactions",
|
||||
"gift_balance_after",
|
||||
existing_type=sa.Numeric(20, 8),
|
||||
nullable=True,
|
||||
)
|
||||
|
||||
if _check_constraint_exists("wallet_transactions", _WALLET_TX_AFTER_CHECK):
|
||||
op.drop_constraint(_WALLET_TX_AFTER_CHECK, "wallet_transactions", type_="check")
|
||||
if _check_constraint_exists("wallet_transactions", _WALLET_TX_BEFORE_CHECK):
|
||||
op.drop_constraint(_WALLET_TX_BEFORE_CHECK, "wallet_transactions", type_="check")
|
||||
@@ -0,0 +1,80 @@
|
||||
"""add missing foreign key indexes for cascade delete performance
|
||||
|
||||
Revision ID: 2d932114930d
|
||||
Revises: 8e71f2a4c9b0
|
||||
Create Date: 2026-03-07 16:28:48.633531+00:00
|
||||
|
||||
"""
|
||||
from alembic import op
|
||||
from sqlalchemy import inspect
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = '2d932114930d'
|
||||
down_revision = '8e71f2a4c9b0'
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def _index_exists(table_name: str, index_name: str) -> bool:
|
||||
bind = op.get_bind()
|
||||
insp = inspect(bind)
|
||||
return any(idx["name"] == index_name for idx in insp.get_indexes(table_name))
|
||||
|
||||
|
||||
def _create_index_if_not_exists(index_name: str, table_name: str, columns: list[str]) -> None:
|
||||
if not _index_exists(table_name, index_name):
|
||||
op.create_index(op.f(index_name), table_name, columns, unique=False)
|
||||
|
||||
|
||||
def _drop_index_if_exists(index_name: str, table_name: str) -> None:
|
||||
if _index_exists(table_name, index_name):
|
||||
op.drop_index(op.f(index_name), table_name=table_name)
|
||||
|
||||
|
||||
# (index_name, table_name, columns)
|
||||
_INDEXES = [
|
||||
# api_keys.user_id (CASCADE -> users.id)
|
||||
('ix_api_keys_user_id', 'api_keys', ['user_id']),
|
||||
# usage: wallet_id, provider_endpoint_id, provider_api_key_id (SET NULL)
|
||||
('ix_usage_wallet_id', 'usage', ['wallet_id']),
|
||||
('ix_usage_provider_endpoint_id', 'usage', ['provider_endpoint_id']),
|
||||
('ix_usage_provider_api_key_id', 'usage', ['provider_api_key_id']),
|
||||
# wallet_transactions.operator_id (SET NULL -> users.id)
|
||||
('ix_wallet_transactions_operator_id', 'wallet_transactions', ['operator_id']),
|
||||
# payment_callbacks.payment_order_id (SET NULL -> payment_orders.id)
|
||||
('ix_payment_callbacks_payment_order_id', 'payment_callbacks', ['payment_order_id']),
|
||||
# refund_requests: payment_order_id, requested_by, approved_by, processed_by (SET NULL)
|
||||
('ix_refund_requests_payment_order_id', 'refund_requests', ['payment_order_id']),
|
||||
('ix_refund_requests_requested_by', 'refund_requests', ['requested_by']),
|
||||
('ix_refund_requests_approved_by', 'refund_requests', ['approved_by']),
|
||||
('ix_refund_requests_processed_by', 'refund_requests', ['processed_by']),
|
||||
# proxy_nodes.registered_by (SET NULL -> users.id)
|
||||
('ix_proxy_nodes_registered_by', 'proxy_nodes', ['registered_by']),
|
||||
# video_tasks: api_key_id, provider_id, endpoint_id, key_id, remixed_from_task_id
|
||||
('ix_video_tasks_api_key_id', 'video_tasks', ['api_key_id']),
|
||||
('ix_video_tasks_provider_id', 'video_tasks', ['provider_id']),
|
||||
('ix_video_tasks_endpoint_id', 'video_tasks', ['endpoint_id']),
|
||||
('ix_video_tasks_key_id', 'video_tasks', ['key_id']),
|
||||
('ix_video_tasks_remixed_from_task_id', 'video_tasks', ['remixed_from_task_id']),
|
||||
# user_preferences.default_provider_id (-> providers.id)
|
||||
('ix_user_preferences_default_provider_id', 'user_preferences', ['default_provider_id']),
|
||||
# announcements.author_id (SET NULL -> users.id)
|
||||
('ix_announcements_author_id', 'announcements', ['author_id']),
|
||||
# announcement_reads.announcement_id (-> announcements.id)
|
||||
('ix_announcement_reads_announcement_id', 'announcement_reads', ['announcement_id']),
|
||||
# request_candidates: user_id, api_key_id, endpoint_id, key_id (CASCADE)
|
||||
('ix_request_candidates_user_id', 'request_candidates', ['user_id']),
|
||||
('ix_request_candidates_api_key_id', 'request_candidates', ['api_key_id']),
|
||||
('ix_request_candidates_endpoint_id', 'request_candidates', ['endpoint_id']),
|
||||
('ix_request_candidates_key_id', 'request_candidates', ['key_id']),
|
||||
]
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
for index_name, table_name, columns in _INDEXES:
|
||||
_create_index_if_not_exists(index_name, table_name, columns)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
for index_name, table_name, _columns in reversed(_INDEXES):
|
||||
_drop_index_if_exists(index_name, table_name)
|
||||
@@ -0,0 +1,224 @@
|
||||
"""usage stats retention: SET NULL on delete and add name snapshots
|
||||
|
||||
Revision ID: 45b118150a78
|
||||
Revises: 2d932114930d
|
||||
Create Date: 2026-03-08 03:48:49.622091+00:00
|
||||
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlalchemy as sa
|
||||
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "45b118150a78"
|
||||
down_revision = "2d932114930d"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
_TABLES = ["usage", "stats_user_daily", "stats_daily_api_key"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Inline helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _SchemaCache:
|
||||
def __init__(self) -> None:
|
||||
self._columns: dict[str, dict[str, str]] = {}
|
||||
self._fk_rules: dict[tuple[str, str], str] = {}
|
||||
self._fk_loaded_tables: set[str] = set()
|
||||
|
||||
def load_columns(self, tables: list[str]) -> None:
|
||||
need = [t for t in tables if t not in self._columns]
|
||||
if not need:
|
||||
return
|
||||
bind = op.get_bind()
|
||||
rows = bind.execute(
|
||||
sa.text(
|
||||
"SELECT table_name, column_name, data_type "
|
||||
"FROM information_schema.columns "
|
||||
"WHERE table_name = ANY(:tables) "
|
||||
" AND table_schema = current_schema()"
|
||||
),
|
||||
{"tables": need},
|
||||
).fetchall()
|
||||
for t in need:
|
||||
self._columns.setdefault(t, {})
|
||||
for table, col, dtype in rows:
|
||||
self._columns[table][col] = dtype
|
||||
|
||||
def load_fk_rules(self, tables: list[str]) -> None:
|
||||
need = [t for t in tables if t not in self._fk_loaded_tables]
|
||||
if not need:
|
||||
return
|
||||
bind = op.get_bind()
|
||||
rows = bind.execute(
|
||||
sa.text(
|
||||
"SELECT tc.table_name, tc.constraint_name, rc.delete_rule "
|
||||
"FROM information_schema.referential_constraints rc "
|
||||
"JOIN information_schema.table_constraints tc "
|
||||
" ON rc.constraint_name = tc.constraint_name "
|
||||
" AND rc.constraint_schema = tc.constraint_schema "
|
||||
"WHERE tc.table_name = ANY(:tables) "
|
||||
" AND tc.table_schema = current_schema()"
|
||||
),
|
||||
{"tables": need},
|
||||
).fetchall()
|
||||
for table, name, rule in rows:
|
||||
self._fk_rules[(table, name)] = rule
|
||||
self._fk_loaded_tables.update(need)
|
||||
|
||||
def column_exists(self, table: str, column: str) -> bool:
|
||||
return column in self._columns.get(table, {})
|
||||
|
||||
def fk_ondelete(self, table: str, constraint: str) -> str | None:
|
||||
return self._fk_rules.get((table, constraint))
|
||||
|
||||
|
||||
def _fk_exists(constraint_name: str, table_name: str) -> bool:
|
||||
bind = op.get_bind()
|
||||
result = bind.execute(
|
||||
sa.text(
|
||||
"SELECT 1 FROM pg_constraint c "
|
||||
"JOIN pg_class r ON c.conrelid = r.oid "
|
||||
"JOIN pg_namespace n ON r.relnamespace = n.oid "
|
||||
"WHERE c.conname = :name AND r.relname = :table "
|
||||
" AND n.nspname = current_schema() AND c.contype = 'f'"
|
||||
),
|
||||
{"name": constraint_name, "table": table_name},
|
||||
)
|
||||
return result.scalar() is not None
|
||||
|
||||
|
||||
def _replace_fk_if_needed(
|
||||
cache: _SchemaCache,
|
||||
constraint_name: str,
|
||||
table_name: str,
|
||||
ref_table: str,
|
||||
local_cols: list[str],
|
||||
remote_cols: list[str],
|
||||
desired_ondelete: str,
|
||||
) -> None:
|
||||
current = cache.fk_ondelete(table_name, constraint_name)
|
||||
if current and current.upper() == desired_ondelete.upper():
|
||||
return
|
||||
if current or _fk_exists(constraint_name, table_name):
|
||||
op.drop_constraint(constraint_name, table_name, type_="foreignkey")
|
||||
op.create_foreign_key(
|
||||
constraint_name,
|
||||
table_name,
|
||||
ref_table,
|
||||
local_cols,
|
||||
remote_cols,
|
||||
ondelete=desired_ondelete,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
c = _SchemaCache()
|
||||
c.load_columns(_TABLES)
|
||||
c.load_fk_rules(["stats_user_daily", "stats_daily_api_key"])
|
||||
|
||||
# --- Usage: add name snapshot columns ---
|
||||
if not c.column_exists("usage", "username"):
|
||||
op.add_column(
|
||||
"usage", sa.Column("username", sa.String(100), nullable=True, comment="用户名快照")
|
||||
)
|
||||
if not c.column_exists("usage", "api_key_name"):
|
||||
op.add_column(
|
||||
"usage",
|
||||
sa.Column("api_key_name", sa.String(200), nullable=True, comment="API Key 名称快照"),
|
||||
)
|
||||
|
||||
# --- StatsUserDaily: CASCADE -> SET NULL, add username snapshot ---
|
||||
_replace_fk_if_needed(
|
||||
c,
|
||||
"stats_user_daily_user_id_fkey",
|
||||
"stats_user_daily",
|
||||
"users",
|
||||
["user_id"],
|
||||
["id"],
|
||||
"SET NULL",
|
||||
)
|
||||
op.alter_column("stats_user_daily", "user_id", existing_type=sa.String(36), nullable=True)
|
||||
if not c.column_exists("stats_user_daily", "username"):
|
||||
op.add_column(
|
||||
"stats_user_daily",
|
||||
sa.Column(
|
||||
"username",
|
||||
sa.String(100),
|
||||
nullable=True,
|
||||
comment="用户名快照(删除用户后仍可追溯)",
|
||||
),
|
||||
)
|
||||
|
||||
# --- StatsDailyApiKey: CASCADE -> SET NULL, add api_key_name snapshot ---
|
||||
_replace_fk_if_needed(
|
||||
c,
|
||||
"stats_daily_api_key_api_key_id_fkey",
|
||||
"stats_daily_api_key",
|
||||
"api_keys",
|
||||
["api_key_id"],
|
||||
["id"],
|
||||
"SET NULL",
|
||||
)
|
||||
op.alter_column("stats_daily_api_key", "api_key_id", existing_type=sa.String(36), nullable=True)
|
||||
if not c.column_exists("stats_daily_api_key", "api_key_name"):
|
||||
op.add_column(
|
||||
"stats_daily_api_key",
|
||||
sa.Column(
|
||||
"api_key_name",
|
||||
sa.String(200),
|
||||
nullable=True,
|
||||
comment="API Key 名称快照(删除 Key 后仍可追溯)",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
c = _SchemaCache()
|
||||
c.load_columns(["stats_daily_api_key", "stats_user_daily", "usage"])
|
||||
c.load_fk_rules(["stats_daily_api_key", "stats_user_daily"])
|
||||
|
||||
# --- Remove snapshot columns ---
|
||||
if c.column_exists("stats_daily_api_key", "api_key_name"):
|
||||
op.drop_column("stats_daily_api_key", "api_key_name")
|
||||
if c.column_exists("stats_user_daily", "username"):
|
||||
op.drop_column("stats_user_daily", "username")
|
||||
if c.column_exists("usage", "api_key_name"):
|
||||
op.drop_column("usage", "api_key_name")
|
||||
if c.column_exists("usage", "username"):
|
||||
op.drop_column("usage", "username")
|
||||
|
||||
# --- StatsDailyApiKey: SET NULL -> CASCADE ---
|
||||
_replace_fk_if_needed(
|
||||
c,
|
||||
"stats_daily_api_key_api_key_id_fkey",
|
||||
"stats_daily_api_key",
|
||||
"api_keys",
|
||||
["api_key_id"],
|
||||
["id"],
|
||||
"CASCADE",
|
||||
)
|
||||
op.alter_column(
|
||||
"stats_daily_api_key", "api_key_id", existing_type=sa.String(36), nullable=False
|
||||
)
|
||||
|
||||
# --- StatsUserDaily: SET NULL -> CASCADE ---
|
||||
_replace_fk_if_needed(
|
||||
c,
|
||||
"stats_user_daily_user_id_fkey",
|
||||
"stats_user_daily",
|
||||
"users",
|
||||
["user_id"],
|
||||
["id"],
|
||||
"CASCADE",
|
||||
)
|
||||
op.alter_column("stats_user_daily", "user_id", existing_type=sa.String(36), nullable=False)
|
||||
+257
@@ -0,0 +1,257 @@
|
||||
"""request_candidates/video_tasks retention: SET NULL and add snapshots
|
||||
|
||||
Revision ID: 13a4c8f6d9e0
|
||||
Revises: 45b118150a78
|
||||
Create Date: 2026-03-08 12:15:00.000000+00:00
|
||||
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlalchemy as sa
|
||||
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "13a4c8f6d9e0"
|
||||
down_revision = "45b118150a78"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
_TABLES = ["request_candidates", "video_tasks"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Inline helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _SchemaCache:
|
||||
def __init__(self) -> None:
|
||||
self._columns: dict[str, dict[str, str]] = {}
|
||||
self._fk_rules: dict[tuple[str, str], str] = {}
|
||||
self._fk_loaded_tables: set[str] = set()
|
||||
|
||||
def load_columns(self, tables: list[str]) -> None:
|
||||
need = [t for t in tables if t not in self._columns]
|
||||
if not need:
|
||||
return
|
||||
bind = op.get_bind()
|
||||
rows = bind.execute(
|
||||
sa.text(
|
||||
"SELECT table_name, column_name, data_type "
|
||||
"FROM information_schema.columns "
|
||||
"WHERE table_name = ANY(:tables) "
|
||||
" AND table_schema = current_schema()"
|
||||
),
|
||||
{"tables": need},
|
||||
).fetchall()
|
||||
for t in need:
|
||||
self._columns.setdefault(t, {})
|
||||
for table, col, dtype in rows:
|
||||
self._columns[table][col] = dtype
|
||||
|
||||
def load_fk_rules(self, tables: list[str]) -> None:
|
||||
need = [t for t in tables if t not in self._fk_loaded_tables]
|
||||
if not need:
|
||||
return
|
||||
bind = op.get_bind()
|
||||
rows = bind.execute(
|
||||
sa.text(
|
||||
"SELECT tc.table_name, tc.constraint_name, rc.delete_rule "
|
||||
"FROM information_schema.referential_constraints rc "
|
||||
"JOIN information_schema.table_constraints tc "
|
||||
" ON rc.constraint_name = tc.constraint_name "
|
||||
" AND rc.constraint_schema = tc.constraint_schema "
|
||||
"WHERE tc.table_name = ANY(:tables) "
|
||||
" AND tc.table_schema = current_schema()"
|
||||
),
|
||||
{"tables": need},
|
||||
).fetchall()
|
||||
for table, name, rule in rows:
|
||||
self._fk_rules[(table, name)] = rule
|
||||
self._fk_loaded_tables.update(need)
|
||||
|
||||
def column_exists(self, table: str, column: str) -> bool:
|
||||
return column in self._columns.get(table, {})
|
||||
|
||||
def fk_ondelete(self, table: str, constraint: str) -> str | None:
|
||||
return self._fk_rules.get((table, constraint))
|
||||
|
||||
|
||||
def _fk_exists(constraint_name: str, table_name: str) -> bool:
|
||||
bind = op.get_bind()
|
||||
result = bind.execute(
|
||||
sa.text(
|
||||
"SELECT 1 FROM pg_constraint c "
|
||||
"JOIN pg_class r ON c.conrelid = r.oid "
|
||||
"JOIN pg_namespace n ON r.relnamespace = n.oid "
|
||||
"WHERE c.conname = :name AND r.relname = :table "
|
||||
" AND n.nspname = current_schema() AND c.contype = 'f'"
|
||||
),
|
||||
{"name": constraint_name, "table": table_name},
|
||||
)
|
||||
return result.scalar() is not None
|
||||
|
||||
|
||||
def _replace_fk_if_needed(
|
||||
cache: _SchemaCache,
|
||||
constraint_name: str,
|
||||
table_name: str,
|
||||
ref_table: str,
|
||||
local_cols: list[str],
|
||||
remote_cols: list[str],
|
||||
desired_ondelete: str,
|
||||
) -> None:
|
||||
current = cache.fk_ondelete(table_name, constraint_name)
|
||||
if current and current.upper() == desired_ondelete.upper():
|
||||
return
|
||||
if current or _fk_exists(constraint_name, table_name):
|
||||
op.drop_constraint(constraint_name, table_name, type_="foreignkey")
|
||||
op.create_foreign_key(
|
||||
constraint_name,
|
||||
table_name,
|
||||
ref_table,
|
||||
local_cols,
|
||||
remote_cols,
|
||||
ondelete=desired_ondelete,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
c = _SchemaCache()
|
||||
c.load_columns(_TABLES)
|
||||
c.load_fk_rules(_TABLES)
|
||||
|
||||
# --- request_candidates: add snapshot columns ---
|
||||
if not c.column_exists("request_candidates", "username"):
|
||||
op.add_column(
|
||||
"request_candidates",
|
||||
sa.Column("username", sa.String(length=100), nullable=True, comment="用户名快照"),
|
||||
)
|
||||
if not c.column_exists("request_candidates", "api_key_name"):
|
||||
op.add_column(
|
||||
"request_candidates",
|
||||
sa.Column(
|
||||
"api_key_name",
|
||||
sa.String(length=200),
|
||||
nullable=True,
|
||||
comment="API Key 名称快照",
|
||||
),
|
||||
)
|
||||
|
||||
# --- request_candidates: CASCADE -> SET NULL ---
|
||||
_replace_fk_if_needed(
|
||||
c,
|
||||
"request_candidates_user_id_fkey",
|
||||
"request_candidates",
|
||||
"users",
|
||||
["user_id"],
|
||||
["id"],
|
||||
"SET NULL",
|
||||
)
|
||||
_replace_fk_if_needed(
|
||||
c,
|
||||
"request_candidates_api_key_id_fkey",
|
||||
"request_candidates",
|
||||
"api_keys",
|
||||
["api_key_id"],
|
||||
["id"],
|
||||
"SET NULL",
|
||||
)
|
||||
|
||||
# --- video_tasks: add snapshot columns ---
|
||||
if not c.column_exists("video_tasks", "username"):
|
||||
op.add_column(
|
||||
"video_tasks",
|
||||
sa.Column("username", sa.String(length=100), nullable=True, comment="用户名快照"),
|
||||
)
|
||||
if not c.column_exists("video_tasks", "api_key_name"):
|
||||
op.add_column(
|
||||
"video_tasks",
|
||||
sa.Column(
|
||||
"api_key_name",
|
||||
sa.String(length=200),
|
||||
nullable=True,
|
||||
comment="API Key 名称快照",
|
||||
),
|
||||
)
|
||||
|
||||
# --- video_tasks: CASCADE -> SET NULL, user_id nullable ---
|
||||
op.alter_column("video_tasks", "user_id", existing_type=sa.String(length=36), nullable=True)
|
||||
_replace_fk_if_needed(
|
||||
c,
|
||||
"video_tasks_user_id_fkey",
|
||||
"video_tasks",
|
||||
"users",
|
||||
["user_id"],
|
||||
["id"],
|
||||
"SET NULL",
|
||||
)
|
||||
_replace_fk_if_needed(
|
||||
c,
|
||||
"video_tasks_api_key_id_fkey",
|
||||
"video_tasks",
|
||||
"api_keys",
|
||||
["api_key_id"],
|
||||
["id"],
|
||||
"SET NULL",
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
c = _SchemaCache()
|
||||
c.load_columns(_TABLES)
|
||||
c.load_fk_rules(_TABLES)
|
||||
|
||||
# --- video_tasks: SET NULL -> default (no action), restore NOT NULL ---
|
||||
_replace_fk_if_needed(
|
||||
c,
|
||||
"video_tasks_api_key_id_fkey",
|
||||
"video_tasks",
|
||||
"api_keys",
|
||||
["api_key_id"],
|
||||
["id"],
|
||||
"NO ACTION",
|
||||
)
|
||||
_replace_fk_if_needed(
|
||||
c,
|
||||
"video_tasks_user_id_fkey",
|
||||
"video_tasks",
|
||||
"users",
|
||||
["user_id"],
|
||||
["id"],
|
||||
"NO ACTION",
|
||||
)
|
||||
op.alter_column("video_tasks", "user_id", existing_type=sa.String(length=36), nullable=False)
|
||||
if c.column_exists("video_tasks", "api_key_name"):
|
||||
op.drop_column("video_tasks", "api_key_name")
|
||||
if c.column_exists("video_tasks", "username"):
|
||||
op.drop_column("video_tasks", "username")
|
||||
|
||||
# --- request_candidates: SET NULL -> CASCADE ---
|
||||
_replace_fk_if_needed(
|
||||
c,
|
||||
"request_candidates_api_key_id_fkey",
|
||||
"request_candidates",
|
||||
"api_keys",
|
||||
["api_key_id"],
|
||||
["id"],
|
||||
"CASCADE",
|
||||
)
|
||||
_replace_fk_if_needed(
|
||||
c,
|
||||
"request_candidates_user_id_fkey",
|
||||
"request_candidates",
|
||||
"users",
|
||||
["user_id"],
|
||||
["id"],
|
||||
"CASCADE",
|
||||
)
|
||||
if c.column_exists("request_candidates", "api_key_name"):
|
||||
op.drop_column("request_candidates", "api_key_name")
|
||||
if c.column_exists("request_candidates", "username"):
|
||||
op.drop_column("request_candidates", "username")
|
||||
@@ -0,0 +1,230 @@
|
||||
"""cost fields: Float -> Numeric(20,8) + provider_api_keys composite index
|
||||
|
||||
Revision ID: 2053ab8ed764
|
||||
Revises: 13a4c8f6d9e0
|
||||
Create Date: 2026-03-08 15:30:00.000000+00:00
|
||||
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from collections import defaultdict
|
||||
from collections.abc import Callable
|
||||
|
||||
import sqlalchemy as sa
|
||||
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "2053ab8ed764"
|
||||
down_revision = "13a4c8f6d9e0"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
# (table_name, column_name, nullable, server_default)
|
||||
_COST_COLUMNS: list[tuple[str, str, bool, str | None]] = [
|
||||
# api_keys
|
||||
("api_keys", "total_cost_usd", True, "0.0"),
|
||||
# usage
|
||||
("usage", "input_cost_usd", True, "0.0"),
|
||||
("usage", "output_cost_usd", True, "0.0"),
|
||||
("usage", "cache_cost_usd", True, "0.0"),
|
||||
("usage", "cache_creation_cost_usd", True, "0.0"),
|
||||
("usage", "cache_read_cost_usd", True, "0.0"),
|
||||
("usage", "request_cost_usd", True, "0.0"),
|
||||
("usage", "total_cost_usd", True, "0.0"),
|
||||
("usage", "actual_input_cost_usd", True, "0.0"),
|
||||
("usage", "actual_output_cost_usd", True, "0.0"),
|
||||
("usage", "actual_cache_creation_cost_usd", True, "0.0"),
|
||||
("usage", "actual_cache_read_cost_usd", True, "0.0"),
|
||||
("usage", "actual_request_cost_usd", True, "0.0"),
|
||||
("usage", "actual_total_cost_usd", True, "0.0"),
|
||||
("usage", "rate_multiplier", True, "1.0"),
|
||||
("usage", "input_price_per_1m", True, None),
|
||||
("usage", "output_price_per_1m", True, None),
|
||||
("usage", "cache_creation_price_per_1m", True, None),
|
||||
("usage", "cache_read_price_per_1m", True, None),
|
||||
("usage", "price_per_request", True, None),
|
||||
# providers
|
||||
("providers", "monthly_quota_usd", True, None),
|
||||
("providers", "monthly_used_usd", True, "0.0"),
|
||||
# global_models
|
||||
("global_models", "default_price_per_request", True, None),
|
||||
# models
|
||||
("models", "price_per_request", True, None),
|
||||
# stats_hourly
|
||||
("stats_hourly", "total_cost", False, "0.0"),
|
||||
("stats_hourly", "actual_total_cost", False, "0.0"),
|
||||
# stats_hourly_user
|
||||
("stats_hourly_user", "total_cost", False, "0.0"),
|
||||
# stats_hourly_model
|
||||
("stats_hourly_model", "total_cost", False, "0.0"),
|
||||
# stats_hourly_provider
|
||||
("stats_hourly_provider", "total_cost", False, "0.0"),
|
||||
# stats_daily
|
||||
("stats_daily", "total_cost", False, "0.0"),
|
||||
("stats_daily", "actual_total_cost", False, "0.0"),
|
||||
("stats_daily", "input_cost", False, "0.0"),
|
||||
("stats_daily", "output_cost", False, "0.0"),
|
||||
("stats_daily", "cache_creation_cost", False, "0.0"),
|
||||
("stats_daily", "cache_read_cost", False, "0.0"),
|
||||
# stats_daily_model
|
||||
("stats_daily_model", "total_cost", False, "0.0"),
|
||||
# stats_daily_provider
|
||||
("stats_daily_provider", "total_cost", False, "0.0"),
|
||||
# stats_daily_api_key
|
||||
("stats_daily_api_key", "total_cost", False, "0.0"),
|
||||
# stats_summary
|
||||
("stats_summary", "all_time_cost", False, "0.0"),
|
||||
("stats_summary", "all_time_actual_cost", False, "0.0"),
|
||||
# stats_user_daily
|
||||
("stats_user_daily", "total_cost", False, "0.0"),
|
||||
]
|
||||
|
||||
_ALL_TABLES = list({t for t, *_ in _COST_COLUMNS})
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Inline helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _SchemaCache:
|
||||
def __init__(self) -> None:
|
||||
self._columns: dict[str, dict[str, str]] = {}
|
||||
|
||||
def load_columns(self, tables: list[str]) -> None:
|
||||
need = [t for t in tables if t not in self._columns]
|
||||
if not need:
|
||||
return
|
||||
bind = op.get_bind()
|
||||
rows = bind.execute(
|
||||
sa.text(
|
||||
"SELECT table_name, column_name, data_type "
|
||||
"FROM information_schema.columns "
|
||||
"WHERE table_name = ANY(:tables) "
|
||||
" AND table_schema = current_schema()"
|
||||
),
|
||||
{"tables": need},
|
||||
).fetchall()
|
||||
for t in need:
|
||||
self._columns.setdefault(t, {})
|
||||
for table, col, dtype in rows:
|
||||
self._columns[table][col] = dtype
|
||||
|
||||
def column_exists(self, table: str, column: str) -> bool:
|
||||
return column in self._columns.get(table, {})
|
||||
|
||||
def column_type(self, table: str, column: str) -> str | None:
|
||||
return self._columns.get(table, {}).get(column)
|
||||
|
||||
def is_numeric(self, table: str, column: str) -> bool:
|
||||
return self.column_type(table, column) == "numeric"
|
||||
|
||||
|
||||
def _index_exists(index_name: str) -> bool:
|
||||
bind = op.get_bind()
|
||||
result = bind.execute(
|
||||
sa.text(
|
||||
"SELECT 1 FROM pg_indexes "
|
||||
"WHERE indexname = :name AND schemaname = current_schema()::text"
|
||||
),
|
||||
{"name": index_name},
|
||||
)
|
||||
return result.scalar() is not None
|
||||
|
||||
|
||||
def _numeric_max(type_spec: str) -> float | None:
|
||||
m = re.match(r"NUMERIC\((\d+),(\d+)\)", type_spec, re.IGNORECASE)
|
||||
if not m:
|
||||
return None
|
||||
precision, scale = int(m.group(1)), int(m.group(2))
|
||||
return 10 ** (precision - scale) - 10 ** (-scale)
|
||||
|
||||
|
||||
def _batch_alter_type(
|
||||
cache: _SchemaCache,
|
||||
columns: list[tuple[str, str, bool, str | None]],
|
||||
cast_suffix: str,
|
||||
type_fn: Callable[[str], str],
|
||||
) -> None:
|
||||
by_table: dict[str, list[tuple[str, str]]] = defaultdict(list)
|
||||
for table, col, _nullable, _default in columns:
|
||||
if not cache.column_exists(table, col):
|
||||
continue
|
||||
by_table[table].append((col, type_fn(col)))
|
||||
|
||||
bind = op.get_bind()
|
||||
for table, col_types in by_table.items():
|
||||
for col, target in col_types:
|
||||
cap = _numeric_max(target)
|
||||
if cap is not None:
|
||||
bind.execute(
|
||||
sa.text(
|
||||
f"UPDATE {table} SET {col} = :cap "
|
||||
f"WHERE {col} IS NOT NULL AND abs({col}) > :cap"
|
||||
),
|
||||
{"cap": cap},
|
||||
)
|
||||
parts = [
|
||||
f"ALTER COLUMN {col} TYPE {target} USING {col}::{cast_suffix}"
|
||||
for col, target in col_types
|
||||
]
|
||||
if parts:
|
||||
bind.execute(sa.text(f"ALTER TABLE {table} " + ", ".join(parts)))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _type_spec(col: str) -> str:
|
||||
"""Return the SQL type literal for a given column name."""
|
||||
return "NUMERIC(10,6)" if col == "rate_multiplier" else "NUMERIC(20,8)"
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
c = _SchemaCache()
|
||||
c.load_columns(_ALL_TABLES)
|
||||
|
||||
# -- 1. cost fields: Float -> Numeric (batched per table)
|
||||
cols_to_convert = [
|
||||
(t, col, n, d)
|
||||
for t, col, n, d in _COST_COLUMNS
|
||||
if c.column_exists(t, col) and not c.is_numeric(t, col)
|
||||
]
|
||||
_batch_alter_type(c, cols_to_convert, cast_suffix="numeric", type_fn=_type_spec)
|
||||
|
||||
# -- 2. provider_api_keys composite index
|
||||
if not _index_exists("idx_provider_api_keys_provider_active"):
|
||||
op.create_index(
|
||||
"idx_provider_api_keys_provider_active",
|
||||
"provider_api_keys",
|
||||
["provider_id", "is_active"],
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
# -- 2. drop composite index
|
||||
if _index_exists("idx_provider_api_keys_provider_active"):
|
||||
op.drop_index(
|
||||
"idx_provider_api_keys_provider_active",
|
||||
table_name="provider_api_keys",
|
||||
)
|
||||
|
||||
# -- 1. Numeric -> Float (batched per table)
|
||||
c = _SchemaCache()
|
||||
c.load_columns(_ALL_TABLES)
|
||||
|
||||
cols_to_revert = [
|
||||
(t, col, n, d)
|
||||
for t, col, n, d in _COST_COLUMNS
|
||||
if c.column_exists(t, col) and c.is_numeric(t, col)
|
||||
]
|
||||
_batch_alter_type(
|
||||
c,
|
||||
cols_to_revert,
|
||||
cast_suffix="double precision",
|
||||
type_fn=lambda _col: "DOUBLE PRECISION",
|
||||
)
|
||||
@@ -0,0 +1,127 @@
|
||||
"""video_tasks.key_id: add ondelete SET NULL
|
||||
|
||||
Revision ID: d7649c1f8e21
|
||||
Revises: 2053ab8ed764
|
||||
Create Date: 2026-03-09 01:00:00.000000+00:00
|
||||
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlalchemy as sa
|
||||
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "d7649c1f8e21"
|
||||
down_revision = "2053ab8ed764"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
_TABLE = "video_tasks"
|
||||
_FK_NAME = "video_tasks_key_id_fkey"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Inline helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _SchemaCache:
|
||||
def __init__(self) -> None:
|
||||
self._fk_rules: dict[tuple[str, str], str] = {}
|
||||
self._fk_loaded_tables: set[str] = set()
|
||||
|
||||
def load_fk_rules(self, tables: list[str]) -> None:
|
||||
need = [t for t in tables if t not in self._fk_loaded_tables]
|
||||
if not need:
|
||||
return
|
||||
bind = op.get_bind()
|
||||
rows = bind.execute(
|
||||
sa.text(
|
||||
"SELECT tc.table_name, tc.constraint_name, rc.delete_rule "
|
||||
"FROM information_schema.referential_constraints rc "
|
||||
"JOIN information_schema.table_constraints tc "
|
||||
" ON rc.constraint_name = tc.constraint_name "
|
||||
" AND rc.constraint_schema = tc.constraint_schema "
|
||||
"WHERE tc.table_name = ANY(:tables) "
|
||||
" AND tc.table_schema = current_schema()"
|
||||
),
|
||||
{"tables": need},
|
||||
).fetchall()
|
||||
for table, name, rule in rows:
|
||||
self._fk_rules[(table, name)] = rule
|
||||
self._fk_loaded_tables.update(need)
|
||||
|
||||
def fk_ondelete(self, table: str, constraint: str) -> str | None:
|
||||
return self._fk_rules.get((table, constraint))
|
||||
|
||||
|
||||
def _fk_exists(constraint_name: str, table_name: str) -> bool:
|
||||
bind = op.get_bind()
|
||||
result = bind.execute(
|
||||
sa.text(
|
||||
"SELECT 1 FROM pg_constraint c "
|
||||
"JOIN pg_class r ON c.conrelid = r.oid "
|
||||
"JOIN pg_namespace n ON r.relnamespace = n.oid "
|
||||
"WHERE c.conname = :name AND r.relname = :table "
|
||||
" AND n.nspname = current_schema() AND c.contype = 'f'"
|
||||
),
|
||||
{"name": constraint_name, "table": table_name},
|
||||
)
|
||||
return result.scalar() is not None
|
||||
|
||||
|
||||
def _replace_fk_if_needed(
|
||||
cache: _SchemaCache,
|
||||
constraint_name: str,
|
||||
table_name: str,
|
||||
ref_table: str,
|
||||
local_cols: list[str],
|
||||
remote_cols: list[str],
|
||||
desired_ondelete: str,
|
||||
) -> None:
|
||||
current = cache.fk_ondelete(table_name, constraint_name)
|
||||
if current and current.upper() == desired_ondelete.upper():
|
||||
return
|
||||
if current or _fk_exists(constraint_name, table_name):
|
||||
op.drop_constraint(constraint_name, table_name, type_="foreignkey")
|
||||
op.create_foreign_key(
|
||||
constraint_name,
|
||||
table_name,
|
||||
ref_table,
|
||||
local_cols,
|
||||
remote_cols,
|
||||
ondelete=desired_ondelete,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
c = _SchemaCache()
|
||||
c.load_fk_rules([_TABLE])
|
||||
_replace_fk_if_needed(
|
||||
c,
|
||||
_FK_NAME,
|
||||
_TABLE,
|
||||
"provider_api_keys",
|
||||
["key_id"],
|
||||
["id"],
|
||||
"SET NULL",
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
c = _SchemaCache()
|
||||
c.load_fk_rules([_TABLE])
|
||||
_replace_fk_if_needed(
|
||||
c,
|
||||
_FK_NAME,
|
||||
_TABLE,
|
||||
"provider_api_keys",
|
||||
["key_id"],
|
||||
["id"],
|
||||
"NO ACTION",
|
||||
)
|
||||
@@ -0,0 +1,42 @@
|
||||
"""Strip request_results_window from health_by_format JSON.
|
||||
|
||||
This data is now maintained in process memory only, no longer persisted to DB.
|
||||
|
||||
Revision ID: a3f1b7c9d2e4
|
||||
Revises: d7649c1f8e21
|
||||
Create Date: 2026-03-10 12:00:00.000000+00:00
|
||||
|
||||
"""
|
||||
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "a3f1b7c9d2e4"
|
||||
down_revision = "d7649c1f8e21"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.execute("""
|
||||
UPDATE provider_api_keys
|
||||
SET health_by_format = (
|
||||
SELECT jsonb_object_agg(
|
||||
fmt_key,
|
||||
fmt_value - 'request_results_window'
|
||||
)
|
||||
FROM jsonb_each(health_by_format) AS x(fmt_key, fmt_value)
|
||||
)
|
||||
WHERE health_by_format IS NOT NULL
|
||||
AND health_by_format != '{}'::jsonb
|
||||
AND EXISTS (
|
||||
SELECT 1
|
||||
FROM jsonb_each(health_by_format) AS x(fmt_key, fmt_value)
|
||||
WHERE fmt_value ? 'request_results_window'
|
||||
)
|
||||
""")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
# No-op: window data is rebuilt from scratch on process start
|
||||
pass
|
||||
@@ -0,0 +1,159 @@
|
||||
"""tighten usage billing state machine
|
||||
|
||||
Revision ID: 9e4f1a2b3c4d
|
||||
Revises: a3f1b7c9d2e4
|
||||
Create Date: 2026-03-11 19:00:00.000000+00:00
|
||||
|
||||
This migration does two things:
|
||||
1. Change new `usage.billing_status` default from `settled` to `pending`.
|
||||
2. Repair only the clearly-safe inconsistent historical rows for production:
|
||||
- failed/cancelled zero-cost rows that were marked settled are converted to void
|
||||
- terminal rows missing finalized_at are backfilled from created_at
|
||||
|
||||
Ambiguous positive-cost settled rows are intentionally left untouched for manual audit.
|
||||
|
||||
All data updates are batched (10000 rows per iteration) to avoid long-held locks
|
||||
and excessive WAL generation on large usage tables.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy import inspect
|
||||
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "9e4f1a2b3c4d"
|
||||
down_revision: str | None = "a3f1b7c9d2e4"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
BATCH_SIZE = 10000
|
||||
|
||||
|
||||
def _table_exists(table_name: str) -> bool:
|
||||
bind = op.get_bind()
|
||||
insp = inspect(bind)
|
||||
insp.clear_cache()
|
||||
return table_name in insp.get_table_names()
|
||||
|
||||
|
||||
def _column_exists(table_name: str, column_name: str) -> bool:
|
||||
bind = op.get_bind()
|
||||
insp = inspect(bind)
|
||||
insp.clear_cache()
|
||||
return column_name in [col["name"] for col in insp.get_columns(table_name)]
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
if not _table_exists("usage"):
|
||||
return
|
||||
|
||||
if _column_exists("usage", "billing_status"):
|
||||
op.alter_column(
|
||||
"usage",
|
||||
"billing_status",
|
||||
existing_type=sa.String(length=20),
|
||||
server_default="pending",
|
||||
existing_nullable=False,
|
||||
)
|
||||
|
||||
required_columns = {
|
||||
"billing_status",
|
||||
"status",
|
||||
"total_cost_usd",
|
||||
"request_cost_usd",
|
||||
"actual_total_cost_usd",
|
||||
"actual_request_cost_usd",
|
||||
"wallet_balance_after",
|
||||
"finalized_at",
|
||||
"created_at",
|
||||
}
|
||||
if not required_columns.issubset(
|
||||
{col for col in required_columns if _column_exists("usage", col)}
|
||||
):
|
||||
return
|
||||
|
||||
conn = op.get_bind()
|
||||
|
||||
# Step 1: billing_status IS NULL -> 'pending' (batched)
|
||||
while True:
|
||||
result = conn.execute(
|
||||
sa.text("""
|
||||
WITH batch AS (
|
||||
SELECT id FROM usage
|
||||
WHERE billing_status IS NULL
|
||||
LIMIT :batch_size
|
||||
FOR UPDATE SKIP LOCKED
|
||||
)
|
||||
UPDATE usage
|
||||
SET billing_status = 'pending'
|
||||
FROM batch WHERE usage.id = batch.id
|
||||
"""),
|
||||
{"batch_size": BATCH_SIZE},
|
||||
)
|
||||
if result.rowcount < BATCH_SIZE:
|
||||
break
|
||||
|
||||
# Step 2: failed/cancelled zero-cost settled -> void (batched)
|
||||
while True:
|
||||
result = conn.execute(
|
||||
sa.text("""
|
||||
WITH batch AS (
|
||||
SELECT id FROM usage
|
||||
WHERE billing_status = 'settled'
|
||||
AND status IN ('failed', 'cancelled')
|
||||
AND COALESCE(total_cost_usd, 0) = 0
|
||||
AND wallet_balance_after IS NULL
|
||||
LIMIT :batch_size
|
||||
FOR UPDATE SKIP LOCKED
|
||||
)
|
||||
UPDATE usage
|
||||
SET billing_status = 'void',
|
||||
finalized_at = COALESCE(usage.finalized_at, usage.created_at),
|
||||
total_cost_usd = 0,
|
||||
request_cost_usd = 0,
|
||||
actual_total_cost_usd = 0,
|
||||
actual_request_cost_usd = 0
|
||||
FROM batch WHERE usage.id = batch.id
|
||||
"""),
|
||||
{"batch_size": BATCH_SIZE},
|
||||
)
|
||||
if result.rowcount < BATCH_SIZE:
|
||||
break
|
||||
|
||||
# Step 3: backfill finalized_at for terminal rows (batched)
|
||||
while True:
|
||||
result = conn.execute(
|
||||
sa.text("""
|
||||
WITH batch AS (
|
||||
SELECT id FROM usage
|
||||
WHERE billing_status IN ('settled', 'void')
|
||||
AND finalized_at IS NULL
|
||||
LIMIT :batch_size
|
||||
FOR UPDATE SKIP LOCKED
|
||||
)
|
||||
UPDATE usage
|
||||
SET finalized_at = COALESCE(usage.finalized_at, usage.created_at)
|
||||
FROM batch WHERE usage.id = batch.id
|
||||
"""),
|
||||
{"batch_size": BATCH_SIZE},
|
||||
)
|
||||
if result.rowcount < BATCH_SIZE:
|
||||
break
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
if not _table_exists("usage") or not _column_exists("usage", "billing_status"):
|
||||
return
|
||||
|
||||
op.alter_column(
|
||||
"usage",
|
||||
"billing_status",
|
||||
existing_type=sa.String(length=20),
|
||||
server_default="settled",
|
||||
existing_nullable=False,
|
||||
)
|
||||
@@ -0,0 +1,113 @@
|
||||
"""add wallet daily usage ledgers
|
||||
|
||||
Revision ID: d4e5f6a7b8c9
|
||||
Revises: 9e4f1a2b3c4d
|
||||
Create Date: 2026-03-11 21:00:00.000000+00:00
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy import inspect
|
||||
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "d4e5f6a7b8c9"
|
||||
down_revision: str | None = "9e4f1a2b3c4d"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _table_exists(table_name: str) -> bool:
|
||||
bind = op.get_bind()
|
||||
insp = inspect(bind)
|
||||
insp.clear_cache()
|
||||
return table_name in insp.get_table_names()
|
||||
|
||||
|
||||
def _column_exists(table_name: str, column_name: str) -> bool:
|
||||
bind = op.get_bind()
|
||||
insp = inspect(bind)
|
||||
insp.clear_cache()
|
||||
return column_name in [col["name"] for col in insp.get_columns(table_name)]
|
||||
|
||||
|
||||
def _index_exists(table_name: str, index_name: str) -> bool:
|
||||
bind = op.get_bind()
|
||||
insp = inspect(bind)
|
||||
insp.clear_cache()
|
||||
return any(idx["name"] == index_name for idx in insp.get_indexes(table_name))
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
if not _table_exists("wallet_daily_usage_ledgers"):
|
||||
op.create_table(
|
||||
"wallet_daily_usage_ledgers",
|
||||
sa.Column("id", sa.String(length=36), nullable=False),
|
||||
sa.Column("wallet_id", sa.String(length=36), nullable=False),
|
||||
sa.Column("billing_date", sa.Date(), nullable=False),
|
||||
sa.Column("billing_timezone", sa.String(length=64), nullable=False),
|
||||
sa.Column("total_cost_usd", sa.Numeric(20, 8), nullable=False, server_default="0"),
|
||||
sa.Column("total_requests", sa.Integer(), nullable=False, server_default="0"),
|
||||
sa.Column("input_tokens", sa.BigInteger(), nullable=False, server_default="0"),
|
||||
sa.Column("output_tokens", sa.BigInteger(), nullable=False, server_default="0"),
|
||||
sa.Column("cache_creation_tokens", sa.BigInteger(), nullable=False, server_default="0"),
|
||||
sa.Column("cache_read_tokens", sa.BigInteger(), nullable=False, server_default="0"),
|
||||
sa.Column("first_finalized_at", sa.DateTime(timezone=True), nullable=True),
|
||||
sa.Column("last_finalized_at", sa.DateTime(timezone=True), nullable=True),
|
||||
sa.Column("aggregated_at", sa.DateTime(timezone=True), nullable=False),
|
||||
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
|
||||
sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False),
|
||||
sa.ForeignKeyConstraint(["wallet_id"], ["wallets.id"], ondelete="CASCADE"),
|
||||
sa.PrimaryKeyConstraint("id"),
|
||||
sa.UniqueConstraint(
|
||||
"wallet_id",
|
||||
"billing_date",
|
||||
"billing_timezone",
|
||||
name="uq_wallet_daily_usage_ledgers_wallet_date_tz",
|
||||
),
|
||||
)
|
||||
|
||||
if not _index_exists("wallet_daily_usage_ledgers", "idx_wallet_daily_usage_wallet_date"):
|
||||
op.create_index(
|
||||
"idx_wallet_daily_usage_wallet_date",
|
||||
"wallet_daily_usage_ledgers",
|
||||
["wallet_id", "billing_date"],
|
||||
)
|
||||
if not _index_exists("wallet_daily_usage_ledgers", "idx_wallet_daily_usage_date"):
|
||||
op.create_index(
|
||||
"idx_wallet_daily_usage_date",
|
||||
"wallet_daily_usage_ledgers",
|
||||
["billing_date"],
|
||||
)
|
||||
|
||||
if (
|
||||
_table_exists("usage")
|
||||
and all(
|
||||
_column_exists("usage", col) for col in ["billing_status", "finalized_at", "wallet_id"]
|
||||
)
|
||||
and not _index_exists("usage", "idx_usage_billing_finalized_wallet")
|
||||
):
|
||||
op.create_index(
|
||||
"idx_usage_billing_finalized_wallet",
|
||||
"usage",
|
||||
["billing_status", "finalized_at", "wallet_id"],
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
if _table_exists("usage") and _index_exists("usage", "idx_usage_billing_finalized_wallet"):
|
||||
op.drop_index("idx_usage_billing_finalized_wallet", table_name="usage")
|
||||
|
||||
if _table_exists("wallet_daily_usage_ledgers"):
|
||||
if _index_exists("wallet_daily_usage_ledgers", "idx_wallet_daily_usage_date"):
|
||||
op.drop_index("idx_wallet_daily_usage_date", table_name="wallet_daily_usage_ledgers")
|
||||
if _index_exists("wallet_daily_usage_ledgers", "idx_wallet_daily_usage_wallet_date"):
|
||||
op.drop_index(
|
||||
"idx_wallet_daily_usage_wallet_date",
|
||||
table_name="wallet_daily_usage_ledgers",
|
||||
)
|
||||
op.drop_table("wallet_daily_usage_ledgers")
|
||||
@@ -0,0 +1,59 @@
|
||||
"""add provider_api_keys usage total columns
|
||||
|
||||
Revision ID: 9b7c6d5e4f3a
|
||||
Revises: d4e5f6a7b8c9
|
||||
Create Date: 2026-03-11 22:00:00.000000+00:00
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy import inspect
|
||||
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "9b7c6d5e4f3a"
|
||||
down_revision: str | None = "d4e5f6a7b8c9"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def column_exists(table_name: str, column_name: str) -> bool:
|
||||
bind = op.get_bind()
|
||||
inspector = inspect(bind)
|
||||
columns = [c["name"] for c in inspector.get_columns(table_name)]
|
||||
return column_name in columns
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
if not column_exists("provider_api_keys", "total_tokens"):
|
||||
op.add_column(
|
||||
"provider_api_keys",
|
||||
sa.Column("total_tokens", sa.BigInteger(), nullable=False, server_default="0"),
|
||||
)
|
||||
if column_exists("provider_api_keys", "total_tokens"):
|
||||
op.alter_column("provider_api_keys", "total_tokens", server_default=None)
|
||||
|
||||
if not column_exists("provider_api_keys", "total_cost_usd"):
|
||||
op.add_column(
|
||||
"provider_api_keys",
|
||||
sa.Column(
|
||||
"total_cost_usd",
|
||||
sa.Numeric(20, 8),
|
||||
nullable=False,
|
||||
server_default="0.0",
|
||||
),
|
||||
)
|
||||
if column_exists("provider_api_keys", "total_cost_usd"):
|
||||
op.alter_column("provider_api_keys", "total_cost_usd", server_default=None)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
if column_exists("provider_api_keys", "total_cost_usd"):
|
||||
op.drop_column("provider_api_keys", "total_cost_usd")
|
||||
|
||||
if column_exists("provider_api_keys", "total_tokens"):
|
||||
op.drop_column("provider_api_keys", "total_tokens")
|
||||
@@ -0,0 +1,116 @@
|
||||
"""cleanup stale provider references after provider deletion
|
||||
|
||||
Revision ID: c1d2e3f4a5b6
|
||||
Revises: 9b7c6d5e4f3a
|
||||
Create Date: 2026-03-11 23:00:00.000000+00:00
|
||||
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlalchemy as sa
|
||||
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "c1d2e3f4a5b6"
|
||||
down_revision = "9b7c6d5e4f3a"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
_users = sa.table(
|
||||
"users",
|
||||
sa.column("id", sa.String(36)),
|
||||
sa.column("allowed_providers", sa.JSON()),
|
||||
)
|
||||
_api_keys = sa.table(
|
||||
"api_keys",
|
||||
sa.column("id", sa.String(36)),
|
||||
sa.column("allowed_providers", sa.JSON()),
|
||||
)
|
||||
_user_preferences = sa.table(
|
||||
"user_preferences",
|
||||
sa.column("id", sa.String(36)),
|
||||
sa.column("default_provider_id", sa.String(36)),
|
||||
)
|
||||
_video_tasks = sa.table(
|
||||
"video_tasks",
|
||||
sa.column("id", sa.String(36)),
|
||||
sa.column("provider_id", sa.String(36)),
|
||||
sa.column("endpoint_id", sa.String(36)),
|
||||
)
|
||||
_providers = sa.table("providers", sa.column("id", sa.String(36)))
|
||||
_provider_endpoints = sa.table("provider_endpoints", sa.column("id", sa.String(36)))
|
||||
|
||||
|
||||
def _load_valid_ids(conn: sa.Connection, table: sa.Table) -> set[str]:
|
||||
return {str(row[0]) for row in conn.execute(sa.select(table.c.id)).fetchall() if row[0]}
|
||||
|
||||
|
||||
def _cleanup_allowed_providers(
|
||||
conn: sa.Connection,
|
||||
table: sa.Table,
|
||||
valid_provider_ids: set[str],
|
||||
) -> None:
|
||||
rows = conn.execute(
|
||||
sa.select(table.c.id, table.c.allowed_providers).where(
|
||||
table.c.allowed_providers.isnot(None)
|
||||
)
|
||||
).fetchall()
|
||||
for row_id, allowed_providers in rows:
|
||||
if not isinstance(allowed_providers, list):
|
||||
continue
|
||||
filtered = [
|
||||
provider_id for provider_id in allowed_providers if provider_id in valid_provider_ids
|
||||
]
|
||||
if filtered == allowed_providers:
|
||||
continue
|
||||
conn.execute(table.update().where(table.c.id == row_id).values(allowed_providers=filtered))
|
||||
|
||||
|
||||
def _nullify_missing_fk(
|
||||
conn: sa.Connection,
|
||||
table: sa.Table,
|
||||
id_column: sa.ColumnElement[str],
|
||||
fk_column: sa.ColumnElement[str],
|
||||
valid_ids: set[str],
|
||||
) -> None:
|
||||
rows = conn.execute(sa.select(id_column, fk_column).where(fk_column.isnot(None))).fetchall()
|
||||
invalid_row_ids = [row_id for row_id, fk_value in rows if fk_value not in valid_ids]
|
||||
if not invalid_row_ids:
|
||||
return
|
||||
conn.execute(table.update().where(id_column.in_(invalid_row_ids)).values({fk_column.key: None}))
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
valid_provider_ids = _load_valid_ids(conn, _providers)
|
||||
valid_endpoint_ids = _load_valid_ids(conn, _provider_endpoints)
|
||||
|
||||
_cleanup_allowed_providers(conn, _users, valid_provider_ids)
|
||||
_cleanup_allowed_providers(conn, _api_keys, valid_provider_ids)
|
||||
_nullify_missing_fk(
|
||||
conn,
|
||||
_user_preferences,
|
||||
_user_preferences.c.id,
|
||||
_user_preferences.c.default_provider_id,
|
||||
valid_provider_ids,
|
||||
)
|
||||
_nullify_missing_fk(
|
||||
conn,
|
||||
_video_tasks,
|
||||
_video_tasks.c.id,
|
||||
_video_tasks.c.provider_id,
|
||||
valid_provider_ids,
|
||||
)
|
||||
_nullify_missing_fk(
|
||||
conn,
|
||||
_video_tasks,
|
||||
_video_tasks.c.id,
|
||||
_video_tasks.c.endpoint_id,
|
||||
valid_endpoint_ids,
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
pass
|
||||
@@ -0,0 +1,64 @@
|
||||
"""decouple request_candidates.key_id foreign key from provider_api_keys lifecycle
|
||||
|
||||
Revision ID: b7c8d9e0f1a2
|
||||
Revises: c1d2e3f4a5b6
|
||||
Create Date: 2026-03-12 19:15:00.000000+00:00
|
||||
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlalchemy as sa
|
||||
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "b7c8d9e0f1a2"
|
||||
down_revision = "c1d2e3f4a5b6"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def _fk_exists(constraint_name: str, table_name: str) -> bool:
|
||||
bind = op.get_bind()
|
||||
result = bind.execute(
|
||||
sa.text(
|
||||
"SELECT 1 FROM pg_constraint c "
|
||||
"JOIN pg_class r ON c.conrelid = r.oid "
|
||||
"JOIN pg_namespace n ON r.relnamespace = n.oid "
|
||||
"WHERE c.conname = :name AND r.relname = :table "
|
||||
" AND n.nspname = current_schema() AND c.contype = 'f'"
|
||||
),
|
||||
{"name": constraint_name, "table": table_name},
|
||||
)
|
||||
return result.scalar() is not None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
if _fk_exists("request_candidates_key_id_fkey", "request_candidates"):
|
||||
op.drop_constraint(
|
||||
"request_candidates_key_id_fkey", "request_candidates", type_="foreignkey"
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
bind = op.get_bind()
|
||||
bind.execute(
|
||||
sa.text(
|
||||
"UPDATE request_candidates rc "
|
||||
"SET key_id = NULL "
|
||||
"WHERE key_id IS NOT NULL "
|
||||
" AND NOT EXISTS ("
|
||||
" SELECT 1 FROM provider_api_keys pak WHERE pak.id = rc.key_id"
|
||||
" )"
|
||||
)
|
||||
)
|
||||
if not _fk_exists("request_candidates_key_id_fkey", "request_candidates"):
|
||||
op.create_foreign_key(
|
||||
"request_candidates_key_id_fkey",
|
||||
"request_candidates",
|
||||
"provider_api_keys",
|
||||
["key_id"],
|
||||
["id"],
|
||||
ondelete="CASCADE",
|
||||
)
|
||||
+54
@@ -0,0 +1,54 @@
|
||||
"""add user rate_limit and backfill normal api key limits
|
||||
|
||||
Revision ID: b7e8f9a0c1d2
|
||||
Revises: b7c8d9e0f1a2
|
||||
Create Date: 2026-03-13 12:00:00.000000+00:00
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy import inspect
|
||||
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "b7e8f9a0c1d2"
|
||||
down_revision: str | None = "b7c8d9e0f1a2"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def column_exists(table_name: str, column_name: str) -> bool:
|
||||
bind = op.get_bind()
|
||||
inspector = inspect(bind)
|
||||
columns = [c["name"] for c in inspector.get_columns(table_name)]
|
||||
return column_name in columns
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
if not column_exists("users", "rate_limit"):
|
||||
op.add_column("users", sa.Column("rate_limit", sa.Integer(), nullable=True))
|
||||
|
||||
# 普通 Key 新语义不再允许 NULL;存量 NULL 统一回填为 0(不限制)。
|
||||
op.execute(sa.text("""
|
||||
UPDATE api_keys
|
||||
SET rate_limit = 0
|
||||
WHERE is_standalone = FALSE
|
||||
AND rate_limit IS NULL
|
||||
"""))
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
# 恢复普通 Key 的 rate_limit 为 NULL(与 upgrade 中回填 0 对应)
|
||||
op.execute(sa.text("""
|
||||
UPDATE api_keys
|
||||
SET rate_limit = NULL
|
||||
WHERE is_standalone = FALSE
|
||||
AND rate_limit = 0
|
||||
"""))
|
||||
|
||||
if column_exists("users", "rate_limit"):
|
||||
op.drop_column("users", "rate_limit")
|
||||
@@ -0,0 +1,87 @@
|
||||
"""add user sessions table for device-level auth
|
||||
|
||||
Revision ID: f6e7d8c9b0a1
|
||||
Revises: b7e8f9a0c1d2
|
||||
Create Date: 2026-03-15 12:00:00.000000+00:00
|
||||
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlalchemy as sa
|
||||
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "f6e7d8c9b0a1"
|
||||
down_revision = "b7e8f9a0c1d2"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
bind = op.get_bind()
|
||||
inspector = sa.inspect(bind)
|
||||
if "user_sessions" in inspector.get_table_names():
|
||||
return
|
||||
|
||||
op.create_table(
|
||||
"user_sessions",
|
||||
sa.Column("id", sa.String(length=36), nullable=False),
|
||||
sa.Column("user_id", sa.String(length=36), nullable=False),
|
||||
sa.Column("client_device_id", sa.String(length=128), nullable=False),
|
||||
sa.Column("device_label", sa.String(length=120), nullable=True),
|
||||
sa.Column("device_type", sa.String(length=20), nullable=False, server_default="unknown"),
|
||||
sa.Column("browser_name", sa.String(length=50), nullable=True),
|
||||
sa.Column("browser_version", sa.String(length=50), nullable=True),
|
||||
sa.Column("os_name", sa.String(length=50), nullable=True),
|
||||
sa.Column("os_version", sa.String(length=50), nullable=True),
|
||||
sa.Column("device_model", sa.String(length=100), nullable=True),
|
||||
sa.Column("ip_address", sa.String(length=45), nullable=True),
|
||||
sa.Column("user_agent", sa.String(length=1000), nullable=True),
|
||||
sa.Column("client_hints", sa.JSON(), nullable=True),
|
||||
sa.Column("refresh_token_hash", sa.String(length=64), nullable=False),
|
||||
sa.Column("prev_refresh_token_hash", sa.String(length=64), nullable=True),
|
||||
sa.Column("rotated_at", sa.DateTime(timezone=True), nullable=True),
|
||||
sa.Column(
|
||||
"last_seen_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now()
|
||||
),
|
||||
sa.Column("expires_at", sa.DateTime(timezone=True), nullable=False),
|
||||
sa.Column("revoked_at", sa.DateTime(timezone=True), nullable=True),
|
||||
sa.Column("revoke_reason", sa.String(length=100), nullable=True),
|
||||
sa.Column(
|
||||
"created_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now()
|
||||
),
|
||||
sa.Column(
|
||||
"updated_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.func.now()
|
||||
),
|
||||
sa.ForeignKeyConstraint(["user_id"], ["users.id"], ondelete="CASCADE"),
|
||||
sa.PrimaryKeyConstraint("id"),
|
||||
)
|
||||
op.create_index("ix_user_sessions_user_id", "user_sessions", ["user_id"], unique=False)
|
||||
op.create_index(
|
||||
"ix_user_sessions_client_device_id",
|
||||
"user_sessions",
|
||||
["client_device_id"],
|
||||
unique=False,
|
||||
)
|
||||
op.create_index(
|
||||
"idx_user_sessions_user_active",
|
||||
"user_sessions",
|
||||
["user_id", "revoked_at", "expires_at"],
|
||||
unique=False,
|
||||
)
|
||||
op.create_index(
|
||||
"idx_user_sessions_user_device",
|
||||
"user_sessions",
|
||||
["user_id", "client_device_id"],
|
||||
unique=False,
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_index("idx_user_sessions_user_device", table_name="user_sessions")
|
||||
op.drop_index("idx_user_sessions_user_active", table_name="user_sessions")
|
||||
op.drop_index("ix_user_sessions_client_device_id", table_name="user_sessions")
|
||||
op.drop_index("ix_user_sessions_user_id", table_name="user_sessions")
|
||||
op.drop_table("user_sessions")
|
||||
+41
@@ -0,0 +1,41 @@
|
||||
"""add status_snapshot column to provider_api_keys
|
||||
|
||||
Revision ID: c9d8e7f6a5b4
|
||||
Revises: f6e7d8c9b0a1
|
||||
Create Date: 2026-03-20 12:00:00.000000
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy import inspect
|
||||
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "c9d8e7f6a5b4"
|
||||
down_revision: str | None = "f6e7d8c9b0a1"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def column_exists(table_name: str, column_name: str) -> bool:
|
||||
bind = op.get_bind()
|
||||
inspector = inspect(bind)
|
||||
columns = [c["name"] for c in inspector.get_columns(table_name)]
|
||||
return column_name in columns
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
if not column_exists("provider_api_keys", "status_snapshot"):
|
||||
op.add_column(
|
||||
"provider_api_keys",
|
||||
sa.Column("status_snapshot", sa.JSON(), nullable=True),
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
if column_exists("provider_api_keys", "status_snapshot"):
|
||||
op.drop_column("provider_api_keys", "status_snapshot")
|
||||
@@ -0,0 +1,317 @@
|
||||
"""usage token semantics v2
|
||||
|
||||
Revision ID: c3d4e5f6a7b8
|
||||
Revises: c9d8e7f6a5b4
|
||||
Create Date: 2026-03-24 14:00:00.000000+00:00
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy import inspect
|
||||
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "c3d4e5f6a7b8"
|
||||
down_revision: str | None = "c9d8e7f6a5b4"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
BACKFILL_BATCH_SIZE = 2000
|
||||
|
||||
# 最大批次数,防止因数据异常导致死循环(2000 * 500000 = 10亿行上限)
|
||||
_MAX_BATCHES = 500000
|
||||
|
||||
# 使用子查询中间层展开 input_output_total_tokens 的计算,
|
||||
# 确保 total_tokens 引用的是本次 SET 后的新值而非旧值。
|
||||
_UPGRADE_BACKFILL_SQL = sa.text(
|
||||
"""
|
||||
UPDATE usage
|
||||
SET
|
||||
input_output_total_tokens = src.new_iot,
|
||||
input_context_tokens = src.new_ict,
|
||||
total_tokens = src.new_total,
|
||||
cache_creation_cost_usd_5m = src.new_cc5m,
|
||||
cache_creation_cost_usd_1h = src.new_cc1h,
|
||||
actual_cache_creation_cost_usd_5m = src.new_acc5m,
|
||||
actual_cache_creation_cost_usd_1h = src.new_acc1h,
|
||||
actual_cache_cost_usd = src.new_accu,
|
||||
cache_creation_price_per_1m_5m = src.new_cp5m,
|
||||
cache_creation_price_per_1m_1h = src.new_cp1h,
|
||||
cache_cost_usd = src.new_ccu
|
||||
FROM (
|
||||
SELECT
|
||||
id,
|
||||
COALESCE(input_output_total_tokens,
|
||||
COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0))
|
||||
AS new_iot,
|
||||
COALESCE(input_tokens, 0) + COALESCE(cache_read_input_tokens, 0)
|
||||
AS new_ict,
|
||||
/* total_tokens 引用本行计算出的 new_iot,避免依赖 SET 顺序 */
|
||||
COALESCE(input_output_total_tokens,
|
||||
COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0))
|
||||
+ COALESCE(cache_creation_input_tokens, 0)
|
||||
+ COALESCE(cache_read_input_tokens, 0)
|
||||
AS new_total,
|
||||
CASE
|
||||
WHEN COALESCE(cache_creation_input_tokens_5m, 0) > 0
|
||||
AND COALESCE(cache_creation_input_tokens_1h, 0) = 0
|
||||
THEN COALESCE(cache_creation_cost_usd, 0)
|
||||
WHEN COALESCE(cache_creation_input_tokens_5m, 0) > 0
|
||||
AND COALESCE(cache_creation_input_tokens, 0) > 0
|
||||
THEN COALESCE(cache_creation_cost_usd, 0)
|
||||
* (COALESCE(cache_creation_input_tokens_5m, 0) * 1.0
|
||||
/ GREATEST(COALESCE(cache_creation_input_tokens, 0), 1))
|
||||
ELSE 0
|
||||
END AS new_cc5m,
|
||||
CASE
|
||||
WHEN COALESCE(cache_creation_input_tokens_1h, 0) > 0
|
||||
AND COALESCE(cache_creation_input_tokens_5m, 0) = 0
|
||||
THEN COALESCE(cache_creation_cost_usd, 0)
|
||||
WHEN COALESCE(cache_creation_input_tokens_1h, 0) > 0
|
||||
AND COALESCE(cache_creation_input_tokens, 0) > 0
|
||||
THEN COALESCE(cache_creation_cost_usd, 0)
|
||||
* (COALESCE(cache_creation_input_tokens_1h, 0) * 1.0
|
||||
/ GREATEST(COALESCE(cache_creation_input_tokens, 0), 1))
|
||||
ELSE 0
|
||||
END AS new_cc1h,
|
||||
CASE
|
||||
WHEN COALESCE(cache_creation_input_tokens_5m, 0) > 0
|
||||
AND COALESCE(cache_creation_input_tokens_1h, 0) = 0
|
||||
THEN COALESCE(actual_cache_creation_cost_usd, 0)
|
||||
WHEN COALESCE(cache_creation_input_tokens_5m, 0) > 0
|
||||
AND COALESCE(cache_creation_input_tokens, 0) > 0
|
||||
THEN COALESCE(actual_cache_creation_cost_usd, 0)
|
||||
* (COALESCE(cache_creation_input_tokens_5m, 0) * 1.0
|
||||
/ GREATEST(COALESCE(cache_creation_input_tokens, 0), 1))
|
||||
ELSE 0
|
||||
END AS new_acc5m,
|
||||
CASE
|
||||
WHEN COALESCE(cache_creation_input_tokens_1h, 0) > 0
|
||||
AND COALESCE(cache_creation_input_tokens_5m, 0) = 0
|
||||
THEN COALESCE(actual_cache_creation_cost_usd, 0)
|
||||
WHEN COALESCE(cache_creation_input_tokens_1h, 0) > 0
|
||||
AND COALESCE(cache_creation_input_tokens, 0) > 0
|
||||
THEN COALESCE(actual_cache_creation_cost_usd, 0)
|
||||
* (COALESCE(cache_creation_input_tokens_1h, 0) * 1.0
|
||||
/ GREATEST(COALESCE(cache_creation_input_tokens, 0), 1))
|
||||
ELSE 0
|
||||
END AS new_acc1h,
|
||||
COALESCE(actual_cache_creation_cost_usd, 0)
|
||||
+ COALESCE(actual_cache_read_cost_usd, 0) AS new_accu,
|
||||
CASE
|
||||
WHEN COALESCE(cache_creation_input_tokens_5m, 0) > 0
|
||||
AND COALESCE(cache_creation_input_tokens_1h, 0) = 0
|
||||
THEN cache_creation_price_per_1m
|
||||
ELSE NULL
|
||||
END AS new_cp5m,
|
||||
CASE
|
||||
WHEN COALESCE(cache_creation_input_tokens_1h, 0) > 0
|
||||
AND COALESCE(cache_creation_input_tokens_5m, 0) = 0
|
||||
THEN cache_creation_price_per_1m
|
||||
ELSE NULL
|
||||
END AS new_cp1h,
|
||||
COALESCE(cache_creation_cost_usd, 0)
|
||||
+ COALESCE(cache_read_cost_usd, 0) AS new_ccu
|
||||
FROM usage
|
||||
WHERE id IN (
|
||||
SELECT id FROM usage
|
||||
WHERE
|
||||
input_output_total_tokens IS DISTINCT FROM
|
||||
COALESCE(input_output_total_tokens,
|
||||
COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0))
|
||||
OR input_context_tokens IS DISTINCT FROM
|
||||
COALESCE(input_tokens, 0) + COALESCE(cache_read_input_tokens, 0)
|
||||
OR total_tokens IS DISTINCT FROM (
|
||||
COALESCE(input_output_total_tokens,
|
||||
COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0))
|
||||
+ COALESCE(cache_creation_input_tokens, 0)
|
||||
+ COALESCE(cache_read_input_tokens, 0)
|
||||
)
|
||||
OR cache_creation_cost_usd_5m IS DISTINCT FROM (
|
||||
CASE
|
||||
WHEN COALESCE(cache_creation_input_tokens_5m, 0) > 0
|
||||
AND COALESCE(cache_creation_input_tokens_1h, 0) = 0
|
||||
THEN COALESCE(cache_creation_cost_usd, 0)
|
||||
WHEN COALESCE(cache_creation_input_tokens_5m, 0) > 0
|
||||
AND COALESCE(cache_creation_input_tokens, 0) > 0
|
||||
THEN COALESCE(cache_creation_cost_usd, 0)
|
||||
* (COALESCE(cache_creation_input_tokens_5m, 0) * 1.0
|
||||
/ GREATEST(COALESCE(cache_creation_input_tokens, 0), 1))
|
||||
ELSE 0
|
||||
END
|
||||
)
|
||||
OR cache_creation_cost_usd_1h IS DISTINCT FROM (
|
||||
CASE
|
||||
WHEN COALESCE(cache_creation_input_tokens_1h, 0) > 0
|
||||
AND COALESCE(cache_creation_input_tokens_5m, 0) = 0
|
||||
THEN COALESCE(cache_creation_cost_usd, 0)
|
||||
WHEN COALESCE(cache_creation_input_tokens_1h, 0) > 0
|
||||
AND COALESCE(cache_creation_input_tokens, 0) > 0
|
||||
THEN COALESCE(cache_creation_cost_usd, 0)
|
||||
* (COALESCE(cache_creation_input_tokens_1h, 0) * 1.0
|
||||
/ GREATEST(COALESCE(cache_creation_input_tokens, 0), 1))
|
||||
ELSE 0
|
||||
END
|
||||
)
|
||||
OR actual_cache_creation_cost_usd_5m IS DISTINCT FROM (
|
||||
CASE
|
||||
WHEN COALESCE(cache_creation_input_tokens_5m, 0) > 0
|
||||
AND COALESCE(cache_creation_input_tokens_1h, 0) = 0
|
||||
THEN COALESCE(actual_cache_creation_cost_usd, 0)
|
||||
WHEN COALESCE(cache_creation_input_tokens_5m, 0) > 0
|
||||
AND COALESCE(cache_creation_input_tokens, 0) > 0
|
||||
THEN COALESCE(actual_cache_creation_cost_usd, 0)
|
||||
* (COALESCE(cache_creation_input_tokens_5m, 0) * 1.0
|
||||
/ GREATEST(COALESCE(cache_creation_input_tokens, 0), 1))
|
||||
ELSE 0
|
||||
END
|
||||
)
|
||||
OR actual_cache_creation_cost_usd_1h IS DISTINCT FROM (
|
||||
CASE
|
||||
WHEN COALESCE(cache_creation_input_tokens_1h, 0) > 0
|
||||
AND COALESCE(cache_creation_input_tokens_5m, 0) = 0
|
||||
THEN COALESCE(actual_cache_creation_cost_usd, 0)
|
||||
WHEN COALESCE(cache_creation_input_tokens_1h, 0) > 0
|
||||
AND COALESCE(cache_creation_input_tokens, 0) > 0
|
||||
THEN COALESCE(actual_cache_creation_cost_usd, 0)
|
||||
* (COALESCE(cache_creation_input_tokens_1h, 0) * 1.0
|
||||
/ GREATEST(COALESCE(cache_creation_input_tokens, 0), 1))
|
||||
ELSE 0
|
||||
END
|
||||
)
|
||||
OR actual_cache_cost_usd IS DISTINCT FROM (
|
||||
COALESCE(actual_cache_creation_cost_usd, 0)
|
||||
+ COALESCE(actual_cache_read_cost_usd, 0)
|
||||
)
|
||||
OR cache_creation_price_per_1m_5m IS DISTINCT FROM (
|
||||
CASE
|
||||
WHEN COALESCE(cache_creation_input_tokens_5m, 0) > 0
|
||||
AND COALESCE(cache_creation_input_tokens_1h, 0) = 0
|
||||
THEN cache_creation_price_per_1m
|
||||
ELSE NULL
|
||||
END
|
||||
)
|
||||
OR cache_creation_price_per_1m_1h IS DISTINCT FROM (
|
||||
CASE
|
||||
WHEN COALESCE(cache_creation_input_tokens_1h, 0) > 0
|
||||
AND COALESCE(cache_creation_input_tokens_5m, 0) = 0
|
||||
THEN cache_creation_price_per_1m
|
||||
ELSE NULL
|
||||
END
|
||||
)
|
||||
OR cache_cost_usd IS DISTINCT FROM (
|
||||
COALESCE(cache_creation_cost_usd, 0) + COALESCE(cache_read_cost_usd, 0)
|
||||
)
|
||||
ORDER BY id
|
||||
LIMIT :batch_size
|
||||
)
|
||||
) AS src
|
||||
WHERE usage.id = src.id
|
||||
"""
|
||||
)
|
||||
|
||||
_DOWNGRADE_BACKFILL_SQL = sa.text(
|
||||
"""
|
||||
UPDATE usage
|
||||
SET total_tokens = COALESCE(input_output_total_tokens, COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0))
|
||||
WHERE id IN (
|
||||
SELECT id
|
||||
FROM usage
|
||||
WHERE total_tokens IS DISTINCT FROM
|
||||
COALESCE(input_output_total_tokens, COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0))
|
||||
ORDER BY id
|
||||
LIMIT :batch_size
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def column_exists(table_name: str, column_name: str) -> bool:
|
||||
bind = op.get_bind()
|
||||
inspector = inspect(bind)
|
||||
columns = [c["name"] for c in inspector.get_columns(table_name)]
|
||||
return column_name in columns
|
||||
|
||||
|
||||
def run_backfill_in_batches(sql: sa.TextClause, batch_size: int = BACKFILL_BATCH_SIZE) -> None:
|
||||
context = op.get_context()
|
||||
for _ in range(_MAX_BATCHES):
|
||||
# Commit the preceding schema transaction before each batch so PostgreSQL
|
||||
# does not keep ALTER TABLE locks for the entire data backfill.
|
||||
with context.autocommit_block():
|
||||
rowcount = op.get_bind().execute(sql, {"batch_size": batch_size}).rowcount
|
||||
if rowcount == 0:
|
||||
break
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"Backfill did not converge after {_MAX_BATCHES} batches "
|
||||
f"(batch_size={batch_size}). Possible infinite loop due to data anomaly."
|
||||
)
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
if column_exists("usage", "total_tokens") and not column_exists("usage", "input_output_total_tokens"):
|
||||
with op.batch_alter_table("usage") as batch_op:
|
||||
batch_op.alter_column(
|
||||
"total_tokens",
|
||||
new_column_name="input_output_total_tokens",
|
||||
existing_type=sa.Integer(),
|
||||
existing_nullable=True,
|
||||
)
|
||||
|
||||
with op.batch_alter_table("usage") as batch_op:
|
||||
if not column_exists("usage", "input_context_tokens"):
|
||||
batch_op.add_column(sa.Column("input_context_tokens", sa.Integer(), nullable=False, server_default="0"))
|
||||
if not column_exists("usage", "total_tokens"):
|
||||
batch_op.add_column(sa.Column("total_tokens", sa.Integer(), nullable=False, server_default="0"))
|
||||
if not column_exists("usage", "cache_creation_cost_usd_5m"):
|
||||
batch_op.add_column(sa.Column("cache_creation_cost_usd_5m", sa.Numeric(20, 8), nullable=False, server_default="0"))
|
||||
if not column_exists("usage", "cache_creation_cost_usd_1h"):
|
||||
batch_op.add_column(sa.Column("cache_creation_cost_usd_1h", sa.Numeric(20, 8), nullable=False, server_default="0"))
|
||||
if not column_exists("usage", "actual_cache_creation_cost_usd_5m"):
|
||||
batch_op.add_column(sa.Column("actual_cache_creation_cost_usd_5m", sa.Numeric(20, 8), nullable=False, server_default="0"))
|
||||
if not column_exists("usage", "actual_cache_creation_cost_usd_1h"):
|
||||
batch_op.add_column(sa.Column("actual_cache_creation_cost_usd_1h", sa.Numeric(20, 8), nullable=False, server_default="0"))
|
||||
if not column_exists("usage", "actual_cache_cost_usd"):
|
||||
batch_op.add_column(sa.Column("actual_cache_cost_usd", sa.Numeric(20, 8), nullable=False, server_default="0"))
|
||||
if not column_exists("usage", "cache_creation_price_per_1m_5m"):
|
||||
batch_op.add_column(sa.Column("cache_creation_price_per_1m_5m", sa.Numeric(20, 8), nullable=True))
|
||||
if not column_exists("usage", "cache_creation_price_per_1m_1h"):
|
||||
batch_op.add_column(sa.Column("cache_creation_price_per_1m_1h", sa.Numeric(20, 8), nullable=True))
|
||||
|
||||
run_backfill_in_batches(_UPGRADE_BACKFILL_SQL)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_backfill_in_batches(_DOWNGRADE_BACKFILL_SQL)
|
||||
|
||||
with op.batch_alter_table("usage") as batch_op:
|
||||
if column_exists("usage", "cache_creation_price_per_1m_1h"):
|
||||
batch_op.drop_column("cache_creation_price_per_1m_1h")
|
||||
if column_exists("usage", "cache_creation_price_per_1m_5m"):
|
||||
batch_op.drop_column("cache_creation_price_per_1m_5m")
|
||||
if column_exists("usage", "actual_cache_cost_usd"):
|
||||
batch_op.drop_column("actual_cache_cost_usd")
|
||||
if column_exists("usage", "actual_cache_creation_cost_usd_1h"):
|
||||
batch_op.drop_column("actual_cache_creation_cost_usd_1h")
|
||||
if column_exists("usage", "actual_cache_creation_cost_usd_5m"):
|
||||
batch_op.drop_column("actual_cache_creation_cost_usd_5m")
|
||||
if column_exists("usage", "cache_creation_cost_usd_1h"):
|
||||
batch_op.drop_column("cache_creation_cost_usd_1h")
|
||||
if column_exists("usage", "cache_creation_cost_usd_5m"):
|
||||
batch_op.drop_column("cache_creation_cost_usd_5m")
|
||||
if column_exists("usage", "input_context_tokens"):
|
||||
batch_op.drop_column("input_context_tokens")
|
||||
if column_exists("usage", "total_tokens"):
|
||||
batch_op.drop_column("total_tokens")
|
||||
if column_exists("usage", "input_output_total_tokens"):
|
||||
batch_op.alter_column(
|
||||
"input_output_total_tokens",
|
||||
new_column_name="total_tokens",
|
||||
existing_type=sa.Integer(),
|
||||
existing_nullable=True,
|
||||
)
|
||||
@@ -15,7 +15,7 @@
|
||||
### 用户系统
|
||||
- **users**: 用户账户管理
|
||||
- **api_keys**: API 密钥管理
|
||||
- **user_quotas**: 用户配额管理
|
||||
- **wallets**: 统一钱包账户(充值余额/赠款余额/无限制模式)
|
||||
- **user_preferences**: 用户偏好设置
|
||||
|
||||
### Provider 三层架构
|
||||
|
||||
@@ -3,45 +3,131 @@
|
||||
#
|
||||
# 用法:
|
||||
# 部署/更新: ./deploy.sh (自动检测所有变化)
|
||||
# 指定 Hub 版本: ./deploy.sh --hub-tag hub-v0.1.0
|
||||
# 更新 Hub: ./deploy.sh --update-hub
|
||||
# GitHub 镜像: ./deploy.sh --mirror https://ghfast.top
|
||||
# 强制重建: ./deploy.sh --rebuild-base
|
||||
# 强制全部重建: ./deploy.sh --force
|
||||
|
||||
set -e
|
||||
set -euo pipefail
|
||||
cd "$(dirname "$0")"
|
||||
|
||||
# 兼容 docker-compose 和 docker compose
|
||||
if command -v docker-compose &> /dev/null; then
|
||||
DC="docker-compose -f docker-compose.build.yml"
|
||||
USE_LEGACY_COMPOSE=true
|
||||
else
|
||||
DC="docker compose -f docker-compose.build.yml"
|
||||
USE_LEGACY_COMPOSE=false
|
||||
fi
|
||||
|
||||
compose_up() {
|
||||
if [ "$USE_LEGACY_COMPOSE" = true ]; then
|
||||
$DC up -d --no-build "$@"
|
||||
else
|
||||
$DC up -d --no-build --pull never "$@"
|
||||
fi
|
||||
}
|
||||
|
||||
# 缓存文件
|
||||
HASH_FILE=".deps-hash"
|
||||
CODE_HASH_FILE=".code-hash"
|
||||
MIGRATION_HASH_FILE=".migration-hash"
|
||||
|
||||
# 提取 pyproject.toml 中"会影响运行时依赖安装"的最小指纹(与 CI 保持一致):
|
||||
# - [build-system] requires / build-backend
|
||||
# - [project] requires-python / dependencies
|
||||
# 使用 Python tomllib 解析,不受 TOML 格式变化影响。
|
||||
pyproject_deps_fingerprint() {
|
||||
python3 - <<'PY'
|
||||
import json, pathlib, tomllib
|
||||
# Hub release 配置
|
||||
GITHUB_REPO="fawney19/Aether"
|
||||
HUB_TAG_STATE_FILE=".hub-tag"
|
||||
|
||||
data = tomllib.loads(pathlib.Path("pyproject.toml").read_text("utf-8"))
|
||||
project = data.get("project") or {}
|
||||
build = data.get("build-system") or {}
|
||||
usage() {
|
||||
cat <<'EOF'
|
||||
Usage: ./deploy.sh [options]
|
||||
|
||||
fingerprint = {
|
||||
"requires-python": project.get("requires-python"),
|
||||
"dependencies": sorted(project.get("dependencies") or []),
|
||||
"build-backend": build.get("build-backend"),
|
||||
"build-requires": sorted(build.get("requires") or []),
|
||||
Options:
|
||||
--hub-tag <hub-vX.Y.Z> 指定 Hub Release tag(例如 hub-v0.1.0)
|
||||
--update-hub 强制刷新 Hub 版本标记(下次构建会重新下载)
|
||||
--mirror <url> GitHub 下载镜像(例如 https://ghfast.top)
|
||||
--rebuild-base, -r 仅重建 base 镜像
|
||||
--force, -f 强制重建全部(hub/base/app)并重启
|
||||
-h, --help 显示帮助
|
||||
EOF
|
||||
}
|
||||
|
||||
print(json.dumps(fingerprint, sort_keys=True, separators=(",", ":")))
|
||||
PY
|
||||
FORCE_REBUILD_ALL=false
|
||||
REBUILD_BASE_ONLY=false
|
||||
FORCE_UPDATE_HUB=false
|
||||
HUB_TAG="${HUB_TAG:-}"
|
||||
GITHUB_MIRROR="${GITHUB_MIRROR:-}"
|
||||
RESOLVED_HUB_TAG=""
|
||||
|
||||
while [ $# -gt 0 ]; do
|
||||
case "$1" in
|
||||
--hub-tag)
|
||||
if [ $# -lt 2 ]; then
|
||||
echo "❌ --hub-tag 需要一个值,例如 hub-v0.1.0"
|
||||
exit 1
|
||||
fi
|
||||
HUB_TAG="$2"
|
||||
shift 2
|
||||
;;
|
||||
--update-hub)
|
||||
FORCE_UPDATE_HUB=true
|
||||
shift
|
||||
;;
|
||||
--mirror)
|
||||
if [ $# -lt 2 ]; then
|
||||
echo "ERROR: --mirror needs a URL, e.g. https://ghfast.top"
|
||||
exit 1
|
||||
fi
|
||||
GITHUB_MIRROR="$2"
|
||||
shift 2
|
||||
;;
|
||||
--rebuild-base|-r)
|
||||
REBUILD_BASE_ONLY=true
|
||||
shift
|
||||
;;
|
||||
--force|-f)
|
||||
FORCE_REBUILD_ALL=true
|
||||
shift
|
||||
;;
|
||||
-h|--help)
|
||||
usage
|
||||
exit 0
|
||||
;;
|
||||
*)
|
||||
echo "❌ 未知参数: $1"
|
||||
usage
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
if [ -n "$HUB_TAG" ]; then
|
||||
case "$HUB_TAG" in
|
||||
hub-v*) ;;
|
||||
*) echo "❌ --hub-tag 格式应为 hub-vX.Y.Z,例如 hub-v0.1.0"; exit 1 ;;
|
||||
esac
|
||||
fi
|
||||
|
||||
# 提取 pyproject.toml 中会影响运行时依赖安装的字段指纹(纯 shell,无需 Python)
|
||||
# 用 sed 提取 dependencies / requires 数组块和单值字段,排序后输出稳定文本
|
||||
pyproject_deps_fingerprint() {
|
||||
local file="pyproject.toml"
|
||||
# 提取 "key = [..." 多行数组块(从 key 行到 ] 行)
|
||||
extract_array() {
|
||||
sed -n "/^$1[[:space:]]*=[[:space:]]*\[/,/\]/p" "$file" | grep '"' | sed 's/.*"\(.*\)".*/\1/' | sort
|
||||
}
|
||||
# 提取 "key = "value"" 单行值
|
||||
extract_value() {
|
||||
grep -m1 "^$1[[:space:]]*=" "$file" 2>/dev/null | sed 's/.*"\(.*\)".*/\1/'
|
||||
}
|
||||
{
|
||||
echo "requires-python=$(extract_value requires-python)"
|
||||
echo "build-backend=$(extract_value build-backend)"
|
||||
echo "dependencies:"
|
||||
extract_array dependencies
|
||||
echo "build-requires:"
|
||||
extract_array requires
|
||||
}
|
||||
}
|
||||
|
||||
# 计算依赖文件的哈希值(包含 Dockerfile.base.local)
|
||||
@@ -63,6 +149,66 @@ calc_code_hash() {
|
||||
} | md5sum | cut -d' ' -f1
|
||||
}
|
||||
|
||||
# 获取最新 hub release tag
|
||||
# 支持 GITHUB_TOKEN 环境变量以避免未认证 API 限流(60 次/小时 -> 5000 次/小时)
|
||||
get_latest_hub_tag() {
|
||||
local auth_args=()
|
||||
if [ -n "${GITHUB_TOKEN:-}" ]; then
|
||||
auth_args=(-H "Authorization: token ${GITHUB_TOKEN}")
|
||||
fi
|
||||
curl -sL "${auth_args[@]}" "https://api.github.com/repos/$GITHUB_REPO/releases" | \
|
||||
python3 -c "
|
||||
import json, sys
|
||||
releases = json.load(sys.stdin)
|
||||
for r in releases:
|
||||
tag = r.get('tag_name', '')
|
||||
if tag.startswith('hub-v') and not r.get('draft') and not r.get('prerelease'):
|
||||
print(tag)
|
||||
break
|
||||
" 2>/dev/null
|
||||
}
|
||||
|
||||
# 解析当前应使用的 Hub release tag(优先使用指定值,否则拉取最新)
|
||||
resolve_hub_tag() {
|
||||
local requested_tag="${1:-}"
|
||||
local latest_tag
|
||||
|
||||
if [ -n "$requested_tag" ]; then
|
||||
echo "$requested_tag"
|
||||
return 0
|
||||
fi
|
||||
|
||||
latest_tag="$(get_latest_hub_tag || true)"
|
||||
if [ -n "$latest_tag" ]; then
|
||||
echo "$latest_tag"
|
||||
return 0
|
||||
fi
|
||||
|
||||
if [ -f "$HUB_TAG_STATE_FILE" ]; then
|
||||
echo "⚠️ 无法查询最新 Hub 版本,回退使用本地记录: $(cat "$HUB_TAG_STATE_FILE")" >&2
|
||||
cat "$HUB_TAG_STATE_FILE"
|
||||
return 0
|
||||
fi
|
||||
|
||||
echo "❌ 无法获取 Hub Release tag,请检查网络或手动指定 --hub-tag" >&2
|
||||
exit 1
|
||||
}
|
||||
|
||||
# 确保本次构建的 Hub tag 已解析(默认追踪最新 release,也可通过 --hub-tag 固定版本)
|
||||
ensure_hub_tag() {
|
||||
local requested_tag="${1:-}"
|
||||
RESOLVED_HUB_TAG="$(resolve_hub_tag "$requested_tag")"
|
||||
|
||||
if [ -f "$HUB_TAG_STATE_FILE" ] && [ "$(cat "$HUB_TAG_STATE_FILE")" = "$RESOLVED_HUB_TAG" ]; then
|
||||
echo ">>> Hub 版本未变化: $RESOLVED_HUB_TAG"
|
||||
return 1
|
||||
fi
|
||||
|
||||
echo "$RESOLVED_HUB_TAG" > "$HUB_TAG_STATE_FILE"
|
||||
echo ">>> 使用 Hub 版本: $RESOLVED_HUB_TAG"
|
||||
return 0
|
||||
}
|
||||
|
||||
# 计算迁移文件的哈希值
|
||||
calc_migration_hash() {
|
||||
find alembic/versions -name "*.py" -type f 2>/dev/null | sort | xargs cat 2>/dev/null | md5sum | cut -d' ' -f1
|
||||
@@ -92,6 +238,8 @@ check_code_changed() {
|
||||
return 0
|
||||
}
|
||||
|
||||
|
||||
|
||||
# 检查迁移是否变化
|
||||
check_migration_changed() {
|
||||
local current_hash=$(calc_migration_hash)
|
||||
@@ -112,10 +260,11 @@ save_migration_hash() { calc_migration_hash > "$MIGRATION_HASH_FILE"; }
|
||||
# 构建基础镜像
|
||||
build_base() {
|
||||
echo ">>> Building base image (dependencies)..."
|
||||
docker build -f Dockerfile.base.local -t aether-base:latest .
|
||||
docker build --pull=false -f Dockerfile.base.local -t aether-base:latest .
|
||||
save_deps_hash
|
||||
}
|
||||
|
||||
|
||||
# 生成版本文件
|
||||
generate_version_file() {
|
||||
# 从 git 获取版本号
|
||||
@@ -137,8 +286,27 @@ EOF
|
||||
# 构建应用镜像
|
||||
build_app() {
|
||||
echo ">>> Building app image (code only)..."
|
||||
if [ -z "${RESOLVED_HUB_TAG:-}" ]; then
|
||||
echo ">>> RESOLVED_HUB_TAG 为空,无法构建 app 镜像"
|
||||
exit 1
|
||||
fi
|
||||
echo ">>> Build args: HUB_TAG=$RESOLVED_HUB_TAG"
|
||||
generate_version_file
|
||||
docker build -f Dockerfile.app.local -t aether-app:latest .
|
||||
local token_args=()
|
||||
if [ -n "${GITHUB_TOKEN:-}" ]; then
|
||||
token_args=(--build-arg "GITHUB_TOKEN=${GITHUB_TOKEN}")
|
||||
fi
|
||||
local mirror_args=()
|
||||
if [ -n "${GITHUB_MIRROR:-}" ]; then
|
||||
mirror_args=(--build-arg "GITHUB_MIRROR=${GITHUB_MIRROR}")
|
||||
fi
|
||||
docker build --pull=false \
|
||||
--build-arg HUB_RELEASE_REPO="$GITHUB_REPO" \
|
||||
--build-arg HUB_TAG="$RESOLVED_HUB_TAG" \
|
||||
"${token_args[@]}" \
|
||||
"${mirror_args[@]}" \
|
||||
-f Dockerfile.app.local \
|
||||
-t aether-app:latest .
|
||||
save_code_hash
|
||||
}
|
||||
|
||||
@@ -186,11 +354,15 @@ print('Old version cleared')
|
||||
}
|
||||
|
||||
# 强制全部重建
|
||||
if [ "$1" = "--force" ] || [ "$1" = "-f" ]; then
|
||||
if [ "$FORCE_REBUILD_ALL" = true ]; then
|
||||
echo ">>> Force rebuilding everything..."
|
||||
if [ "$FORCE_UPDATE_HUB" = true ]; then
|
||||
rm -f "$HUB_TAG_STATE_FILE"
|
||||
fi
|
||||
ensure_hub_tag "$HUB_TAG" || true
|
||||
build_base
|
||||
build_app
|
||||
$DC up -d --force-recreate
|
||||
compose_up --force-recreate
|
||||
sleep 3
|
||||
run_migration
|
||||
docker image prune -f
|
||||
@@ -200,19 +372,25 @@ if [ "$1" = "--force" ] || [ "$1" = "-f" ]; then
|
||||
fi
|
||||
|
||||
# 强制重建基础镜像
|
||||
if [ "$1" = "--rebuild-base" ] || [ "$1" = "-r" ]; then
|
||||
if [ "$REBUILD_BASE_ONLY" = true ]; then
|
||||
build_base
|
||||
echo ">>> Base image rebuilt. Run ./deploy.sh to deploy."
|
||||
exit 0
|
||||
fi
|
||||
|
||||
# 拉取最新代码
|
||||
echo ">>> Pulling latest code..."
|
||||
git pull
|
||||
# 更新 Hub 版本标记
|
||||
if [ "$FORCE_UPDATE_HUB" = true ]; then
|
||||
rm -f "$HUB_TAG_STATE_FILE"
|
||||
ensure_hub_tag "$HUB_TAG" || true
|
||||
echo ">>> Hub tag updated: $RESOLVED_HUB_TAG"
|
||||
echo ">>> Run ./deploy.sh to build app image with the new Hub release."
|
||||
exit 0
|
||||
fi
|
||||
|
||||
# 标记是否需要重启
|
||||
NEED_RESTART=false
|
||||
BASE_REBUILT=false
|
||||
HUB_UPDATED=false
|
||||
|
||||
# 检查基础镜像是否存在,或依赖是否变化
|
||||
if ! docker image inspect aether-base:latest >/dev/null 2>&1; then
|
||||
@@ -229,6 +407,14 @@ else
|
||||
echo ">>> Dependencies unchanged."
|
||||
fi
|
||||
|
||||
# 解析/检查 Hub 版本(构建时由 Dockerfile 从 GitHub Release 下载)
|
||||
if ensure_hub_tag "$HUB_TAG"; then
|
||||
HUB_UPDATED=true
|
||||
NEED_RESTART=true
|
||||
else
|
||||
echo ">>> Hub version unchanged."
|
||||
fi
|
||||
|
||||
# 检查代码或迁移是否变化,或者 base 重建了(app 依赖 base)
|
||||
# 注意:迁移文件打包在镜像中,所以迁移变化也需要重建 app 镜像
|
||||
MIGRATION_CHANGED=false
|
||||
@@ -244,6 +430,10 @@ elif [ "$BASE_REBUILT" = true ]; then
|
||||
echo ">>> Base image rebuilt, rebuilding app image..."
|
||||
build_app
|
||||
NEED_RESTART=true
|
||||
elif [ "$HUB_UPDATED" = true ]; then
|
||||
echo ">>> Hub version updated, rebuilding app image..."
|
||||
build_app
|
||||
NEED_RESTART=true
|
||||
elif check_code_changed; then
|
||||
echo ">>> Code changed, rebuilding app image..."
|
||||
build_app
|
||||
@@ -265,10 +455,10 @@ fi
|
||||
# 有变化时重启,或容器未运行时启动
|
||||
if [ "$NEED_RESTART" = true ]; then
|
||||
echo ">>> Restarting services..."
|
||||
$DC up -d
|
||||
compose_up
|
||||
elif [ "$CONTAINERS_RUNNING" = false ]; then
|
||||
echo ">>> Containers not running, starting services..."
|
||||
$DC up -d
|
||||
compose_up
|
||||
else
|
||||
echo ">>> No changes detected, skipping restart."
|
||||
fi
|
||||
|
||||
@@ -11,10 +11,16 @@ set +a
|
||||
export DATABASE_URL="postgresql://${DB_USER:-postgres}:${DB_PASSWORD}@${DB_HOST:-localhost}:${DB_PORT:-5432}/${DB_NAME:-aether}"
|
||||
export REDIS_URL=redis://:${REDIS_PASSWORD}@${REDIS_HOST:-localhost}:${REDIS_PORT:-6379}/0
|
||||
|
||||
# 启动 uvicorn(热重载模式)
|
||||
echo "🚀 启动本地开发服务器..."
|
||||
echo "📍 后端地址: http://localhost:8084"
|
||||
echo "📊 数据库: ${DATABASE_URL}"
|
||||
# 开发环境连接池低配(节省内存)
|
||||
export DB_POOL_SIZE=${DB_POOL_SIZE:-5}
|
||||
export DB_MAX_OVERFLOW=${DB_MAX_OVERFLOW:-5}
|
||||
export HTTP_MAX_CONNECTIONS=${HTTP_MAX_CONNECTIONS:-20}
|
||||
export HTTP_KEEPALIVE_CONNECTIONS=${HTTP_KEEPALIVE_CONNECTIONS:-5}
|
||||
|
||||
# 启动 uvicorn(热重载模式,只监视 src 目录)
|
||||
echo "=> 启动本地开发服务器..."
|
||||
echo "=> 后端地址: http://localhost:8084"
|
||||
echo "=> 数据库: ${DATABASE_URL}"
|
||||
echo ""
|
||||
|
||||
uv run uvicorn src.main:app --reload --port 8084
|
||||
uv run uvicorn src.main:app --reload --reload-dir src --port 8084
|
||||
|
||||
@@ -1,7 +1,10 @@
|
||||
# Aether 部署配置 - 本地构建
|
||||
# 使用方法:
|
||||
# 首次构建 base: docker build -f Dockerfile.base -t aether-base:latest .
|
||||
# Hub 二进制: cd aether-hub && ./build.sh(默认 binary 模式)供发布;deploy.sh 构建 app 时会自动下载
|
||||
# Hub 镜像发布: cd aether-hub && ./build.sh --image --tag <tag> --push(可选)
|
||||
# 启动服务: docker compose -f docker-compose.build.yml up -d --build
|
||||
# 或使用: ./deploy.sh
|
||||
|
||||
services:
|
||||
postgres:
|
||||
@@ -26,9 +29,7 @@ services:
|
||||
redis:
|
||||
image: redis:7-alpine
|
||||
container_name: aether-redis
|
||||
command: redis-server --appendonly yes --requirepass ${REDIS_PASSWORD}
|
||||
volumes:
|
||||
- redis_data:/data
|
||||
command: redis-server --appendonly no --save "" --requirepass ${REDIS_PASSWORD}
|
||||
ports:
|
||||
- "${REDIS_PORT:-6379}:6379"
|
||||
healthcheck:
|
||||
@@ -42,6 +43,8 @@ services:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile.app.local
|
||||
args:
|
||||
GITHUB_TOKEN: ${GITHUB_TOKEN:-}
|
||||
image: aether-app:latest
|
||||
container_name: aether-app
|
||||
env_file:
|
||||
@@ -67,4 +70,3 @@ services:
|
||||
|
||||
volumes:
|
||||
postgres_data:
|
||||
redis_data:
|
||||
|
||||
+4
-5
@@ -22,9 +22,7 @@ services:
|
||||
redis:
|
||||
image: redis:7-alpine
|
||||
container_name: aether-redis
|
||||
command: redis-server --appendonly yes --requirepass ${REDIS_PASSWORD}
|
||||
volumes:
|
||||
- redis_data:/data
|
||||
command: redis-server --appendonly no --save "" --requirepass ${REDIS_PASSWORD}
|
||||
healthcheck:
|
||||
test: [ "CMD", "redis-cli", "--raw", "incr", "ping" ]
|
||||
interval: 5s
|
||||
@@ -33,7 +31,7 @@ services:
|
||||
restart: unless-stopped
|
||||
|
||||
app:
|
||||
image: ghcr.io/fawney19/aether:latest
|
||||
image: ${APP_IMAGE:-ghcr.io/fawney19/aether:latest}
|
||||
container_name: aether-app
|
||||
env_file:
|
||||
- .env
|
||||
@@ -58,4 +56,5 @@ services:
|
||||
|
||||
volumes:
|
||||
postgres_data:
|
||||
redis_data:
|
||||
|
||||
|
||||
|
||||
@@ -1,6 +1,26 @@
|
||||
#!/bin/bash
|
||||
set -e
|
||||
|
||||
# Wait for PostgreSQL to be ready
|
||||
MAX_ATTEMPTS=30
|
||||
ATTEMPT=0
|
||||
until python -c "
|
||||
from sqlalchemy import create_engine, text
|
||||
import os
|
||||
engine = create_engine(os.environ['DATABASE_URL'])
|
||||
with engine.connect() as conn:
|
||||
conn.execute(text('SELECT 1'))
|
||||
" 2>/dev/null; do
|
||||
ATTEMPT=$((ATTEMPT + 1))
|
||||
if [ "$ATTEMPT" -ge "$MAX_ATTEMPTS" ]; then
|
||||
echo "Database not ready after $MAX_ATTEMPTS attempts, exiting."
|
||||
exit 1
|
||||
fi
|
||||
echo "Waiting for database... (attempt $ATTEMPT/$MAX_ATTEMPTS)"
|
||||
sleep 2
|
||||
done
|
||||
echo "Database is ready."
|
||||
|
||||
echo "Running database migrations..."
|
||||
alembic upgrade head
|
||||
|
||||
|
||||
Generated
+7
@@ -24,6 +24,7 @@
|
||||
"marked": "^16.0.0",
|
||||
"otpauth": "^9.5.0",
|
||||
"pinia": "^3.0.3",
|
||||
"pinyin-pro": "^3.28.0",
|
||||
"radix-vue": "^1.9.17",
|
||||
"tailwind-merge": "^3.3.1",
|
||||
"three": "^0.180.0",
|
||||
@@ -4718,6 +4719,12 @@
|
||||
}
|
||||
}
|
||||
},
|
||||
"node_modules/pinyin-pro": {
|
||||
"version": "3.28.0",
|
||||
"resolved": "https://registry.npmmirror.com/pinyin-pro/-/pinyin-pro-3.28.0.tgz",
|
||||
"integrity": "sha512-mMRty6RisoyYNphJrTo3pnvp3w8OMZBrXm9YSWkxhAfxKj1KZk2y8T2PDIZlDDRsvZ0No+Hz6FI4sZpA6Ey25g==",
|
||||
"license": "MIT"
|
||||
},
|
||||
"node_modules/pirates": {
|
||||
"version": "4.0.7",
|
||||
"resolved": "https://registry.npmjs.org/pirates/-/pirates-4.0.7.tgz",
|
||||
|
||||
@@ -32,6 +32,7 @@
|
||||
"marked": "^16.0.0",
|
||||
"otpauth": "^9.5.0",
|
||||
"pinia": "^3.0.3",
|
||||
"pinyin-pro": "^3.28.0",
|
||||
"radix-vue": "^1.9.17",
|
||||
"tailwind-merge": "^3.3.1",
|
||||
"three": "^0.180.0",
|
||||
|
||||
+66
-2
@@ -5,12 +5,14 @@
|
||||
</template>
|
||||
|
||||
<script setup lang="ts">
|
||||
import { onMounted, onErrorCaptured } from 'vue'
|
||||
import { onMounted, onErrorCaptured, onUnmounted } from 'vue'
|
||||
import { useAuthStore } from '@/stores/auth'
|
||||
import ToastContainer from '@/components/ToastContainer.vue'
|
||||
import ConfirmContainer from '@/components/ConfirmContainer.vue'
|
||||
import apiClient from '@/api/client'
|
||||
import apiClient, { AUTH_STATE_CHANGE_EVENT } from '@/api/client'
|
||||
import { NETWORK_CONFIG, AUTH_CONFIG } from '@/config/constants'
|
||||
import router from '@/router'
|
||||
import { hasAuthIdentityChanged } from '@/utils/authToken'
|
||||
import { log } from '@/utils/logger'
|
||||
|
||||
const authStore = useAuthStore()
|
||||
@@ -86,7 +88,59 @@ if (typeof window !== 'undefined') {
|
||||
})
|
||||
}
|
||||
|
||||
async function syncExternalAuthState(nextToken: string | null): Promise<void> {
|
||||
const previousToken = authStore.token
|
||||
const previousUser = authStore.user
|
||||
? {
|
||||
id: authStore.user.id,
|
||||
role: authStore.user.role,
|
||||
}
|
||||
: null
|
||||
|
||||
authStore.syncToken()
|
||||
|
||||
if (!nextToken) {
|
||||
if (previousToken || previousUser) {
|
||||
authStore.applyExternalLogout()
|
||||
await router.replace('/')
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
const identityChanged = hasAuthIdentityChanged(previousToken, nextToken, previousUser)
|
||||
if (!identityChanged && previousUser) {
|
||||
return
|
||||
}
|
||||
|
||||
const user = await authStore.fetchCurrentUser()
|
||||
if (!user) {
|
||||
return
|
||||
}
|
||||
|
||||
if (router.currentRoute.value.path.startsWith('/admin') && user.role !== 'admin') {
|
||||
await router.replace('/dashboard')
|
||||
}
|
||||
}
|
||||
|
||||
function handleAuthStorageChange(event: StorageEvent): void {
|
||||
if (event.key !== 'access_token') {
|
||||
return
|
||||
}
|
||||
|
||||
syncExternalAuthState(event.newValue).catch((err) => log.error('syncExternalAuthState failed', err))
|
||||
}
|
||||
|
||||
function handleLocalAuthStateChange(event: Event): void {
|
||||
const authEvent = event as CustomEvent<{ token: string | null }>
|
||||
syncExternalAuthState(authEvent.detail?.token ?? apiClient.getToken()).catch((err) => log.error('syncExternalAuthState failed', err))
|
||||
}
|
||||
|
||||
onMounted(async () => {
|
||||
if (typeof window !== 'undefined') {
|
||||
window.addEventListener('storage', handleAuthStorageChange)
|
||||
window.addEventListener(AUTH_STATE_CHANGE_EVENT, handleLocalAuthStateChange as (event: Event) => void)
|
||||
}
|
||||
|
||||
// 延迟检查认证状态,让页面先加载
|
||||
setTimeout(async () => {
|
||||
try {
|
||||
@@ -97,4 +151,14 @@ onMounted(async () => {
|
||||
}
|
||||
}, AUTH_CONFIG.TOKEN_REFRESH_INTERVAL)
|
||||
})
|
||||
|
||||
onUnmounted(() => {
|
||||
if (typeof window !== 'undefined') {
|
||||
window.removeEventListener('storage', handleAuthStorageChange)
|
||||
window.removeEventListener(
|
||||
AUTH_STATE_CHANGE_EVENT,
|
||||
handleLocalAuthStateChange as (event: Event) => void,
|
||||
)
|
||||
}
|
||||
})
|
||||
</script>
|
||||
|
||||
@@ -0,0 +1,31 @@
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
|
||||
|
||||
import apiClient, { AUTH_STATE_CHANGE_EVENT } from '@/api/client'
|
||||
|
||||
describe('apiClient auth state change event', () => {
|
||||
beforeEach(() => {
|
||||
localStorage.clear()
|
||||
apiClient.clearAuth()
|
||||
})
|
||||
|
||||
afterEach(() => {
|
||||
localStorage.clear()
|
||||
apiClient.clearAuth()
|
||||
})
|
||||
|
||||
it('dispatches a same-tab auth change event when clearing auth', () => {
|
||||
const handler = vi.fn()
|
||||
window.addEventListener(AUTH_STATE_CHANGE_EVENT, handler as EventListener)
|
||||
|
||||
apiClient.setToken('access-token')
|
||||
apiClient.clearAuth()
|
||||
|
||||
expect(localStorage.getItem('access_token')).toBeNull()
|
||||
expect(handler).toHaveBeenCalledTimes(1)
|
||||
|
||||
const event = handler.mock.calls[0][0] as CustomEvent<{ token: string | null }>
|
||||
expect(event.detail).toEqual({ token: null })
|
||||
|
||||
window.removeEventListener(AUTH_STATE_CHANGE_EVENT, handler as EventListener)
|
||||
})
|
||||
})
|
||||
@@ -0,0 +1,93 @@
|
||||
import apiClient from './client'
|
||||
import type { PaymentOrder } from './wallet'
|
||||
|
||||
export interface PaymentCallbackRecord {
|
||||
id: string
|
||||
payment_order_id: string | null
|
||||
payment_method: string
|
||||
callback_key: string
|
||||
order_no: string | null
|
||||
gateway_order_id: string | null
|
||||
payload_hash: string | null
|
||||
signature_valid: boolean
|
||||
status: string
|
||||
payload: Record<string, unknown> | null
|
||||
error_message: string | null
|
||||
created_at: string
|
||||
processed_at: string | null
|
||||
}
|
||||
|
||||
export interface AdminPaymentOrderListResponse {
|
||||
items: PaymentOrder[]
|
||||
total: number
|
||||
limit: number
|
||||
offset: number
|
||||
}
|
||||
|
||||
export interface AdminPaymentCallbacksResponse {
|
||||
items: PaymentCallbackRecord[]
|
||||
total: number
|
||||
limit: number
|
||||
offset: number
|
||||
}
|
||||
|
||||
export interface AdminPaymentCreditRequest {
|
||||
gateway_order_id?: string
|
||||
pay_amount?: number
|
||||
pay_currency?: string
|
||||
exchange_rate?: number
|
||||
gateway_response?: Record<string, unknown>
|
||||
}
|
||||
|
||||
export const adminPaymentsApi = {
|
||||
async listOrders(params?: {
|
||||
status?: string
|
||||
payment_method?: string
|
||||
limit?: number
|
||||
offset?: number
|
||||
}): Promise<AdminPaymentOrderListResponse> {
|
||||
const response = await apiClient.get<AdminPaymentOrderListResponse>('/api/admin/payments/orders', { params })
|
||||
return response.data
|
||||
},
|
||||
|
||||
async getOrder(orderId: string): Promise<{ order: PaymentOrder }> {
|
||||
const response = await apiClient.get<{ order: PaymentOrder }>(`/api/admin/payments/orders/${orderId}`)
|
||||
return response.data
|
||||
},
|
||||
|
||||
async expireOrder(orderId: string): Promise<{ order: PaymentOrder; expired: boolean }> {
|
||||
const response = await apiClient.post<{ order: PaymentOrder; expired: boolean }>(
|
||||
`/api/admin/payments/orders/${orderId}/expire`,
|
||||
{}
|
||||
)
|
||||
return response.data
|
||||
},
|
||||
|
||||
async failOrder(orderId: string): Promise<{ order: PaymentOrder }> {
|
||||
const response = await apiClient.post<{ order: PaymentOrder }>(
|
||||
`/api/admin/payments/orders/${orderId}/fail`,
|
||||
{}
|
||||
)
|
||||
return response.data
|
||||
},
|
||||
|
||||
async creditOrder(
|
||||
orderId: string,
|
||||
payload: AdminPaymentCreditRequest
|
||||
): Promise<{ order: PaymentOrder; credited: boolean }> {
|
||||
const response = await apiClient.post<{ order: PaymentOrder; credited: boolean }>(
|
||||
`/api/admin/payments/orders/${orderId}/credit`,
|
||||
payload
|
||||
)
|
||||
return response.data
|
||||
},
|
||||
|
||||
async listCallbacks(params?: {
|
||||
payment_method?: string
|
||||
limit?: number
|
||||
offset?: number
|
||||
}): Promise<AdminPaymentCallbacksResponse> {
|
||||
const response = await apiClient.get<AdminPaymentCallbacksResponse>('/api/admin/payments/callbacks', { params })
|
||||
return response.data
|
||||
},
|
||||
}
|
||||
@@ -0,0 +1,247 @@
|
||||
import apiClient from './client'
|
||||
import type { RefundRequest, WalletSummary, WalletTransaction } from './wallet'
|
||||
|
||||
export interface AdminWallet extends WalletSummary {
|
||||
user_id: string | null
|
||||
api_key_id: string | null
|
||||
owner_type: 'user' | 'api_key'
|
||||
owner_name: string | null
|
||||
created_at: string
|
||||
}
|
||||
|
||||
export interface AdminWalletListResponse {
|
||||
items: AdminWallet[]
|
||||
total: number
|
||||
limit: number
|
||||
offset: number
|
||||
}
|
||||
|
||||
export interface AdminWalletDetailResponse extends AdminWallet {
|
||||
pending_refund_count: number
|
||||
}
|
||||
|
||||
export interface AdminWalletTransactionsResponse {
|
||||
wallet: AdminWallet
|
||||
items: WalletTransaction[]
|
||||
total: number
|
||||
limit: number
|
||||
offset: number
|
||||
}
|
||||
|
||||
export interface AdminWalletRefundsResponse {
|
||||
wallet: AdminWallet
|
||||
items: RefundRequest[]
|
||||
total: number
|
||||
limit: number
|
||||
offset: number
|
||||
}
|
||||
|
||||
export interface AdminLedgerTransaction extends WalletTransaction {
|
||||
wallet_id: string
|
||||
owner_type: 'user' | 'api_key'
|
||||
owner_name: string | null
|
||||
wallet_status?: string | null
|
||||
}
|
||||
|
||||
export interface AdminGlobalRefund extends RefundRequest {
|
||||
wallet_id: string
|
||||
owner_type: 'user' | 'api_key'
|
||||
owner_name: string | null
|
||||
wallet_status?: string | null
|
||||
}
|
||||
|
||||
export interface AdminLedgerResponse {
|
||||
items: AdminLedgerTransaction[]
|
||||
total: number
|
||||
limit: number
|
||||
offset: number
|
||||
}
|
||||
|
||||
export interface AdminGlobalRefundsListResponse {
|
||||
items: AdminGlobalRefund[]
|
||||
total: number
|
||||
limit: number
|
||||
offset: number
|
||||
}
|
||||
|
||||
export interface ManualRechargeRequest {
|
||||
amount_usd: number
|
||||
payment_method?: string
|
||||
description?: string
|
||||
}
|
||||
|
||||
export interface WalletAdjustRequest {
|
||||
amount_usd: number
|
||||
balance_type?: 'recharge' | 'gift'
|
||||
description?: string
|
||||
}
|
||||
|
||||
export interface RefundFailRequest {
|
||||
reason: string
|
||||
}
|
||||
|
||||
export interface RefundCompleteRequest {
|
||||
gateway_refund_id?: string
|
||||
payout_reference?: string
|
||||
payout_proof?: Record<string, unknown>
|
||||
}
|
||||
|
||||
export const adminWalletApi = {
|
||||
async listWallets(params?: {
|
||||
status?: string
|
||||
limit?: number
|
||||
offset?: number
|
||||
}): Promise<AdminWalletListResponse> {
|
||||
const response = await apiClient.get<AdminWalletListResponse>('/api/admin/wallets', { params })
|
||||
return response.data
|
||||
},
|
||||
|
||||
async listAllWallets(params?: {
|
||||
status?: string
|
||||
}): Promise<AdminWallet[]> {
|
||||
const items: AdminWallet[] = []
|
||||
const limit = 200
|
||||
const maxPages = 200
|
||||
let offset = 0
|
||||
let page = 0
|
||||
|
||||
while (page < maxPages) {
|
||||
const response = await apiClient.get<AdminWalletListResponse>('/api/admin/wallets', {
|
||||
params: {
|
||||
...params,
|
||||
limit,
|
||||
offset,
|
||||
},
|
||||
})
|
||||
const data = response.data
|
||||
items.push(...data.items)
|
||||
|
||||
if (items.length >= data.total || data.items.length < limit) {
|
||||
break
|
||||
}
|
||||
|
||||
const nextOffset = offset + data.items.length
|
||||
if (nextOffset <= offset) {
|
||||
throw new Error('分页游标未前进,终止全量钱包拉取以避免死循环')
|
||||
}
|
||||
offset = nextOffset
|
||||
page += 1
|
||||
}
|
||||
|
||||
if (page >= maxPages) {
|
||||
throw new Error(`钱包列表分页超过最大页数 ${maxPages},已中止请求`)
|
||||
}
|
||||
|
||||
return items
|
||||
},
|
||||
|
||||
async getWalletDetail(walletId: string): Promise<AdminWalletDetailResponse> {
|
||||
const response = await apiClient.get<AdminWalletDetailResponse>(`/api/admin/wallets/${walletId}`)
|
||||
return response.data
|
||||
},
|
||||
|
||||
async listLedger(params?: {
|
||||
category?: string
|
||||
reason_code?: string
|
||||
owner_type?: string
|
||||
limit?: number
|
||||
offset?: number
|
||||
}): Promise<AdminLedgerResponse> {
|
||||
const response = await apiClient.get<AdminLedgerResponse>('/api/admin/wallets/ledger', { params })
|
||||
return response.data
|
||||
},
|
||||
|
||||
async listGlobalRefunds(params?: {
|
||||
status?: string
|
||||
owner_type?: string
|
||||
limit?: number
|
||||
offset?: number
|
||||
}): Promise<AdminGlobalRefundsListResponse> {
|
||||
const response = await apiClient.get<AdminGlobalRefundsListResponse>('/api/admin/wallets/refund-requests', {
|
||||
params,
|
||||
})
|
||||
return response.data
|
||||
},
|
||||
|
||||
async getWalletTransactions(
|
||||
walletId: string,
|
||||
params?: { limit?: number; offset?: number }
|
||||
): Promise<AdminWalletTransactionsResponse> {
|
||||
const response = await apiClient.get<AdminWalletTransactionsResponse>(
|
||||
`/api/admin/wallets/${walletId}/transactions`,
|
||||
{ params }
|
||||
)
|
||||
return response.data
|
||||
},
|
||||
|
||||
async getWalletRefunds(
|
||||
walletId: string,
|
||||
params?: { limit?: number; offset?: number }
|
||||
): Promise<AdminWalletRefundsResponse> {
|
||||
const response = await apiClient.get<AdminWalletRefundsResponse>(
|
||||
`/api/admin/wallets/${walletId}/refunds`,
|
||||
{ params }
|
||||
)
|
||||
return response.data
|
||||
},
|
||||
|
||||
async rechargeWallet(walletId: string, payload: ManualRechargeRequest): Promise<{
|
||||
wallet: AdminWallet
|
||||
payment_order: {
|
||||
id: string
|
||||
order_no: string
|
||||
amount_usd: number
|
||||
payment_method: string
|
||||
status: string
|
||||
created_at: string
|
||||
credited_at: string | null
|
||||
}
|
||||
}> {
|
||||
const response = await apiClient.post(`/api/admin/wallets/${walletId}/recharge`, payload)
|
||||
return response.data
|
||||
},
|
||||
|
||||
async adjustWallet(walletId: string, payload: WalletAdjustRequest): Promise<{
|
||||
wallet: AdminWallet
|
||||
transaction: WalletTransaction
|
||||
}> {
|
||||
const response = await apiClient.post(`/api/admin/wallets/${walletId}/adjust`, payload)
|
||||
return response.data
|
||||
},
|
||||
|
||||
async processRefund(walletId: string, refundId: string): Promise<{
|
||||
wallet: AdminWallet
|
||||
refund: RefundRequest
|
||||
transaction: WalletTransaction
|
||||
}> {
|
||||
const response = await apiClient.post(
|
||||
`/api/admin/wallets/${walletId}/refunds/${refundId}/process`,
|
||||
{}
|
||||
)
|
||||
return response.data
|
||||
},
|
||||
|
||||
async failRefund(walletId: string, refundId: string, payload: RefundFailRequest): Promise<{
|
||||
wallet: AdminWallet
|
||||
refund: RefundRequest
|
||||
transaction: WalletTransaction | null
|
||||
}> {
|
||||
const response = await apiClient.post(
|
||||
`/api/admin/wallets/${walletId}/refunds/${refundId}/fail`,
|
||||
payload
|
||||
)
|
||||
return response.data
|
||||
},
|
||||
|
||||
async completeRefund(
|
||||
walletId: string,
|
||||
refundId: string,
|
||||
payload: RefundCompleteRequest
|
||||
): Promise<{ refund: RefundRequest }> {
|
||||
const response = await apiClient.post(
|
||||
`/api/admin/wallets/${walletId}/refunds/${refundId}/complete`,
|
||||
payload
|
||||
)
|
||||
return response.data
|
||||
},
|
||||
}
|
||||
+61
-282
@@ -1,4 +1,6 @@
|
||||
import apiClient from './client'
|
||||
import { cachedRequest, buildCacheKey } from '@/utils/cache'
|
||||
import type { BillingSummary } from './auth'
|
||||
|
||||
// LDAP 配置导出结构
|
||||
export interface LDAPConfigExport {
|
||||
@@ -33,14 +35,38 @@ export interface OAuthProviderExport {
|
||||
is_enabled?: boolean
|
||||
}
|
||||
|
||||
export interface SystemConfigExport {
|
||||
key: string
|
||||
value: unknown
|
||||
description?: string | null
|
||||
}
|
||||
|
||||
// 配置导出数据结构
|
||||
export interface ConfigExportData {
|
||||
version: string
|
||||
exported_at: string
|
||||
global_models: GlobalModelExport[]
|
||||
providers: ProviderExport[]
|
||||
proxy_nodes?: ProxyNodeExport[]
|
||||
ldap_config?: LDAPConfigExport | null
|
||||
oauth_providers?: OAuthProviderExport[]
|
||||
system_configs?: SystemConfigExport[]
|
||||
}
|
||||
|
||||
export interface ProxyNodeExport {
|
||||
id: string
|
||||
name: string
|
||||
ip: string
|
||||
port: number
|
||||
region?: string | null
|
||||
is_manual: boolean
|
||||
proxy_url?: string | null
|
||||
proxy_username?: string | null
|
||||
proxy_password?: string | null
|
||||
tunnel_mode: boolean
|
||||
heartbeat_interval: number
|
||||
remote_config?: Record<string, unknown> | null
|
||||
config_version: number
|
||||
}
|
||||
|
||||
// 用户导出数据结构
|
||||
@@ -53,31 +79,31 @@ export interface UsersExportData {
|
||||
|
||||
export interface UserExport {
|
||||
email: string
|
||||
email_verified?: boolean
|
||||
username: string
|
||||
password_hash: string
|
||||
role: string
|
||||
allowed_providers?: string[] | null
|
||||
allowed_api_formats?: string[] | null
|
||||
allowed_models?: string[] | null
|
||||
rate_limit?: number | null // null = 跟随系统默认,0 = 不限制
|
||||
model_capability_settings?: Record<string, Record<string, boolean>>
|
||||
quota_usd?: number | null
|
||||
used_usd?: number
|
||||
total_usd?: number
|
||||
unlimited?: boolean
|
||||
wallet?: BillingSummary | null
|
||||
is_active: boolean
|
||||
api_keys: UserApiKeyExport[]
|
||||
}
|
||||
|
||||
export interface UserApiKeyExport {
|
||||
key?: string | null
|
||||
key_hash: string
|
||||
key_encrypted?: string | null
|
||||
name?: string | null
|
||||
is_standalone: boolean
|
||||
balance_used_usd?: number
|
||||
current_balance_usd?: number | null
|
||||
allowed_providers?: string[] | null
|
||||
allowed_api_formats?: string[] | null
|
||||
allowed_models?: string[] | null
|
||||
rate_limit?: number | null // null = 无限制
|
||||
rate_limit?: number | null // legacy/null 兼容;1.3+ standalone null = 跟随系统默认
|
||||
concurrent_limit?: number | null
|
||||
force_capabilities?: Record<string, boolean>
|
||||
is_active: boolean
|
||||
@@ -104,14 +130,18 @@ export interface ProviderExport {
|
||||
name: string
|
||||
description?: string | null
|
||||
website?: string | null
|
||||
provider_type?: string
|
||||
billing_type?: string | null
|
||||
monthly_quota_usd?: number | null
|
||||
quota_reset_day?: number
|
||||
rpm_limit?: number | null
|
||||
provider_priority?: number
|
||||
keep_priority_on_conversion?: boolean
|
||||
enable_format_conversion?: boolean
|
||||
is_active: boolean
|
||||
concurrent_limit?: number | null
|
||||
max_retries?: number | null
|
||||
stream_first_byte_timeout?: number | null
|
||||
request_timeout?: number | null
|
||||
proxy?: Record<string, unknown>
|
||||
config?: Record<string, unknown>
|
||||
endpoints: EndpointExport[]
|
||||
@@ -122,19 +152,24 @@ export interface ProviderExport {
|
||||
export interface EndpointExport {
|
||||
api_format: string
|
||||
base_url: string
|
||||
headers?: Record<string, unknown>
|
||||
header_rules?: Record<string, unknown>[] | null
|
||||
body_rules?: Record<string, unknown>[] | null
|
||||
max_retries?: number
|
||||
is_active: boolean
|
||||
custom_path?: string | null
|
||||
config?: Record<string, unknown>
|
||||
format_acceptance_config?: Record<string, unknown> | null
|
||||
proxy?: Record<string, unknown>
|
||||
}
|
||||
|
||||
export interface ProviderKeyExport {
|
||||
api_key: string
|
||||
auth_type?: string
|
||||
auth_config?: string | Record<string, unknown> | null
|
||||
name?: string | null
|
||||
note?: string | null
|
||||
api_formats: string[]
|
||||
supported_endpoints?: string[]
|
||||
rate_multipliers?: Record<string, number> | null
|
||||
internal_priority?: number
|
||||
global_priority_by_format?: Record<string, number> | null
|
||||
@@ -143,7 +178,13 @@ export interface ProviderKeyExport {
|
||||
capabilities?: Record<string, boolean>
|
||||
cache_ttl_minutes?: number
|
||||
max_probe_interval_minutes?: number
|
||||
auto_fetch_models?: boolean
|
||||
locked_models?: string[] | null
|
||||
model_include_patterns?: string[] | null
|
||||
model_exclude_patterns?: string[] | null
|
||||
is_active: boolean
|
||||
proxy?: Record<string, unknown> | null
|
||||
fingerprint?: Record<string, unknown> | null
|
||||
}
|
||||
|
||||
export interface ModelExport {
|
||||
@@ -285,6 +326,7 @@ export interface ConfigImportResponse {
|
||||
message: string
|
||||
stats: {
|
||||
global_models: { created: number; updated: number; skipped: number }
|
||||
proxy_nodes?: { created: number; updated: number; skipped: number }
|
||||
providers: { created: number; updated: number; skipped: number }
|
||||
endpoints: { created: number; updated: number; skipped: number }
|
||||
keys: { created: number; updated: number; skipped: number }
|
||||
@@ -304,14 +346,11 @@ export interface AdminApiKey {
|
||||
name?: string
|
||||
key_display?: string // 脱敏后的密钥显示
|
||||
is_active: boolean
|
||||
is_locked: boolean // 管理员锁定标志
|
||||
is_standalone: boolean // 是否为独立余额Key
|
||||
balance_used_usd?: number // 已使用余额(仅独立Key)
|
||||
current_balance_usd?: number | null // 当前余额(独立Key预付费模式,null表示无限制)
|
||||
total_requests?: number
|
||||
total_tokens?: number
|
||||
total_cost_usd?: number
|
||||
rate_limit?: number | null // null = 无限制
|
||||
rate_limit?: number | null // null = 跟随系统默认,0 = 不限制
|
||||
allowed_providers?: string[] | null // 允许的提供商列表
|
||||
allowed_api_formats?: string[] | null // 允许的 API 格式列表
|
||||
allowed_models?: string[] | null // 允许的模型列表
|
||||
@@ -327,9 +366,10 @@ export interface CreateStandaloneApiKeyRequest {
|
||||
allowed_providers?: string[] | null
|
||||
allowed_api_formats?: string[] | null
|
||||
allowed_models?: string[] | null
|
||||
rate_limit?: number | null // null = 无限制
|
||||
rate_limit?: number | null // null = 跟随系统默认,0 = 不限制
|
||||
expires_at?: string | null // ISO 日期字符串,如 "2025-12-31",null = 永不过期
|
||||
initial_balance_usd: number // 初始余额,必须设置
|
||||
initial_balance_usd: number | null // 初始余额,null = 无限制
|
||||
unlimited_balance?: boolean | null // 编辑时仅切换额度模式,不调整余额数值
|
||||
auto_delete_on_expiry?: boolean // 过期后是否自动删除
|
||||
}
|
||||
|
||||
@@ -340,102 +380,6 @@ export interface AdminApiKeysResponse {
|
||||
skip: number
|
||||
}
|
||||
|
||||
export interface LeaderboardItem {
|
||||
rank: number
|
||||
id: string
|
||||
name: string
|
||||
value: number
|
||||
requests: number
|
||||
tokens: number
|
||||
cost: number
|
||||
}
|
||||
|
||||
export interface LeaderboardResponse {
|
||||
items: LeaderboardItem[]
|
||||
total: number
|
||||
metric: string
|
||||
start_date?: string | null
|
||||
end_date?: string | null
|
||||
}
|
||||
|
||||
export interface CostForecastResponse {
|
||||
history: Array<{ date: string; total_cost: number }>
|
||||
forecast: Array<{ date: string; total_cost: number }>
|
||||
slope: number
|
||||
intercept: number
|
||||
start_date: string
|
||||
end_date: string
|
||||
}
|
||||
|
||||
export interface CostSavingsResponse {
|
||||
cache_read_tokens: number
|
||||
cache_read_cost: number
|
||||
cache_creation_cost: number
|
||||
estimated_full_cost: number
|
||||
cache_savings: number
|
||||
}
|
||||
|
||||
export interface QuotaUsageProvider {
|
||||
id: string
|
||||
name: string
|
||||
quota_usd: number
|
||||
used_usd: number
|
||||
remaining_usd: number
|
||||
usage_percent: number
|
||||
quota_expires_at?: string | null
|
||||
estimated_exhaust_at?: string | null
|
||||
}
|
||||
|
||||
export interface QuotaUsageResponse {
|
||||
providers: QuotaUsageProvider[]
|
||||
}
|
||||
|
||||
export interface PercentileItem {
|
||||
date: string
|
||||
p50_response_time_ms?: number | null
|
||||
p90_response_time_ms?: number | null
|
||||
p99_response_time_ms?: number | null
|
||||
p50_first_byte_time_ms?: number | null
|
||||
p90_first_byte_time_ms?: number | null
|
||||
p99_first_byte_time_ms?: number | null
|
||||
}
|
||||
|
||||
export interface ErrorDistributionItem {
|
||||
category: string
|
||||
count: number
|
||||
}
|
||||
|
||||
export interface ErrorTrendItem {
|
||||
date: string
|
||||
total: number
|
||||
categories: Record<string, number>
|
||||
}
|
||||
|
||||
export interface ErrorDistributionResponse {
|
||||
distribution: ErrorDistributionItem[]
|
||||
trend: ErrorTrendItem[]
|
||||
}
|
||||
|
||||
export interface ComparisonMetric {
|
||||
total_requests: number
|
||||
total_tokens: number
|
||||
total_cost: number
|
||||
actual_total_cost: number
|
||||
avg_response_time_ms: number
|
||||
error_requests: number
|
||||
}
|
||||
|
||||
export interface ComparisonResponse {
|
||||
current: ComparisonMetric
|
||||
comparison: ComparisonMetric
|
||||
change_percent: Record<string, number | null>
|
||||
current_start: string
|
||||
current_end: string
|
||||
comparison_start: string
|
||||
comparison_end: string
|
||||
}
|
||||
|
||||
|
||||
export interface ApiKeyToggleResponse {
|
||||
id: string // UUID
|
||||
is_active: boolean
|
||||
@@ -477,7 +421,10 @@ export const adminApi = {
|
||||
},
|
||||
|
||||
// 更新独立余额Key
|
||||
async updateApiKey(keyId: string, data: Partial<CreateStandaloneApiKeyRequest>): Promise<AdminApiKey & { message: string }> {
|
||||
async updateApiKey(
|
||||
keyId: string,
|
||||
data: Partial<CreateStandaloneApiKeyRequest>
|
||||
): Promise<AdminApiKey & { message: string }> {
|
||||
const response = await apiClient.put<AdminApiKey & { message: string }>(
|
||||
`/api/admin/api-keys/${keyId}`,
|
||||
data
|
||||
@@ -501,27 +448,10 @@ export const adminApi = {
|
||||
return response.data
|
||||
},
|
||||
|
||||
// 切换API密钥锁定状态(锁定/解锁)
|
||||
async toggleLockApiKey(keyId: string): Promise<ApiKeyLockResponse> {
|
||||
// 切换用户普通 API Key 锁定状态(锁定/解锁)
|
||||
async toggleUserApiKeyLock(userId: string, keyId: string): Promise<ApiKeyLockResponse> {
|
||||
const response = await apiClient.patch<ApiKeyLockResponse>(
|
||||
`/api/admin/api-keys/${keyId}/lock`
|
||||
)
|
||||
return response.data
|
||||
},
|
||||
|
||||
// 为独立余额Key调整余额
|
||||
async addApiKeyBalance(keyId: string, amountUsd: number): Promise<AdminApiKey & { message: string }> {
|
||||
const response = await apiClient.patch<AdminApiKey & { message: string }>(
|
||||
`/api/admin/api-keys/${keyId}/balance`,
|
||||
{ amount_usd: amountUsd }
|
||||
)
|
||||
return response.data
|
||||
},
|
||||
|
||||
// 重置独立余额Key的已使用额度
|
||||
async resetApiKeyUsage(keyId: string): Promise<AdminApiKey & { message: string }> {
|
||||
const response = await apiClient.patch<AdminApiKey & { message: string }>(
|
||||
`/api/admin/api-keys/${keyId}/reset-usage`
|
||||
`/api/admin/users/${userId}/api-keys/${keyId}/lock`
|
||||
)
|
||||
return response.data
|
||||
},
|
||||
@@ -727,142 +657,6 @@ export const adminApi = {
|
||||
return response.data
|
||||
},
|
||||
|
||||
// Stats / Leaderboards
|
||||
async getLeaderboardUsers(params?: {
|
||||
start_date?: string
|
||||
end_date?: string
|
||||
preset?: string
|
||||
timezone?: string
|
||||
tz_offset_minutes?: number
|
||||
metric?: 'requests' | 'tokens' | 'cost'
|
||||
order?: 'asc' | 'desc'
|
||||
limit?: number
|
||||
offset?: number
|
||||
provider_name?: string
|
||||
model?: string
|
||||
include_inactive?: boolean
|
||||
exclude_admin?: boolean
|
||||
}): Promise<LeaderboardResponse> {
|
||||
const response = await apiClient.get<LeaderboardResponse>('/api/admin/stats/leaderboard/users', {
|
||||
params
|
||||
})
|
||||
return response.data
|
||||
},
|
||||
|
||||
async getLeaderboardApiKeys(params?: {
|
||||
start_date?: string
|
||||
end_date?: string
|
||||
preset?: string
|
||||
timezone?: string
|
||||
tz_offset_minutes?: number
|
||||
metric?: 'requests' | 'tokens' | 'cost'
|
||||
order?: 'asc' | 'desc'
|
||||
limit?: number
|
||||
offset?: number
|
||||
provider_name?: string
|
||||
model?: string
|
||||
include_inactive?: boolean
|
||||
exclude_admin?: boolean
|
||||
}): Promise<LeaderboardResponse> {
|
||||
const response = await apiClient.get<LeaderboardResponse>('/api/admin/stats/leaderboard/api-keys', {
|
||||
params
|
||||
})
|
||||
return response.data
|
||||
},
|
||||
|
||||
async getLeaderboardModels(params?: {
|
||||
start_date?: string
|
||||
end_date?: string
|
||||
preset?: string
|
||||
timezone?: string
|
||||
tz_offset_minutes?: number
|
||||
metric?: 'requests' | 'tokens' | 'cost'
|
||||
order?: 'asc' | 'desc'
|
||||
limit?: number
|
||||
offset?: number
|
||||
provider_name?: string
|
||||
model?: string
|
||||
}): Promise<LeaderboardResponse> {
|
||||
const response = await apiClient.get<LeaderboardResponse>('/api/admin/stats/leaderboard/models', {
|
||||
params
|
||||
})
|
||||
return response.data
|
||||
},
|
||||
|
||||
async getCostForecast(params?: {
|
||||
start_date?: string
|
||||
end_date?: string
|
||||
preset?: string
|
||||
timezone?: string
|
||||
tz_offset_minutes?: number
|
||||
days?: number
|
||||
forecast_days?: number
|
||||
}): Promise<CostForecastResponse> {
|
||||
const response = await apiClient.get<CostForecastResponse>('/api/admin/stats/cost/forecast', {
|
||||
params
|
||||
})
|
||||
return response.data
|
||||
},
|
||||
|
||||
async getCostSavings(params?: {
|
||||
start_date?: string
|
||||
end_date?: string
|
||||
preset?: string
|
||||
timezone?: string
|
||||
tz_offset_minutes?: number
|
||||
provider_name?: string
|
||||
model?: string
|
||||
}): Promise<CostSavingsResponse> {
|
||||
const response = await apiClient.get<CostSavingsResponse>('/api/admin/stats/cost/savings', {
|
||||
params
|
||||
})
|
||||
return response.data
|
||||
},
|
||||
|
||||
async getQuotaUsage(): Promise<QuotaUsageResponse> {
|
||||
const response = await apiClient.get<QuotaUsageResponse>('/api/admin/stats/providers/quota-usage')
|
||||
return response.data
|
||||
},
|
||||
|
||||
async getPercentiles(params?: {
|
||||
start_date?: string
|
||||
end_date?: string
|
||||
preset?: string
|
||||
timezone?: string
|
||||
tz_offset_minutes?: number
|
||||
}): Promise<PercentileItem[]> {
|
||||
const response = await apiClient.get<PercentileItem[]>('/api/admin/stats/performance/percentiles', {
|
||||
params
|
||||
})
|
||||
return response.data
|
||||
},
|
||||
|
||||
async getErrorDistribution(params?: {
|
||||
start_date?: string
|
||||
end_date?: string
|
||||
preset?: string
|
||||
timezone?: string
|
||||
tz_offset_minutes?: number
|
||||
}): Promise<ErrorDistributionResponse> {
|
||||
const response = await apiClient.get<ErrorDistributionResponse>('/api/admin/stats/errors/distribution', {
|
||||
params
|
||||
})
|
||||
return response.data
|
||||
},
|
||||
|
||||
async getComparison(params: {
|
||||
current_start: string
|
||||
current_end: string
|
||||
comparison_type?: 'period' | 'year'
|
||||
timezone?: string
|
||||
tz_offset_minutes?: number
|
||||
}): Promise<ComparisonResponse> {
|
||||
const response = await apiClient.get<ComparisonResponse>('/api/admin/stats/comparison', {
|
||||
params
|
||||
})
|
||||
return response.data
|
||||
},
|
||||
|
||||
// 数据清空
|
||||
purgeConfig: () => purge<{ message: string; deleted: Record<string, number> }>('config'),
|
||||
purgeUsers: () => purge<{ message: string; deleted: Record<string, number> }>('users'),
|
||||
@@ -871,19 +665,4 @@ export const adminApi = {
|
||||
purgeRequestBodies: () => purge<{ message: string; cleaned: Record<string, number> }>('request-bodies'),
|
||||
purgeStats: () => purge<{ message: string }>('stats'),
|
||||
|
||||
async getTimeSeries(params?: {
|
||||
start_date?: string
|
||||
end_date?: string
|
||||
preset?: string
|
||||
granularity?: 'hour' | 'day' | 'week' | 'month'
|
||||
timezone?: string
|
||||
tz_offset_minutes?: number
|
||||
user_id?: string
|
||||
model?: string
|
||||
provider_name?: string
|
||||
}): Promise<Array<Record<string, unknown>>> {
|
||||
const response = await apiClient.get<Array<Record<string, unknown>>>('/api/admin/stats/time-series', { params })
|
||||
return response.data
|
||||
},
|
||||
|
||||
}
|
||||
|
||||
@@ -0,0 +1,396 @@
|
||||
import apiClient from './client'
|
||||
import { cachedRequest, buildCacheKey } from '@/utils/cache'
|
||||
import type { ActivityHeatmap } from '@/types/activity'
|
||||
|
||||
export type AnalyticsScopeKind = 'global' | 'me' | 'user' | 'api_key'
|
||||
export type AnalyticsGranularity = 'hour' | 'day' | 'week' | 'month'
|
||||
export type AnalyticsBreakdownDimension = 'model' | 'provider' | 'api_format' | 'api_key' | 'user'
|
||||
export type AnalyticsBreakdownMetric = 'requests_total' | 'total_tokens' | 'total_cost_usd' | 'actual_total_cost_usd'
|
||||
export type AnalyticsLeaderboardEntity = 'user' | 'api_key'
|
||||
export type AnalyticsLeaderboardMetric = 'requests_total' | 'total_tokens' | 'total_cost_usd'
|
||||
|
||||
export interface AnalyticsScope {
|
||||
kind: AnalyticsScopeKind
|
||||
user_id?: string | null
|
||||
api_key_id?: string | null
|
||||
}
|
||||
|
||||
export interface AnalyticsTimeRange {
|
||||
start_date?: string
|
||||
end_date?: string
|
||||
preset?: string
|
||||
granularity?: AnalyticsGranularity
|
||||
timezone?: string | null
|
||||
tz_offset_minutes?: number
|
||||
}
|
||||
|
||||
export interface AnalyticsFilters {
|
||||
user_ids?: string[]
|
||||
provider_names?: string[]
|
||||
models?: string[]
|
||||
target_models?: string[]
|
||||
api_key_ids?: string[]
|
||||
api_formats?: string[]
|
||||
request_types?: string[]
|
||||
statuses?: string[]
|
||||
error_categories?: string[]
|
||||
is_stream?: boolean | null
|
||||
has_format_conversion?: boolean | null
|
||||
}
|
||||
|
||||
export interface AnalyticsBaseRequest {
|
||||
scope: AnalyticsScope
|
||||
time_range: AnalyticsTimeRange
|
||||
filters?: AnalyticsFilters
|
||||
}
|
||||
|
||||
export interface AnalyticsSummary {
|
||||
requests_total: number
|
||||
requests_success: number
|
||||
requests_error: number
|
||||
requests_stream: number
|
||||
success_rate: number
|
||||
input_tokens: number
|
||||
output_tokens: number
|
||||
input_output_total_tokens: number
|
||||
cache_creation_input_tokens: number
|
||||
cache_creation_input_tokens_5m: number
|
||||
cache_creation_input_tokens_1h: number
|
||||
cache_read_input_tokens: number
|
||||
input_context_tokens: number
|
||||
total_tokens: number
|
||||
cache_hit_rate: number
|
||||
input_cost_usd: number
|
||||
output_cost_usd: number
|
||||
cache_creation_cost_usd: number
|
||||
cache_creation_cost_usd_5m: number
|
||||
cache_creation_cost_usd_1h: number
|
||||
cache_read_cost_usd: number
|
||||
cache_cost_usd: number
|
||||
request_cost_usd: number
|
||||
total_cost_usd: number
|
||||
actual_total_cost_usd: number
|
||||
actual_cache_cost_usd: number
|
||||
avg_response_time_ms: number
|
||||
avg_first_byte_time_ms: number
|
||||
format_conversion_count: number
|
||||
models_used_count: number
|
||||
}
|
||||
|
||||
export interface AnalyticsCompositionSegment {
|
||||
key: string
|
||||
value: number
|
||||
percentage: number
|
||||
}
|
||||
|
||||
export interface AnalyticsOverviewResponse {
|
||||
query_context: {
|
||||
scope: AnalyticsScope
|
||||
time_range: AnalyticsTimeRange
|
||||
}
|
||||
summary: AnalyticsSummary
|
||||
composition: {
|
||||
token_segments: AnalyticsCompositionSegment[]
|
||||
cost_segments: AnalyticsCompositionSegment[]
|
||||
}
|
||||
}
|
||||
|
||||
export interface AnalyticsTimeseriesBucket extends AnalyticsSummary {
|
||||
bucket_start: string
|
||||
bucket_end: string
|
||||
}
|
||||
|
||||
export interface AnalyticsTimeseriesResponse {
|
||||
buckets: AnalyticsTimeseriesBucket[]
|
||||
}
|
||||
|
||||
export interface AnalyticsBreakdownRequest extends AnalyticsBaseRequest {
|
||||
dimension: AnalyticsBreakdownDimension
|
||||
metric?: AnalyticsBreakdownMetric
|
||||
limit?: number
|
||||
}
|
||||
|
||||
export interface AnalyticsBreakdownRow extends AnalyticsSummary {
|
||||
key: string
|
||||
label: string
|
||||
share_of_total_cost: number
|
||||
share_of_total_tokens: number
|
||||
}
|
||||
|
||||
export interface AnalyticsBreakdownResponse {
|
||||
dimension: AnalyticsBreakdownDimension
|
||||
metric: AnalyticsBreakdownMetric
|
||||
rows: AnalyticsBreakdownRow[]
|
||||
}
|
||||
|
||||
export interface AnalyticsRecordsRequest extends AnalyticsBaseRequest {
|
||||
search?: {
|
||||
text?: string | null
|
||||
request_id?: string | null
|
||||
}
|
||||
pagination?: {
|
||||
limit?: number
|
||||
offset?: number
|
||||
}
|
||||
}
|
||||
|
||||
export interface AnalyticsRecord {
|
||||
id: string
|
||||
request_id: string
|
||||
created_at: string | null
|
||||
user_id: string | null
|
||||
username: string | null
|
||||
api_key_id: string | null
|
||||
api_key_name: string | null
|
||||
provider_api_key_name: string | null
|
||||
provider_name: string | null
|
||||
model: string
|
||||
target_model: string | null
|
||||
api_format: string | null
|
||||
request_type: string | null
|
||||
status: string
|
||||
billing_status: string
|
||||
is_stream: boolean
|
||||
has_format_conversion: boolean | null
|
||||
has_fallback?: boolean
|
||||
has_retry?: boolean
|
||||
status_code: number | null
|
||||
error_message: string | null
|
||||
error_category: string | null
|
||||
response_time_ms: number | null
|
||||
first_byte_time_ms: number | null
|
||||
input_tokens: number
|
||||
output_tokens: number
|
||||
input_output_total_tokens: number
|
||||
cache_creation_input_tokens: number
|
||||
cache_creation_input_tokens_5m: number
|
||||
cache_creation_input_tokens_1h: number
|
||||
cache_read_input_tokens: number
|
||||
input_context_tokens: number
|
||||
total_tokens: number
|
||||
input_cost_usd: number
|
||||
output_cost_usd: number
|
||||
cache_creation_cost_usd: number
|
||||
cache_creation_cost_usd_5m: number
|
||||
cache_creation_cost_usd_1h: number
|
||||
cache_read_cost_usd: number
|
||||
cache_cost_usd: number
|
||||
request_cost_usd: number
|
||||
total_cost_usd: number
|
||||
actual_total_cost_usd: number
|
||||
actual_cache_cost_usd: number
|
||||
rate_multiplier: number
|
||||
}
|
||||
|
||||
export interface AnalyticsRecordsResponse {
|
||||
total: number
|
||||
limit: number
|
||||
offset: number
|
||||
records: AnalyticsRecord[]
|
||||
}
|
||||
|
||||
export interface AnalyticsFilterOption {
|
||||
value: string
|
||||
label: string
|
||||
}
|
||||
|
||||
export interface AnalyticsFilterOptionsResponse {
|
||||
providers: AnalyticsFilterOption[]
|
||||
models: AnalyticsFilterOption[]
|
||||
target_models: AnalyticsFilterOption[]
|
||||
api_formats: AnalyticsFilterOption[]
|
||||
request_types: AnalyticsFilterOption[]
|
||||
error_categories: AnalyticsFilterOption[]
|
||||
statuses: AnalyticsFilterOption[]
|
||||
users?: AnalyticsFilterOption[]
|
||||
api_keys?: AnalyticsFilterOption[]
|
||||
}
|
||||
|
||||
export interface AnalyticsLeaderboardRequest extends AnalyticsBaseRequest {
|
||||
entity: AnalyticsLeaderboardEntity
|
||||
metric: AnalyticsLeaderboardMetric
|
||||
limit?: number
|
||||
}
|
||||
|
||||
export interface AnalyticsLeaderboardItem {
|
||||
rank: number
|
||||
id: string
|
||||
label: string
|
||||
requests_total: number
|
||||
total_tokens: number
|
||||
total_cost_usd: number
|
||||
actual_total_cost_usd: number
|
||||
metric_value: number
|
||||
}
|
||||
|
||||
export interface AnalyticsLeaderboardResponse {
|
||||
entity: AnalyticsLeaderboardEntity
|
||||
metric: AnalyticsLeaderboardMetric
|
||||
items: AnalyticsLeaderboardItem[]
|
||||
}
|
||||
|
||||
export interface AnalyticsPercentilePoint {
|
||||
date: string
|
||||
p50_response_time_ms: number | null
|
||||
p90_response_time_ms: number | null
|
||||
p99_response_time_ms: number | null
|
||||
p50_first_byte_time_ms: number | null
|
||||
p90_first_byte_time_ms: number | null
|
||||
p99_first_byte_time_ms: number | null
|
||||
}
|
||||
|
||||
export interface AnalyticsErrorCategory {
|
||||
category: string
|
||||
label: string
|
||||
count: number
|
||||
}
|
||||
|
||||
export interface AnalyticsErrorTrendItem {
|
||||
date: string
|
||||
total: number
|
||||
}
|
||||
|
||||
export interface AnalyticsProviderHealthItem {
|
||||
provider_name: string
|
||||
requests_total: number
|
||||
success_rate: number
|
||||
error_rate: number
|
||||
avg_response_time_ms: number
|
||||
avg_first_byte_time_ms: number
|
||||
}
|
||||
|
||||
export interface AnalyticsPerformanceResponse {
|
||||
latency: {
|
||||
response_time_ms: {
|
||||
avg: number
|
||||
p50: number | null
|
||||
p90: number | null
|
||||
p99: number | null
|
||||
}
|
||||
first_byte_time_ms: {
|
||||
avg: number
|
||||
p50: number | null
|
||||
p90: number | null
|
||||
p99: number | null
|
||||
}
|
||||
}
|
||||
percentiles: AnalyticsPercentilePoint[]
|
||||
errors: {
|
||||
total: number
|
||||
rate: number
|
||||
categories: AnalyticsErrorCategory[]
|
||||
trend: AnalyticsErrorTrendItem[]
|
||||
}
|
||||
provider_health: AnalyticsProviderHealthItem[]
|
||||
}
|
||||
|
||||
export interface AnalyticsActiveRequest {
|
||||
id: string
|
||||
status: 'pending' | 'streaming' | 'completed' | 'failed' | 'cancelled'
|
||||
input_tokens: number
|
||||
output_tokens: number
|
||||
cache_creation_input_tokens?: number | null
|
||||
cache_read_input_tokens?: number | null
|
||||
cost?: number
|
||||
total_cost_usd?: number
|
||||
actual_cost?: number | null
|
||||
actual_total_cost_usd?: number | null
|
||||
rate_multiplier?: number | null
|
||||
response_time_ms: number | null
|
||||
first_byte_time_ms: number | null
|
||||
provider?: string | null
|
||||
provider_name?: string | null
|
||||
api_key_name?: string | null
|
||||
api_format?: string | null
|
||||
endpoint_api_format?: string | null
|
||||
has_format_conversion?: boolean | null
|
||||
target_model?: string | null
|
||||
}
|
||||
|
||||
export interface AnalyticsActiveRequestsResponse {
|
||||
requests: AnalyticsActiveRequest[]
|
||||
}
|
||||
|
||||
export const analyticsApi = {
|
||||
async getOverview(payload: AnalyticsBaseRequest): Promise<AnalyticsOverviewResponse> {
|
||||
const response = await apiClient.post<AnalyticsOverviewResponse>('/api/analytics/overview', payload)
|
||||
return response.data
|
||||
},
|
||||
|
||||
async getTimeseries(payload: AnalyticsBaseRequest): Promise<AnalyticsTimeseriesResponse> {
|
||||
const response = await apiClient.post<AnalyticsTimeseriesResponse>('/api/analytics/timeseries', payload)
|
||||
return response.data
|
||||
},
|
||||
|
||||
async getBreakdown(payload: AnalyticsBreakdownRequest): Promise<AnalyticsBreakdownResponse> {
|
||||
const response = await apiClient.post<AnalyticsBreakdownResponse>('/api/analytics/breakdown', payload)
|
||||
return response.data
|
||||
},
|
||||
|
||||
async getRecords(payload: AnalyticsRecordsRequest): Promise<AnalyticsRecordsResponse> {
|
||||
const response = await apiClient.post<AnalyticsRecordsResponse>('/api/analytics/records', payload)
|
||||
return response.data
|
||||
},
|
||||
|
||||
async getFilterOptions(payload: AnalyticsBaseRequest): Promise<AnalyticsFilterOptionsResponse> {
|
||||
const response = await apiClient.post<AnalyticsFilterOptionsResponse>('/api/analytics/filter-options', payload)
|
||||
return response.data
|
||||
},
|
||||
|
||||
async getLeaderboard(payload: AnalyticsLeaderboardRequest): Promise<AnalyticsLeaderboardResponse> {
|
||||
const response = await apiClient.post<AnalyticsLeaderboardResponse>('/api/analytics/leaderboard', payload)
|
||||
return response.data
|
||||
},
|
||||
|
||||
async getPerformance(payload: AnalyticsBaseRequest): Promise<AnalyticsPerformanceResponse> {
|
||||
const response = await apiClient.post<AnalyticsPerformanceResponse>('/api/analytics/performance', payload)
|
||||
return response.data
|
||||
},
|
||||
|
||||
async getHeatmap(payload: {
|
||||
scope: AnalyticsScope
|
||||
user_id?: string | null
|
||||
api_key_id?: string | null
|
||||
}): Promise<ActivityHeatmap> {
|
||||
const cacheKey = buildCacheKey('analytics:heatmap', payload as Record<string, unknown>)
|
||||
return cachedRequest(
|
||||
cacheKey,
|
||||
async () => {
|
||||
const response = await apiClient.post<ActivityHeatmap>('/api/analytics/heatmap', payload)
|
||||
return response.data
|
||||
},
|
||||
60000,
|
||||
)
|
||||
},
|
||||
|
||||
async getActiveRequests(payload: {
|
||||
scope: AnalyticsScope
|
||||
ids?: string[]
|
||||
}): Promise<AnalyticsActiveRequestsResponse> {
|
||||
const response = await apiClient.post<AnalyticsActiveRequestsResponse>('/api/analytics/active-requests', {
|
||||
...payload,
|
||||
ids: payload.ids ?? [],
|
||||
})
|
||||
return response.data
|
||||
},
|
||||
|
||||
async analyzeCacheAffinityTTL(payload: {
|
||||
scope: AnalyticsScope
|
||||
user_id?: string | null
|
||||
api_key_id?: string | null
|
||||
hours?: number
|
||||
}): Promise<Record<string, unknown>> {
|
||||
const response = await apiClient.post('/api/analytics/cache-affinity/ttl-analysis', payload)
|
||||
return response.data
|
||||
},
|
||||
|
||||
async analyzeCacheAffinityHit(payload: {
|
||||
scope: AnalyticsScope
|
||||
user_id?: string | null
|
||||
api_key_id?: string | null
|
||||
hours?: number
|
||||
}): Promise<Record<string, unknown>> {
|
||||
const response = await apiClient.post('/api/analytics/cache-affinity/hit-analysis', payload)
|
||||
return response.data
|
||||
},
|
||||
}
|
||||
@@ -33,6 +33,203 @@ export interface AuditFilters {
|
||||
offset?: number
|
||||
}
|
||||
|
||||
export type MonitoringMetricStatus = 'ok' | 'warning' | 'danger' | 'degraded' | 'error' | 'unknown'
|
||||
|
||||
export interface MonitoringCpuMetric {
|
||||
status: MonitoringMetricStatus
|
||||
label: string
|
||||
usage_percent: number | null
|
||||
load_percent: number | null
|
||||
core_count: number
|
||||
message?: string | null
|
||||
}
|
||||
|
||||
export interface MonitoringMemoryMetric {
|
||||
status: MonitoringMetricStatus
|
||||
label: string
|
||||
used_percent: number | null
|
||||
used_bytes: number | null
|
||||
available_bytes: number | null
|
||||
total_bytes: number | null
|
||||
message?: string | null
|
||||
}
|
||||
|
||||
export interface MonitoringServiceMetric {
|
||||
status: MonitoringMetricStatus
|
||||
label: string
|
||||
latency_ms: number | null
|
||||
memory_status?: MonitoringMetricStatus
|
||||
memory_label?: string | null
|
||||
used_memory_bytes?: number | null
|
||||
peak_memory_bytes?: number | null
|
||||
maxmemory_bytes?: number | null
|
||||
memory_ceiling_bytes?: number | null
|
||||
memory_source?: 'configured' | 'maxmemory' | 'system' | 'unknown'
|
||||
available_memory_bytes?: number | null
|
||||
memory_percent?: number | null
|
||||
message?: string | null
|
||||
}
|
||||
|
||||
export interface MonitoringPostgresMetric {
|
||||
status: MonitoringMetricStatus
|
||||
label: string
|
||||
usage_percent: number | null
|
||||
pool_usage_percent?: number | null
|
||||
checked_out: number
|
||||
pool_size: number
|
||||
overflow: number
|
||||
max_capacity: number
|
||||
pool_timeout: number
|
||||
server_connections?: number | null
|
||||
server_max_connections?: number | null
|
||||
server_usage_percent?: number | null
|
||||
storage_status?: MonitoringMetricStatus
|
||||
storage_label?: string | null
|
||||
storage_total_bytes?: number | null
|
||||
storage_free_bytes?: number | null
|
||||
storage_free_percent?: number | null
|
||||
database_size_bytes?: number | null
|
||||
storage_message?: string | null
|
||||
message?: string | null
|
||||
}
|
||||
|
||||
export interface MonitoringSystemStatus {
|
||||
timestamp: string
|
||||
users: {
|
||||
total: number
|
||||
active: number
|
||||
}
|
||||
providers: {
|
||||
total: number
|
||||
active: number
|
||||
}
|
||||
api_keys: {
|
||||
total: number
|
||||
active: number
|
||||
}
|
||||
today_stats: {
|
||||
requests: number
|
||||
tokens: number
|
||||
cost_usd: number
|
||||
}
|
||||
recent_errors: number
|
||||
system_metrics?: {
|
||||
cpu: MonitoringCpuMetric
|
||||
memory: MonitoringMemoryMetric
|
||||
redis: MonitoringServiceMetric
|
||||
postgres: MonitoringPostgresMetric
|
||||
}
|
||||
}
|
||||
|
||||
function parseNumericLike(value: unknown): number {
|
||||
if (typeof value === 'number') {
|
||||
return Number.isFinite(value) ? value : 0
|
||||
}
|
||||
if (typeof value === 'string') {
|
||||
const normalized = value.trim().replace(/[$,\s]/g, '')
|
||||
if (!normalized) return 0
|
||||
const parsed = Number(normalized)
|
||||
return Number.isFinite(parsed) ? parsed : 0
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
function parseNullableNumericLike(value: unknown): number | null {
|
||||
if (value == null) {
|
||||
return null
|
||||
}
|
||||
if (typeof value === 'number') {
|
||||
return Number.isFinite(value) ? value : null
|
||||
}
|
||||
if (typeof value === 'string') {
|
||||
const normalized = value.trim().replace(/[$,\s]/g, '')
|
||||
if (!normalized) return null
|
||||
const parsed = Number(normalized)
|
||||
return Number.isFinite(parsed) ? parsed : null
|
||||
}
|
||||
return null
|
||||
}
|
||||
|
||||
function normalizeMonitoringSystemStatus(data: MonitoringSystemStatus): MonitoringSystemStatus {
|
||||
return {
|
||||
...data,
|
||||
users: {
|
||||
total: parseNumericLike(data.users?.total),
|
||||
active: parseNumericLike(data.users?.active),
|
||||
},
|
||||
providers: {
|
||||
total: parseNumericLike(data.providers?.total),
|
||||
active: parseNumericLike(data.providers?.active),
|
||||
},
|
||||
api_keys: {
|
||||
total: parseNumericLike(data.api_keys?.total),
|
||||
active: parseNumericLike(data.api_keys?.active),
|
||||
},
|
||||
today_stats: {
|
||||
requests: parseNumericLike(data.today_stats?.requests),
|
||||
tokens: parseNumericLike(data.today_stats?.tokens),
|
||||
cost_usd: parseNumericLike(data.today_stats?.cost_usd),
|
||||
},
|
||||
recent_errors: parseNumericLike(data.recent_errors),
|
||||
system_metrics: data.system_metrics ? {
|
||||
cpu: {
|
||||
status: data.system_metrics.cpu?.status ?? 'unknown',
|
||||
label: data.system_metrics.cpu?.label ?? '未知',
|
||||
usage_percent: parseNullableNumericLike(data.system_metrics.cpu?.usage_percent),
|
||||
load_percent: parseNullableNumericLike(data.system_metrics.cpu?.load_percent),
|
||||
core_count: parseNumericLike(data.system_metrics.cpu?.core_count),
|
||||
message: data.system_metrics.cpu?.message ?? null,
|
||||
},
|
||||
memory: {
|
||||
status: data.system_metrics.memory?.status ?? 'unknown',
|
||||
label: data.system_metrics.memory?.label ?? '未知',
|
||||
used_percent: parseNullableNumericLike(data.system_metrics.memory?.used_percent),
|
||||
used_bytes: parseNullableNumericLike(data.system_metrics.memory?.used_bytes),
|
||||
available_bytes: parseNullableNumericLike(data.system_metrics.memory?.available_bytes),
|
||||
total_bytes: parseNullableNumericLike(data.system_metrics.memory?.total_bytes),
|
||||
message: data.system_metrics.memory?.message ?? null,
|
||||
},
|
||||
redis: {
|
||||
status: data.system_metrics.redis?.status ?? 'unknown',
|
||||
label: data.system_metrics.redis?.label ?? '未知',
|
||||
latency_ms: parseNullableNumericLike(data.system_metrics.redis?.latency_ms),
|
||||
memory_status: data.system_metrics.redis?.memory_status ?? 'unknown',
|
||||
memory_label: data.system_metrics.redis?.memory_label ?? '未知',
|
||||
used_memory_bytes: parseNullableNumericLike(data.system_metrics.redis?.used_memory_bytes),
|
||||
peak_memory_bytes: parseNullableNumericLike(data.system_metrics.redis?.peak_memory_bytes),
|
||||
maxmemory_bytes: parseNullableNumericLike(data.system_metrics.redis?.maxmemory_bytes),
|
||||
memory_ceiling_bytes: parseNullableNumericLike(data.system_metrics.redis?.memory_ceiling_bytes),
|
||||
memory_source: (data.system_metrics.redis?.memory_source as MonitoringServiceMetric['memory_source']) ?? 'unknown',
|
||||
available_memory_bytes: parseNullableNumericLike(data.system_metrics.redis?.available_memory_bytes),
|
||||
memory_percent: parseNullableNumericLike(data.system_metrics.redis?.memory_percent),
|
||||
message: data.system_metrics.redis?.message ?? null,
|
||||
},
|
||||
postgres: {
|
||||
status: data.system_metrics.postgres?.status ?? 'unknown',
|
||||
label: data.system_metrics.postgres?.label ?? '未知',
|
||||
usage_percent: parseNullableNumericLike(data.system_metrics.postgres?.usage_percent),
|
||||
pool_usage_percent: parseNullableNumericLike(data.system_metrics.postgres?.pool_usage_percent),
|
||||
checked_out: parseNumericLike(data.system_metrics.postgres?.checked_out),
|
||||
pool_size: parseNumericLike(data.system_metrics.postgres?.pool_size),
|
||||
overflow: parseNumericLike(data.system_metrics.postgres?.overflow),
|
||||
max_capacity: parseNumericLike(data.system_metrics.postgres?.max_capacity),
|
||||
pool_timeout: parseNumericLike(data.system_metrics.postgres?.pool_timeout),
|
||||
server_connections: parseNullableNumericLike(data.system_metrics.postgres?.server_connections),
|
||||
server_max_connections: parseNullableNumericLike(data.system_metrics.postgres?.server_max_connections),
|
||||
server_usage_percent: parseNullableNumericLike(data.system_metrics.postgres?.server_usage_percent),
|
||||
storage_status: data.system_metrics.postgres?.storage_status ?? 'unknown',
|
||||
storage_label: data.system_metrics.postgres?.storage_label ?? '未知',
|
||||
storage_total_bytes: parseNullableNumericLike(data.system_metrics.postgres?.storage_total_bytes),
|
||||
storage_free_bytes: parseNullableNumericLike(data.system_metrics.postgres?.storage_free_bytes),
|
||||
storage_free_percent: parseNullableNumericLike(data.system_metrics.postgres?.storage_free_percent),
|
||||
database_size_bytes: parseNullableNumericLike(data.system_metrics.postgres?.database_size_bytes),
|
||||
storage_message: data.system_metrics.postgres?.storage_message ?? null,
|
||||
message: data.system_metrics.postgres?.message ?? null,
|
||||
},
|
||||
} : undefined,
|
||||
}
|
||||
}
|
||||
|
||||
function normalizeAuditResponse(data: Record<string, unknown>): AuditLogsResponse {
|
||||
const items: AuditLog[] = (data.items ?? data.logs ?? []) as AuditLog[]
|
||||
const meta: PaginationMeta = (data.meta as PaginationMeta) ?? {
|
||||
@@ -78,6 +275,11 @@ export const auditApi = {
|
||||
return response.data
|
||||
},
|
||||
|
||||
async getSystemStatus(): Promise<MonitoringSystemStatus> {
|
||||
const response = await apiClient.get<MonitoringSystemStatus>('/api/admin/monitoring/system-status')
|
||||
return normalizeMonitoringSystemStatus(response.data)
|
||||
},
|
||||
|
||||
// 分析用户行为 (管理员)
|
||||
async analyzeUserBehavior(userId: number, days: number = 7): Promise<{
|
||||
analysis: Record<string, unknown>
|
||||
|
||||
+21
-26
@@ -9,7 +9,6 @@ export interface LoginRequest {
|
||||
|
||||
export interface LoginResponse {
|
||||
access_token: string
|
||||
refresh_token?: string
|
||||
token_type?: string
|
||||
expires_in?: number
|
||||
user_id?: string // UUID
|
||||
@@ -25,13 +24,6 @@ export interface UserPreferences {
|
||||
[key: string]: unknown // 允许扩展其他偏好设置
|
||||
}
|
||||
|
||||
export interface UserStats {
|
||||
total_requests?: number
|
||||
total_cost?: number
|
||||
last_request_at?: string
|
||||
[key: string]: unknown // 允许扩展其他统计数据
|
||||
}
|
||||
|
||||
export interface SendVerificationCodeRequest {
|
||||
email: string
|
||||
}
|
||||
@@ -81,6 +73,7 @@ export interface RegistrationSettingsResponse {
|
||||
enable_registration: boolean
|
||||
require_email_verification: boolean
|
||||
email_configured: boolean
|
||||
password_policy_level: string
|
||||
}
|
||||
|
||||
export interface AuthSettingsResponse {
|
||||
@@ -89,32 +82,39 @@ export interface AuthSettingsResponse {
|
||||
ldap_exclusive: boolean
|
||||
}
|
||||
|
||||
export interface BillingSummary {
|
||||
id?: string | null
|
||||
balance: number
|
||||
recharge_balance: number
|
||||
gift_balance: number
|
||||
refundable_balance: number
|
||||
currency: string
|
||||
status: string
|
||||
limit_mode: 'finite' | 'unlimited'
|
||||
unlimited: boolean
|
||||
total_recharged: number
|
||||
total_consumed: number
|
||||
total_refunded: number
|
||||
total_adjusted: number
|
||||
updated_at?: string | null
|
||||
}
|
||||
|
||||
export interface User {
|
||||
id: string // UUID
|
||||
username: string
|
||||
email?: string
|
||||
role: string // 'admin' or 'user'
|
||||
is_active: boolean
|
||||
quota_usd?: number | null
|
||||
used_usd?: number
|
||||
total_usd?: number
|
||||
allowed_providers?: string[] | null // 允许使用的提供商 ID 列表
|
||||
allowed_api_formats?: string[] | null // 允许使用的 API 格式列表
|
||||
allowed_models?: string[] | null // 允许使用的模型名称列表
|
||||
billing?: BillingSummary
|
||||
created_at: string
|
||||
last_login_at?: string
|
||||
preferences?: UserPreferences
|
||||
stats?: UserStats
|
||||
}
|
||||
|
||||
export const authApi = {
|
||||
async login(credentials: LoginRequest): Promise<LoginResponse> {
|
||||
const response = await apiClient.post<LoginResponse>('/api/auth/login', credentials)
|
||||
apiClient.setToken(response.data.access_token)
|
||||
// 后端暂时没有返回 refresh_token
|
||||
if (response.data.refresh_token) {
|
||||
localStorage.setItem('refresh_token', response.data.refresh_token)
|
||||
}
|
||||
return response.data
|
||||
},
|
||||
|
||||
@@ -136,14 +136,9 @@ export const authApi = {
|
||||
return response.data
|
||||
},
|
||||
|
||||
async refreshToken(refreshToken: string): Promise<LoginResponse> {
|
||||
const response = await apiClient.post<LoginResponse>('/api/auth/refresh', {
|
||||
refresh_token: refreshToken
|
||||
})
|
||||
async refreshToken(): Promise<LoginResponse> {
|
||||
const response = await apiClient.post<LoginResponse>('/api/auth/refresh', {})
|
||||
apiClient.setToken(response.data.access_token)
|
||||
if (response.data.refresh_token) {
|
||||
localStorage.setItem('refresh_token', response.data.refresh_token)
|
||||
}
|
||||
return response.data
|
||||
},
|
||||
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
*/
|
||||
|
||||
import api from './client'
|
||||
import { cachedRequest, buildCacheKey } from '@/utils/cache'
|
||||
|
||||
export interface CacheStats {
|
||||
scheduler: string
|
||||
@@ -287,7 +288,12 @@ export const cacheAnalysisApi = {
|
||||
api_key_id?: string
|
||||
hours?: number
|
||||
}): Promise<TTLAnalysisResponse> {
|
||||
const response = await api.get('/api/admin/usage/cache-affinity/ttl-analysis', { params })
|
||||
const response = await api.post('/api/analytics/cache-affinity/ttl-analysis', {
|
||||
scope: { kind: 'global' },
|
||||
user_id: params?.user_id ?? null,
|
||||
api_key_id: params?.api_key_id ?? null,
|
||||
hours: params?.hours,
|
||||
})
|
||||
return response.data
|
||||
},
|
||||
|
||||
@@ -299,7 +305,12 @@ export const cacheAnalysisApi = {
|
||||
api_key_id?: string
|
||||
hours?: number
|
||||
}): Promise<CacheHitAnalysisResponse> {
|
||||
const response = await api.get('/api/admin/usage/cache-affinity/hit-analysis', { params })
|
||||
const response = await api.post('/api/analytics/cache-affinity/hit-analysis', {
|
||||
scope: { kind: 'global' },
|
||||
user_id: params?.user_id ?? null,
|
||||
api_key_id: params?.api_key_id ?? null,
|
||||
hours: params?.hours,
|
||||
})
|
||||
return response.data
|
||||
},
|
||||
|
||||
@@ -314,8 +325,21 @@ export const cacheAnalysisApi = {
|
||||
user_id?: string
|
||||
include_user_info?: boolean
|
||||
}): Promise<IntervalTimelineResponse> {
|
||||
const response = await api.get('/api/admin/usage/cache-affinity/interval-timeline', { params })
|
||||
const cacheKey = buildCacheKey('cache-affinity:interval-timeline', params as Record<string, unknown> | undefined)
|
||||
return cachedRequest(
|
||||
cacheKey,
|
||||
async () => {
|
||||
const response = await api.post('/api/analytics/interval-timeline', {
|
||||
scope: params?.user_id ? { kind: 'user', user_id: params.user_id } : { kind: 'global' },
|
||||
user_id: params?.user_id ?? null,
|
||||
hours: params?.hours,
|
||||
limit: params?.limit,
|
||||
include_user_info: params?.include_user_info ?? false,
|
||||
})
|
||||
return response.data
|
||||
},
|
||||
30000
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+87
-48
@@ -3,10 +3,13 @@ import type { AxiosInstance, AxiosRequestConfig, AxiosResponse, InternalAxiosReq
|
||||
import { NETWORK_CONFIG, AUTH_CONFIG } from '@/config/constants'
|
||||
import { isDemoMode } from '@/config/demo'
|
||||
import { handleMockRequest, setMockUserToken } from '@/mocks'
|
||||
import { getClientDeviceId } from '@/utils/deviceId'
|
||||
import { CrossTabRefreshCoordinator } from '@/utils/crossTabRefresh'
|
||||
import { log } from '@/utils/logger'
|
||||
|
||||
// 在开发环境下使用代理,生产环境使用环境变量
|
||||
const API_BASE_URL = import.meta.env.VITE_API_URL || ''
|
||||
export const AUTH_STATE_CHANGE_EVENT = 'aether-auth-state-change'
|
||||
|
||||
/**
|
||||
* 判断请求是否为公共端点
|
||||
@@ -31,17 +34,15 @@ function isAuthRequest(url?: string): boolean {
|
||||
}
|
||||
|
||||
/**
|
||||
* 判断是否为可刷新的认证错误
|
||||
* 判断 403 错误是否表示用户账号级别的问题(需要清除认证并跳转)
|
||||
*/
|
||||
function isRefreshableAuthError(errorDetail: string): boolean {
|
||||
const nonRefreshableErrors = [
|
||||
function isAccountLevelForbidden(status: number, errorDetail: string): boolean {
|
||||
if (status !== 403) return false
|
||||
const accountErrors = [
|
||||
'用户不存在或已禁用',
|
||||
'需要管理员权限',
|
||||
'权限不足',
|
||||
'用户已禁用',
|
||||
]
|
||||
|
||||
return !nonRefreshableErrors.some((msg) => errorDetail.includes(msg))
|
||||
return accountErrors.some((msg) => errorDetail.includes(msg))
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -83,12 +84,21 @@ class ApiClient {
|
||||
private client: AxiosInstance
|
||||
private token: string | null = null
|
||||
private isRefreshing = false
|
||||
private refreshPromise: Promise<AxiosResponse> | null = null
|
||||
private refreshPromise: Promise<string> | null = null
|
||||
private readonly refreshCoordinator = new CrossTabRefreshCoordinator()
|
||||
|
||||
private readonly onStorageSync = (event: StorageEvent): void => {
|
||||
if (event.key !== 'access_token') {
|
||||
return
|
||||
}
|
||||
this.syncTokenState(event.newValue)
|
||||
}
|
||||
|
||||
constructor() {
|
||||
this.client = axios.create({
|
||||
baseURL: API_BASE_URL,
|
||||
timeout: NETWORK_CONFIG.API_TIMEOUT,
|
||||
withCredentials: true,
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
},
|
||||
@@ -99,6 +109,7 @@ class ApiClient {
|
||||
this.client.defaults.adapter = createDemoAdapter(defaultAdapter)
|
||||
|
||||
this.setupInterceptors()
|
||||
this.setupCrossTabAuthSync()
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -108,6 +119,10 @@ class ApiClient {
|
||||
// 请求拦截器 - 仅处理认证
|
||||
this.client.interceptors.request.use(
|
||||
(config) => {
|
||||
if (config.url?.includes('/api/')) {
|
||||
config.headers['X-Client-Device-Id'] = getClientDeviceId()
|
||||
}
|
||||
|
||||
const requiresAuth = !isPublicEndpoint(config.url, config.method) &&
|
||||
config.url?.includes('/api/')
|
||||
|
||||
@@ -129,6 +144,23 @@ class ApiClient {
|
||||
)
|
||||
}
|
||||
|
||||
private setupCrossTabAuthSync(): void {
|
||||
if (typeof window !== 'undefined') {
|
||||
window.addEventListener('storage', this.onStorageSync)
|
||||
}
|
||||
}
|
||||
|
||||
private emitAuthStateChange(token: string | null): void {
|
||||
if (typeof window === 'undefined') {
|
||||
return
|
||||
}
|
||||
window.dispatchEvent(
|
||||
new CustomEvent<{ token: string | null }>(AUTH_STATE_CHANGE_EVENT, {
|
||||
detail: { token },
|
||||
})
|
||||
)
|
||||
}
|
||||
|
||||
/**
|
||||
* 处理响应错误
|
||||
*/
|
||||
@@ -155,8 +187,22 @@ class ApiClient {
|
||||
return Promise.reject(error)
|
||||
}
|
||||
|
||||
const status = error.response?.status ?? 0
|
||||
|
||||
// 处理 403 用户账号级别错误(被禁用/删除)
|
||||
if (status === 403) {
|
||||
const rawDetail = (error.response?.data as Record<string, unknown>)?.detail
|
||||
const errorDetail = typeof rawDetail === 'string' ? rawDetail : ''
|
||||
if (isAccountLevelForbidden(status, errorDetail)) {
|
||||
log.info('User account issue detected, clearing auth', { errorDetail })
|
||||
this.clearAuth()
|
||||
window.location.href = '/'
|
||||
return Promise.reject(error)
|
||||
}
|
||||
}
|
||||
|
||||
// 处理401错误
|
||||
if (error.response?.status === 401) {
|
||||
if (status === 401) {
|
||||
return this.handle401Error(error, originalRequest)
|
||||
}
|
||||
|
||||
@@ -177,25 +223,7 @@ class ApiClient {
|
||||
return Promise.reject(error)
|
||||
}
|
||||
|
||||
const errorDetail = (error.response?.data as Record<string, unknown>)?.detail as string || ''
|
||||
log.debug('Got 401 error, attempting token refresh', { errorDetail })
|
||||
|
||||
// 检查是否为业务相关的401错误(用户被禁用/删除等)
|
||||
if (!isRefreshableAuthError(errorDetail)) {
|
||||
log.info('User account issue detected, logging out and redirecting to home', { errorDetail })
|
||||
this.clearAuth()
|
||||
// 跳转到首页
|
||||
window.location.href = '/'
|
||||
return Promise.reject(error)
|
||||
}
|
||||
|
||||
// 获取refresh token
|
||||
const refreshToken = localStorage.getItem('refresh_token')
|
||||
if (!refreshToken) {
|
||||
log.info('No refresh token available, clearing invalid token')
|
||||
this.clearAuth()
|
||||
return Promise.reject(error)
|
||||
}
|
||||
log.debug('Got 401 error, attempting token refresh')
|
||||
|
||||
// 标记为已重试
|
||||
originalRequest._retry = true
|
||||
@@ -210,8 +238,8 @@ class ApiClient {
|
||||
// 如果正在刷新,等待刷新完成
|
||||
if (this.isRefreshing) {
|
||||
try {
|
||||
await this.refreshPromise
|
||||
originalRequest.headers.Authorization = `Bearer ${this.getToken()}`
|
||||
const accessToken = await this.refreshPromise
|
||||
originalRequest.headers.Authorization = `Bearer ${accessToken}`
|
||||
return this.client.request(originalRequest)
|
||||
} catch {
|
||||
return Promise.reject(error)
|
||||
@@ -219,29 +247,27 @@ class ApiClient {
|
||||
}
|
||||
|
||||
// 开始刷新token
|
||||
return this.refreshTokenAndRetry(refreshToken, originalRequest, error)
|
||||
return this.refreshTokenAndRetry(originalRequest, error)
|
||||
}
|
||||
|
||||
/**
|
||||
* 刷新token并重试原始请求
|
||||
*/
|
||||
private async refreshTokenAndRetry(
|
||||
refreshToken: string,
|
||||
originalRequest: InternalAxiosRequestConfig,
|
||||
originalError: import('axios').AxiosError
|
||||
): Promise<AxiosResponse> {
|
||||
this.isRefreshing = true
|
||||
this.refreshPromise = this.refreshToken(refreshToken)
|
||||
this.refreshPromise = this.coordinatedRefresh()
|
||||
|
||||
try {
|
||||
const response = await this.refreshPromise
|
||||
this.setToken(response.data.access_token)
|
||||
localStorage.setItem('refresh_token', response.data.refresh_token)
|
||||
const accessToken = await this.refreshPromise
|
||||
this.setToken(accessToken)
|
||||
this.isRefreshing = false
|
||||
this.refreshPromise = null
|
||||
|
||||
// 重试原始请求
|
||||
originalRequest.headers.Authorization = `Bearer ${response.data.access_token}`
|
||||
originalRequest.headers.Authorization = `Bearer ${accessToken}`
|
||||
return this.client.request(originalRequest)
|
||||
} catch (refreshError: unknown) {
|
||||
log.error('Token refresh failed', refreshError instanceof Error ? refreshError.message : String(refreshError))
|
||||
@@ -252,15 +278,29 @@ class ApiClient {
|
||||
}
|
||||
}
|
||||
|
||||
setToken(token: string): void {
|
||||
private async coordinatedRefresh(): Promise<string> {
|
||||
return this.refreshCoordinator.run(async () => {
|
||||
const response = await this.refreshToken()
|
||||
const accessToken = response.data.access_token
|
||||
if (!accessToken) {
|
||||
throw new Error('Refresh response missing access token')
|
||||
}
|
||||
return accessToken
|
||||
})
|
||||
}
|
||||
|
||||
private syncTokenState(token: string | null): void {
|
||||
this.token = token
|
||||
localStorage.setItem('access_token', token)
|
||||
// 同步到 mock handler
|
||||
if (isDemoMode()) {
|
||||
setMockUserToken(token)
|
||||
}
|
||||
}
|
||||
|
||||
setToken(token: string): void {
|
||||
this.syncTokenState(token)
|
||||
localStorage.setItem('access_token', token)
|
||||
}
|
||||
|
||||
getToken(): string | null {
|
||||
if (!this.token) {
|
||||
this.token = localStorage.getItem('access_token')
|
||||
@@ -273,18 +313,17 @@ class ApiClient {
|
||||
}
|
||||
|
||||
clearAuth(): void {
|
||||
this.token = null
|
||||
const hadAuth = this.token !== null || localStorage.getItem('access_token') !== null
|
||||
this.syncTokenState(null)
|
||||
localStorage.removeItem('access_token')
|
||||
localStorage.removeItem('refresh_token')
|
||||
// 同步清除 mock token
|
||||
if (isDemoMode()) {
|
||||
setMockUserToken(null)
|
||||
// 同标签页内清理认证状态时不会触发 storage 事件,这里主动广播一次。
|
||||
if (hadAuth) {
|
||||
this.emitAuthStateChange(null)
|
||||
}
|
||||
}
|
||||
|
||||
async refreshToken(refreshToken: string): Promise<AxiosResponse> {
|
||||
// refreshToken 会通过 adapter 处理 Demo 模式
|
||||
return this.client.post('/api/auth/refresh', { refresh_token: refreshToken })
|
||||
async refreshToken(): Promise<AxiosResponse> {
|
||||
return this.client.post('/api/auth/refresh', {})
|
||||
}
|
||||
|
||||
// 以下方法直接委托给 axios client,Demo 模式由 adapter 统一处理
|
||||
|
||||
@@ -1,339 +0,0 @@
|
||||
import apiClient from './client'
|
||||
|
||||
export interface DashboardStat {
|
||||
name: string
|
||||
value: string
|
||||
subValue?: string
|
||||
change?: string
|
||||
changeType?: 'increase' | 'decrease' | 'neutral'
|
||||
extraBadge?: string
|
||||
icon: string
|
||||
}
|
||||
|
||||
export interface RecentRequest {
|
||||
id: string // UUID
|
||||
user: string
|
||||
model: string
|
||||
tokens: number
|
||||
time: string
|
||||
}
|
||||
|
||||
export interface ProviderStatus {
|
||||
name: string
|
||||
status: 'active' | 'inactive'
|
||||
requests: number
|
||||
}
|
||||
|
||||
// 系统健康指标(管理员专用)
|
||||
export interface SystemHealth {
|
||||
avg_response_time: number
|
||||
error_rate: number
|
||||
error_requests: number
|
||||
fallback_count: number
|
||||
total_requests: number
|
||||
}
|
||||
|
||||
// 成本统计(管理员专用)
|
||||
export interface CostStats {
|
||||
total_cost: number
|
||||
total_actual_cost: number
|
||||
cost_savings: number
|
||||
}
|
||||
|
||||
// 缓存统计
|
||||
export interface CacheStats {
|
||||
cache_creation_tokens: number
|
||||
cache_read_tokens: number
|
||||
cache_creation_cost?: number
|
||||
cache_read_cost?: number
|
||||
cache_hit_rate?: number
|
||||
total_cache_tokens: number
|
||||
}
|
||||
|
||||
// 用户统计(管理员专用)
|
||||
export interface UserStats {
|
||||
total: number
|
||||
active: number
|
||||
}
|
||||
|
||||
// Token 详细分类
|
||||
export interface TokenBreakdown {
|
||||
input: number
|
||||
output: number
|
||||
cache_creation: number
|
||||
cache_read: number
|
||||
}
|
||||
|
||||
export interface DashboardStatsResponse {
|
||||
stats: DashboardStat[]
|
||||
today?: {
|
||||
requests: number
|
||||
tokens: number
|
||||
cost: number
|
||||
actual_cost?: number
|
||||
cache_creation_tokens?: number
|
||||
cache_read_tokens?: number
|
||||
}
|
||||
api_keys?: {
|
||||
total: number
|
||||
active: number
|
||||
}
|
||||
tokens?: {
|
||||
month: number
|
||||
}
|
||||
// 管理员专用字段
|
||||
system_health?: SystemHealth
|
||||
cost_stats?: CostStats
|
||||
cache_stats?: CacheStats
|
||||
users?: UserStats
|
||||
token_breakdown?: TokenBreakdown
|
||||
// 普通用户专用字段
|
||||
monthly_cost?: number
|
||||
}
|
||||
|
||||
export interface RecentRequestsResponse {
|
||||
requests: RecentRequest[]
|
||||
}
|
||||
|
||||
export interface ProviderStatusResponse {
|
||||
providers: ProviderStatus[]
|
||||
}
|
||||
|
||||
// 视频/图像/音频计费信息
|
||||
export interface VideoBilling {
|
||||
task_type: 'video' | 'image' | 'audio'
|
||||
duration_seconds?: number // 视频时长(秒)
|
||||
resolution?: string // 分辨率
|
||||
video_price_per_second?: number // 每秒单价
|
||||
video_cost?: number // 视频费用
|
||||
cost?: number // 总费用
|
||||
rule_name?: string // 计费规则名称
|
||||
expression?: string // 计费公式
|
||||
status?: string // 计费状态
|
||||
}
|
||||
|
||||
export interface RequestDetail {
|
||||
id: string // UUID
|
||||
request_id: string
|
||||
user: {
|
||||
id: string // UUID
|
||||
username: string
|
||||
email: string
|
||||
}
|
||||
api_key: {
|
||||
id: string // UUID
|
||||
name: string
|
||||
display: string
|
||||
}
|
||||
provider: string
|
||||
api_format?: string
|
||||
model: string
|
||||
target_model?: string | null // 映射后的目标模型名
|
||||
tokens: {
|
||||
input: number
|
||||
output: number
|
||||
total: number
|
||||
}
|
||||
cost: {
|
||||
input: number
|
||||
output: number
|
||||
total: number
|
||||
}
|
||||
// Additional token fields
|
||||
input_tokens?: number
|
||||
output_tokens?: number
|
||||
total_tokens?: number
|
||||
cache_creation_input_tokens?: number
|
||||
cache_read_input_tokens?: number
|
||||
// Additional cost fields
|
||||
input_cost?: number
|
||||
output_cost?: number
|
||||
total_cost?: number
|
||||
cache_creation_cost?: number
|
||||
cache_read_cost?: number
|
||||
request_cost?: number // 按次计费费用
|
||||
// Historical pricing fields (per 1M tokens)
|
||||
input_price_per_1m?: number
|
||||
output_price_per_1m?: number
|
||||
cache_creation_price_per_1m?: number
|
||||
cache_read_price_per_1m?: number
|
||||
price_per_request?: number // 按次计费价格
|
||||
request_type: string
|
||||
is_stream: boolean
|
||||
status_code: number
|
||||
status?: string // pending, streaming, completed, failed, cancelled
|
||||
error_message?: string
|
||||
response_time_ms: number
|
||||
created_at: string
|
||||
request_headers?: Record<string, unknown>
|
||||
request_body?: Record<string, unknown>
|
||||
provider_request_headers?: Record<string, unknown>
|
||||
provider_request_body?: Record<string, unknown>
|
||||
response_headers?: Record<string, unknown>
|
||||
client_response_headers?: Record<string, unknown>
|
||||
response_body?: Record<string, unknown>
|
||||
client_response_body?: Record<string, unknown>
|
||||
metadata?: Record<string, unknown>
|
||||
// 阶梯计费信息
|
||||
tiered_pricing?: {
|
||||
total_input_context: number // 总输入上下文 (input + cache_read)
|
||||
tier_index: number // 命中的阶梯索引 (0-based)
|
||||
tier_count: number // 阶梯总数
|
||||
source?: 'provider' | 'global' // 定价来源: 提供商或全局
|
||||
current_tier: { // 当前命中的阶梯配置
|
||||
up_to?: number | null
|
||||
input_price_per_1m: number
|
||||
output_price_per_1m: number
|
||||
cache_creation_price_per_1m?: number
|
||||
cache_read_price_per_1m?: number
|
||||
cache_ttl_pricing?: Array<{
|
||||
ttl_minutes: number
|
||||
cache_read_price_per_1m: number
|
||||
}>
|
||||
}
|
||||
tiers: Array<{ // 完整阶梯配置列表
|
||||
up_to?: number | null
|
||||
input_price_per_1m: number
|
||||
output_price_per_1m: number
|
||||
cache_creation_price_per_1m?: number
|
||||
cache_read_price_per_1m?: number
|
||||
cache_ttl_pricing?: Array<{
|
||||
ttl_minutes: number
|
||||
cache_read_price_per_1m: number
|
||||
}>
|
||||
}>
|
||||
} | null
|
||||
// 视频/图像/音频计费信息
|
||||
video_billing?: VideoBilling | null
|
||||
}
|
||||
|
||||
export interface CurlData {
|
||||
url: string
|
||||
method: string
|
||||
headers: Record<string, string>
|
||||
body: Record<string, unknown>
|
||||
curl: string
|
||||
}
|
||||
|
||||
export interface ReplayRequest {
|
||||
provider_id?: string
|
||||
endpoint_id?: string
|
||||
api_key_id?: string
|
||||
body_override?: Record<string, unknown>
|
||||
}
|
||||
|
||||
export interface ReplayResponse {
|
||||
url: string
|
||||
provider: string
|
||||
status_code: number
|
||||
response_headers: Record<string, string>
|
||||
response_body: Record<string, unknown>
|
||||
response_time_ms: number
|
||||
}
|
||||
|
||||
export interface ModelBreakdown {
|
||||
model: string
|
||||
requests: number
|
||||
tokens: number
|
||||
cost: number
|
||||
}
|
||||
|
||||
export interface ModelSummary {
|
||||
model: string
|
||||
requests: number
|
||||
tokens: number
|
||||
cost: number
|
||||
avg_response_time: number
|
||||
cost_per_request: number
|
||||
tokens_per_request: number
|
||||
}
|
||||
|
||||
export interface ProviderSummary {
|
||||
provider: string
|
||||
requests: number
|
||||
tokens: number
|
||||
cost: number
|
||||
}
|
||||
|
||||
export interface DailyStat {
|
||||
date: string // ISO date string
|
||||
requests: number
|
||||
tokens: number
|
||||
cost: number
|
||||
avg_response_time: number // in seconds
|
||||
unique_models: number
|
||||
unique_providers?: number // 仅管理员返回
|
||||
model_breakdown: ModelBreakdown[]
|
||||
}
|
||||
|
||||
export interface DailyStatsResponse {
|
||||
daily_stats: DailyStat[]
|
||||
model_summary: ModelSummary[]
|
||||
provider_summary?: ProviderSummary[] // 仅管理员返回
|
||||
period: {
|
||||
start_date: string
|
||||
end_date: string
|
||||
days: number
|
||||
}
|
||||
}
|
||||
|
||||
export interface TimeRangeParams {
|
||||
start_date?: string
|
||||
end_date?: string
|
||||
preset?: string
|
||||
granularity?: 'hour' | 'day' | 'week' | 'month'
|
||||
timezone?: string
|
||||
tz_offset_minutes?: number
|
||||
}
|
||||
|
||||
export const dashboardApi = {
|
||||
// 获取仪表盘统计数据
|
||||
async getStats(params?: TimeRangeParams): Promise<DashboardStatsResponse> {
|
||||
const response = await apiClient.get<DashboardStatsResponse>('/api/dashboard/stats', { params })
|
||||
return response.data
|
||||
},
|
||||
|
||||
// 获取最近的请求记录
|
||||
async getRecentRequests(limit: number = 10): Promise<RecentRequest[]> {
|
||||
const response = await apiClient.get<RecentRequestsResponse>('/api/dashboard/recent-requests', {
|
||||
params: { limit }
|
||||
})
|
||||
return response.data.requests
|
||||
},
|
||||
|
||||
// 获取提供商状态
|
||||
async getProviderStatus(): Promise<ProviderStatus[]> {
|
||||
const response = await apiClient.get<ProviderStatusResponse>('/api/dashboard/provider-status')
|
||||
return response.data.providers
|
||||
},
|
||||
|
||||
// 获取请求详情
|
||||
// NOTE: This method now calls the new RESTful API at /api/admin/usage/{id}
|
||||
async getRequestDetail(requestId: string): Promise<RequestDetail> {
|
||||
const response = await apiClient.get<RequestDetail>(`/api/admin/usage/${requestId}`)
|
||||
return response.data
|
||||
},
|
||||
|
||||
// 获取每日统计数据
|
||||
async getDailyStats(params?: TimeRangeParams & { days?: number }): Promise<DailyStatsResponse> {
|
||||
const response = await apiClient.get<DailyStatsResponse>('/api/dashboard/daily-stats', {
|
||||
params
|
||||
})
|
||||
return response.data
|
||||
},
|
||||
|
||||
// 获取 cURL 命令数据(含明文 API Key)
|
||||
async getCurlData(requestId: string): Promise<CurlData> {
|
||||
const response = await apiClient.get<CurlData>(`/api/admin/usage/${requestId}/curl`)
|
||||
return response.data
|
||||
},
|
||||
|
||||
// 回放请求到提供商
|
||||
async replayRequest(requestId: string, params?: ReplayRequest): Promise<ReplayResponse> {
|
||||
const response = await apiClient.post<ReplayResponse>(
|
||||
`/api/admin/usage/${requestId}/replay`,
|
||||
params || {}
|
||||
)
|
||||
return response.data
|
||||
}
|
||||
}
|
||||
@@ -68,3 +68,13 @@ export async function deleteEndpoint(endpointId: string): Promise<{ message: str
|
||||
const response = await client.delete(`/api/admin/endpoints/${endpointId}`)
|
||||
return response.data
|
||||
}
|
||||
|
||||
/**
|
||||
* 获取指定 API 格式的默认请求体规则
|
||||
*/
|
||||
export async function getDefaultBodyRules(apiFormat: string, providerType?: string): Promise<{ api_format: string; body_rules: BodyRule[] }> {
|
||||
const params: Record<string, string> = {}
|
||||
if (providerType) params.provider_type = providerType
|
||||
const response = await client.get(`/api/admin/endpoints/defaults/${encodeURIComponent(apiFormat)}/body-rules`, { params })
|
||||
return response.data
|
||||
}
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import client from '../client'
|
||||
import { dedupedRequest, buildCacheKey } from '@/utils/cache'
|
||||
import type {
|
||||
GlobalModelCreate,
|
||||
GlobalModelUpdate,
|
||||
@@ -27,16 +28,21 @@ export async function getGlobalModels(params?: {
|
||||
is_active?: boolean
|
||||
search?: string
|
||||
}): Promise<GlobalModelListResponse> {
|
||||
const key = buildCacheKey('global-models:list', params as Record<string, unknown> | undefined)
|
||||
return dedupedRequest(key, async () => {
|
||||
const response = await client.get('/api/admin/models/global', { params })
|
||||
return response.data
|
||||
})
|
||||
}
|
||||
|
||||
/**
|
||||
* 获取单个 GlobalModel 详情
|
||||
*/
|
||||
export async function getGlobalModel(id: string): Promise<GlobalModelWithStats> {
|
||||
return dedupedRequest(`global-models:detail:${id}`, async () => {
|
||||
const response = await client.get(`/api/admin/models/global/${id}`)
|
||||
return response.data
|
||||
})
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -112,10 +118,12 @@ export async function getGlobalModelProviders(globalModelId: string): Promise<{
|
||||
providers: ModelCatalogProviderDetail[]
|
||||
total: number
|
||||
}> {
|
||||
return dedupedRequest(`global-models:providers:${globalModelId}`, async () => {
|
||||
const response = await client.get(
|
||||
`/api/admin/models/global/${globalModelId}/providers`
|
||||
)
|
||||
return response.data
|
||||
})
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
@@ -4,6 +4,34 @@ import type { EndpointAPIKey, AllowedModels } from './types'
|
||||
// Re-export types for convenience
|
||||
export type { EndpointAPIKey, AllowedModels }
|
||||
|
||||
export interface GroupedFormatKey {
|
||||
id: string
|
||||
provider_id: string
|
||||
name: string
|
||||
auth_type?: string
|
||||
api_key_masked: string
|
||||
internal_priority: number
|
||||
global_priority_by_format: Record<string, number> | null
|
||||
format_priority: number | null
|
||||
rate_multipliers: Record<string, number> | null
|
||||
is_active: boolean
|
||||
provider_active: boolean
|
||||
pool_enabled: boolean
|
||||
circuit_breaker_open: boolean
|
||||
provider_name: string
|
||||
api_format: string
|
||||
api_formats: string[]
|
||||
capabilities: string[]
|
||||
health_score: number | null
|
||||
success_rate: number | null
|
||||
avg_response_time_ms: number | null
|
||||
request_count: number
|
||||
}
|
||||
|
||||
function toNumberOrNull(value: unknown): number | null {
|
||||
return typeof value === 'number' && Number.isFinite(value) ? value : null
|
||||
}
|
||||
|
||||
/**
|
||||
* 能力定义类型
|
||||
*/
|
||||
@@ -56,7 +84,7 @@ export async function getModelCapabilities(modelName: string): Promise<ModelCapa
|
||||
* 获取完整的 API Key(用于查看和复制)
|
||||
*/
|
||||
export interface RevealKeyResult {
|
||||
auth_type: 'api_key' | 'vertex_ai' | 'oauth'
|
||||
auth_type: 'api_key' | 'service_account' | 'oauth'
|
||||
api_key?: string
|
||||
refresh_token?: string
|
||||
auth_config?: string | Record<string, unknown>
|
||||
@@ -83,6 +111,20 @@ export async function deleteEndpointKey(keyId: string): Promise<{ message: strin
|
||||
return response.data
|
||||
}
|
||||
|
||||
/**
|
||||
* 批量删除 Keys
|
||||
*/
|
||||
export interface BatchDeleteKeysResult {
|
||||
success_count: number
|
||||
failed_count: number
|
||||
failed: Array<{ id: string; error: string }>
|
||||
}
|
||||
|
||||
export async function batchDeleteEndpointKeys(ids: string[]): Promise<BatchDeleteKeysResult> {
|
||||
const response = await client.post('/api/admin/endpoints/keys/batch-delete', { ids })
|
||||
return response.data
|
||||
}
|
||||
|
||||
|
||||
// ========== Provider 级别的 Keys API ==========
|
||||
|
||||
@@ -111,6 +153,51 @@ export async function getProviderKeys(providerId: string): Promise<EndpointAPIKe
|
||||
return allKeys
|
||||
}
|
||||
|
||||
/**
|
||||
* 获取按 API 格式分组的 Key 列表
|
||||
*/
|
||||
export async function getKeysGroupedByFormat(): Promise<Record<string, GroupedFormatKey[]>> {
|
||||
const response = await client.get('/api/admin/endpoints/keys/grouped-by-format')
|
||||
const grouped = response.data as Record<string, Array<Record<string, unknown>>>
|
||||
const result: Record<string, GroupedFormatKey[]> = {}
|
||||
|
||||
for (const [apiFormat, keys] of Object.entries(grouped || {})) {
|
||||
if (!Array.isArray(keys)) continue
|
||||
|
||||
result[apiFormat] = keys.map((key) => ({
|
||||
id: String(key.id || ''),
|
||||
provider_id: String(key.provider_id || ''),
|
||||
name: String(key.name || 'Unnamed Key'),
|
||||
auth_type: typeof key.auth_type === 'string' ? key.auth_type : undefined,
|
||||
api_key_masked: String(key.api_key_masked || '***'),
|
||||
internal_priority: toNumberOrNull(key.internal_priority) ?? 0,
|
||||
global_priority_by_format:
|
||||
key.global_priority_by_format && typeof key.global_priority_by_format === 'object'
|
||||
? (key.global_priority_by_format as Record<string, number>)
|
||||
: null,
|
||||
format_priority: toNumberOrNull(key.format_priority),
|
||||
rate_multipliers:
|
||||
key.rate_multipliers && typeof key.rate_multipliers === 'object'
|
||||
? (key.rate_multipliers as Record<string, number>)
|
||||
: null,
|
||||
is_active: key.is_active !== false,
|
||||
provider_active: key.provider_active !== false,
|
||||
pool_enabled: key.pool_enabled === true,
|
||||
circuit_breaker_open: key.circuit_breaker_open === true,
|
||||
provider_name: String(key.provider_name || 'Unknown Provider'),
|
||||
api_format: typeof key.api_format === 'string' ? key.api_format : apiFormat,
|
||||
api_formats: Array.isArray(key.api_formats) ? key.api_formats.map(item => String(item)) : [apiFormat],
|
||||
capabilities: Array.isArray(key.capabilities) ? key.capabilities.map(item => String(item)) : [],
|
||||
health_score: toNumberOrNull(key.health_score),
|
||||
success_rate: toNumberOrNull(key.success_rate),
|
||||
avg_response_time_ms: toNumberOrNull(key.avg_response_time_ms),
|
||||
request_count: toNumberOrNull(key.request_count) ?? 0,
|
||||
}))
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
/**
|
||||
* 为 Provider 添加 Key
|
||||
*/
|
||||
@@ -119,7 +206,7 @@ export async function addProviderKey(
|
||||
data: {
|
||||
api_formats: string[] // 支持的 API 格式列表(必填)
|
||||
api_key: string
|
||||
auth_type?: 'api_key' | 'vertex_ai' | 'oauth' // 认证类型
|
||||
auth_type?: 'api_key' | 'service_account' | 'oauth' // 认证类型
|
||||
auth_config?: Record<string, unknown> // 认证配置(Vertex AI Service Account JSON)
|
||||
name: string
|
||||
rate_multipliers?: Record<string, number> | null // 按 API 格式的成本倍率
|
||||
@@ -147,7 +234,7 @@ export async function updateProviderKey(
|
||||
data: Partial<{
|
||||
api_formats: string[] // 支持的 API 格式列表
|
||||
api_key: string
|
||||
auth_type: 'api_key' | 'vertex_ai' | 'oauth' // 认证类型
|
||||
auth_type: 'api_key' | 'service_account' | 'oauth' // 认证类型
|
||||
auth_config: Record<string, unknown> // 认证配置(Vertex AI Service Account JSON)
|
||||
name: string
|
||||
rate_multipliers: Record<string, number> | null // 按 API 格式的成本倍率
|
||||
@@ -197,8 +284,16 @@ export interface RefreshQuotaResult {
|
||||
}>
|
||||
}
|
||||
|
||||
export async function refreshProviderQuota(providerId: string): Promise<RefreshQuotaResult> {
|
||||
const response = await client.post(`/api/admin/endpoints/providers/${providerId}/refresh-quota`)
|
||||
export async function refreshProviderQuota(
|
||||
providerId: string,
|
||||
keyIds?: string[],
|
||||
): Promise<RefreshQuotaResult> {
|
||||
const body = keyIds && keyIds.length > 0 ? { key_ids: keyIds } : undefined
|
||||
const response = await client.post(
|
||||
`/api/admin/endpoints/providers/${providerId}/refresh-quota`,
|
||||
body,
|
||||
{ timeout: 5 * 60 * 1000 },
|
||||
)
|
||||
return response.data
|
||||
}
|
||||
|
||||
|
||||
@@ -1,4 +1,9 @@
|
||||
import client from '../client'
|
||||
import { dedupedRequest } from '@/utils/cache'
|
||||
import type { AllowedModels, OAuthOrganizationInfo, ProxyConfig } from './types/provider'
|
||||
import type { ProviderKeyStatusSnapshot } from './types/statusSnapshot'
|
||||
|
||||
const POOL_BATCH_ACTION_TIMEOUT_MS = 5 * 60 * 1000
|
||||
|
||||
export interface PoolKeyStatus {
|
||||
key_id: string
|
||||
@@ -25,7 +30,7 @@ export interface PoolStatusResponse {
|
||||
* 获取 Provider 的号池状态
|
||||
*/
|
||||
export async function getPoolStatus(providerId: string): Promise<PoolStatusResponse> {
|
||||
const response = await client.get(`/api/admin/providers/${providerId}/pool-status`)
|
||||
const response = await client.get<PoolStatusResponse>(`/api/admin/providers/${providerId}/pool-status`)
|
||||
return response.data
|
||||
}
|
||||
|
||||
@@ -36,7 +41,7 @@ export async function clearPoolCooldown(
|
||||
providerId: string,
|
||||
keyId: string,
|
||||
): Promise<{ message: string }> {
|
||||
const response = await client.post(
|
||||
const response = await client.post<{ message: string }>(
|
||||
`/api/admin/providers/${providerId}/pool/clear-cooldown/${keyId}`,
|
||||
)
|
||||
return response.data
|
||||
@@ -49,7 +54,7 @@ export async function resetPoolCost(
|
||||
providerId: string,
|
||||
keyId: string,
|
||||
): Promise<{ message: string }> {
|
||||
const response = await client.post(
|
||||
const response = await client.post<{ message: string }>(
|
||||
`/api/admin/providers/${providerId}/pool/reset-cost/${keyId}`,
|
||||
)
|
||||
return response.data
|
||||
@@ -73,19 +78,95 @@ export interface PoolOverviewResponse {
|
||||
items: PoolOverviewItem[]
|
||||
}
|
||||
|
||||
export interface PoolPresetModeMeta {
|
||||
value: string
|
||||
label: string
|
||||
}
|
||||
|
||||
export interface PoolPresetMeta {
|
||||
name: string
|
||||
label: string
|
||||
description: string
|
||||
providers: string[]
|
||||
modes?: PoolPresetModeMeta[] | null
|
||||
default_mode?: string | null
|
||||
mutex_group?: string | null
|
||||
evidence_hint?: string | null
|
||||
}
|
||||
|
||||
export interface PoolKeyDetail {
|
||||
key_id: string
|
||||
key_name: string
|
||||
is_active: boolean
|
||||
auth_type: string
|
||||
oauth_expires_at?: number | null
|
||||
oauth_invalid_at?: number | null // 兼容字段;优先使用 status_snapshot.oauth
|
||||
oauth_invalid_reason?: string | null // 兼容字段;优先使用 status_snapshot.oauth
|
||||
oauth_plan_type?: string | null
|
||||
oauth_account_id?: string | null
|
||||
oauth_account_user_id?: string | null
|
||||
oauth_account_name?: string | null
|
||||
oauth_organizations?: OAuthOrganizationInfo[] | null
|
||||
account_status_code?: string | null // 兼容字段;优先使用 status_snapshot.account
|
||||
account_status_label?: string | null // 兼容字段;优先使用 status_snapshot.account
|
||||
account_status_reason?: string | null // 兼容字段;优先使用 status_snapshot.account
|
||||
account_status_blocked?: boolean // 兼容字段;优先使用 status_snapshot.account
|
||||
account_status_recoverable?: boolean // 兼容字段;优先使用 status_snapshot.account
|
||||
account_status_source?: string | null // 兼容字段;优先使用 status_snapshot.account
|
||||
status_snapshot?: ProviderKeyStatusSnapshot | null
|
||||
quota_updated_at?: number | null
|
||||
health_score?: number
|
||||
circuit_breaker_open?: boolean
|
||||
api_formats?: string[]
|
||||
rate_multipliers?: Record<string, number> | null
|
||||
internal_priority?: number
|
||||
rpm_limit?: number | null
|
||||
cache_ttl_minutes?: number
|
||||
max_probe_interval_minutes?: number
|
||||
note?: string | null
|
||||
allowed_models?: AllowedModels
|
||||
capabilities?: Record<string, boolean> | null
|
||||
auto_fetch_models?: boolean
|
||||
locked_models?: string[] | null
|
||||
model_include_patterns?: string[] | null
|
||||
model_exclude_patterns?: string[] | null
|
||||
proxy?: ProxyConfig | null
|
||||
account_quota: string | null
|
||||
cooldown_reason: string | null
|
||||
cooldown_ttl_seconds: number | null
|
||||
cost_window_usage: number
|
||||
cost_limit: number | null
|
||||
request_count: number
|
||||
total_tokens: number
|
||||
total_cost_usd: string
|
||||
sticky_sessions: number
|
||||
lru_score: number | null
|
||||
created_at: string | null
|
||||
last_used_at: string | null
|
||||
scheduling_status?: 'available' | 'degraded' | 'blocked'
|
||||
scheduling_reason?:
|
||||
| 'available'
|
||||
| 'manual_disabled'
|
||||
| 'cooldown'
|
||||
| 'circuit_open'
|
||||
| 'cost_exhausted'
|
||||
| 'cost_soft'
|
||||
| 'cost'
|
||||
| 'health_low'
|
||||
| 'health_degraded'
|
||||
| 'health'
|
||||
| string
|
||||
scheduling_label?: string
|
||||
scheduling_reasons?: PoolSchedulingReason[]
|
||||
}
|
||||
|
||||
export interface PoolSchedulingReason {
|
||||
code: string
|
||||
label: string
|
||||
blocking: boolean
|
||||
source: 'manual' | 'pool' | 'health' | 'policy' | string
|
||||
ttl_seconds?: number | null
|
||||
detail?: string | null
|
||||
}
|
||||
|
||||
export interface PoolKeysPageResponse {
|
||||
@@ -100,50 +181,118 @@ export interface PoolKeysQuery {
|
||||
page_size?: number
|
||||
search?: string
|
||||
status?: 'all' | 'active' | 'cooldown' | 'inactive'
|
||||
quick_selectors?: string[]
|
||||
search_scope?: 'name' | 'full'
|
||||
}
|
||||
|
||||
export interface PoolKeyImportItem {
|
||||
name: string
|
||||
api_key: string
|
||||
auth_type?: string
|
||||
export interface PoolKeySelectionRequest {
|
||||
search?: string
|
||||
quick_selectors?: string[]
|
||||
}
|
||||
|
||||
export interface BatchImportResponse {
|
||||
imported: number
|
||||
skipped: number
|
||||
errors: { index: number; reason: string }[]
|
||||
export interface PoolKeySelectionItem {
|
||||
key_id: string
|
||||
key_name: string
|
||||
auth_type: string
|
||||
}
|
||||
|
||||
export interface PoolKeySelectionResponse {
|
||||
total: number
|
||||
items: PoolKeySelectionItem[]
|
||||
}
|
||||
|
||||
export interface PoolBatchAction {
|
||||
key_ids: string[]
|
||||
action: 'enable' | 'disable' | 'delete' | 'clear_cooldown' | 'reset_cost'
|
||||
action:
|
||||
| 'enable'
|
||||
| 'disable'
|
||||
| 'delete'
|
||||
| 'clear_cooldown'
|
||||
| 'reset_cost'
|
||||
| 'regenerate_fingerprint'
|
||||
| 'clear_proxy'
|
||||
| 'set_proxy'
|
||||
payload?: Record<string, unknown> | null
|
||||
}
|
||||
|
||||
export async function getPoolOverview(): Promise<PoolOverviewResponse> {
|
||||
const response = await client.get('/api/admin/pool/overview')
|
||||
return dedupedRequest('pool:overview', async () => {
|
||||
const response = await client.get<PoolOverviewResponse>('/api/admin/pool/overview')
|
||||
return response.data
|
||||
})
|
||||
}
|
||||
|
||||
export async function getPoolSchedulingPresets(): Promise<PoolPresetMeta[]> {
|
||||
return dedupedRequest('pool:scheduling-presets', async () => {
|
||||
const response = await client.get<PoolPresetMeta[]>('/api/admin/pool/scheduling-presets')
|
||||
return response.data
|
||||
})
|
||||
}
|
||||
|
||||
export async function listPoolKeys(
|
||||
providerId: string,
|
||||
params: PoolKeysQuery = {},
|
||||
): Promise<PoolKeysPageResponse> {
|
||||
const response = await client.get(`/api/admin/pool/${providerId}/keys`, { params })
|
||||
const normalizedParams = {
|
||||
...params,
|
||||
quick_selectors: params.quick_selectors?.length ? params.quick_selectors.join(',') : undefined,
|
||||
}
|
||||
const key = `pool:keys:${providerId}|${normalizedParams.page ?? ''}|${normalizedParams.page_size ?? ''}|${normalizedParams.search ?? ''}|${normalizedParams.status ?? ''}|${normalizedParams.quick_selectors ?? ''}|${normalizedParams.search_scope ?? ''}`
|
||||
return dedupedRequest(key, async () => {
|
||||
const response = await client.get<PoolKeysPageResponse>(`/api/admin/pool/${providerId}/keys`, { params: normalizedParams })
|
||||
return response.data
|
||||
})
|
||||
}
|
||||
|
||||
export async function batchImportPoolKeys(
|
||||
export async function resolvePoolKeySelection(
|
||||
providerId: string,
|
||||
keys: PoolKeyImportItem[],
|
||||
): Promise<BatchImportResponse> {
|
||||
const response = await client.post(`/api/admin/pool/${providerId}/keys/batch-import`, { keys })
|
||||
body: PoolKeySelectionRequest,
|
||||
): Promise<PoolKeySelectionResponse> {
|
||||
const response = await client.post<PoolKeySelectionResponse>(
|
||||
`/api/admin/pool/${providerId}/keys/resolve-selection`,
|
||||
body,
|
||||
{ timeout: POOL_BATCH_ACTION_TIMEOUT_MS },
|
||||
)
|
||||
return response.data
|
||||
}
|
||||
|
||||
export async function batchActionPoolKeys(
|
||||
providerId: string,
|
||||
body: PoolBatchAction,
|
||||
): Promise<{ affected: number; message: string; task_id?: string }> {
|
||||
const response = await client.post(
|
||||
`/api/admin/pool/${providerId}/keys/batch-action`,
|
||||
body,
|
||||
{ timeout: POOL_BATCH_ACTION_TIMEOUT_MS },
|
||||
)
|
||||
return response.data
|
||||
}
|
||||
|
||||
export interface BatchDeleteTaskStatus {
|
||||
task_id: string
|
||||
status: 'pending' | 'running' | 'completed' | 'failed'
|
||||
total: number
|
||||
deleted: number
|
||||
message: string
|
||||
}
|
||||
|
||||
export async function getPoolBatchDeleteTask(
|
||||
providerId: string,
|
||||
taskId: string,
|
||||
): Promise<BatchDeleteTaskStatus> {
|
||||
const response = await client.get<BatchDeleteTaskStatus>(
|
||||
`/api/admin/pool/${providerId}/keys/batch-delete-task/${taskId}`,
|
||||
)
|
||||
return response.data
|
||||
}
|
||||
|
||||
export async function cleanupBannedPoolKeys(
|
||||
providerId: string,
|
||||
): Promise<{ affected: number; message: string }> {
|
||||
const response = await client.post(`/api/admin/pool/${providerId}/keys/batch-action`, body)
|
||||
const response = await client.post(
|
||||
`/api/admin/pool/${providerId}/keys/cleanup-banned`,
|
||||
undefined,
|
||||
{ timeout: POOL_BATCH_ACTION_TIMEOUT_MS },
|
||||
)
|
||||
return response.data
|
||||
}
|
||||
|
||||
@@ -18,6 +18,8 @@ export interface ProviderOAuthCompleteResponse {
|
||||
expires_at?: number | null
|
||||
has_refresh_token: boolean
|
||||
email?: string | null
|
||||
account_state_recheck_attempted?: boolean
|
||||
account_state_recheck_error?: string | null
|
||||
}
|
||||
|
||||
export interface ProviderOAuthCompleteResponseWithKey {
|
||||
@@ -28,6 +30,48 @@ export interface ProviderOAuthCompleteResponseWithKey {
|
||||
email?: string | null
|
||||
}
|
||||
|
||||
export interface OAuthBatchImportResultItem {
|
||||
index: number
|
||||
status: 'success' | 'error'
|
||||
key_id?: string
|
||||
key_name?: string
|
||||
auth_method?: string
|
||||
error?: string
|
||||
replaced?: boolean
|
||||
}
|
||||
|
||||
export type OAuthBatchImportTaskStatus = 'submitted' | 'processing' | 'completed' | 'failed'
|
||||
|
||||
export interface OAuthBatchImportTaskStartResponse {
|
||||
task_id: string
|
||||
status: OAuthBatchImportTaskStatus
|
||||
total: number
|
||||
processed: number
|
||||
success: number
|
||||
failed: number
|
||||
progress_percent: number
|
||||
message?: string | null
|
||||
}
|
||||
|
||||
export interface OAuthBatchImportTaskStatusResponse {
|
||||
task_id: string
|
||||
provider_id: string
|
||||
provider_type: string
|
||||
status: OAuthBatchImportTaskStatus
|
||||
total: number
|
||||
processed: number
|
||||
success: number
|
||||
failed: number
|
||||
progress_percent: number
|
||||
message?: string | null
|
||||
error?: string | null
|
||||
error_samples: OAuthBatchImportResultItem[]
|
||||
created_at: number
|
||||
started_at?: number | null
|
||||
finished_at?: number | null
|
||||
updated_at: number
|
||||
}
|
||||
|
||||
export async function refreshProviderOAuth(keyId: string): Promise<ProviderOAuthCompleteResponse> {
|
||||
const resp = await client.post(`/api/admin/provider-oauth/keys/${keyId}/refresh`)
|
||||
return resp.data
|
||||
@@ -50,12 +94,32 @@ export async function completeProviderLevelOAuth(
|
||||
|
||||
export async function importProviderRefreshToken(
|
||||
providerId: string,
|
||||
data: { refresh_token: string; name?: string; proxy_node_id?: string }
|
||||
data: { refresh_token?: string; access_token?: string; name?: string; proxy_node_id?: string }
|
||||
): Promise<ProviderOAuthCompleteResponseWithKey> {
|
||||
const resp = await client.post(`/api/admin/provider-oauth/providers/${providerId}/import-refresh-token`, data)
|
||||
return resp.data
|
||||
}
|
||||
|
||||
export async function startBatchImportOAuthTask(
|
||||
providerId: string,
|
||||
credentials: string,
|
||||
proxyNodeId?: string
|
||||
): Promise<OAuthBatchImportTaskStartResponse> {
|
||||
const resp = await client.post(`/api/admin/provider-oauth/providers/${providerId}/batch-import/tasks`, {
|
||||
credentials,
|
||||
proxy_node_id: proxyNodeId || undefined,
|
||||
})
|
||||
return resp.data
|
||||
}
|
||||
|
||||
export async function getBatchImportOAuthTaskStatus(
|
||||
providerId: string,
|
||||
taskId: string
|
||||
): Promise<OAuthBatchImportTaskStatusResponse> {
|
||||
const resp = await client.get(`/api/admin/provider-oauth/providers/${providerId}/batch-import/tasks/${taskId}`)
|
||||
return resp.data
|
||||
}
|
||||
|
||||
// Device Authorization (AWS SSO OIDC)
|
||||
|
||||
export interface DeviceAuthorizeRequest {
|
||||
|
||||
@@ -1,16 +1,39 @@
|
||||
import client from '../client'
|
||||
import { dedupedRequest } from '@/utils/cache'
|
||||
import type {
|
||||
ClaudeCodeAdvancedConfig,
|
||||
FailoverRulesConfig,
|
||||
PoolAdvancedConfig,
|
||||
ProviderWithEndpointsSummary,
|
||||
ProxyConfig,
|
||||
} from './types'
|
||||
|
||||
/**
|
||||
* 获取 Providers 摘要(包含 Endpoints 统计)
|
||||
* 获取 Providers 摘要(分页)
|
||||
*/
|
||||
export async function getProvidersSummary(): Promise<ProviderWithEndpointsSummary[]> {
|
||||
const response = await client.get('/api/admin/providers/summary')
|
||||
export interface ProviderSummaryQuery {
|
||||
page?: number
|
||||
page_size?: number
|
||||
search?: string
|
||||
status?: string
|
||||
api_format?: string
|
||||
model_id?: string
|
||||
}
|
||||
|
||||
export interface ProviderSummaryPageResponse {
|
||||
total: number
|
||||
page: number
|
||||
page_size: number
|
||||
items: ProviderWithEndpointsSummary[]
|
||||
}
|
||||
|
||||
export async function getProvidersSummary(
|
||||
params: ProviderSummaryQuery = {},
|
||||
): Promise<ProviderSummaryPageResponse> {
|
||||
const response = await client.get<ProviderSummaryPageResponse>(
|
||||
'/api/admin/providers/summary',
|
||||
{ params },
|
||||
)
|
||||
return response.data
|
||||
}
|
||||
|
||||
@@ -18,8 +41,10 @@ export async function getProvidersSummary(): Promise<ProviderWithEndpointsSummar
|
||||
* 获取单个 Provider 的详细信息
|
||||
*/
|
||||
export async function getProvider(providerId: string): Promise<ProviderWithEndpointsSummary> {
|
||||
const response = await client.get(`/api/admin/providers/${providerId}/summary`)
|
||||
return dedupedRequest(`providers:detail:${providerId}`, async () => {
|
||||
const response = await client.get<ProviderWithEndpointsSummary>(`/api/admin/providers/${providerId}/summary`)
|
||||
return response.data
|
||||
})
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -29,10 +54,11 @@ export async function updateProvider(
|
||||
providerId: string,
|
||||
data: Partial<{
|
||||
name: string
|
||||
provider_type: 'custom' | 'claude_code' | 'codex' | 'gemini_cli' | 'antigravity' | 'kiro'
|
||||
provider_type: 'custom' | 'vertex_ai' | 'claude_code' | 'codex' | 'gemini_cli' | 'antigravity' | 'kiro'
|
||||
description: string | null
|
||||
website: string
|
||||
provider_priority: number
|
||||
keep_priority_on_conversion: boolean
|
||||
billing_type: 'monthly_quota' | 'pay_as_you_go' | 'free_tier'
|
||||
monthly_quota_usd: number
|
||||
quota_reset_day: number
|
||||
@@ -48,6 +74,7 @@ export async function updateProvider(
|
||||
is_active: boolean
|
||||
claude_code_advanced: ClaudeCodeAdvancedConfig | null
|
||||
pool_advanced: PoolAdvancedConfig | null
|
||||
failover_rules: FailoverRulesConfig | null
|
||||
}>
|
||||
): Promise<ProviderWithEndpointsSummary> {
|
||||
const response = await client.patch(`/api/admin/providers/${providerId}`, data)
|
||||
@@ -60,7 +87,7 @@ export async function updateProvider(
|
||||
export async function createProvider(
|
||||
data: {
|
||||
name: string
|
||||
provider_type?: 'custom' | 'claude_code' | 'codex' | 'gemini_cli' | 'antigravity' | 'kiro'
|
||||
provider_type?: 'custom' | 'vertex_ai' | 'claude_code' | 'codex' | 'gemini_cli' | 'antigravity' | 'kiro'
|
||||
description?: string
|
||||
website?: string
|
||||
billing_type?: 'monthly_quota' | 'pay_as_you_go' | 'free_tier'
|
||||
@@ -77,6 +104,7 @@ export async function createProvider(
|
||||
proxy?: ProxyConfig | null
|
||||
claude_code_advanced?: ClaudeCodeAdvancedConfig | null
|
||||
pool_advanced?: PoolAdvancedConfig | null
|
||||
failover_rules?: FailoverRulesConfig | null
|
||||
}
|
||||
): Promise<{ id: string; name: string; message?: string }> {
|
||||
const response = await client.post('/api/admin/providers/', data)
|
||||
@@ -86,8 +114,36 @@ export async function createProvider(
|
||||
/**
|
||||
* 删除 Provider
|
||||
*/
|
||||
export async function deleteProvider(providerId: string): Promise<{ message: string }> {
|
||||
const response = await client.delete(`/api/admin/providers/${providerId}`)
|
||||
export interface ProviderDeleteSubmitResponse {
|
||||
task_id: string
|
||||
status: string
|
||||
message: string
|
||||
}
|
||||
|
||||
export interface ProviderDeleteTaskResponse {
|
||||
task_id: string
|
||||
provider_id: string
|
||||
status: string
|
||||
stage: string
|
||||
total_keys: number
|
||||
deleted_keys: number
|
||||
total_endpoints: number
|
||||
deleted_endpoints: number
|
||||
message: string
|
||||
}
|
||||
|
||||
export async function deleteProvider(providerId: string): Promise<ProviderDeleteSubmitResponse> {
|
||||
const response = await client.delete<ProviderDeleteSubmitResponse>(`/api/admin/providers/${providerId}`)
|
||||
return response.data
|
||||
}
|
||||
|
||||
export async function getProviderDeleteTask(
|
||||
providerId: string,
|
||||
taskId: string,
|
||||
): Promise<ProviderDeleteTaskResponse> {
|
||||
const response = await client.get<ProviderDeleteTaskResponse>(
|
||||
`/api/admin/providers/${providerId}/delete-task/${taskId}`,
|
||||
)
|
||||
return response.data
|
||||
}
|
||||
|
||||
@@ -122,7 +178,68 @@ export interface TestModelResponse {
|
||||
}
|
||||
|
||||
export async function testModel(data: TestModelRequest): Promise<TestModelResponse> {
|
||||
const response = await client.post('/api/admin/provider-query/test-model', data)
|
||||
const response = await client.post('/api/admin/provider-query/test-model', data, {
|
||||
timeout: 10 * 60 * 1000,
|
||||
})
|
||||
return response.data
|
||||
}
|
||||
|
||||
/**
|
||||
* 带故障转移的模型测试
|
||||
*/
|
||||
export interface TestModelFailoverRequest {
|
||||
provider_id: string
|
||||
mode: 'global' | 'direct'
|
||||
model_name: string
|
||||
api_format?: string
|
||||
endpoint_id?: string
|
||||
message?: string
|
||||
request_headers?: Record<string, unknown>
|
||||
request_body?: Record<string, unknown>
|
||||
request_id?: string
|
||||
concurrency?: number
|
||||
}
|
||||
|
||||
export interface TestAttemptDetail {
|
||||
candidate_index: number
|
||||
retry_index?: number
|
||||
endpoint_api_format: string
|
||||
endpoint_base_url: string
|
||||
key_name: string | null
|
||||
key_id: string
|
||||
auth_type: string
|
||||
effective_model?: string | null
|
||||
status: 'success' | 'failed' | 'skipped' | 'cancelled' | 'pending' | 'streaming' | 'stream_interrupted' | 'available' | 'unused'
|
||||
skip_reason?: string | null
|
||||
error_message?: string | null
|
||||
status_code?: number | null
|
||||
latency_ms?: number | null
|
||||
request_url?: string | null
|
||||
request_headers?: Record<string, unknown> | null
|
||||
request_body?: unknown
|
||||
response_headers?: Record<string, unknown> | null
|
||||
response_body?: unknown
|
||||
}
|
||||
|
||||
export interface TestModelFailoverResponse {
|
||||
success: boolean
|
||||
model: string
|
||||
provider: { id: string; name: string }
|
||||
attempts: TestAttemptDetail[]
|
||||
total_candidates: number
|
||||
total_attempts: number
|
||||
data?: Record<string, unknown> | null
|
||||
error?: string | null
|
||||
}
|
||||
|
||||
export async function testModelFailover(
|
||||
data: TestModelFailoverRequest,
|
||||
options: { signal?: AbortSignal } = {}
|
||||
): Promise<TestModelFailoverResponse> {
|
||||
const response = await client.post('/api/admin/provider-query/test-model-failover', data, {
|
||||
timeout: 10 * 60 * 1000,
|
||||
signal: options.signal,
|
||||
})
|
||||
return response.data
|
||||
}
|
||||
|
||||
@@ -169,6 +286,8 @@ export interface ProviderMappingPreviewResponse {
|
||||
export async function getProviderMappingPreview(
|
||||
providerId: string
|
||||
): Promise<ProviderMappingPreviewResponse> {
|
||||
const response = await client.get(`/api/admin/providers/${providerId}/mapping-preview`)
|
||||
return dedupedRequest(`providers:mapping-preview:${providerId}`, async () => {
|
||||
const response = await client.get<ProviderMappingPreviewResponse>(`/api/admin/providers/${providerId}/mapping-preview`)
|
||||
return response.data
|
||||
})
|
||||
}
|
||||
|
||||
@@ -5,6 +5,7 @@ export const API_FORMATS = {
|
||||
CLAUDE_CLI: 'claude:cli',
|
||||
OPENAI: 'openai:chat',
|
||||
OPENAI_CLI: 'openai:cli',
|
||||
OPENAI_COMPACT: 'openai:compact',
|
||||
OPENAI_VIDEO: 'openai:video',
|
||||
GEMINI: 'gemini:chat',
|
||||
GEMINI_CLI: 'gemini:cli',
|
||||
@@ -19,6 +20,7 @@ export const API_FORMAT_LABELS: Record<string, string> = {
|
||||
[API_FORMATS.CLAUDE_CLI]: 'Claude CLI',
|
||||
[API_FORMATS.OPENAI]: 'OpenAI Chat',
|
||||
[API_FORMATS.OPENAI_CLI]: 'OpenAI CLI',
|
||||
[API_FORMATS.OPENAI_COMPACT]: 'OpenAI Compact',
|
||||
[API_FORMATS.OPENAI_VIDEO]: 'OpenAI Video',
|
||||
[API_FORMATS.GEMINI]: 'Gemini Chat',
|
||||
[API_FORMATS.GEMINI_CLI]: 'Gemini CLI',
|
||||
@@ -28,6 +30,7 @@ export const API_FORMAT_LABELS: Record<string, string> = {
|
||||
CLAUDE_CLI: 'Claude CLI',
|
||||
OPENAI: 'OpenAI Chat',
|
||||
OPENAI_CLI: 'OpenAI CLI',
|
||||
OPENAI_COMPACT: 'OpenAI Compact',
|
||||
OPENAI_VIDEO: 'OpenAI Video',
|
||||
GEMINI: 'Gemini Chat',
|
||||
GEMINI_CLI: 'Gemini CLI',
|
||||
@@ -38,6 +41,7 @@ export const API_FORMAT_LABELS: Record<string, string> = {
|
||||
export const API_FORMAT_SHORT: Record<string, string> = {
|
||||
[API_FORMATS.OPENAI]: 'O',
|
||||
[API_FORMATS.OPENAI_CLI]: 'OC',
|
||||
[API_FORMATS.OPENAI_COMPACT]: 'OCP',
|
||||
[API_FORMATS.OPENAI_VIDEO]: 'OV',
|
||||
[API_FORMATS.CLAUDE]: 'C',
|
||||
[API_FORMATS.CLAUDE_CLI]: 'CC',
|
||||
@@ -47,6 +51,7 @@ export const API_FORMAT_SHORT: Record<string, string> = {
|
||||
// legacy 兼容(仅用于展示历史数据)
|
||||
OPENAI: 'O',
|
||||
OPENAI_CLI: 'OC',
|
||||
OPENAI_COMPACT: 'OCP',
|
||||
OPENAI_VIDEO: 'OV',
|
||||
CLAUDE: 'C',
|
||||
CLAUDE_CLI: 'CC',
|
||||
@@ -59,6 +64,7 @@ export const API_FORMAT_SHORT: Record<string, string> = {
|
||||
export const API_FORMAT_ORDER: string[] = [
|
||||
API_FORMATS.OPENAI,
|
||||
API_FORMATS.OPENAI_CLI,
|
||||
API_FORMATS.OPENAI_COMPACT,
|
||||
API_FORMATS.OPENAI_VIDEO,
|
||||
API_FORMATS.CLAUDE,
|
||||
API_FORMATS.CLAUDE_CLI,
|
||||
@@ -67,6 +73,62 @@ export const API_FORMAT_ORDER: string[] = [
|
||||
API_FORMATS.GEMINI_VIDEO,
|
||||
]
|
||||
|
||||
// Family 显示名称映射
|
||||
export const API_FORMAT_FAMILY_LABELS: Record<string, string> = {
|
||||
openai: 'OpenAI',
|
||||
claude: 'Claude',
|
||||
gemini: 'Gemini',
|
||||
}
|
||||
|
||||
// Kind 显示名称映射
|
||||
export const API_FORMAT_KIND_LABELS: Record<string, string> = {
|
||||
chat: 'Chat',
|
||||
cli: 'CLI',
|
||||
compact: 'Compact',
|
||||
video: 'Video',
|
||||
}
|
||||
|
||||
// Family 排序顺序
|
||||
const FAMILY_ORDER = ['openai', 'claude', 'gemini']
|
||||
|
||||
// 工具函数:从 API 格式中提取 family 和 kind
|
||||
export function parseApiFormat(format: string): { family: string; kind: string } {
|
||||
const idx = format.indexOf(':')
|
||||
if (idx === -1) return { family: format.toLowerCase(), kind: '' }
|
||||
return { family: format.slice(0, idx).toLowerCase(), kind: format.slice(idx + 1).toLowerCase() }
|
||||
}
|
||||
|
||||
// 工具函数:按 family 分组并排序 API 格式数组
|
||||
export interface ApiFormatGroup {
|
||||
family: string
|
||||
label: string
|
||||
formats: string[]
|
||||
}
|
||||
|
||||
export function groupApiFormats(formats: string[]): ApiFormatGroup[] {
|
||||
const sorted = sortApiFormats(formats)
|
||||
const groups = new Map<string, string[]>()
|
||||
for (const f of sorted) {
|
||||
const { family } = parseApiFormat(f)
|
||||
if (!groups.has(family)) groups.set(family, [])
|
||||
groups.get(family)?.push(f)
|
||||
}
|
||||
return [...groups.entries()]
|
||||
.sort(([a], [b]) => {
|
||||
const ai = FAMILY_ORDER.indexOf(a)
|
||||
const bi = FAMILY_ORDER.indexOf(b)
|
||||
if (ai === -1 && bi === -1) return 0
|
||||
if (ai === -1) return 1
|
||||
if (bi === -1) return -1
|
||||
return ai - bi
|
||||
})
|
||||
.map(([family, fmts]) => ({
|
||||
family,
|
||||
label: API_FORMAT_FAMILY_LABELS[family] || family,
|
||||
formats: fmts,
|
||||
}))
|
||||
}
|
||||
|
||||
// 工具函数:将 API 格式签名转为友好显示名称
|
||||
export function formatApiFormat(format: string | null | undefined): string {
|
||||
if (!format) return '-'
|
||||
|
||||
@@ -1,3 +1,5 @@
|
||||
import type { ProviderKeyStatusSnapshot } from './statusSnapshot'
|
||||
|
||||
/**
|
||||
* 代理配置类型
|
||||
* 支持两种模式:
|
||||
@@ -12,6 +14,13 @@ export interface ProxyConfig {
|
||||
enabled?: boolean // 是否启用代理(false 时保留配置但不使用)
|
||||
}
|
||||
|
||||
export interface OAuthOrganizationInfo {
|
||||
id?: string | null
|
||||
title?: string | null
|
||||
is_default?: boolean | null
|
||||
role?: string | null
|
||||
}
|
||||
|
||||
/**
|
||||
* 请求头规则类型
|
||||
* - set: 设置/覆盖请求头
|
||||
@@ -35,8 +44,6 @@ export interface HeaderRuleRename {
|
||||
to: string
|
||||
}
|
||||
|
||||
export type HeaderRule = HeaderRuleSet | HeaderRuleDrop | HeaderRuleRename
|
||||
|
||||
/**
|
||||
* 请求体规则类型
|
||||
* - set: 设置/覆盖字段
|
||||
@@ -129,10 +136,28 @@ export type BodyRuleConditionOp =
|
||||
| 'exists' | 'not_exists'
|
||||
| 'in' | 'type_is'
|
||||
|
||||
export interface BodyRuleCondition {
|
||||
export interface BodyRuleConditionLeaf {
|
||||
path: string
|
||||
op: BodyRuleConditionOp
|
||||
value?: unknown // exists / not_exists 不需要 value
|
||||
source?: 'original' | 'current'
|
||||
}
|
||||
|
||||
export interface BodyRuleConditionAll {
|
||||
all: BodyRuleCondition[]
|
||||
}
|
||||
|
||||
export interface BodyRuleConditionAny {
|
||||
any: BodyRuleCondition[]
|
||||
}
|
||||
|
||||
export type BodyRuleCondition =
|
||||
| BodyRuleConditionLeaf
|
||||
| BodyRuleConditionAll
|
||||
| BodyRuleConditionAny
|
||||
|
||||
export type HeaderRule = (HeaderRuleSet | HeaderRuleDrop | HeaderRuleRename) & {
|
||||
condition?: BodyRuleCondition
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -204,7 +229,7 @@ export interface EndpointAPIKey {
|
||||
api_formats: string[] // 支持的 endpoint signature 列表(如 "openai:chat")
|
||||
api_key_masked: string
|
||||
api_key_plain?: string | null
|
||||
auth_type: 'api_key' | 'vertex_ai' | 'oauth' // 认证类型(必返回)
|
||||
auth_type: 'api_key' | 'service_account' | 'oauth' // 认证类型(必返回)
|
||||
name: string // 密钥名称(必填,用于识别)
|
||||
rate_multipliers?: Record<string, number> | null // 按 endpoint signature 的成本倍率
|
||||
internal_priority: number // Key 内部优先级
|
||||
@@ -264,8 +289,13 @@ export interface EndpointAPIKey {
|
||||
oauth_email?: string | null // OAuth 授权的邮箱
|
||||
oauth_plan_type?: string | null // Codex 订阅类型: plus/free/team/enterprise
|
||||
oauth_account_id?: string | null // Codex ChatGPT 账号 ID
|
||||
oauth_invalid_at?: number | null // OAuth Token 失效时间(Unix 时间戳)
|
||||
oauth_invalid_reason?: string | null // OAuth Token 失效原因
|
||||
oauth_account_user_id?: string | null // Codex ChatGPT account-user 联合 ID
|
||||
oauth_account_name?: string | null
|
||||
oauth_organizations?: OAuthOrganizationInfo[] | null // OAuth 关联组织/工作区摘要
|
||||
oauth_temporary?: boolean | null // 是否为仅 Access Token 导入的临时 OAuth 账号
|
||||
oauth_invalid_at?: number | null // 兼容字段;优先使用 status_snapshot.oauth
|
||||
oauth_invalid_reason?: string | null // 兼容字段;优先使用 status_snapshot.oauth
|
||||
status_snapshot?: ProviderKeyStatusSnapshot | null
|
||||
// 上游元数据(由上游响应采集,如 Codex 额度信息 / Antigravity 配额信息)
|
||||
upstream_metadata?: UpstreamMetadata | null
|
||||
// Key 级别代理配置(覆盖 Provider 级别代理)
|
||||
@@ -284,10 +314,6 @@ export interface CodexUpstreamMetadata {
|
||||
secondary_reset_seconds?: number // 5H限额重置剩余秒数
|
||||
secondary_reset_at?: number // 5H限额重置时间(Unix 时间戳)
|
||||
secondary_window_minutes?: number // 5H限额窗口大小(分钟)
|
||||
code_review_used_percent?: number // 代码审查限额使用百分比
|
||||
code_review_reset_seconds?: number // 代码审查限额重置剩余秒数
|
||||
code_review_reset_at?: number // 代码审查限额重置时间(Unix 时间戳)
|
||||
code_review_window_minutes?: number // 代码审查限额窗口大小(分钟)
|
||||
has_credits?: boolean // 是否有积分
|
||||
credits_balance?: number // 积分余额
|
||||
}
|
||||
@@ -351,7 +377,7 @@ export interface EndpointAPIKeyUpdate {
|
||||
api_formats?: string[] // 支持的 API 格式列表
|
||||
name?: string
|
||||
api_key?: string // 仅在需要更新时提供
|
||||
auth_type?: 'api_key' | 'vertex_ai' | 'oauth' // 认证类型
|
||||
auth_type?: 'api_key' | 'service_account' | 'oauth' // 认证类型
|
||||
auth_config?: Record<string, unknown> // 认证配置(Vertex AI Service Account JSON)
|
||||
rate_multipliers?: Record<string, number> | null // 按 API 格式的成本倍率
|
||||
internal_priority?: number
|
||||
@@ -440,14 +466,12 @@ export interface PublicEndpointStatusMonitorResponse {
|
||||
formats: PublicEndpointStatusMonitor[]
|
||||
}
|
||||
|
||||
export type ProviderType = 'custom' | 'claude_code' | 'codex' | 'gemini_cli' | 'antigravity' | 'kiro'
|
||||
export type ProviderType = 'custom' | 'claude_code' | 'codex' | 'gemini_cli' | 'antigravity' | 'kiro' | 'vertex_ai'
|
||||
|
||||
export interface ClaudeCodeAdvancedConfig {
|
||||
// 会话数量控制:null/undefined 表示不限制
|
||||
max_sessions?: number | null
|
||||
session_idle_timeout_minutes?: number | null
|
||||
// TLS 指纹模拟(模拟 Node.js/Claude Code 客户端指纹)
|
||||
enable_tls_fingerprint?: boolean
|
||||
// 会话 ID 伪装(固定 metadata.user_id 中 session 片段)
|
||||
session_id_masking_enabled?: boolean
|
||||
// Cache TTL 统一(强制所有 cache_control 使用相同 TTL 类型)
|
||||
@@ -457,10 +481,29 @@ export interface ClaudeCodeAdvancedConfig {
|
||||
cli_only_enabled?: boolean
|
||||
}
|
||||
|
||||
export interface SchedulingPresetItem {
|
||||
preset: string
|
||||
enabled: boolean
|
||||
mode?: string | null
|
||||
}
|
||||
|
||||
export interface PoolAdvancedConfig {
|
||||
global_priority?: number | null
|
||||
sticky_session_ttl_seconds?: number | null
|
||||
load_threshold_percent?: number | null
|
||||
// 旧字段(兼容读取)
|
||||
lru_enabled?: boolean
|
||||
scheduling_mode?: 'lru' | 'multi_score' | null
|
||||
// 新格式:对象列表;旧格式:字符串列表
|
||||
scheduling_presets?: SchedulingPresetItem[] | string[] | null
|
||||
scoring_weights?: {
|
||||
lru?: number
|
||||
latency?: number
|
||||
health?: number
|
||||
cost_remaining?: number
|
||||
} | null
|
||||
latency_window_seconds?: number | null
|
||||
latency_sample_limit?: number | null
|
||||
cost_window_seconds?: number | null
|
||||
cost_limit_per_key_tokens?: number | null
|
||||
cost_soft_threshold_percent?: number | null
|
||||
@@ -469,6 +512,21 @@ export interface PoolAdvancedConfig {
|
||||
proactive_refresh_seconds?: number | null
|
||||
health_policy_enabled?: boolean
|
||||
unschedulable_rules?: Array<Record<string, unknown>> | null
|
||||
batch_concurrency?: number | null
|
||||
probing_enabled?: boolean
|
||||
probing_interval_minutes?: number | null
|
||||
auto_remove_banned_keys?: boolean
|
||||
}
|
||||
|
||||
export interface FailoverRuleItem {
|
||||
pattern: string
|
||||
description?: string
|
||||
status_codes?: number[]
|
||||
}
|
||||
|
||||
export interface FailoverRulesConfig {
|
||||
success_failover_patterns: FailoverRuleItem[]
|
||||
error_stop_patterns: FailoverRuleItem[]
|
||||
}
|
||||
|
||||
export interface ProviderWithEndpointsSummary {
|
||||
@@ -506,6 +564,7 @@ export interface ProviderWithEndpointsSummary {
|
||||
endpoint_health_details: EndpointHealthDetail[]
|
||||
claude_code_advanced?: ClaudeCodeAdvancedConfig | null
|
||||
pool_advanced?: PoolAdvancedConfig | null
|
||||
failover_rules?: FailoverRulesConfig | null
|
||||
ops_configured: boolean // 是否配置了扩展操作(余额监控等)
|
||||
ops_architecture_id?: string // 扩展操作使用的架构 ID(如 cubence, anyrouter)
|
||||
created_at: string
|
||||
@@ -536,6 +595,7 @@ export interface HealthSummary {
|
||||
total: number
|
||||
active: number
|
||||
unhealthy: number
|
||||
circuit_open?: number
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,36 @@
|
||||
export interface OAuthStatusSnapshot {
|
||||
code: 'none' | 'valid' | 'expiring' | 'expired' | 'invalid' | 'check_failed'
|
||||
label?: string | null
|
||||
reason?: string | null
|
||||
expires_at?: number | null
|
||||
invalid_at?: number | null
|
||||
source?: string | null
|
||||
requires_reauth?: boolean
|
||||
expiring_soon?: boolean
|
||||
}
|
||||
|
||||
export interface AccountStatusSnapshot {
|
||||
code: string
|
||||
label?: string | null
|
||||
reason?: string | null
|
||||
blocked: boolean
|
||||
source?: string | null
|
||||
recoverable?: boolean
|
||||
}
|
||||
|
||||
export interface QuotaStatusSnapshot {
|
||||
code: 'unknown' | 'ok' | 'exhausted'
|
||||
label?: string | null
|
||||
reason?: string | null
|
||||
exhausted: boolean
|
||||
usage_ratio?: number | null
|
||||
updated_at?: number | null
|
||||
reset_seconds?: number | null
|
||||
plan_type?: string | null
|
||||
}
|
||||
|
||||
export interface ProviderKeyStatusSnapshot {
|
||||
oauth: OAuthStatusSnapshot
|
||||
account: AccountStatusSnapshot
|
||||
quota: QuotaStatusSnapshot
|
||||
}
|
||||
+38
-165
@@ -1,6 +1,9 @@
|
||||
import apiClient from './client'
|
||||
import type { ActivityHeatmap } from '@/types/activity'
|
||||
import type { TieredPricingConfig } from './endpoints/types'
|
||||
import type { BillingSummary } from './auth'
|
||||
import type { UserSession } from '@/types/session'
|
||||
|
||||
export type { UserSession }
|
||||
|
||||
export interface Profile {
|
||||
id: string // UUID
|
||||
@@ -8,9 +11,7 @@ export interface Profile {
|
||||
username: string
|
||||
role: string
|
||||
is_active: boolean
|
||||
quota_usd: number | null
|
||||
used_usd: number
|
||||
total_usd?: number // 累积消费总额
|
||||
billing: BillingSummary
|
||||
created_at: string
|
||||
updated_at?: string
|
||||
last_login_at?: string
|
||||
@@ -22,8 +23,6 @@ export interface Profile {
|
||||
export interface UserPreferences {
|
||||
avatar_url?: string
|
||||
bio?: string
|
||||
default_provider_id?: string // UUID
|
||||
default_provider?: Record<string, unknown>
|
||||
theme: string
|
||||
language: string
|
||||
timezone?: string
|
||||
@@ -34,88 +33,6 @@ export interface UserPreferences {
|
||||
}
|
||||
}
|
||||
|
||||
// 提供商配置接口
|
||||
export interface ProviderConfig {
|
||||
provider_id: string
|
||||
priority: number // 优先级(越高越优先)
|
||||
weight: number // 负载均衡权重
|
||||
enabled: boolean // 是否启用
|
||||
}
|
||||
|
||||
// 使用记录接口
|
||||
export interface UsageRecordDetail {
|
||||
id: string
|
||||
provider: string
|
||||
model: string
|
||||
input_tokens: number
|
||||
output_tokens: number
|
||||
total_tokens: number
|
||||
cost: number // 官方费率
|
||||
actual_cost?: number // 倍率消耗(仅管理员可见)
|
||||
rate_multiplier?: number // 成本倍率(仅管理员可见)
|
||||
response_time_ms?: number
|
||||
is_stream: boolean
|
||||
created_at: string
|
||||
cache_creation_input_tokens?: number
|
||||
cache_read_input_tokens?: number
|
||||
status_code: number
|
||||
error_message?: string
|
||||
input_price_per_1m: number
|
||||
output_price_per_1m: number
|
||||
cache_creation_price_per_1m?: number
|
||||
cache_read_price_per_1m?: number
|
||||
price_per_request?: number // 按次计费价格
|
||||
api_key?: {
|
||||
id: string
|
||||
name: string
|
||||
display: string
|
||||
}
|
||||
}
|
||||
|
||||
// 模型统计接口
|
||||
export interface ModelSummary {
|
||||
model: string
|
||||
requests: number
|
||||
input_tokens: number
|
||||
output_tokens: number
|
||||
total_tokens: number
|
||||
total_cost_usd: number
|
||||
actual_total_cost_usd?: number // 倍率消耗(仅管理员可见)
|
||||
}
|
||||
|
||||
// 提供商统计接口
|
||||
export interface ProviderSummary {
|
||||
provider: string
|
||||
requests: number
|
||||
total_tokens: number
|
||||
total_cost_usd: number
|
||||
success_rate: number | null
|
||||
avg_response_time_ms: number | null
|
||||
}
|
||||
|
||||
// 使用统计响应接口
|
||||
export interface UsageResponse {
|
||||
total_requests: number
|
||||
total_input_tokens: number
|
||||
total_output_tokens: number
|
||||
total_tokens: number
|
||||
total_cost: number // 官方费率
|
||||
total_actual_cost?: number // 倍率消耗(仅管理员可见)
|
||||
avg_response_time: number
|
||||
quota_usd: number | null
|
||||
used_usd: number
|
||||
summary_by_model: ModelSummary[]
|
||||
summary_by_provider?: ProviderSummary[]
|
||||
pagination?: {
|
||||
total: number
|
||||
limit: number
|
||||
offset: number
|
||||
has_more: boolean
|
||||
}
|
||||
records: UsageRecordDetail[]
|
||||
activity_heatmap?: ActivityHeatmap | null
|
||||
}
|
||||
|
||||
export interface ApiKey {
|
||||
id: string // UUID
|
||||
name: string
|
||||
@@ -127,7 +44,7 @@ export interface ApiKey {
|
||||
created_at: string
|
||||
total_requests?: number
|
||||
total_cost_usd?: number
|
||||
allowed_providers?: ProviderConfig[]
|
||||
rate_limit?: number | null
|
||||
force_capabilities?: Record<string, boolean> | null // 强制能力配置
|
||||
}
|
||||
|
||||
@@ -160,14 +77,36 @@ export const meApi = {
|
||||
return response.data
|
||||
},
|
||||
|
||||
async listSessions(): Promise<UserSession[]> {
|
||||
const response = await apiClient.get<UserSession[]>('/api/users/me/sessions')
|
||||
return response.data
|
||||
},
|
||||
|
||||
async updateSessionLabel(sessionId: string, deviceLabel: string): Promise<UserSession> {
|
||||
const response = await apiClient.patch<UserSession>(`/api/users/me/sessions/${sessionId}`, {
|
||||
device_label: deviceLabel,
|
||||
})
|
||||
return response.data
|
||||
},
|
||||
|
||||
async revokeSession(sessionId: string): Promise<{ message: string }> {
|
||||
const response = await apiClient.delete(`/api/users/me/sessions/${sessionId}`)
|
||||
return response.data
|
||||
},
|
||||
|
||||
async revokeOtherSessions(): Promise<{ message: string; revoked_count: number }> {
|
||||
const response = await apiClient.delete('/api/users/me/sessions/others')
|
||||
return response.data
|
||||
},
|
||||
|
||||
// API密钥管理
|
||||
async getApiKeys(): Promise<ApiKey[]> {
|
||||
const response = await apiClient.get<ApiKey[]>('/api/users/me/api-keys')
|
||||
return response.data
|
||||
},
|
||||
|
||||
async createApiKey(name: string): Promise<ApiKey> {
|
||||
const response = await apiClient.post<ApiKey>('/api/users/me/api-keys', { name })
|
||||
async createApiKey(data: { name: string; rate_limit?: number }): Promise<ApiKey> {
|
||||
const response = await apiClient.post<ApiKey>('/api/users/me/api-keys', data)
|
||||
return response.data
|
||||
},
|
||||
|
||||
@@ -197,48 +136,14 @@ export const meApi = {
|
||||
return response.data
|
||||
},
|
||||
|
||||
// 使用统计
|
||||
async getUsage(params?: {
|
||||
start_date?: string
|
||||
end_date?: string
|
||||
preset?: string
|
||||
timezone?: string
|
||||
tz_offset_minutes?: number
|
||||
search?: string // 通用搜索:密钥名、模型名
|
||||
limit?: number
|
||||
offset?: number
|
||||
}): Promise<UsageResponse> {
|
||||
const response = await apiClient.get<UsageResponse>('/api/users/me/usage', { params })
|
||||
return response.data
|
||||
},
|
||||
|
||||
// 获取活跃请求状态(用于轮询更新)
|
||||
async getActiveRequests(ids?: string): Promise<{
|
||||
requests: Array<{
|
||||
id: string
|
||||
status: 'pending' | 'streaming' | 'completed' | 'failed' | 'cancelled'
|
||||
input_tokens: number
|
||||
output_tokens: number
|
||||
cache_creation_input_tokens?: number | null
|
||||
cache_read_input_tokens?: number | null
|
||||
cost: number
|
||||
actual_cost?: number | null
|
||||
rate_multiplier?: number | null
|
||||
response_time_ms: number | null
|
||||
first_byte_time_ms: number | null
|
||||
api_format?: string | null
|
||||
endpoint_api_format?: string | null
|
||||
has_format_conversion?: boolean | null
|
||||
}>
|
||||
}> {
|
||||
const params = ids ? { ids } : {}
|
||||
const response = await apiClient.get('/api/users/me/usage/active', { params })
|
||||
return response.data
|
||||
},
|
||||
|
||||
// 获取可用的提供商
|
||||
async getAvailableProviders(): Promise<Array<Record<string, unknown>>> {
|
||||
const response = await apiClient.get('/api/users/me/providers')
|
||||
async updateApiKey(
|
||||
keyId: string,
|
||||
data: { name?: string; rate_limit?: number | null }
|
||||
): Promise<ApiKey & { message: string }> {
|
||||
const response = await apiClient.put<ApiKey & { message: string }>(
|
||||
`/api/users/me/api-keys/${keyId}`,
|
||||
data
|
||||
)
|
||||
return response.data
|
||||
},
|
||||
|
||||
@@ -282,16 +187,6 @@ export const meApi = {
|
||||
return response.data
|
||||
},
|
||||
|
||||
// 提供商绑定管理相关方法已移除,改为直接从可用提供商中选择
|
||||
|
||||
// API密钥提供商关联
|
||||
async updateApiKeyProviders(keyId: string, data: {
|
||||
allowed_providers?: ProviderConfig[]
|
||||
}): Promise<{ message: string }> {
|
||||
const response = await apiClient.put(`/api/users/me/api-keys/${keyId}/providers`, data)
|
||||
return response.data
|
||||
},
|
||||
|
||||
// API密钥能力配置
|
||||
async updateApiKeyCapabilities(keyId: string, data: {
|
||||
force_capabilities?: Record<string, boolean> | null
|
||||
@@ -318,26 +213,4 @@ export const meApi = {
|
||||
return response.data
|
||||
},
|
||||
|
||||
// 获取请求间隔时间线(用于散点图)
|
||||
async getIntervalTimeline(params?: {
|
||||
hours?: number
|
||||
limit?: number
|
||||
}): Promise<{
|
||||
analysis_period_hours: number
|
||||
total_points: number
|
||||
points: Array<{ x: string; y: number; model?: string }>
|
||||
models?: string[]
|
||||
}> {
|
||||
const response = await apiClient.get('/api/users/me/usage/interval-timeline', { params })
|
||||
return response.data
|
||||
},
|
||||
|
||||
/**
|
||||
* 获取活跃度热力图数据(用户)
|
||||
* 后端已缓存5分钟
|
||||
*/
|
||||
async getActivityHeatmap(): Promise<ActivityHeatmap> {
|
||||
const response = await apiClient.get<ActivityHeatmap>('/api/users/me/usage/heatmap')
|
||||
return response.data
|
||||
}
|
||||
}
|
||||
|
||||
@@ -5,6 +5,7 @@ export interface ProxyNodeRemoteConfig {
|
||||
allowed_ports?: number[]
|
||||
log_level?: string
|
||||
heartbeat_interval?: number
|
||||
upgrade_to?: string | null
|
||||
}
|
||||
|
||||
export interface ProxyNode {
|
||||
@@ -37,6 +38,7 @@ export interface ProxyNode {
|
||||
failed_requests: number
|
||||
dns_failures: number
|
||||
stream_errors: number
|
||||
proxy_metadata: Record<string, unknown> | null
|
||||
created_at: string
|
||||
updated_at: string
|
||||
}
|
||||
@@ -109,6 +111,14 @@ export const proxyNodesApi = {
|
||||
return response.data
|
||||
},
|
||||
|
||||
async batchUpgrade(version: string): Promise<{ version: string; updated: number; skipped: number; node_ids: string[] }> {
|
||||
const response = await apiClient.post<{ version: string; updated: number; skipped: number; node_ids: string[] }>(
|
||||
'/api/admin/proxy-nodes/upgrade',
|
||||
{ version }
|
||||
)
|
||||
return response.data
|
||||
},
|
||||
|
||||
async testProxyUrl(data: { proxy_url: string; username?: string; password?: string }): Promise<ProxyNodeTestResult> {
|
||||
const response = await apiClient.post<ProxyNodeTestResult>('/api/admin/proxy-nodes/test-url', data)
|
||||
return response.data
|
||||
|
||||
@@ -0,0 +1,161 @@
|
||||
import apiClient from './client'
|
||||
|
||||
export interface VideoBilling {
|
||||
task_type: 'video' | 'image' | 'audio'
|
||||
duration_seconds?: number
|
||||
resolution?: string
|
||||
video_price_per_second?: number
|
||||
video_cost?: number
|
||||
cost?: number
|
||||
rule_name?: string
|
||||
expression?: string
|
||||
status?: string
|
||||
}
|
||||
|
||||
export interface RequestDetail {
|
||||
id: string
|
||||
request_id: string
|
||||
user: {
|
||||
id: string | null
|
||||
username: string | null
|
||||
email: string | null
|
||||
}
|
||||
api_key: {
|
||||
id: string | null
|
||||
name: string | null
|
||||
display: string | null
|
||||
}
|
||||
provider_api_key: {
|
||||
id: string | null
|
||||
name: string | null
|
||||
}
|
||||
provider: string
|
||||
api_format?: string
|
||||
model: string
|
||||
target_model?: string | null
|
||||
tokens: {
|
||||
input: number
|
||||
output: number
|
||||
total: number
|
||||
}
|
||||
cost: {
|
||||
input: number
|
||||
output: number
|
||||
total: number
|
||||
}
|
||||
input_tokens?: number
|
||||
output_tokens?: number
|
||||
total_tokens?: number
|
||||
cache_creation_input_tokens?: number
|
||||
cache_creation_input_tokens_5m?: number
|
||||
cache_creation_input_tokens_1h?: number
|
||||
cache_read_input_tokens?: number
|
||||
input_cost?: number
|
||||
output_cost?: number
|
||||
total_cost?: number
|
||||
cache_creation_cost?: number
|
||||
cache_read_cost?: number
|
||||
request_cost?: number
|
||||
input_price_per_1m?: number
|
||||
output_price_per_1m?: number
|
||||
cache_creation_price_per_1m?: number
|
||||
cache_read_price_per_1m?: number
|
||||
price_per_request?: number
|
||||
request_type: string
|
||||
is_stream: boolean
|
||||
status_code: number
|
||||
status?: string
|
||||
error_message?: string
|
||||
response_time_ms: number
|
||||
created_at: string
|
||||
request_headers?: Record<string, unknown>
|
||||
request_body?: Record<string, unknown>
|
||||
provider_request_headers?: Record<string, unknown>
|
||||
provider_request_body?: Record<string, unknown>
|
||||
response_headers?: Record<string, unknown>
|
||||
client_response_headers?: Record<string, unknown>
|
||||
response_body?: Record<string, unknown>
|
||||
client_response_body?: Record<string, unknown>
|
||||
has_request_body?: boolean
|
||||
has_provider_request_body?: boolean
|
||||
has_response_body?: boolean
|
||||
has_client_response_body?: boolean
|
||||
metadata?: Record<string, unknown>
|
||||
tiered_pricing?: {
|
||||
total_input_context: number
|
||||
tier_index: number
|
||||
tier_count: number
|
||||
source?: 'provider' | 'global'
|
||||
current_tier: {
|
||||
up_to?: number | null
|
||||
input_price_per_1m: number
|
||||
output_price_per_1m: number
|
||||
cache_creation_price_per_1m?: number
|
||||
cache_read_price_per_1m?: number
|
||||
cache_ttl_pricing?: Array<{
|
||||
ttl_minutes: number
|
||||
cache_creation_price_per_1m?: number
|
||||
cache_read_price_per_1m?: number
|
||||
}>
|
||||
}
|
||||
tiers: Array<{
|
||||
up_to?: number | null
|
||||
input_price_per_1m: number
|
||||
output_price_per_1m: number
|
||||
cache_creation_price_per_1m?: number
|
||||
cache_read_price_per_1m?: number
|
||||
cache_ttl_pricing?: Array<{
|
||||
ttl_minutes: number
|
||||
cache_creation_price_per_1m?: number
|
||||
cache_read_price_per_1m?: number
|
||||
}>
|
||||
}>
|
||||
} | null
|
||||
video_billing?: VideoBilling | null
|
||||
}
|
||||
|
||||
export interface CurlData {
|
||||
url: string
|
||||
method: string
|
||||
headers: Record<string, string>
|
||||
body: Record<string, unknown>
|
||||
curl: string
|
||||
}
|
||||
|
||||
export interface ReplayRequest {
|
||||
provider_id?: string
|
||||
endpoint_id?: string
|
||||
api_key_id?: string
|
||||
body_override?: Record<string, unknown>
|
||||
}
|
||||
|
||||
export interface ReplayResponse {
|
||||
url: string
|
||||
provider: string
|
||||
status_code: number
|
||||
response_headers: Record<string, string>
|
||||
response_body: Record<string, unknown>
|
||||
response_time_ms: number
|
||||
}
|
||||
|
||||
export const requestDetailsApi = {
|
||||
async getRequestDetail(requestId: string, options: { includeBodies?: boolean } = {}): Promise<RequestDetail> {
|
||||
const response = await apiClient.get<RequestDetail>(`/api/admin/usage/${requestId}`, {
|
||||
params: { include_bodies: options.includeBodies ?? true },
|
||||
})
|
||||
return response.data
|
||||
},
|
||||
|
||||
async getCurlData(requestId: string): Promise<CurlData> {
|
||||
const response = await apiClient.get<CurlData>(`/api/admin/usage/${requestId}/curl`)
|
||||
return response.data
|
||||
},
|
||||
|
||||
async replayRequest(requestId: string, params?: ReplayRequest): Promise<ReplayResponse> {
|
||||
const response = await apiClient.post<ReplayResponse>(
|
||||
`/api/admin/usage/${requestId}/replay`,
|
||||
params || {},
|
||||
)
|
||||
return response.data
|
||||
},
|
||||
}
|
||||
@@ -12,8 +12,9 @@ export interface CandidateRecord {
|
||||
endpoint_name?: string // 端点显示名称(api_format)
|
||||
key_id?: string
|
||||
key_name?: string // 密钥名称
|
||||
key_account_label?: string // 更适合展示的测试账号标签(优先 OAuth 邮箱)
|
||||
key_preview?: string // 密钥脱敏预览(如 sk-***abc),OAuth 类型不返回
|
||||
key_auth_type?: string // 密钥认证类型(api_key, oauth, vertex_ai 等)
|
||||
key_auth_type?: string // 密钥认证类型(api_key, service_account, oauth 等)
|
||||
key_oauth_plan_type?: string // OAuth 账号套餐类型(free/plus/team/enterprise)
|
||||
key_capabilities?: Record<string, boolean> | null // Key 支持的能力
|
||||
required_capabilities?: Record<string, boolean> | null // 请求实际需要的能力标签
|
||||
@@ -55,8 +56,14 @@ export const requestTraceApi = {
|
||||
/**
|
||||
* 获取特定请求的完整追踪信息
|
||||
*/
|
||||
async getRequestTrace(requestId: string): Promise<RequestTrace> {
|
||||
const response = await apiClient.get<RequestTrace>(`/api/admin/monitoring/trace/${requestId}`)
|
||||
async getRequestTrace(
|
||||
requestId: string,
|
||||
options: { attemptedOnly?: boolean } = {},
|
||||
): Promise<RequestTrace> {
|
||||
const attemptedOnly = options.attemptedOnly ?? false
|
||||
const response = await apiClient.get<RequestTrace>(`/api/admin/monitoring/trace/${requestId}`, {
|
||||
params: { attempted_only: attemptedOnly },
|
||||
})
|
||||
return response.data
|
||||
},
|
||||
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user