mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-09 16:06:27 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
bf1c6211b8 | ||
|
|
d388fa344a | ||
|
|
eabb61b08c | ||
|
|
a5c3438927 | ||
|
|
b6948542bb | ||
|
|
0242bd3193 | ||
|
|
229ef48470 | ||
|
|
19c35ba901 | ||
|
|
3c06ae97e8 | ||
|
|
f34966e4c7 | ||
|
|
e9ebf45640 | ||
|
|
c00e5bd243 | ||
|
|
ab5f18a437 | ||
|
|
330b222359 | ||
|
|
d81cb2d4f0 | ||
|
|
8e2edb680d | ||
|
|
00879be47d | ||
|
|
34656a3f07 | ||
|
|
27cf037309 | ||
|
|
2198fa69ab | ||
|
|
95f6434443 | ||
|
|
8734b0d3d5 | ||
|
|
790d0d9628 | ||
|
|
3c562563c5 | ||
|
|
f19785d0e1 | ||
|
|
a3d098a87c | ||
|
|
449d7c4218 | ||
|
|
d199141a91 | ||
|
|
df83ffdafb | ||
|
|
cf03c56cfa | ||
|
|
3d1e33558a | ||
|
|
f314ebb3a2 | ||
|
|
31d6ef9585 | ||
|
|
fdfb77c4a2 | ||
|
|
a93e3853f7 |
@@ -76,7 +76,7 @@ jobs:
|
||||
${{ hashFiles('trivy-results.sarif') != '' &&
|
||||
(github.event_name != 'pull_request' ||
|
||||
github.event.pull_request.head.repo.full_name == github.repository) }}
|
||||
uses: github/codeql-action/upload-sarif@e4fba868fa4b1b91e1fdab776edc8cfbe6e9fb81 # v4.37.3
|
||||
uses: github/codeql-action/upload-sarif@f205ea1c3313d32999d8d6a48b4f6530d4437b38 # v4.37.4
|
||||
with:
|
||||
sarif_file: trivy-results.sarif
|
||||
category: runtime-container
|
||||
|
||||
@@ -1,9 +1,6 @@
|
||||
name: Release
|
||||
|
||||
on:
|
||||
push:
|
||||
tags:
|
||||
- 'v*'
|
||||
workflow_dispatch:
|
||||
inputs:
|
||||
tag:
|
||||
@@ -15,7 +12,7 @@ permissions:
|
||||
contents: read
|
||||
|
||||
concurrency:
|
||||
group: release-${{ github.event_name == 'workflow_dispatch' && inputs.tag || github.ref_name }}
|
||||
group: release-${{ inputs.tag }}
|
||||
cancel-in-progress: false
|
||||
|
||||
env:
|
||||
@@ -43,14 +40,14 @@ jobs:
|
||||
- name: Checkout
|
||||
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
|
||||
with:
|
||||
ref: ${{ github.event_name == 'workflow_dispatch' && inputs.tag || github.ref }}
|
||||
ref: ${{ inputs.tag }}
|
||||
fetch-depth: 0
|
||||
|
||||
- name: Read release metadata
|
||||
id: version
|
||||
shell: bash
|
||||
env:
|
||||
RELEASE_TAG: ${{ github.event_name == 'workflow_dispatch' && inputs.tag || github.ref_name }}
|
||||
RELEASE_TAG: ${{ inputs.tag }}
|
||||
run: |
|
||||
ruby <<'RUBY' >> "$GITHUB_OUTPUT"
|
||||
require "json"
|
||||
@@ -188,9 +185,9 @@ jobs:
|
||||
~/.cargo/registry
|
||||
~/.cargo/git
|
||||
backend/target
|
||||
key: cargo-${{ runner.os }}-${{ matrix.target }}-${{ hashFiles('backend/Cargo.lock') }}
|
||||
key: cargo-v2-${{ runner.os }}-${{ matrix.target }}-${{ hashFiles('backend/Cargo.lock', 'backend/rust-toolchain.toml', 'backend/**/Cargo.toml', 'backend/.cargo/config.toml') }}
|
||||
restore-keys: |
|
||||
cargo-${{ runner.os }}-${{ matrix.target }}-
|
||||
cargo-v2-${{ runner.os }}-${{ matrix.target }}-
|
||||
|
||||
- name: Build binary
|
||||
shell: bash
|
||||
@@ -253,9 +250,9 @@ jobs:
|
||||
~/.cargo/registry
|
||||
~/.cargo/git
|
||||
backend/target
|
||||
key: cargo-${{ runner.os }}-${{ matrix.target }}-${{ hashFiles('backend/Cargo.lock') }}
|
||||
key: cargo-v2-${{ runner.os }}-${{ matrix.target }}-${{ hashFiles('backend/Cargo.lock', 'backend/rust-toolchain.toml', 'backend/**/Cargo.toml', 'backend/.cargo/config.toml') }}
|
||||
restore-keys: |
|
||||
cargo-${{ runner.os }}-${{ matrix.target }}-
|
||||
cargo-v2-${{ runner.os }}-${{ matrix.target }}-
|
||||
|
||||
- name: Build binary
|
||||
shell: bash
|
||||
@@ -378,7 +375,7 @@ jobs:
|
||||
PY
|
||||
|
||||
- name: Generate release asset provenance
|
||||
uses: actions/attest@f7c74d28b9d84cb8768d0b8ca14a4bac6ef463e6 # v4.2.0
|
||||
uses: actions/attest@508db95dd578ae2727ebd6217d5ba78e4fbda05d # v4.2.1
|
||||
with:
|
||||
subject-path: dist/release/*
|
||||
|
||||
@@ -454,7 +451,7 @@ jobs:
|
||||
uses: docker/setup-buildx-action@bb05f3f5519dd87d3ba754cc423b652a5edd6d2c # v4.2.0
|
||||
|
||||
- name: Login to GHCR
|
||||
uses: docker/login-action@abd2ef45e78c5afb21d64d4ca52ee8550d9572c7 # v4.5.1
|
||||
uses: docker/login-action@dbcb813823bdd20940b903addbd779551569679f # v4.6.0
|
||||
with:
|
||||
registry: ${{ env.REGISTRY }}
|
||||
username: ${{ github.actor }}
|
||||
@@ -471,8 +468,8 @@ jobs:
|
||||
push: true
|
||||
build-args: |
|
||||
CPR_PREBUILT_DIR=dist/docker
|
||||
cache-from: type=gha,scope=runtime-prebuilt-${{ matrix.os }}-${{ matrix.arch }}
|
||||
cache-to: type=gha,mode=max,scope=runtime-prebuilt-${{ matrix.os }}-${{ matrix.arch }}
|
||||
cache-from: type=registry,ref=${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:buildcache-${{ matrix.os }}-${{ matrix.arch }}
|
||||
cache-to: type=registry,ref=${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:buildcache-${{ matrix.os }}-${{ matrix.arch }},mode=min
|
||||
provenance: mode=max
|
||||
sbom: true
|
||||
labels: |
|
||||
@@ -525,7 +522,7 @@ jobs:
|
||||
uses: sigstore/cosign-installer@6f9f17788090df1f26f669e9d70d6ae9567deba6 # v4.1.2
|
||||
|
||||
- name: Login to GHCR
|
||||
uses: docker/login-action@abd2ef45e78c5afb21d64d4ca52ee8550d9572c7 # v4.5.1
|
||||
uses: docker/login-action@dbcb813823bdd20940b903addbd779551569679f # v4.6.0
|
||||
with:
|
||||
registry: ${{ env.REGISTRY }}
|
||||
username: ${{ github.actor }}
|
||||
@@ -584,7 +581,7 @@ jobs:
|
||||
echo "digest=${digest}" >> "$GITHUB_OUTPUT"
|
||||
|
||||
- name: Generate container provenance attestation
|
||||
uses: actions/attest@f7c74d28b9d84cb8768d0b8ca14a4bac6ef463e6 # v4.2.0
|
||||
uses: actions/attest@508db95dd578ae2727ebd6217d5ba78e4fbda05d # v4.2.1
|
||||
with:
|
||||
subject-name: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
|
||||
subject-digest: ${{ steps.manifest.outputs.digest }}
|
||||
|
||||
@@ -182,9 +182,10 @@ upgrade 与真实客户端 IP;不要暴露 PostgreSQL 或 Redis。
|
||||
release/publish <version>
|
||||
```
|
||||
|
||||
脚本会更新 `release/version.yaml`、创建约定提交与带注释的 `v<version>` tag,并原子推送分支和 tag。
|
||||
不要手工修改版本后单独打 tag。该脚本只发布仓库版本和触发 Release 构建,不会登录服务器,也不会
|
||||
改变任何运行实例;源码提交、仓库发版和实例升级是三个独立状态。
|
||||
脚本依赖已登录的 GitHub CLI(`gh auth login`)。它会更新 `release/version.yaml`、创建约定提交与带
|
||||
注释的 `v<version>` tag,原子推送分支和 tag,再从发布分支触发 Release 构建,使构建缓存可以跨版本
|
||||
复用。不要手工修改版本后单独打 tag。该脚本不会登录服务器,也不会改变任何运行实例;源码提交、
|
||||
仓库发版和实例升级是三个独立状态。
|
||||
|
||||
## 文档
|
||||
|
||||
|
||||
Generated
+115
-101
@@ -239,7 +239,7 @@ dependencies = [
|
||||
"bytes",
|
||||
"fastrand",
|
||||
"hex",
|
||||
"hmac 0.13.0",
|
||||
"hmac",
|
||||
"http 0.2.12",
|
||||
"http 1.4.2",
|
||||
"http-body 1.0.1",
|
||||
@@ -265,7 +265,7 @@ dependencies = [
|
||||
"bytes",
|
||||
"form_urlencoded",
|
||||
"hex",
|
||||
"hmac 0.13.0",
|
||||
"hmac",
|
||||
"http 0.2.12",
|
||||
"http 1.4.2",
|
||||
"percent-encoding",
|
||||
@@ -566,9 +566,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "base16ct"
|
||||
version = "0.2.0"
|
||||
version = "1.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4c7f02d4ea65f2c1853089ffd8d2787bdbc63de2f0d29dedbcf8ccdfa0ccd4cf"
|
||||
checksum = "fd307490d624467aa6f74b0eabb77633d1f758a7b25f12bceb0b22e08d9726f6"
|
||||
|
||||
[[package]]
|
||||
name = "base64"
|
||||
@@ -793,12 +793,6 @@ dependencies = [
|
||||
"yaml-rust2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "const-oid"
|
||||
version = "0.9.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c2459377285ad874054d797f3ccebf984978aa39129f6eafde5cdc8315b612f8"
|
||||
|
||||
[[package]]
|
||||
name = "const-oid"
|
||||
version = "0.10.2"
|
||||
@@ -850,6 +844,12 @@ version = "0.8.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b"
|
||||
|
||||
[[package]]
|
||||
name = "cpubits"
|
||||
version = "0.1.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "15b85f9c39137c3a891689859392b1bd49812121d0d61c9caf00d46ed5ce06ae"
|
||||
|
||||
[[package]]
|
||||
name = "cpufeatures"
|
||||
version = "0.2.17"
|
||||
@@ -940,12 +940,16 @@ checksum = "61803da095bee82a81bb1a452ecc25d3b2f1416d1897eb86430c6159ef717c17"
|
||||
|
||||
[[package]]
|
||||
name = "crypto-bigint"
|
||||
version = "0.5.5"
|
||||
version = "0.7.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0dc92fb57ca44df6db8059111ab3af99a63d5d0f8375d9972e319a379c6bab76"
|
||||
checksum = "1a52aa3fcda4e6302a9f48734f234d35d4721b96f8fe07d073f07ce9df4f0271"
|
||||
dependencies = [
|
||||
"generic-array",
|
||||
"rand_core 0.6.4",
|
||||
"cpubits",
|
||||
"ctutils",
|
||||
"getrandom 0.4.3",
|
||||
"hybrid-array",
|
||||
"num-traits",
|
||||
"rand_core 0.10.1",
|
||||
"subtle",
|
||||
"zeroize",
|
||||
]
|
||||
@@ -967,7 +971,9 @@ version = "0.2.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ce6e4c961d6cd6c9a86db418387425e8bdeaf05b3c8bc1411e6dca4c252f1453"
|
||||
dependencies = [
|
||||
"getrandom 0.4.3",
|
||||
"hybrid-array",
|
||||
"rand_core 0.10.1",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1007,6 +1013,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7d5515a3834141de9eafb9717ad39eea8247b5674e6066c404e8c4b365d2a29e"
|
||||
dependencies = [
|
||||
"cmov",
|
||||
"subtle",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1075,24 +1082,14 @@ version = "0.1.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "092966b41edc516079bdf31ec78a2e0588d1d0c08f78b91d8307215928642b2b"
|
||||
|
||||
[[package]]
|
||||
name = "der"
|
||||
version = "0.7.10"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e7c1832837b905bbfb5101e07cc24c8deddf52f93225eee6ead5f4d63d53ddcb"
|
||||
dependencies = [
|
||||
"const-oid 0.9.6",
|
||||
"pem-rfc7468",
|
||||
"zeroize",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "der"
|
||||
version = "0.8.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a69dedd701da44b0536442edf09c81a64b0ab97a7a4a5e3d1971f00027cbc63d"
|
||||
dependencies = [
|
||||
"const-oid 0.10.2",
|
||||
"const-oid",
|
||||
"pem-rfc7468",
|
||||
"zeroize",
|
||||
]
|
||||
|
||||
@@ -1109,7 +1106,6 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292"
|
||||
dependencies = [
|
||||
"block-buffer 0.10.4",
|
||||
"const-oid 0.9.6",
|
||||
"crypto-common 0.1.7",
|
||||
"subtle",
|
||||
]
|
||||
@@ -1121,7 +1117,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f1dd6dbb5841937940781866fa1281a1ff7bd3bf827091440879f9994983d5c2"
|
||||
dependencies = [
|
||||
"block-buffer 0.12.1",
|
||||
"const-oid 0.10.2",
|
||||
"const-oid",
|
||||
"crypto-common 0.2.2",
|
||||
"ctutils",
|
||||
]
|
||||
@@ -1160,16 +1156,17 @@ checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813"
|
||||
|
||||
[[package]]
|
||||
name = "ecdsa"
|
||||
version = "0.16.9"
|
||||
version = "0.17.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ee27f32b5c5292967d2d4a9d7f1e0b0aed2c15daded5a60300e4abb9d8020bca"
|
||||
checksum = "c0681a4fc24c767085329728d8dfba959af91228aa4610cca4f8ce317ba46ae0"
|
||||
dependencies = [
|
||||
"der 0.7.10",
|
||||
"digest 0.10.7",
|
||||
"der",
|
||||
"digest 0.11.3",
|
||||
"elliptic-curve",
|
||||
"rfc6979",
|
||||
"signature 2.2.0",
|
||||
"spki 0.7.3",
|
||||
"signature 3.0.0",
|
||||
"spki",
|
||||
"zeroize",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1178,7 +1175,7 @@ version = "3.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "29fcf32e6c73d1079f83ab4d782de2d81620346a5f38c6237a86a22f8368980a"
|
||||
dependencies = [
|
||||
"pkcs8 0.11.0",
|
||||
"pkcs8",
|
||||
"signature 3.0.0",
|
||||
]
|
||||
|
||||
@@ -1208,19 +1205,20 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "elliptic-curve"
|
||||
version = "0.13.8"
|
||||
version = "0.14.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b5e6043086bf7973472e0c7dff2142ea0b680d30e18d9cc40f267efbf222bd47"
|
||||
checksum = "9d65aa39b3a5c1c9c1b745c9a019234bb7a21b77abcb4f4d266d706e2d577d65"
|
||||
dependencies = [
|
||||
"base16ct",
|
||||
"crypto-bigint",
|
||||
"digest 0.10.7",
|
||||
"crypto-common 0.2.2",
|
||||
"digest 0.11.3",
|
||||
"ff",
|
||||
"generic-array",
|
||||
"group",
|
||||
"hybrid-array",
|
||||
"pem-rfc7468",
|
||||
"pkcs8 0.10.2",
|
||||
"rand_core 0.6.4",
|
||||
"pkcs8",
|
||||
"rand_core 0.10.1",
|
||||
"sec1",
|
||||
"subtle",
|
||||
"zeroize",
|
||||
@@ -1289,11 +1287,11 @@ checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6"
|
||||
|
||||
[[package]]
|
||||
name = "ff"
|
||||
version = "0.13.1"
|
||||
version = "0.14.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c0b50bfb653653f9ca9095b427bed08ab8d75a137839d9ad64eb11810d5b6393"
|
||||
checksum = "a1f686ab92a9fb0eaf188f6c6c87b89490baa6fdb0db4544ba4dc47f7942489f"
|
||||
dependencies = [
|
||||
"rand_core 0.6.4",
|
||||
"rand_core 0.10.1",
|
||||
"subtle",
|
||||
]
|
||||
|
||||
@@ -1677,12 +1675,12 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "group"
|
||||
version = "0.13.0"
|
||||
version = "0.14.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f0f9ef7462f7c099f518d754361858f86d8a07af53ba9af0fe635bbccb151a63"
|
||||
checksum = "7fd1a1c7a5206c5b7a3f5a0d7ccd3ff85d0c8f5133d62a02680255b0004af5f4"
|
||||
dependencies = [
|
||||
"ff",
|
||||
"rand_core 0.6.4",
|
||||
"rand_core 0.10.1",
|
||||
"subtle",
|
||||
]
|
||||
|
||||
@@ -1779,16 +1777,7 @@ version = "0.13.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4aaa26c720c68b866f2c96ef5c1264b3e6f473fe5d4ce61cd44bbe913e553018"
|
||||
dependencies = [
|
||||
"hmac 0.13.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hmac"
|
||||
version = "0.12.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6c49c37c09c17a53d937dfbb742eb3a961d65a994e6bcdcf37e7399d0cc8ab5e"
|
||||
dependencies = [
|
||||
"digest 0.10.7",
|
||||
"hmac",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1879,7 +1868,9 @@ version = "0.4.13"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "818356c5132c1fede50f837ca96afbe78ff42413047f4abb886217845e1b6c8c"
|
||||
dependencies = [
|
||||
"subtle",
|
||||
"typenum",
|
||||
"zeroize",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -2359,14 +2350,15 @@ checksum = "1a80800c0488c3a21695ea981a54918fbb37abf04f4d0720c453632255e2ff0e"
|
||||
|
||||
[[package]]
|
||||
name = "p256"
|
||||
version = "0.13.2"
|
||||
version = "0.14.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c9863ad85fa8f4460f9c48cb909d38a0d689dba1f6f6988a5e3e0d31071bcd4b"
|
||||
checksum = "d2c9239b2dbc807adbbe147e8cf72ea7450c3a0aabe62cb8e75ff4ec22e1f72a"
|
||||
dependencies = [
|
||||
"ecdsa",
|
||||
"elliptic-curve",
|
||||
"primefield",
|
||||
"primeorder",
|
||||
"sha2 0.10.9",
|
||||
"sha2 0.11.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -2427,9 +2419,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "pem-rfc7468"
|
||||
version = "0.7.0"
|
||||
version = "1.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "88b39c9bfcfc231068454382784bb460aae594343fb030d46e9f50a645418412"
|
||||
checksum = "a6305423e0e7738146434843d1694d621cce767262b2a86910beab705e4493d9"
|
||||
dependencies = [
|
||||
"base64ct",
|
||||
]
|
||||
@@ -2512,24 +2504,14 @@ version = "0.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184"
|
||||
|
||||
[[package]]
|
||||
name = "pkcs8"
|
||||
version = "0.10.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f950b2377845cebe5cf8b5165cb3cc1a5e0fa5cfa3e1f7f55707d8fd82e0a7b7"
|
||||
dependencies = [
|
||||
"der 0.7.10",
|
||||
"spki 0.7.3",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pkcs8"
|
||||
version = "0.11.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "451913da69c775a56034ea8d9003d27ee8948e12443eae7c038ba100a4f21cb7"
|
||||
dependencies = [
|
||||
"der 0.8.1",
|
||||
"spki 0.8.0",
|
||||
"der",
|
||||
"spki",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -2574,12 +2556,30 @@ dependencies = [
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "primeorder"
|
||||
version = "0.13.6"
|
||||
name = "primefield"
|
||||
version = "0.14.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "353e1ca18966c16d9deb1c69278edbc5f194139612772bd9537af60ac231e1e6"
|
||||
checksum = "c555a6e4eb7d4e158fcb028c835c3b8642206ddc279b5c6b202ef9a8bdb592f4"
|
||||
dependencies = [
|
||||
"crypto-bigint",
|
||||
"crypto-common 0.2.2",
|
||||
"ff",
|
||||
"rand_core 0.10.1",
|
||||
"subtle",
|
||||
"zeroize",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "primeorder"
|
||||
version = "0.14.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5c9f42978c78a00e3d68f69fc03e57a234debae69da4020a4fb588fcdcd07b06"
|
||||
dependencies = [
|
||||
"elliptic-curve",
|
||||
"once_cell",
|
||||
"primefield",
|
||||
"serdect",
|
||||
"wnaf",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -2833,9 +2833,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "redis"
|
||||
version = "1.4.1"
|
||||
version = "1.5.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b0b9503711b03773e43b31668c7b5bd279ee7cd9b7d18cff7c23a42cc1d08e5a"
|
||||
checksum = "3257df217f7eab0044627a268c9cc6cdb60c0c421c88f83ac41c4e31520b6b84"
|
||||
dependencies = [
|
||||
"arc-swap",
|
||||
"arcstr",
|
||||
@@ -2948,12 +2948,12 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "rfc6979"
|
||||
version = "0.4.0"
|
||||
version = "0.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f8dd2a808d456c4a54e300a23e9f5a67e122c3024119acbfd73e3bf664491cb2"
|
||||
checksum = "b4a459cddafb3fe76b31fd8f1108007566c40301feb64dc7b54656eb7388172b"
|
||||
dependencies = [
|
||||
"hmac 0.12.1",
|
||||
"subtle",
|
||||
"crypto-bigint",
|
||||
"hmac",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -3095,14 +3095,14 @@ checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49"
|
||||
|
||||
[[package]]
|
||||
name = "sec1"
|
||||
version = "0.7.3"
|
||||
version = "0.8.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d3e97a565f76233a6003f9f5c54be1d9c5bdfa3eccfb189469f11ec4901c47dc"
|
||||
checksum = "d56d437c2f19203ce5f7122e507831de96f3d2d4d3be5af44a0b0a09d8a80e4d"
|
||||
dependencies = [
|
||||
"base16ct",
|
||||
"der 0.7.10",
|
||||
"generic-array",
|
||||
"pkcs8 0.10.2",
|
||||
"ctutils",
|
||||
"der",
|
||||
"hybrid-array",
|
||||
"subtle",
|
||||
"zeroize",
|
||||
]
|
||||
@@ -3213,6 +3213,16 @@ dependencies = [
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serdect"
|
||||
version = "0.4.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "66cf8fedced2fcf12406bcb34223dffb92eaf34908ede12fed414c82b7f00b3e"
|
||||
dependencies = [
|
||||
"base16ct",
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "sha1"
|
||||
version = "0.10.7"
|
||||
@@ -3294,7 +3304,6 @@ version = "2.2.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "77549399552de45a898a580c1b41d445bf730df867cc44e6c0233bbc4b8329de"
|
||||
dependencies = [
|
||||
"digest 0.10.7",
|
||||
"rand_core 0.6.4",
|
||||
]
|
||||
|
||||
@@ -3303,6 +3312,10 @@ name = "signature"
|
||||
version = "3.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "28d567dcbaf0049cb8ac2608a76cd95ff9e4412e1899d389ee400918ca7537f5"
|
||||
dependencies = [
|
||||
"digest 0.11.3",
|
||||
"rand_core 0.10.1",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "simd-adler32"
|
||||
@@ -3368,16 +3381,6 @@ version = "0.10.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "023a211cb3138dbc438680b32560ad89f699977624c9f8dbb95a47d5b4c07dd3"
|
||||
|
||||
[[package]]
|
||||
name = "spki"
|
||||
version = "0.7.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d91ed6c858b01f942cd56b37a94b3e0a1798290327d1236e4d9cf4eaca44d29d"
|
||||
dependencies = [
|
||||
"base64ct",
|
||||
"der 0.7.10",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "spki"
|
||||
version = "0.8.0"
|
||||
@@ -3385,7 +3388,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1d9efca8738c78ee9484207732f728b1ef517bbb1833d6fc0879ca898a522f6f"
|
||||
dependencies = [
|
||||
"base64ct",
|
||||
"der 0.8.1",
|
||||
"der",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -3523,7 +3526,7 @@ dependencies = [
|
||||
"futures-util",
|
||||
"hex",
|
||||
"hkdf",
|
||||
"hmac 0.13.0",
|
||||
"hmac",
|
||||
"itoa",
|
||||
"log",
|
||||
"md-5",
|
||||
@@ -4493,6 +4496,17 @@ version = "0.57.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e"
|
||||
|
||||
[[package]]
|
||||
name = "wnaf"
|
||||
version = "0.14.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ab12e7090f27e2ffd9322651492942d50c2926094af30601e1964337db39daf1"
|
||||
dependencies = [
|
||||
"ff",
|
||||
"group",
|
||||
"hybrid-array",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "writeable"
|
||||
version = "0.6.3"
|
||||
|
||||
@@ -134,7 +134,7 @@ fn config_loader_should_reject_removed_fingerprint_section() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn config_loader_should_reject_invalid_codex_core_version() {
|
||||
fn config_loader_should_reject_invalid_codex_cli_version() {
|
||||
assert_rejected(valid_config().replace("codex_version: '0.146.0'", "codex_version: 'latest'"));
|
||||
}
|
||||
|
||||
|
||||
@@ -220,7 +220,8 @@ impl BackupTask {
|
||||
sha256: Some(artifact.sha256.clone()),
|
||||
..Default::default()
|
||||
};
|
||||
self.repository
|
||||
let Some(uploading_record) = self
|
||||
.repository
|
||||
.transition_status(
|
||||
&record.id,
|
||||
BackupStatus::Dumping,
|
||||
@@ -229,9 +230,12 @@ impl BackupTask {
|
||||
now,
|
||||
)
|
||||
.await
|
||||
.map_err(repo_error)?;
|
||||
.map_err(repo_error)?
|
||||
else {
|
||||
return Ok(());
|
||||
};
|
||||
self.upload_and_verify(
|
||||
record,
|
||||
&uploading_record,
|
||||
&artifact.path,
|
||||
artifact.size_bytes,
|
||||
&artifact.sha256,
|
||||
@@ -394,7 +398,8 @@ impl BackupTask {
|
||||
sha256: Some(artifact.sha256.clone()),
|
||||
..Default::default()
|
||||
};
|
||||
self.repository
|
||||
let Some(uploading_record) = self
|
||||
.repository
|
||||
.transition_status(
|
||||
&record.id,
|
||||
BackupStatus::Dumping,
|
||||
@@ -403,10 +408,13 @@ impl BackupTask {
|
||||
now,
|
||||
)
|
||||
.await
|
||||
.map_err(repo_error)?;
|
||||
.map_err(repo_error)?
|
||||
else {
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
self.upload_and_verify(
|
||||
record,
|
||||
&uploading_record,
|
||||
&artifact.path,
|
||||
artifact.size_bytes,
|
||||
&artifact.sha256,
|
||||
@@ -448,8 +456,12 @@ impl BackupTask {
|
||||
.await
|
||||
{
|
||||
let _ = store.delete_object(&config, &record.object_key).await;
|
||||
self.fail_task(record, error.code(), "上传到对象存储失败")
|
||||
.await?;
|
||||
self.fail_task(
|
||||
record,
|
||||
error.code(),
|
||||
format!("上传到对象存储失败:{}", error.message()),
|
||||
)
|
||||
.await?;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
@@ -475,19 +487,46 @@ impl BackupTask {
|
||||
);
|
||||
}
|
||||
}
|
||||
Ok(_) => {
|
||||
Ok(Some(remote)) => {
|
||||
let size_matches = remote.size_bytes == size_bytes;
|
||||
let sha256_matches = remote.sha256 == sha256;
|
||||
warn!(
|
||||
backup_id = %record.id,
|
||||
expected_size_bytes = size_bytes,
|
||||
remote_size_bytes = remote.size_bytes,
|
||||
sha256_matches,
|
||||
"备份远端对象校验不一致"
|
||||
);
|
||||
let _ = store.delete_object(&config, &record.object_key).await;
|
||||
self.fail_task(
|
||||
record,
|
||||
code::REMOTE_VERIFICATION_FAILED,
|
||||
"远端对象大小或校验值不一致",
|
||||
match (size_matches, sha256_matches) {
|
||||
(false, true) => "远端对象大小不一致",
|
||||
(true, false) => "远端对象 sha256 校验值不一致",
|
||||
(false, false) => "远端对象大小和 sha256 校验值均不一致",
|
||||
(true, true) => "远端对象校验状态异常",
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
Ok(None) => {
|
||||
warn!(backup_id = %record.id, "上传完成后未找到远端对象");
|
||||
self.fail_task(
|
||||
record,
|
||||
code::REMOTE_VERIFICATION_FAILED,
|
||||
"上传完成后未找到远端对象",
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
Err(error) => {
|
||||
let _ = store.delete_object(&config, &record.object_key).await;
|
||||
self.fail_task(record, error.code(), "远端对象校验失败")
|
||||
.await?;
|
||||
self.fail_task(
|
||||
record,
|
||||
error.code(),
|
||||
format!("远端对象校验失败:{}", error.message()),
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
@@ -500,10 +539,17 @@ impl BackupTask {
|
||||
error_code: &'static str,
|
||||
message: impl Into<String>,
|
||||
) -> Result<(), WorkerTaskError> {
|
||||
let message = message.into();
|
||||
warn!(
|
||||
backup_id = %record.id,
|
||||
error_code,
|
||||
error_message = %message,
|
||||
"备份任务失败"
|
||||
);
|
||||
let now = Utc::now();
|
||||
let update = StatusTransitionUpdate {
|
||||
error_code: Some(error_code.to_string()),
|
||||
error_message: Some(message.into()),
|
||||
error_message: Some(message),
|
||||
completed_at: Some(now),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
//! 多 Provider 账号目录与连接测试的公共事实。
|
||||
|
||||
use std::pin::Pin;
|
||||
use std::{pin::Pin, time::SystemTime};
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use futures::Stream;
|
||||
@@ -13,13 +13,146 @@ use super::{PageSize, Revision, observability::TimeRange};
|
||||
pub use gateway_core::engine::credential::AccountAvailability;
|
||||
|
||||
/// 管理页使用的归一化账号状态。
|
||||
///
|
||||
/// 只回答两个问题:是否可调度、是否需要人工干预。具体错误原因不进入此枚举,
|
||||
/// 由 [`AccountErrorReason`] 承载;新增原因只扩展 reason,不改变分类。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum AccountStatus {
|
||||
Active,
|
||||
Expired,
|
||||
/// 正常,可调度。
|
||||
Normal,
|
||||
/// 配额耗尽(持久化终态),等待额度恢复,无需干预。
|
||||
QuotaExhausted,
|
||||
/// 限流中(quota 快照级信号),自动恢复,无需干预。
|
||||
RateLimited,
|
||||
/// 管理员手动停用(`enabled = false`)。
|
||||
Disabled,
|
||||
/// 错误(过期 / 失效 / 封禁等),需要人工干预。具体原因见 [`AccountErrorReason`]。
|
||||
Error,
|
||||
}
|
||||
|
||||
impl AccountStatus {
|
||||
/// 返回稳定 wire 值。
|
||||
#[must_use]
|
||||
pub const fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::Normal => "normal",
|
||||
Self::QuotaExhausted => "quota_exhausted",
|
||||
Self::RateLimited => "rate_limited",
|
||||
Self::Disabled => "disabled",
|
||||
Self::Error => "error",
|
||||
}
|
||||
}
|
||||
|
||||
/// 解析稳定 wire 值。
|
||||
#[must_use]
|
||||
pub fn parse(value: &str) -> Option<Self> {
|
||||
match value {
|
||||
"normal" => Some(Self::Normal),
|
||||
"quota_exhausted" => Some(Self::QuotaExhausted),
|
||||
"rate_limited" => Some(Self::RateLimited),
|
||||
"disabled" => Some(Self::Disabled),
|
||||
"error" => Some(Self::Error),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 状态派生所需的最小事实;由 store / use_case 各自构造,派生规则唯一。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AccountStatusSignals {
|
||||
pub enabled: bool,
|
||||
pub availability: AccountAvailability,
|
||||
pub access_token_expired: bool,
|
||||
pub quota_limit_reached: bool,
|
||||
/// 429 临时限流(Redis 冷却)到期时间;`Some` 表示账号当前处于临时限流。
|
||||
pub rate_limited_until: Option<SystemTime>,
|
||||
}
|
||||
|
||||
impl AccountStatusSignals {
|
||||
/// 派生产出运营分类。终态优先于瞬时信号:停用 > 配额耗尽 > 错误 > 限流。
|
||||
#[must_use]
|
||||
pub const fn derive(self) -> AccountStatus {
|
||||
if !self.enabled {
|
||||
AccountStatus::Disabled
|
||||
} else if matches!(self.availability, AccountAvailability::QuotaExhausted) {
|
||||
AccountStatus::QuotaExhausted
|
||||
} else if self.access_token_expired
|
||||
|| matches!(
|
||||
self.availability,
|
||||
AccountAvailability::Expired
|
||||
| AccountAvailability::Invalid
|
||||
| AccountAvailability::Banned
|
||||
| AccountAvailability::Unknown
|
||||
)
|
||||
{
|
||||
AccountStatus::Error
|
||||
} else if self.rate_limited_until.is_some() {
|
||||
AccountStatus::RateLimited
|
||||
} else if self.quota_limit_reached {
|
||||
AccountStatus::QuotaExhausted
|
||||
} else {
|
||||
AccountStatus::Normal
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 错误分类下的具体原因;仅在 `AccountStatus::Error` 时存在。
|
||||
///
|
||||
/// 受控枚举用于展示与聚合,扩展不改变状态分类。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum AccountErrorReason {
|
||||
Expired,
|
||||
TokenExpired,
|
||||
Invalid,
|
||||
Banned,
|
||||
Unknown,
|
||||
}
|
||||
|
||||
impl AccountErrorReason {
|
||||
/// 返回稳定 wire 值。
|
||||
#[must_use]
|
||||
pub const fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::Expired => "expired",
|
||||
Self::TokenExpired => "token_expired",
|
||||
Self::Invalid => "invalid",
|
||||
Self::Banned => "banned",
|
||||
Self::Unknown => "unknown",
|
||||
}
|
||||
}
|
||||
|
||||
/// 解析稳定 wire 值。
|
||||
#[must_use]
|
||||
pub fn parse(value: &str) -> Option<Self> {
|
||||
match value {
|
||||
"expired" => Some(Self::Expired),
|
||||
"token_expired" => Some(Self::TokenExpired),
|
||||
"invalid" => Some(Self::Invalid),
|
||||
"banned" => Some(Self::Banned),
|
||||
"unknown" => Some(Self::Unknown),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 由信号派生错误原因;非错误分类(正常 / 配额耗尽 / 限流 / 停用)返回 `None`。
|
||||
///
|
||||
/// 与 [`AccountStatusSignals::derive`] 保持同一判定:`Error` 当且仅当此处返回 `Some`。
|
||||
#[must_use]
|
||||
pub fn derive_error_reason(signals: &AccountStatusSignals) -> Option<AccountErrorReason> {
|
||||
if !signals.enabled {
|
||||
return None;
|
||||
}
|
||||
match signals.availability {
|
||||
AccountAvailability::QuotaExhausted => None,
|
||||
AccountAvailability::Ready => signals
|
||||
.access_token_expired
|
||||
.then_some(AccountErrorReason::TokenExpired),
|
||||
AccountAvailability::Expired => Some(AccountErrorReason::Expired),
|
||||
AccountAvailability::Invalid => Some(AccountErrorReason::Invalid),
|
||||
AccountAvailability::Banned => Some(AccountErrorReason::Banned),
|
||||
AccountAvailability::Unknown => Some(AccountErrorReason::Unknown),
|
||||
}
|
||||
}
|
||||
|
||||
/// 账号列表排序字段。
|
||||
@@ -75,9 +208,9 @@ pub struct AccountRecord {
|
||||
pub next_refresh_at: Option<DateTime<Utc>>,
|
||||
pub enabled: bool,
|
||||
pub availability: AccountAvailability,
|
||||
pub availability_reason: Option<String>,
|
||||
pub cooldown_until: Option<DateTime<Utc>>,
|
||||
pub availability_observed_at: DateTime<Utc>,
|
||||
/// 最近一次非 ready 状态变更的原始原因(上游错误原文 / 原因码);恢复 ready 后清空。
|
||||
pub last_error_message: Option<String>,
|
||||
pub quota_observed_at: Option<DateTime<Utc>>,
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub updated_at: DateTime<Utc>,
|
||||
@@ -170,13 +303,15 @@ pub struct AccountPage {
|
||||
|
||||
/// 统一账号目录的全局状态计数,不受当前筛选和分页影响。
|
||||
///
|
||||
/// 配额耗尽是可恢复状态,始终单列;`unavailable` 只包含过期、禁用和封禁账号。
|
||||
/// 计数与 [`AccountStatus`] 一一对应,由 store 按派生状态聚合。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct AccountSummary {
|
||||
pub total: u64,
|
||||
pub active: u64,
|
||||
pub normal: u64,
|
||||
pub quota_exhausted: u64,
|
||||
pub unavailable: u64,
|
||||
pub rate_limited: u64,
|
||||
pub disabled: u64,
|
||||
pub error: u64,
|
||||
}
|
||||
|
||||
/// 账号启停写入命令。
|
||||
|
||||
@@ -443,14 +443,17 @@ pub struct RequestMetricPoint {
|
||||
///
|
||||
/// 首页概览的 `unavailable` 表示不可参与调度,包含配额耗尽及过期、禁用和封禁账号。
|
||||
///
|
||||
/// `quota_exhausted` 仍单列,以便状态详情展示和诊断。
|
||||
/// `quota_exhausted`(402 确认耗尽)与 `rate_limited`(429 限流中)单列,
|
||||
/// 以便状态详情展示和诊断。
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct AccountPoolMetrics {
|
||||
pub total: u64,
|
||||
pub enabled: u64,
|
||||
pub unavailable: u64,
|
||||
pub active: u64,
|
||||
pub rate_limited: u64,
|
||||
pub expired: u64,
|
||||
pub invalid: u64,
|
||||
pub quota_exhausted: u64,
|
||||
pub refreshing: Option<u64>,
|
||||
pub disabled: u64,
|
||||
|
||||
@@ -11,7 +11,10 @@ use gateway_core::{
|
||||
|
||||
use super::{
|
||||
AdminError, MutationActor, MutationContext, PageSize, Revision,
|
||||
accounts::{AccountAvailability, AccountRecord, AccountStatus, AccountSummary, AccountUsage},
|
||||
accounts::{
|
||||
AccountAvailability, AccountErrorReason, AccountRecord, AccountStatus, AccountSummary,
|
||||
AccountUsage,
|
||||
},
|
||||
};
|
||||
|
||||
/// Provider-owned JSON;公共层只搬运且 Debug 不输出值。
|
||||
@@ -153,8 +156,6 @@ pub struct PreparedCredentialCreate {
|
||||
pub next_refresh_at: Option<DateTime<Utc>>,
|
||||
pub enabled: bool,
|
||||
pub availability: AccountAvailability,
|
||||
pub availability_reason: Option<String>,
|
||||
pub cooldown_until: Option<DateTime<Utc>>,
|
||||
pub availability_observed_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
@@ -573,6 +574,7 @@ pub struct ProviderQuotaWindow {
|
||||
pub window_seconds: Option<u64>,
|
||||
pub used_percent: Option<f64>,
|
||||
pub reset_at: Option<DateTime<Utc>>,
|
||||
pub limit_reached: bool,
|
||||
pub local_usage: Option<AccountUsage>,
|
||||
pub provider_data: Option<ProviderDocument>,
|
||||
}
|
||||
@@ -591,6 +593,11 @@ pub struct ProviderQuota {
|
||||
pub observed_at: Option<DateTime<Utc>>,
|
||||
pub refresh_token_expires_at: Option<DateTime<Utc>>,
|
||||
pub windows: Vec<ProviderQuotaWindow>,
|
||||
/// 快照级限流事实(顶层或任一窗口触顶);Admin 状态派生用,前端不再遍历窗口猜测。
|
||||
pub limit_reached: bool,
|
||||
/// 429 临时限流(Redis 冷却)到期时间;`Some` 时账号处于临时限流,到期自动恢复。
|
||||
/// 与 `limit_reached`(配额耗尽)是独立维度,不污染额度窗口。
|
||||
pub rate_limited_until: Option<DateTime<Utc>>,
|
||||
pub provider_data: Option<ProviderDocument>,
|
||||
}
|
||||
|
||||
@@ -687,6 +694,8 @@ impl fmt::Debug for ProviderExport {
|
||||
pub struct AccountDirectoryItem {
|
||||
pub account: AccountRecord,
|
||||
pub status: AccountStatus,
|
||||
/// `status == Error` 时的具体原因;其余状态为 `None`。
|
||||
pub error_reason: Option<AccountErrorReason>,
|
||||
pub usage: Option<super::accounts::AccountUsage>,
|
||||
pub quota: ProviderQuota,
|
||||
}
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
//! 统一账号目录与跨 Provider 动态分派。
|
||||
|
||||
use std::{collections::BTreeMap, sync::Arc};
|
||||
use std::{collections::BTreeMap, sync::Arc, time::SystemTime};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{Duration, Utc};
|
||||
@@ -18,7 +18,8 @@ use crate::{
|
||||
AdminError, MutationContext,
|
||||
accounts::{
|
||||
AccountConnectionTestEvent, AccountConnectionTestEventStream, AccountListQuery,
|
||||
AccountRecord, AccountStatus, AccountUsageWindowQuery,
|
||||
AccountRecord, AccountStatus, AccountStatusSignals, AccountUsageWindowQuery,
|
||||
derive_error_reason,
|
||||
},
|
||||
observability::TimeRange,
|
||||
provider_credentials::{
|
||||
@@ -216,8 +217,11 @@ impl DefaultAccountsService {
|
||||
std::slice::from_mut(&mut quota),
|
||||
)
|
||||
.await?;
|
||||
let signals = account_signals(&account, "a, now);
|
||||
let error_reason = derive_error_reason(&signals);
|
||||
Ok(AccountDirectoryItem {
|
||||
status: account_status(&account, now),
|
||||
status: signals.derive(),
|
||||
error_reason,
|
||||
usage: usage.into_iter().next(),
|
||||
account,
|
||||
quota,
|
||||
@@ -260,18 +264,37 @@ impl AccountsService for DefaultAccountsService {
|
||||
let mut quotas = futures::future::join_all(page.items.iter().map(|account| async {
|
||||
let account_id = ProviderAccountId::new(account.id.clone())
|
||||
.map_err(|_| AdminError::invalid("Invalid provider account ID"))?;
|
||||
let provider = self
|
||||
.providers
|
||||
.require(&account.provider_kind)
|
||||
.map_err(|error| map_provider_error(error, "provider quota"))?;
|
||||
provider
|
||||
// 单个账号的 quota 投影失败(Provider 未注册或 quota 读取失败)不拖垮整页:
|
||||
// 该账号降级为空额度投影,其余账号与页面状态照常返回。
|
||||
let provider = match self.providers.require(&account.provider_kind) {
|
||||
Ok(provider) => provider,
|
||||
Err(error) => {
|
||||
tracing::warn!(
|
||||
account_id = %account.id,
|
||||
error = %error,
|
||||
"account directory provider is not registered; showing empty quota"
|
||||
);
|
||||
return Ok(empty_quota());
|
||||
}
|
||||
};
|
||||
match provider
|
||||
.quota(ProviderQuotaRequest {
|
||||
account_id,
|
||||
refresh: false,
|
||||
rolling_usage: rolling_usage.get(&account.id).cloned(),
|
||||
})
|
||||
.await
|
||||
.map_err(|error| map_provider_error(error, "provider quota"))
|
||||
{
|
||||
Ok(quota) => Ok(quota),
|
||||
Err(error) => {
|
||||
tracing::warn!(
|
||||
account_id = %account.id,
|
||||
error = %error,
|
||||
"account directory quota projection failed; showing empty quota"
|
||||
);
|
||||
Ok(empty_quota())
|
||||
}
|
||||
}
|
||||
}))
|
||||
.await
|
||||
.into_iter()
|
||||
@@ -282,11 +305,16 @@ impl AccountsService for DefaultAccountsService {
|
||||
.items
|
||||
.into_iter()
|
||||
.zip(quotas)
|
||||
.map(|(account, quota)| AccountDirectoryItem {
|
||||
status: account_status(&account, now),
|
||||
usage: usage.get(&account.id).cloned(),
|
||||
account,
|
||||
quota,
|
||||
.map(|(account, quota)| {
|
||||
let signals = account_signals(&account, "a, now);
|
||||
let error_reason = derive_error_reason(&signals);
|
||||
AccountDirectoryItem {
|
||||
status: signals.derive(),
|
||||
error_reason,
|
||||
usage: usage.get(&account.id).cloned(),
|
||||
account,
|
||||
quota,
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
Ok(AccountDirectoryPage {
|
||||
@@ -407,7 +435,7 @@ impl AccountsService for DefaultAccountsService {
|
||||
upstream_model: UpstreamModelId,
|
||||
) -> Result<AccountConnectionTestEventStream, AdminError> {
|
||||
let (account, provider) = self.provider_for_account(&account_id).await?;
|
||||
let initial_status = account_status(&account, Utc::now());
|
||||
let initial_status = account_status(&account, false, Utc::now());
|
||||
let model = upstream_model.as_str().to_owned();
|
||||
let operation = provider
|
||||
.connection_test_operation(&upstream_model, CONNECTION_TEST_INPUT)
|
||||
@@ -441,7 +469,7 @@ impl AccountsService for DefaultAccountsService {
|
||||
.ok()
|
||||
.flatten()
|
||||
.map_or(initial_status, |account| {
|
||||
account_status(&account, Utc::now())
|
||||
account_status(&account, false, Utc::now())
|
||||
});
|
||||
match result {
|
||||
Ok(result) => result
|
||||
@@ -486,31 +514,48 @@ impl AccountsService for DefaultAccountsService {
|
||||
}
|
||||
}
|
||||
|
||||
fn account_status(account: &AccountRecord, now: chrono::DateTime<Utc>) -> AccountStatus {
|
||||
if !account.enabled {
|
||||
AccountStatus::Disabled
|
||||
} else if account.availability == crate::model::accounts::AccountAvailability::Banned {
|
||||
AccountStatus::Banned
|
||||
} else if account.availability == crate::model::accounts::AccountAvailability::QuotaExhausted
|
||||
|| (account.availability == crate::model::accounts::AccountAvailability::Cooldown
|
||||
&& account.availability_reason.as_deref() == Some("usage_limit_exhausted")
|
||||
&& account.cooldown_until.is_none_or(|until| until > now))
|
||||
{
|
||||
AccountStatus::QuotaExhausted
|
||||
} else if matches!(
|
||||
account.availability,
|
||||
crate::model::accounts::AccountAvailability::Expired
|
||||
| crate::model::accounts::AccountAvailability::Invalid
|
||||
) || account
|
||||
.access_token_expires_at
|
||||
.is_some_and(|expires_at| expires_at <= now)
|
||||
{
|
||||
AccountStatus::Expired
|
||||
} else {
|
||||
AccountStatus::Active
|
||||
/// 账号目录中单个账号 quota 读取失败时使用的空额度投影。
|
||||
fn empty_quota() -> ProviderQuota {
|
||||
ProviderQuota {
|
||||
observed_at: None,
|
||||
refresh_token_expires_at: None,
|
||||
windows: Vec::new(),
|
||||
limit_reached: false,
|
||||
rate_limited_until: None,
|
||||
provider_data: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// 构造状态派生所需的信号快照。
|
||||
fn account_signals(
|
||||
account: &AccountRecord,
|
||||
quota: &ProviderQuota,
|
||||
now: chrono::DateTime<Utc>,
|
||||
) -> AccountStatusSignals {
|
||||
AccountStatusSignals {
|
||||
enabled: account.enabled,
|
||||
availability: account.availability,
|
||||
access_token_expired: account
|
||||
.access_token_expires_at
|
||||
.is_some_and(|expires_at| expires_at <= now),
|
||||
quota_limit_reached: quota.limit_reached,
|
||||
rate_limited_until: quota.rate_limited_until.map(SystemTime::from),
|
||||
}
|
||||
}
|
||||
|
||||
/// 账号运营分类(互斥展示状态)。
|
||||
fn account_status(
|
||||
account: &AccountRecord,
|
||||
quota_limit_reached: bool,
|
||||
now: chrono::DateTime<Utc>,
|
||||
) -> AccountStatus {
|
||||
let quota = ProviderQuota {
|
||||
limit_reached: quota_limit_reached,
|
||||
..empty_quota()
|
||||
};
|
||||
account_signals(account, "a, now).derive()
|
||||
}
|
||||
|
||||
fn retained_usage_range(now: chrono::DateTime<Utc>, retention_days: u32) -> TimeRange {
|
||||
// 账号留存投影是内部查询,不受外部观测接口 366 天上限约束。
|
||||
TimeRange {
|
||||
|
||||
@@ -13,9 +13,10 @@ use gateway_admin::{
|
||||
MutationActor, MutationContext, Revision,
|
||||
auth::AdminAuditEvent,
|
||||
backup::{
|
||||
BackupObjectMetadata, BackupRecord, BackupRecordListQuery, BackupRecordPage,
|
||||
BackupRecordSeed, BackupSettings, BackupStatus, BackupStorageConfig, BackupTriggerKind,
|
||||
ConnectionTestResult, UpdateBackupScheduleCommand, UpdateBackupStorageCommand,
|
||||
BackupError, BackupObjectMetadata, BackupRecord, BackupRecordListQuery,
|
||||
BackupRecordPage, BackupRecordSeed, BackupSettings, BackupStatus, BackupStorageConfig,
|
||||
BackupTriggerKind, ConnectionTestResult, UpdateBackupScheduleCommand,
|
||||
UpdateBackupStorageCommand,
|
||||
},
|
||||
},
|
||||
ports::{
|
||||
@@ -497,15 +498,23 @@ impl DatabaseDumpPort for FakeDumpPort {
|
||||
/// 内存版对象存储:把上传内容按 key 保存并支持 Head/Delete/预签名。
|
||||
pub(crate) struct FakeObjectStore {
|
||||
pub(crate) objects: Mutex<std::collections::HashMap<String, Vec<u8>>>,
|
||||
upload_error: Option<BackupError>,
|
||||
}
|
||||
|
||||
impl FakeObjectStore {
|
||||
pub(crate) fn new() -> Self {
|
||||
Self {
|
||||
objects: Mutex::new(std::collections::HashMap::new()),
|
||||
upload_error: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// 让上传返回指定的脱敏 S3 错误。
|
||||
pub(crate) fn fail_upload(mut self, code: &'static str, message: impl Into<String>) -> Self {
|
||||
self.upload_error = Some(BackupError::new(code, message.into()));
|
||||
self
|
||||
}
|
||||
|
||||
pub(crate) fn object(&self, key: &str) -> Option<Vec<u8>> {
|
||||
self.objects.lock().expect("objects").get(key).cloned()
|
||||
}
|
||||
@@ -530,6 +539,9 @@ impl BackupObjectStorePort for FakeObjectStore {
|
||||
_config: &BackupStorageConfig,
|
||||
request: UploadObjectRequest,
|
||||
) -> Result<(), gateway_admin::model::backup::BackupError> {
|
||||
if let Some(error) = &self.upload_error {
|
||||
return Err(error.clone());
|
||||
}
|
||||
let content = std::fs::read(&request.source).map_err(|_| {
|
||||
gateway_admin::model::backup::BackupError::new(
|
||||
gateway_admin::model::backup::code::S3_UPLOAD_FAILED,
|
||||
|
||||
@@ -6,9 +6,9 @@ use chrono::{Duration, Utc};
|
||||
use gateway_admin::{
|
||||
backup::task::BackupTask,
|
||||
model::backup::{
|
||||
BackupRecordSeed, BackupStatus, BackupTriggerKind, UpdateBackupScheduleCommand,
|
||||
BackupRecordSeed, BackupStatus, BackupTriggerKind, UpdateBackupScheduleCommand, code,
|
||||
},
|
||||
ports::backup::BackupRepository,
|
||||
ports::backup::{BackupRepository, DatabaseDumpPort, DumpRequest},
|
||||
};
|
||||
use gateway_core::engine::CancellationToken;
|
||||
|
||||
@@ -89,6 +89,78 @@ async fn backup_task_marks_failed_when_dump_port_fails() {
|
||||
assert!(object_store.object(&seed.object_key).is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn backup_task_marks_failed_with_the_upload_error() {
|
||||
let repository = Arc::new(FakeBackupRepository::new(configured_settings()));
|
||||
let dump = Arc::new(FakeDumpPort::new());
|
||||
let object_store =
|
||||
Arc::new(FakeObjectStore::new().fail_upload(code::S3_UPLOAD_FAILED, "对象存储拒绝上传"));
|
||||
let seed = BackupRecordSeed {
|
||||
id: backup_id("upload-failure"),
|
||||
trigger_kind: BackupTriggerKind::Manual,
|
||||
scheduled_at: None,
|
||||
object_key: "codex/2026/08/01/upload-failure.dump".to_owned(),
|
||||
expires_at: None,
|
||||
};
|
||||
repository
|
||||
.insert_backup_record(seed.clone())
|
||||
.await
|
||||
.expect("insert queued");
|
||||
|
||||
let task = BackupTask::new(repository.clone(), dump, object_store);
|
||||
task.run_cycle(&CancellationToken::new())
|
||||
.await
|
||||
.expect("run cycle");
|
||||
|
||||
let records = repository.all_records();
|
||||
let failed = &records[0];
|
||||
assert_eq!(failed.status, BackupStatus::Failed);
|
||||
assert_eq!(failed.error_code.as_deref(), Some(code::S3_UPLOAD_FAILED));
|
||||
assert_eq!(
|
||||
failed.error_message.as_deref(),
|
||||
Some("上传到对象存储失败:对象存储拒绝上传")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn backup_task_recovery_marks_failed_with_the_upload_error() {
|
||||
let repository = Arc::new(FakeBackupRepository::new(configured_settings()));
|
||||
let dump = Arc::new(FakeDumpPort::new());
|
||||
let object_store =
|
||||
Arc::new(FakeObjectStore::new().fail_upload(code::S3_UPLOAD_FAILED, "对象存储拒绝上传"));
|
||||
let seed = BackupRecordSeed {
|
||||
id: backup_id("recovered-upload-failure"),
|
||||
trigger_kind: BackupTriggerKind::Manual,
|
||||
scheduled_at: None,
|
||||
object_key: "codex/2026/08/01/recovered-upload-failure.dump".to_owned(),
|
||||
expires_at: None,
|
||||
};
|
||||
repository
|
||||
.insert_backup_record(seed.clone())
|
||||
.await
|
||||
.expect("insert queued");
|
||||
let cancellation = CancellationToken::new();
|
||||
dump.dump(DumpRequest {
|
||||
backup_id: seed.id.clone(),
|
||||
cancellation: cancellation.clone(),
|
||||
})
|
||||
.await
|
||||
.expect("stage dump");
|
||||
repository.force_status(&seed.id, BackupStatus::Dumping);
|
||||
|
||||
let task = BackupTask::new(repository.clone(), dump, object_store);
|
||||
task.run_cycle(&cancellation).await.expect("run cycle");
|
||||
|
||||
let records = repository.all_records();
|
||||
let failed = &records[0];
|
||||
assert_eq!(failed.status, BackupStatus::Failed);
|
||||
assert_eq!(failed.error_code.as_deref(), Some(code::S3_UPLOAD_FAILED));
|
||||
assert_eq!(
|
||||
failed.error_message.as_deref(),
|
||||
Some("上传到对象存储失败:对象存储拒绝上传")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn retention_cleans_expired_scheduled_backups() {
|
||||
let repository = Arc::new(FakeBackupRepository::new(configured_settings()));
|
||||
|
||||
@@ -127,6 +127,8 @@ fn representative_quota_should_prefer_short_window_and_highest_usage() {
|
||||
quota_window("shortTerm", Some(18_000), Some(30.0)),
|
||||
quota_window("shortTerm", Some(14_400), Some(45.0)),
|
||||
],
|
||||
limit_reached: false,
|
||||
rate_limited_until: None,
|
||||
provider_data: None,
|
||||
};
|
||||
|
||||
@@ -146,6 +148,7 @@ fn quota_window(
|
||||
window_seconds,
|
||||
used_percent,
|
||||
reset_at: None,
|
||||
limit_reached: false,
|
||||
local_usage: None,
|
||||
provider_data: None,
|
||||
}
|
||||
|
||||
@@ -373,7 +373,7 @@ impl ProviderAdmin for FakeProviderAdmin {
|
||||
|
||||
pub(super) struct FakeAccountStore {
|
||||
events: EventLog,
|
||||
account: AccountRecord,
|
||||
accounts: Mutex<Vec<AccountRecord>>,
|
||||
account_after_probe: Mutex<Option<AccountRecord>>,
|
||||
fail_commit: Mutex<bool>,
|
||||
audit_requests: Mutex<Vec<String>>,
|
||||
@@ -388,7 +388,7 @@ impl FakeAccountStore {
|
||||
fn with_account(account: AccountRecord, events: EventLog) -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
events,
|
||||
account,
|
||||
accounts: Mutex::new(vec![account]),
|
||||
account_after_probe: Mutex::new(None),
|
||||
fail_commit: Mutex::new(false),
|
||||
audit_requests: Mutex::new(Vec::new()),
|
||||
@@ -408,6 +408,10 @@ impl FakeAccountStore {
|
||||
*self.quota_window_usage.lock().expect("quota window usage") = usage;
|
||||
}
|
||||
|
||||
pub(super) fn set_accounts(&self, accounts: Vec<AccountRecord>) {
|
||||
*self.accounts.lock().expect("accounts") = accounts;
|
||||
}
|
||||
|
||||
fn set_account_after_probe(&self, account: AccountRecord) {
|
||||
*self
|
||||
.account_after_probe
|
||||
@@ -439,28 +443,41 @@ impl FakeAccountStore {
|
||||
#[async_trait]
|
||||
impl AccountStore for FakeAccountStore {
|
||||
async fn list_accounts(&self, _: AccountListQuery) -> AdminStoreResult<AccountPage> {
|
||||
self.record("store.list_accounts");
|
||||
let accounts = self.accounts.lock().expect("accounts").clone();
|
||||
Ok(AccountPage {
|
||||
config_revision: revision(1),
|
||||
items: vec![self.account.clone()],
|
||||
total: 1,
|
||||
items: accounts.clone(),
|
||||
total: accounts.len() as u64,
|
||||
summary: AccountSummary {
|
||||
total: 1,
|
||||
active: 1,
|
||||
total: accounts.len() as u64,
|
||||
normal: 1,
|
||||
quota_exhausted: 0,
|
||||
unavailable: 0,
|
||||
rate_limited: 0,
|
||||
disabled: 0,
|
||||
error: 0,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
async fn load_account(&self, account_id: &str) -> AdminStoreResult<Option<AccountRecord>> {
|
||||
self.record("store.load_account");
|
||||
// probe 后的账号状态覆盖只对同一 id 生效;其余按账号列表查询。
|
||||
let account = self
|
||||
.account_after_probe
|
||||
.lock()
|
||||
.expect("account after probe")
|
||||
.clone()
|
||||
.unwrap_or_else(|| self.account.clone());
|
||||
Ok((account_id == account.id.as_str()).then_some(account))
|
||||
.filter(|account| account.id == account_id)
|
||||
.or_else(|| {
|
||||
self.accounts
|
||||
.lock()
|
||||
.expect("accounts")
|
||||
.iter()
|
||||
.find(|account| account.id == account_id)
|
||||
.cloned()
|
||||
});
|
||||
Ok(account)
|
||||
}
|
||||
|
||||
async fn load_account_usage(
|
||||
@@ -488,11 +505,12 @@ impl AccountStore for FakeAccountStore {
|
||||
_: CredentialListQuery,
|
||||
) -> AdminStoreResult<CredentialPage> {
|
||||
self.record("store.list_credentials");
|
||||
let accounts = self.accounts.lock().expect("accounts").clone();
|
||||
Ok(CredentialPage {
|
||||
config_revision: revision(1),
|
||||
items: (provider_kind == &self.account.provider_kind)
|
||||
.then(|| self.account.clone())
|
||||
items: accounts
|
||||
.into_iter()
|
||||
.filter(|account| &account.provider_kind == provider_kind)
|
||||
.collect(),
|
||||
next_cursor: None,
|
||||
})
|
||||
@@ -504,14 +522,19 @@ impl AccountStore for FakeAccountStore {
|
||||
account_id: &ProviderAccountId,
|
||||
) -> AdminStoreResult<Option<CredentialDetails>> {
|
||||
self.record("store.credential_details");
|
||||
Ok(
|
||||
(provider_kind == &self.account.provider_kind
|
||||
&& account_id.as_str() == self.account.id)
|
||||
.then(|| CredentialDetails {
|
||||
config_revision: revision(1),
|
||||
credential: self.account.clone(),
|
||||
}),
|
||||
)
|
||||
let account = self
|
||||
.accounts
|
||||
.lock()
|
||||
.expect("accounts")
|
||||
.iter()
|
||||
.find(|account| {
|
||||
&account.provider_kind == provider_kind && account.id == account_id.as_str()
|
||||
})
|
||||
.cloned();
|
||||
Ok(account.map(|credential| CredentialDetails {
|
||||
config_revision: revision(1),
|
||||
credential,
|
||||
}))
|
||||
}
|
||||
|
||||
async fn load_credentials_for_export(
|
||||
@@ -520,21 +543,26 @@ impl AccountStore for FakeAccountStore {
|
||||
account_ids: &[ProviderAccountId],
|
||||
) -> AdminStoreResult<Vec<ProviderExportCredentialInput>> {
|
||||
self.record("store.load_credentials_for_export");
|
||||
if provider_kind != &self.account.provider_kind
|
||||
|| account_ids
|
||||
.iter()
|
||||
.any(|account_id| account_id.as_str() != self.account.id)
|
||||
{
|
||||
let accounts = self.accounts.lock().expect("accounts").clone();
|
||||
if accounts.iter().any(|account| {
|
||||
&account.provider_kind != provider_kind
|
||||
|| !account_ids
|
||||
.iter()
|
||||
.any(|account_id| account_id.as_str() == account.id)
|
||||
}) {
|
||||
return Err(AdminStoreError::new(
|
||||
AdminStoreErrorKind::NotFound,
|
||||
"test account",
|
||||
"credential not found",
|
||||
));
|
||||
}
|
||||
Ok(vec![ProviderExportCredentialInput {
|
||||
account: self.account.clone(),
|
||||
provider_material: document(),
|
||||
}])
|
||||
Ok(accounts
|
||||
.into_iter()
|
||||
.map(|account| ProviderExportCredentialInput {
|
||||
account,
|
||||
provider_material: document(),
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn commit_credential_import(
|
||||
@@ -818,7 +846,7 @@ async fn connection_test_should_probe_unavailable_account() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn connection_test_reloads_resettable_usage_limit_status_after_probe() {
|
||||
async fn connection_test_rate_limited_probe_preserves_normal_account_status() {
|
||||
let events = events();
|
||||
let provider = FakeProviderAdmin::new("xai", events.clone());
|
||||
let store = FakeAccountStore::new("xai", events);
|
||||
@@ -841,7 +869,7 @@ async fn connection_test_reloads_resettable_usage_limit_status_after_probe() {
|
||||
assert!(matches!(
|
||||
events.last(),
|
||||
Some(AccountConnectionTestEvent::Failed {
|
||||
account_status: gateway_admin::model::accounts::AccountStatus::QuotaExhausted,
|
||||
account_status: gateway_admin::model::accounts::AccountStatus::Normal,
|
||||
..
|
||||
})
|
||||
));
|
||||
@@ -880,7 +908,7 @@ async fn connection_test_should_preserve_disabled_account_status() {
|
||||
async fn accounts_export_should_pass_store_loaded_timestamps_and_material_to_provider() {
|
||||
let provider = FakeProviderAdmin::new("openai", events());
|
||||
let store = FakeAccountStore::new("openai", events());
|
||||
let expected = store.account.clone();
|
||||
let expected = store.accounts.lock().expect("accounts")[0].clone();
|
||||
accounts_service(provider.clone(), store)
|
||||
.await
|
||||
.accounts()
|
||||
@@ -922,7 +950,7 @@ async fn accounts_refresh_should_keep_guard_through_store_commit() {
|
||||
assert_eq!(result.account.account.provider_kind.as_str(), "openai");
|
||||
assert_eq!(
|
||||
result.account.status,
|
||||
gateway_admin::model::accounts::AccountStatus::Active
|
||||
gateway_admin::model::accounts::AccountStatus::Normal
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
@@ -958,22 +986,87 @@ async fn accounts_list_should_return_complete_directory_semantics() {
|
||||
.expect("complete account directory");
|
||||
|
||||
assert_eq!(page.summary.total, 1);
|
||||
assert_eq!(page.summary.active, 1);
|
||||
assert_eq!(page.summary.normal, 1);
|
||||
let account = page.items.first().expect("account item");
|
||||
assert_eq!(account.account.provider_kind.as_str(), "openai");
|
||||
assert_eq!(
|
||||
account.status,
|
||||
gateway_admin::model::accounts::AccountStatus::Active
|
||||
gateway_admin::model::accounts::AccountStatus::Normal
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn accounts_list_should_prefer_usage_limit_exhaustion_over_stale_token_expiry() {
|
||||
async fn accounts_list_should_degrade_quota_failure_to_empty_window_without_dropping_page() {
|
||||
let events = events();
|
||||
let openai = FakeProviderAdmin::new("openai", events.clone());
|
||||
openai.set_quota(ProviderQuota {
|
||||
observed_at: Some(Utc::now()),
|
||||
refresh_token_expires_at: None,
|
||||
windows: vec![ProviderQuotaWindow {
|
||||
key: "primary".to_owned(),
|
||||
group: "shortTerm".to_owned(),
|
||||
label: "5小时限额".to_owned(),
|
||||
source: None,
|
||||
window_seconds: Some(5 * 60 * 60),
|
||||
used_percent: Some(97.0),
|
||||
reset_at: Some(Utc::now() + TimeDelta::hours(1)),
|
||||
limit_reached: false,
|
||||
local_usage: None,
|
||||
provider_data: None,
|
||||
}],
|
||||
limit_reached: false,
|
||||
rate_limited_until: None,
|
||||
provider_data: None,
|
||||
});
|
||||
let failing = FakeProviderAdmin::new("xai", events.clone());
|
||||
failing.fail_next_quota(ProviderAdminErrorKind::Invalid);
|
||||
let store = FakeAccountStore::new("openai", events.clone());
|
||||
let mut xai_account = account_record("xai");
|
||||
xai_account.id = "acct_bad_xai".to_owned();
|
||||
store.set_accounts(vec![account_record("openai"), xai_account]);
|
||||
let services = super::AdminHarness::new()
|
||||
.accounts(store.clone())
|
||||
.settings(Arc::new(StaticSettingsStore))
|
||||
.provider(openai)
|
||||
.provider(failing)
|
||||
.probe(Arc::new(SuccessfulAccountProbe))
|
||||
.build()
|
||||
.await;
|
||||
|
||||
let page = services
|
||||
.accounts()
|
||||
.list(AccountListQuery {
|
||||
page: 1,
|
||||
page_size: gateway_admin::model::PageSize::new(20).expect("page size"),
|
||||
provider_kind: None,
|
||||
search: None,
|
||||
status: None,
|
||||
sort: None,
|
||||
})
|
||||
.await
|
||||
.expect("one failing quota must not fail the directory page");
|
||||
|
||||
assert_eq!(page.items.len(), 2);
|
||||
let by_id = page
|
||||
.items
|
||||
.iter()
|
||||
.map(|item| (item.account.id.as_str(), item))
|
||||
.collect::<std::collections::HashMap<_, _>>();
|
||||
let healthy = by_id.get("acct_test").expect("healthy account item");
|
||||
assert_eq!(healthy.quota.windows.len(), 1);
|
||||
assert_eq!(healthy.quota.windows[0].key, "primary");
|
||||
let degraded = by_id.get("acct_bad_xai").expect("degraded account item");
|
||||
assert!(degraded.quota.windows.is_empty());
|
||||
assert!(degraded.quota.observed_at.is_none());
|
||||
assert!(degraded.quota.provider_data.is_none());
|
||||
assert_eq!(page.summary.total, 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn accounts_list_should_prefer_quota_exhaustion_over_stale_token_expiry() {
|
||||
let provider = FakeProviderAdmin::new("openai", events());
|
||||
let mut account = account_record("openai");
|
||||
account.availability = AccountAvailability::Cooldown;
|
||||
account.availability_reason = Some("usage_limit_exhausted".to_owned());
|
||||
account.cooldown_until = Some(Utc::now() + TimeDelta::minutes(5));
|
||||
account.availability = AccountAvailability::QuotaExhausted;
|
||||
account.access_token_expires_at = Some(Utc::now() - TimeDelta::minutes(5));
|
||||
let store = FakeAccountStore::with_account(account, events());
|
||||
|
||||
@@ -997,6 +1090,125 @@ async fn accounts_list_should_prefer_usage_limit_exhaustion_over_stale_token_exp
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn accounts_list_should_map_unknown_availability_to_error_not_normal() {
|
||||
// Unknown 不可调度,Admin 不得显示为 normal。
|
||||
let provider = FakeProviderAdmin::new("openai", events());
|
||||
let mut account = account_record("openai");
|
||||
account.availability = AccountAvailability::Unknown;
|
||||
let store = FakeAccountStore::with_account(account, events());
|
||||
|
||||
let page = accounts_service(provider, store)
|
||||
.await
|
||||
.accounts()
|
||||
.list(AccountListQuery {
|
||||
page: 1,
|
||||
page_size: gateway_admin::model::PageSize::new(20).expect("page size"),
|
||||
provider_kind: None,
|
||||
search: None,
|
||||
status: None,
|
||||
sort: None,
|
||||
})
|
||||
.await
|
||||
.expect("unknown account directory");
|
||||
|
||||
assert_eq!(
|
||||
page.items.first().expect("account item").status,
|
||||
gateway_admin::model::accounts::AccountStatus::Error,
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn accounts_list_should_map_rate_limited_from_quota_limit_reached() {
|
||||
// ready + 快照级 quota.limitReached → quota_exhausted(配额耗尽);
|
||||
// rate_limited 只来自 ProviderQuota.rate_limited_until(429 Redis 冷却)。
|
||||
let provider = FakeProviderAdmin::new("openai", events());
|
||||
provider.set_quota(ProviderQuota {
|
||||
observed_at: Some(Utc::now()),
|
||||
refresh_token_expires_at: None,
|
||||
windows: vec![ProviderQuotaWindow {
|
||||
key: "primary".to_owned(),
|
||||
group: "shortTerm".to_owned(),
|
||||
label: "5小时限额".to_owned(),
|
||||
source: None,
|
||||
window_seconds: Some(5 * 60 * 60),
|
||||
used_percent: Some(100.0),
|
||||
reset_at: Some(Utc::now() + TimeDelta::hours(1)),
|
||||
limit_reached: true,
|
||||
local_usage: None,
|
||||
provider_data: None,
|
||||
}],
|
||||
limit_reached: true,
|
||||
rate_limited_until: Some(Utc::now() + TimeDelta::minutes(5)),
|
||||
provider_data: None,
|
||||
});
|
||||
let store = FakeAccountStore::new("openai", events());
|
||||
|
||||
let page = accounts_service(provider, store)
|
||||
.await
|
||||
.accounts()
|
||||
.list(AccountListQuery {
|
||||
page: 1,
|
||||
page_size: gateway_admin::model::PageSize::new(20).expect("page size"),
|
||||
provider_kind: None,
|
||||
search: None,
|
||||
status: None,
|
||||
sort: None,
|
||||
})
|
||||
.await
|
||||
.expect("rate limited account directory");
|
||||
|
||||
assert_eq!(
|
||||
page.items.first().expect("account item").status,
|
||||
gateway_admin::model::accounts::AccountStatus::RateLimited,
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn accounts_list_should_map_quota_exhausted_when_only_limit_reached() {
|
||||
// ready + 快照级 quota.limitReached、无 429 冷却 → quota_exhausted。
|
||||
let provider = FakeProviderAdmin::new("openai", events());
|
||||
provider.set_quota(ProviderQuota {
|
||||
observed_at: Some(Utc::now()),
|
||||
refresh_token_expires_at: None,
|
||||
windows: vec![ProviderQuotaWindow {
|
||||
key: "primary".to_owned(),
|
||||
group: "shortTerm".to_owned(),
|
||||
label: "5小时限额".to_owned(),
|
||||
source: None,
|
||||
window_seconds: Some(5 * 60 * 60),
|
||||
used_percent: Some(100.0),
|
||||
reset_at: Some(Utc::now() + TimeDelta::hours(1)),
|
||||
limit_reached: true,
|
||||
local_usage: None,
|
||||
provider_data: None,
|
||||
}],
|
||||
limit_reached: true,
|
||||
rate_limited_until: None,
|
||||
provider_data: None,
|
||||
});
|
||||
let store = FakeAccountStore::new("openai", events());
|
||||
|
||||
let page = accounts_service(provider, store)
|
||||
.await
|
||||
.accounts()
|
||||
.list(AccountListQuery {
|
||||
page: 1,
|
||||
page_size: gateway_admin::model::PageSize::new(20).expect("page size"),
|
||||
provider_kind: None,
|
||||
search: None,
|
||||
status: None,
|
||||
sort: None,
|
||||
})
|
||||
.await
|
||||
.expect("quota exhausted account directory");
|
||||
|
||||
assert_eq!(
|
||||
page.items.first().expect("account item").status,
|
||||
gateway_admin::model::accounts::AccountStatus::QuotaExhausted,
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn accounts_list_should_attach_local_usage_to_quota_windows() {
|
||||
let provider = FakeProviderAdmin::new("openai", events());
|
||||
@@ -1012,9 +1224,12 @@ async fn accounts_list_should_attach_local_usage_to_quota_windows() {
|
||||
window_seconds: Some(5 * 60 * 60),
|
||||
used_percent: Some(97.0),
|
||||
reset_at: Some(reset_at),
|
||||
limit_reached: false,
|
||||
local_usage: None,
|
||||
provider_data: None,
|
||||
}],
|
||||
limit_reached: false,
|
||||
rate_limited_until: None,
|
||||
provider_data: None,
|
||||
});
|
||||
let store = FakeAccountStore::new("openai", events());
|
||||
@@ -1122,6 +1337,8 @@ fn empty_quota() -> ProviderQuota {
|
||||
observed_at: None,
|
||||
refresh_token_expires_at: None,
|
||||
windows: Vec::new(),
|
||||
limit_reached: false,
|
||||
rate_limited_until: None,
|
||||
provider_data: None,
|
||||
}
|
||||
}
|
||||
@@ -1166,9 +1383,8 @@ pub(super) fn account_record(kind: &str) -> AccountRecord {
|
||||
next_refresh_at: Some(now + TimeDelta::minutes(30)),
|
||||
enabled: true,
|
||||
availability: AccountAvailability::Ready,
|
||||
availability_reason: None,
|
||||
cooldown_until: None,
|
||||
availability_observed_at: now,
|
||||
last_error_message: None,
|
||||
quota_observed_at: None,
|
||||
created_at: now,
|
||||
updated_at: now,
|
||||
@@ -1204,8 +1420,6 @@ fn prepared_create_with_id(
|
||||
next_refresh_at: Some(now + TimeDelta::minutes(30)),
|
||||
enabled: true,
|
||||
availability: AccountAvailability::Ready,
|
||||
availability_reason: None,
|
||||
cooldown_until: None,
|
||||
availability_observed_at: now,
|
||||
}
|
||||
}
|
||||
@@ -1290,9 +1504,7 @@ impl AccountProbe for FreeModelQuotaProbe {
|
||||
) -> BoxFuture<'_, Result<AccountProbeResult, AccountProbeError>> {
|
||||
Box::pin(async move {
|
||||
let mut account = account_record("xai");
|
||||
account.availability = AccountAvailability::Cooldown;
|
||||
account.availability_reason = Some("usage_limit_exhausted".to_owned());
|
||||
account.cooldown_until = Some(Utc::now() + TimeDelta::hours(24));
|
||||
account.availability = AccountAvailability::Ready;
|
||||
self.store.set_account_after_probe(account);
|
||||
Err(GatewayError::new(
|
||||
GatewayErrorKind::RateLimited,
|
||||
|
||||
@@ -91,7 +91,8 @@ impl ListQuery {
|
||||
let status = match self.status.as_deref().map(str::trim) {
|
||||
None | Some("") => None,
|
||||
Some(value) => Some(
|
||||
parse_account_status(value).ok_or_else(|| WireValidationError::new("status"))?,
|
||||
DomainAccountStatus::parse(&value.to_ascii_lowercase())
|
||||
.ok_or_else(|| WireValidationError::new("status"))?,
|
||||
),
|
||||
};
|
||||
let sort = match (self.sort_by.as_deref(), self.sort_direction.as_deref()) {
|
||||
@@ -127,17 +128,6 @@ fn parse_provider(value: &str) -> Result<Option<ProviderKind>, WireValidationErr
|
||||
.map_err(|_| WireValidationError::new("provider"))
|
||||
}
|
||||
|
||||
fn parse_account_status(value: &str) -> Option<DomainAccountStatus> {
|
||||
match value.trim().to_ascii_lowercase().as_str() {
|
||||
"active" => Some(DomainAccountStatus::Active),
|
||||
"expired" => Some(DomainAccountStatus::Expired),
|
||||
"quota_exhausted" => Some(DomainAccountStatus::QuotaExhausted),
|
||||
"disabled" => Some(DomainAccountStatus::Disabled),
|
||||
"banned" => Some(DomainAccountStatus::Banned),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_sort_field(value: &str) -> Option<AccountSortField> {
|
||||
match value.trim() {
|
||||
"email" => Some(AccountSortField::Email),
|
||||
@@ -172,9 +162,11 @@ pub struct AccountPageData {
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AccountSummaryView {
|
||||
pub total: u64,
|
||||
pub active: u64,
|
||||
pub normal: u64,
|
||||
pub quota_exhausted: u64,
|
||||
pub unavailable: u64,
|
||||
pub rate_limited: u64,
|
||||
pub disabled: u64,
|
||||
pub error: u64,
|
||||
}
|
||||
|
||||
/// 一条安全账号视图。
|
||||
@@ -194,7 +186,10 @@ pub struct AccountView {
|
||||
pub has_refresh_token: bool,
|
||||
pub status: String,
|
||||
pub display_status: String,
|
||||
pub token_refreshing: bool,
|
||||
/// `status == "error"` 时的具体原因;其余状态为 `null`。
|
||||
pub error_reason: Option<String>,
|
||||
/// 最近一次失败的上游错误描述;仅错误状态存在。
|
||||
pub error_message: Option<String>,
|
||||
pub availability: String,
|
||||
pub enabled: bool,
|
||||
pub access_token_expires_at: Option<String>,
|
||||
@@ -214,6 +209,9 @@ pub struct AccountView {
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AccountQuotaView {
|
||||
pub refreshed_at_display: String,
|
||||
pub limit_reached: bool,
|
||||
/// 429 临时限流(Redis 冷却)到期时间展示;非限流中为 `null`。
|
||||
pub rate_limited_until: Option<String>,
|
||||
pub windows: Vec<AccountQuotaWindowView>,
|
||||
}
|
||||
|
||||
@@ -227,6 +225,7 @@ pub struct AccountQuotaWindowView {
|
||||
pub label_display: String,
|
||||
pub used_percent: Option<f64>,
|
||||
pub used_percent_display: String,
|
||||
pub limit_reached: bool,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub local_usage: Option<serde_json::Value>,
|
||||
pub reset_at_display: String,
|
||||
@@ -526,7 +525,7 @@ impl From<DomainConnectionTestEvent> for AccountConnectionTestEvent {
|
||||
DomainConnectionTestEvent::Completed { account_status } => serde_json::json!({
|
||||
"type": "test_complete",
|
||||
"success": true,
|
||||
"accountStatus": account_status_name(account_status)
|
||||
"accountStatus": account_status.as_str()
|
||||
}),
|
||||
DomainConnectionTestEvent::Failed {
|
||||
message,
|
||||
@@ -544,7 +543,7 @@ impl From<DomainConnectionTestEvent> for AccountConnectionTestEvent {
|
||||
"upstreamStatus": upstream_status,
|
||||
"upstreamContentType": upstream_content_type,
|
||||
"upstreamBody": upstream_body,
|
||||
"accountStatus": account_status_name(account_status)
|
||||
"accountStatus": account_status.as_str()
|
||||
}),
|
||||
};
|
||||
Self { data }
|
||||
@@ -1332,9 +1331,11 @@ fn account_page_data(
|
||||
page: PageMeta::new(page, u32::from(page_size), result.total, total_pages),
|
||||
summary: AccountSummaryView {
|
||||
total: result.summary.total,
|
||||
active: result.summary.active,
|
||||
normal: result.summary.normal,
|
||||
quota_exhausted: result.summary.quota_exhausted,
|
||||
unavailable: result.summary.unavailable,
|
||||
rate_limited: result.summary.rate_limited,
|
||||
disabled: result.summary.disabled,
|
||||
error: result.summary.error,
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -1365,10 +1366,11 @@ fn account_view(item: AccountDirectoryItem, now: DateTime<Utc>) -> AccountView {
|
||||
let AccountDirectoryItem {
|
||||
account,
|
||||
status,
|
||||
error_reason,
|
||||
usage,
|
||||
quota,
|
||||
} = item;
|
||||
let status = account_status_name(status).to_owned();
|
||||
let status = status.as_str().to_owned();
|
||||
let expires_at = account.access_token_expires_at.as_ref().map(china_rfc3339);
|
||||
let added_at = china_rfc3339(&account.created_at);
|
||||
let updated_at = china_rfc3339(&account.updated_at);
|
||||
@@ -1387,7 +1389,8 @@ fn account_view(item: AccountDirectoryItem, now: DateTime<Utc>) -> AccountView {
|
||||
has_refresh_token: account.has_refresh_token,
|
||||
status: status.clone(),
|
||||
display_status: status,
|
||||
token_refreshing: false,
|
||||
error_reason: error_reason.map(|reason| reason.as_str().to_owned()),
|
||||
error_message: account.last_error_message,
|
||||
availability: account.availability.as_str().to_owned(),
|
||||
enabled: account.enabled,
|
||||
access_token_expires_at: expires_at,
|
||||
@@ -1417,9 +1420,12 @@ fn account_quota_view(
|
||||
.observed_at
|
||||
.map_or_else(|| "—".to_owned(), |value| relative_time(value, now));
|
||||
let windows = quota.windows.into_iter().map(quota_window_view).collect();
|
||||
let rate_limited_until = quota.rate_limited_until.map(|until| china_datetime(&until));
|
||||
(
|
||||
AccountQuotaView {
|
||||
refreshed_at_display,
|
||||
limit_reached: quota.limit_reached,
|
||||
rate_limited_until,
|
||||
windows,
|
||||
},
|
||||
refresh_token_expires_at,
|
||||
@@ -1435,6 +1441,7 @@ fn quota_window_view(window: ProviderQuotaWindow) -> AccountQuotaWindowView {
|
||||
window_seconds,
|
||||
used_percent,
|
||||
reset_at,
|
||||
limit_reached,
|
||||
local_usage,
|
||||
provider_data: _,
|
||||
} = window;
|
||||
@@ -1446,6 +1453,7 @@ fn quota_window_view(window: ProviderQuotaWindow) -> AccountQuotaWindowView {
|
||||
used_percent,
|
||||
used_percent_display: used_percent
|
||||
.map_or_else(|| "—".to_owned(), |value| format!("{value:.1}%")),
|
||||
limit_reached,
|
||||
local_usage: local_usage.as_ref().map(quota_local_usage),
|
||||
reset_at_display: reset_at.map_or_else(|| "—".to_owned(), |value| china_datetime(&value)),
|
||||
}
|
||||
@@ -1710,16 +1718,6 @@ fn provider_document_value(document: ProviderDocument) -> Value {
|
||||
Value::Object(document.into_provider_data().into_inner())
|
||||
}
|
||||
|
||||
fn account_status_name(status: DomainAccountStatus) -> &'static str {
|
||||
match status {
|
||||
DomainAccountStatus::Active => "active",
|
||||
DomainAccountStatus::Expired => "expired",
|
||||
DomainAccountStatus::QuotaExhausted => "quota_exhausted",
|
||||
DomainAccountStatus::Disabled => "disabled",
|
||||
DomainAccountStatus::Banned => "banned",
|
||||
}
|
||||
}
|
||||
|
||||
fn relative_time(value: DateTime<Utc>, now: DateTime<Utc>) -> String {
|
||||
let elapsed = now.signed_duration_since(value);
|
||||
if elapsed.num_seconds() < 0 {
|
||||
|
||||
@@ -784,7 +784,9 @@ pub(crate) fn dashboard_view(
|
||||
pool_summary: DashboardPoolSummaryView {
|
||||
total: provider_accounts.total,
|
||||
active: provider_accounts.active,
|
||||
rate_limited: provider_accounts.rate_limited,
|
||||
expired: provider_accounts.expired,
|
||||
invalid: provider_accounts.invalid,
|
||||
quota_exhausted: provider_accounts.quota_exhausted,
|
||||
refreshing: provider_accounts.refreshing,
|
||||
disabled: provider_accounts.disabled,
|
||||
|
||||
@@ -348,7 +348,9 @@ pub struct DashboardAccountRequestBucketView {
|
||||
pub struct DashboardPoolSummaryView {
|
||||
pub total: u64,
|
||||
pub active: u64,
|
||||
pub rate_limited: u64,
|
||||
pub expired: u64,
|
||||
pub invalid: u64,
|
||||
pub quota_exhausted: u64,
|
||||
pub refreshing: Option<u64>,
|
||||
pub disabled: u64,
|
||||
|
||||
@@ -10,7 +10,7 @@ mod query {
|
||||
"pageSize": 20,
|
||||
"provider": "xai",
|
||||
"search": " operator ",
|
||||
"status": "active",
|
||||
"status": "normal",
|
||||
"sortBy": "lastUsedAt",
|
||||
"sortDirection": "desc"
|
||||
}))
|
||||
@@ -23,7 +23,7 @@ mod query {
|
||||
Some(ref provider) if provider.as_str() == "xai"
|
||||
));
|
||||
assert_eq!(query.search.as_deref(), Some("operator"));
|
||||
assert_eq!(query.status, Some(AccountStatus::Active));
|
||||
assert_eq!(query.status, Some(AccountStatus::Normal));
|
||||
assert_eq!(
|
||||
query.sort.expect("sort").field,
|
||||
AccountSortField::LastUsedAt
|
||||
@@ -34,6 +34,19 @@ mod query {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn account_query_should_parse_rate_limited_status() {
|
||||
let query: ListQuery = serde_json::from_value(json!({
|
||||
"status": "rate_limited"
|
||||
}))
|
||||
.expect("deserialize account query");
|
||||
|
||||
assert_eq!(
|
||||
query.validate().expect("validate account query").status,
|
||||
Some(AccountStatus::RateLimited)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn account_query_should_reject_unbounded_page_size() {
|
||||
let query: ListQuery =
|
||||
@@ -331,7 +344,7 @@ mod actions {
|
||||
text: "OK".to_owned(),
|
||||
},
|
||||
DomainConnectionTestEvent::Completed {
|
||||
account_status: AccountStatus::Active,
|
||||
account_status: AccountStatus::Normal,
|
||||
},
|
||||
DomainConnectionTestEvent::Failed {
|
||||
message: "upstream unavailable".to_owned(),
|
||||
@@ -365,7 +378,7 @@ mod actions {
|
||||
}
|
||||
}),
|
||||
json!({ "type": "content", "text": "OK" }),
|
||||
json!({ "type": "test_complete", "success": true, "accountStatus": "active" }),
|
||||
json!({ "type": "test_complete", "success": true, "accountStatus": "normal" }),
|
||||
json!({
|
||||
"type": "error",
|
||||
"error": "upstream unavailable",
|
||||
|
||||
@@ -224,7 +224,9 @@ async fn dashboard_summary_should_include_quota_exhaustion_in_unavailable_headli
|
||||
enabled: 804,
|
||||
unavailable: 792,
|
||||
active: 15,
|
||||
rate_limited: 0,
|
||||
expired: 2,
|
||||
invalid: 0,
|
||||
quota_exhausted: 785,
|
||||
refreshing: None,
|
||||
disabled: 3,
|
||||
|
||||
@@ -660,6 +660,9 @@ where
|
||||
ProviderBoundary::Result(result) => match *result {
|
||||
Ok(stream) => stream,
|
||||
Err(error) => {
|
||||
if self.prepare_unavailable_native_continuation_replay(&error) {
|
||||
return Ok(Some(PullOutcome::AttemptDiscarded));
|
||||
}
|
||||
if matches!(
|
||||
error.kind(),
|
||||
ProviderErrorKind::AccountCapacityUnavailable
|
||||
@@ -1096,6 +1099,41 @@ where
|
||||
true
|
||||
}
|
||||
|
||||
/// 原生续写的原账号在本地调度阶段已不可用时,改用完整输入重放。
|
||||
///
|
||||
/// 这不是强制 Smart:下一次 attempt 仍使用请求配置的调度策略。只有网关持有
|
||||
/// 对应 Provider 的可携带会话状态时,才可安全移除 native handle 并换号。
|
||||
fn prepare_unavailable_native_continuation_replay(&mut self, error: &ProviderError) -> bool {
|
||||
if self.account_selection.required_account().is_some()
|
||||
|| self.continuation_attempt != ContinuationAttempt::Native
|
||||
|| !matches!(
|
||||
error.kind(),
|
||||
ProviderErrorKind::NoEligibleAccount
|
||||
| ProviderErrorKind::AccountCapacityUnavailable
|
||||
)
|
||||
{
|
||||
return false;
|
||||
}
|
||||
let Some(pin) = self
|
||||
.continuation
|
||||
.as_ref()
|
||||
.and_then(ContinuationBinding::pinned)
|
||||
else {
|
||||
return false;
|
||||
};
|
||||
if self
|
||||
.operation
|
||||
.provider_session_state(pin.provider().as_str())
|
||||
.is_none()
|
||||
{
|
||||
return false;
|
||||
}
|
||||
|
||||
self.continuation_attempt = ContinuationAttempt::ReplayAny;
|
||||
self.excluded_accounts.insert(pin.account().clone());
|
||||
true
|
||||
}
|
||||
|
||||
fn reset_uncommitted_observations(&mut self) {
|
||||
self.usage = Usage::new();
|
||||
self.cost = CostEstimate::unavailable();
|
||||
|
||||
@@ -188,11 +188,13 @@ impl fmt::Debug for OpaqueProviderData {
|
||||
}
|
||||
|
||||
/// 数据库中固定的账号运行状态。
|
||||
///
|
||||
/// 只承载终态或需要外部干预的状态;临时限流(429)不进入此状态,由
|
||||
/// Provider quota 数据驱动(`AccountQuotaSignals::limit_reached`)。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
||||
pub enum AccountAvailability {
|
||||
Unknown,
|
||||
Ready,
|
||||
Cooldown,
|
||||
QuotaExhausted,
|
||||
Expired,
|
||||
Banned,
|
||||
@@ -206,7 +208,6 @@ impl AccountAvailability {
|
||||
match self {
|
||||
Self::Unknown => "unknown",
|
||||
Self::Ready => "ready",
|
||||
Self::Cooldown => "cooldown",
|
||||
Self::QuotaExhausted => "quota_exhausted",
|
||||
Self::Expired => "expired",
|
||||
Self::Banned => "banned",
|
||||
@@ -220,7 +221,6 @@ impl AccountAvailability {
|
||||
match value {
|
||||
"unknown" => Some(Self::Unknown),
|
||||
"ready" => Some(Self::Ready),
|
||||
"cooldown" => Some(Self::Cooldown),
|
||||
"quota_exhausted" => Some(Self::QuotaExhausted),
|
||||
"expired" => Some(Self::Expired),
|
||||
"banned" => Some(Self::Banned),
|
||||
@@ -244,7 +244,6 @@ pub struct ProviderAccount {
|
||||
revision: CredentialRevision,
|
||||
enabled: bool,
|
||||
availability: AccountAvailability,
|
||||
cooldown_until: Option<SystemTime>,
|
||||
access_token_expires_at: Option<SystemTime>,
|
||||
next_refresh_at: Option<SystemTime>,
|
||||
has_refresh_token: bool,
|
||||
@@ -274,7 +273,6 @@ impl ProviderAccount {
|
||||
revision,
|
||||
enabled: true,
|
||||
availability: AccountAvailability::Unknown,
|
||||
cooldown_until: None,
|
||||
access_token_expires_at,
|
||||
next_refresh_at: None,
|
||||
has_refresh_token: false,
|
||||
@@ -299,11 +297,9 @@ impl ProviderAccount {
|
||||
mut self,
|
||||
enabled: bool,
|
||||
availability: AccountAvailability,
|
||||
cooldown_until: Option<SystemTime>,
|
||||
) -> Self {
|
||||
self.enabled = enabled;
|
||||
self.availability = availability;
|
||||
self.cooldown_until = cooldown_until;
|
||||
self
|
||||
}
|
||||
|
||||
@@ -373,11 +369,6 @@ impl ProviderAccount {
|
||||
self.enabled
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn cooldown_until(&self) -> Option<SystemTime> {
|
||||
self.cooldown_until
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn access_token_expires_at(&self) -> Option<SystemTime> {
|
||||
self.access_token_expires_at
|
||||
@@ -398,7 +389,6 @@ impl ProviderAccount {
|
||||
pub fn is_schedulable(&self, now: SystemTime) -> bool {
|
||||
let available = match self.availability {
|
||||
AccountAvailability::Ready => true,
|
||||
AccountAvailability::Cooldown => self.cooldown_until.is_some_and(|until| until <= now),
|
||||
AccountAvailability::Unknown
|
||||
| AccountAvailability::QuotaExhausted
|
||||
| AccountAvailability::Expired
|
||||
@@ -588,6 +578,8 @@ pub struct QuotaObservation {
|
||||
pub expected_revision: CredentialRevision,
|
||||
pub quota: Option<OpaqueProviderData>,
|
||||
pub observed_at: Option<SystemTime>,
|
||||
/// 快照级限流事实(Provider 解析后物化);None 表示不更新物化列。
|
||||
pub limit_reached: Option<bool>,
|
||||
}
|
||||
|
||||
impl fmt::Debug for QuotaObservation {
|
||||
@@ -614,9 +606,9 @@ pub struct AccountStateChange {
|
||||
pub account_id: ProviderAccountId,
|
||||
pub expected_revision: CredentialRevision,
|
||||
pub availability: AccountAvailability,
|
||||
pub reason: Option<String>,
|
||||
pub cooldown_until: Option<SystemTime>,
|
||||
pub observed_at: SystemTime,
|
||||
/// 上游原因(错误原文 / 原因码);`Some` 写入持久化错误信息,`None` 不更新,恢复 `Ready` 时清空。
|
||||
pub message: Option<String>,
|
||||
}
|
||||
|
||||
/// `provider_accounts` 的数据库中立端口。
|
||||
@@ -756,6 +748,7 @@ pub struct AccountRuntimeSignals {
|
||||
pub last_started_at: Option<SystemTime>,
|
||||
pub quota_reset_at: Option<SystemTime>,
|
||||
pub quota_remaining_rank: Option<u64>,
|
||||
pub quota_limit_reached: bool,
|
||||
pub failure_rate_basis_points: Option<u16>,
|
||||
pub first_output_latency_ms: Option<u64>,
|
||||
}
|
||||
@@ -903,6 +896,7 @@ impl AccountRuntimeSignals {
|
||||
if let Some(quota) = quota {
|
||||
self.quota_reset_at = quota.reset_at;
|
||||
self.quota_remaining_rank = quota.remaining_rank;
|
||||
self.quota_limit_reached = quota.limit_reached;
|
||||
}
|
||||
self
|
||||
}
|
||||
@@ -924,14 +918,20 @@ impl AccountRuntimeSignals {
|
||||
pub struct AccountQuotaSignals {
|
||||
reset_at: Option<SystemTime>,
|
||||
remaining_rank: Option<u64>,
|
||||
limit_reached: bool,
|
||||
}
|
||||
|
||||
impl AccountQuotaSignals {
|
||||
#[must_use]
|
||||
pub const fn new(reset_at: Option<SystemTime>, remaining_rank: Option<u64>) -> Self {
|
||||
pub const fn new(
|
||||
reset_at: Option<SystemTime>,
|
||||
remaining_rank: Option<u64>,
|
||||
limit_reached: bool,
|
||||
) -> Self {
|
||||
Self {
|
||||
reset_at,
|
||||
remaining_rank,
|
||||
limit_reached,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -944,6 +944,12 @@ impl AccountQuotaSignals {
|
||||
pub const fn remaining_rank(self) -> Option<u64> {
|
||||
self.remaining_rank
|
||||
}
|
||||
|
||||
/// 任一计量窗口已触顶(`limit_reached`/`used_percent >= 100`);调度据此排除候选。
|
||||
#[must_use]
|
||||
pub const fn limit_reached(self) -> bool {
|
||||
self.limit_reached
|
||||
}
|
||||
}
|
||||
|
||||
/// 账号持久事实与可重建运行信号的请求级组合。
|
||||
|
||||
@@ -635,6 +635,13 @@ pub trait ProviderCooldownPort: Send + Sync {
|
||||
scope: &'a ProviderCooldownScope,
|
||||
through_revision: CredentialRevision,
|
||||
) -> BoxFuture<'a, Result<bool, ProviderStoreError>>;
|
||||
|
||||
/// 删除账号时清除该账号全部 account/model scope cooldown key。
|
||||
/// 生命周期与 credential revision 无关;不要求调用方持有 revision。
|
||||
fn clear_all<'a>(
|
||||
&'a self,
|
||||
account_id: &'a ProviderAccountId,
|
||||
) -> BoxFuture<'a, Result<bool, ProviderStoreError>>;
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
|
||||
@@ -513,17 +513,27 @@ fn image_stream(image_output_tokens: Option<u64>) -> Vec<Result<GatewayEvent, Pr
|
||||
}
|
||||
|
||||
fn plan(operation: &Operation) -> RoutingPlan {
|
||||
plan_with_policy(
|
||||
operation,
|
||||
AccountSelectionPolicy::new(
|
||||
RotationStrategy::Smart,
|
||||
NonZeroU32::new(2).expect("account concurrency"),
|
||||
Duration::from_millis(50),
|
||||
),
|
||||
)
|
||||
}
|
||||
|
||||
fn plan_with_policy(
|
||||
operation: &Operation,
|
||||
account_selection_policy: AccountSelectionPolicy,
|
||||
) -> RoutingPlan {
|
||||
let provider = ProviderKind::new("openai").expect("provider");
|
||||
let public_model = PublicModelId::new("gpt-5").expect("public model");
|
||||
let capabilities = ModelCapabilities::new(BTreeSet::from([operation.kind()]), Some(32_000))
|
||||
.with_upstream_feature_validation();
|
||||
RuntimeSnapshot::new(
|
||||
ConfigRevision::new(1).expect("config revision"),
|
||||
AccountSelectionPolicy::new(
|
||||
RotationStrategy::Smart,
|
||||
NonZeroU32::new(2).expect("account concurrency"),
|
||||
Duration::from_millis(50),
|
||||
),
|
||||
account_selection_policy,
|
||||
vec![provider.clone()],
|
||||
vec![ProviderModel::new(
|
||||
provider.clone(),
|
||||
@@ -1514,6 +1524,70 @@ fn native_continuation_replays_owner_before_safely_switching_account() {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unavailable_native_continuation_replays_with_the_configured_account_policy() {
|
||||
let Operation::Generate(generate) = generate_operation() else {
|
||||
panic!("generate operation");
|
||||
};
|
||||
let operation = Operation::Generate(generate.with_provider_session_state(
|
||||
ProviderSessionState::new("openai", Map::new()).expect("provider session state"),
|
||||
));
|
||||
let policy = AccountSelectionPolicy::new(
|
||||
RotationStrategy::RoundRobin,
|
||||
NonZeroU32::new(2).expect("account concurrency"),
|
||||
Duration::from_millis(50),
|
||||
);
|
||||
let route_plan = plan_with_policy(&operation, policy);
|
||||
let (coordinator, store, provider) = coordinator(vec![
|
||||
Script::Error(ProviderError::new(
|
||||
ProviderErrorKind::NoEligibleAccount,
|
||||
UpstreamSendState::NotSent,
|
||||
)),
|
||||
Script::Stream {
|
||||
account_id: "acct_two",
|
||||
items: complete_stream(None),
|
||||
},
|
||||
]);
|
||||
let original = ProviderAccountId::new("acct_one").expect("account");
|
||||
let continuation = NativeContinuationPin::new(
|
||||
PreviousResponseId::new("previous-response"),
|
||||
PreviousResponseId::new("upstream-response"),
|
||||
ProviderKind::new("openai").expect("provider"),
|
||||
original.clone(),
|
||||
);
|
||||
|
||||
let mut session = block_on(coordinator.start(
|
||||
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
|
||||
operation,
|
||||
route_plan,
|
||||
None,
|
||||
Some(ContinuationBinding::Pinned(continuation)),
|
||||
CancellationToken::new(),
|
||||
))
|
||||
.expect("start execution");
|
||||
block_on(session.collect_uncommitted()).expect("portable replay succeeds");
|
||||
block_on(session.commit_downstream(Some(200))).expect("commit response");
|
||||
|
||||
let contexts = provider.contexts.lock().expect("contexts lock");
|
||||
assert_eq!(contexts.len(), 2);
|
||||
assert_eq!(
|
||||
contexts[0].continuation_attempt(),
|
||||
ContinuationAttempt::Native
|
||||
);
|
||||
assert_eq!(
|
||||
contexts[1].continuation_attempt(),
|
||||
ContinuationAttempt::ReplayAny
|
||||
);
|
||||
assert_eq!(
|
||||
contexts[1].account_selection_policy().strategy(),
|
||||
RotationStrategy::RoundRobin
|
||||
);
|
||||
assert!(contexts[1].excluded_accounts().contains(&original));
|
||||
let state = store.state.lock().expect("store lock");
|
||||
assert_eq!(state.attempts.len(), 1);
|
||||
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Succeeded);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn diagnostic_account_reaches_provider_and_matching_metadata_succeeds() {
|
||||
let operation = operation(RetrySafety::Idempotent);
|
||||
|
||||
@@ -24,7 +24,7 @@ fn account(id: &str) -> ProviderAccount {
|
||||
CredentialRevision::new(1).expect("valid revision"),
|
||||
Some(SystemTime::now() + Duration::from_secs(3600)),
|
||||
)
|
||||
.with_runtime_state(true, AccountAvailability::Ready, None)
|
||||
.with_runtime_state(true, AccountAvailability::Ready)
|
||||
}
|
||||
|
||||
fn candidate(id: &str, in_flight: u32, remaining: Option<u64>) -> AccountCandidate {
|
||||
@@ -35,6 +35,7 @@ fn candidate(id: &str, in_flight: u32, remaining: Option<u64>) -> AccountCandida
|
||||
last_started_at: None,
|
||||
quota_reset_at: None,
|
||||
quota_remaining_rank: remaining,
|
||||
quota_limit_reached: false,
|
||||
failure_rate_basis_points: None,
|
||||
first_output_latency_ms: None,
|
||||
},
|
||||
@@ -110,31 +111,27 @@ fn availability_should_round_trip_all_database_values() {
|
||||
#[test]
|
||||
fn diagnostic_selection_bypasses_all_local_account_availability() {
|
||||
let exhausted = AccountCandidate {
|
||||
account: account("acct_exhausted").with_runtime_state(
|
||||
true,
|
||||
AccountAvailability::QuotaExhausted,
|
||||
None,
|
||||
),
|
||||
account: account("acct_exhausted")
|
||||
.with_runtime_state(true, AccountAvailability::QuotaExhausted),
|
||||
signals: AccountRuntimeSignals {
|
||||
in_flight: 0,
|
||||
last_started_at: None,
|
||||
quota_reset_at: None,
|
||||
quota_remaining_rank: None,
|
||||
quota_limit_reached: false,
|
||||
failure_rate_basis_points: None,
|
||||
first_output_latency_ms: None,
|
||||
},
|
||||
};
|
||||
let disabled = AccountCandidate {
|
||||
account: account("acct_disabled").with_runtime_state(
|
||||
false,
|
||||
AccountAvailability::QuotaExhausted,
|
||||
None,
|
||||
),
|
||||
account: account("acct_disabled")
|
||||
.with_runtime_state(false, AccountAvailability::QuotaExhausted),
|
||||
signals: AccountRuntimeSignals {
|
||||
in_flight: 0,
|
||||
last_started_at: None,
|
||||
quota_reset_at: None,
|
||||
quota_remaining_rank: None,
|
||||
quota_limit_reached: false,
|
||||
failure_rate_basis_points: None,
|
||||
first_output_latency_ms: None,
|
||||
},
|
||||
@@ -158,33 +155,11 @@ fn diagnostic_selection_bypasses_all_local_account_availability() {
|
||||
|
||||
#[test]
|
||||
fn disabled_account_should_not_be_schedulable() {
|
||||
let account =
|
||||
account("acct_disabled").with_runtime_state(false, AccountAvailability::Ready, None);
|
||||
let account = account("acct_disabled").with_runtime_state(false, AccountAvailability::Ready);
|
||||
|
||||
assert!(!account.is_schedulable(SystemTime::now()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cooldown_account_should_become_schedulable_at_its_deadline() {
|
||||
let now = SystemTime::now();
|
||||
let account =
|
||||
account("acct_cooldown").with_runtime_state(true, AccountAvailability::Cooldown, Some(now));
|
||||
|
||||
assert!(account.is_schedulable(now));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cooldown_account_should_remain_isolated_before_its_deadline() {
|
||||
let now = SystemTime::now();
|
||||
let account = account("acct_cooldown").with_runtime_state(
|
||||
true,
|
||||
AccountAvailability::Cooldown,
|
||||
Some(now + Duration::from_secs(1)),
|
||||
);
|
||||
|
||||
assert!(!account.is_schedulable(now));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn smart_selector_should_prefer_lower_inflight_count() {
|
||||
let candidates = vec![
|
||||
@@ -311,10 +286,15 @@ fn provider_quota_overlay_should_preserve_store_concurrency_facts() {
|
||||
last_started_at: Some(last_started_at),
|
||||
quota_reset_at: None,
|
||||
quota_remaining_rank: None,
|
||||
quota_limit_reached: false,
|
||||
failure_rate_basis_points: None,
|
||||
first_output_latency_ms: None,
|
||||
}
|
||||
.with_provider_quota(Some(AccountQuotaSignals::new(Some(reset_at), Some(75))));
|
||||
.with_provider_quota(Some(AccountQuotaSignals::new(
|
||||
Some(reset_at),
|
||||
Some(75),
|
||||
false,
|
||||
)));
|
||||
|
||||
assert_eq!(
|
||||
(
|
||||
@@ -629,3 +609,48 @@ fn credential_cas_should_reject_refresh_schedule_without_refresh_token() {
|
||||
.is_err()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn enforced_availability_excludes_terminal_states_despite_quota_signals() {
|
||||
// availability 终态(Banned/Expired/Invalid)始终排除,
|
||||
// 不受 quota 信号或 skip_exhausted 放宽影响。
|
||||
let mut banned = candidate("acct_banned", 0, Some(100));
|
||||
banned.account = banned
|
||||
.account
|
||||
.with_runtime_state(true, AccountAvailability::Banned);
|
||||
let mut expired = candidate("acct_expired", 0, Some(100));
|
||||
expired.account = expired
|
||||
.account
|
||||
.with_runtime_state(true, AccountAvailability::Expired);
|
||||
let mut invalid = candidate("acct_invalid", 0, Some(100));
|
||||
invalid.account = invalid
|
||||
.account
|
||||
.with_runtime_state(true, AccountAvailability::Invalid);
|
||||
let healthy = candidate("acct_healthy", 0, Some(100));
|
||||
let candidates = vec![banned, expired, invalid, healthy];
|
||||
|
||||
let selected = AccountSelector
|
||||
.select(&candidates, &context(RotationStrategy::Smart))
|
||||
.expect("healthy candidate available");
|
||||
|
||||
assert_eq!(selected.candidate().account.id().as_str(), "acct_healthy");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn enforced_availability_allows_quota_exhausted_with_skip_exhausted_projection() {
|
||||
// QuotaExhausted 是可恢复额度耗尽,skip_exhausted=false 的
|
||||
// 放宽投影(provider 层)允许其参与调度;core 的 Enforce 策略本身
|
||||
// 只排除终态。这里验证 Enforce 下 QuotaExhausted 仍被排除(严格模式)。
|
||||
let mut exhausted = candidate("acct_exhausted", 0, Some(100));
|
||||
exhausted.account = exhausted
|
||||
.account
|
||||
.with_runtime_state(true, AccountAvailability::QuotaExhausted);
|
||||
let healthy = candidate("acct_healthy", 0, Some(100));
|
||||
let candidates = vec![exhausted, healthy];
|
||||
|
||||
let selected = AccountSelector
|
||||
.select(&candidates, &context(RotationStrategy::Smart))
|
||||
.expect("healthy candidate available");
|
||||
|
||||
assert_eq!(selected.candidate().account.id().as_str(), "acct_healthy");
|
||||
}
|
||||
|
||||
@@ -101,6 +101,7 @@ fn scheduling_state_preserves_provider_neutral_signals() {
|
||||
last_started_at: None,
|
||||
quota_reset_at: None,
|
||||
quota_remaining_rank: Some(7),
|
||||
quota_limit_reached: false,
|
||||
failure_rate_basis_points: Some(125),
|
||||
first_output_latency_ms: Some(250),
|
||||
},
|
||||
|
||||
@@ -37,7 +37,7 @@ uuid = { version = "1.23.3", features = ["v4", "v7"] }
|
||||
|
||||
[dev-dependencies]
|
||||
tempfile = "3.23"
|
||||
tokio = { version = "1.52.3", default-features = false, features = ["macros", "rt-multi-thread", "fs", "io-util", "process", "time"] }
|
||||
tokio = { version = "1.52.3", default-features = false, features = ["macros", "rt-multi-thread", "fs", "io-util", "net", "process", "time"] }
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
|
||||
@@ -385,7 +385,13 @@ async fn upload_parts(
|
||||
BackupError::new(code::S3_UPLOAD_FAILED, "上传并发闸门关闭".to_owned())
|
||||
})?;
|
||||
let offset = (part - 1) * (PART_SIZE as u64);
|
||||
let bytes = read_part(&source, offset, PART_SIZE).await?;
|
||||
let remaining = file_size.checked_sub(offset).ok_or_else(|| {
|
||||
BackupError::new(code::S3_UPLOAD_FAILED, "分片偏移超出归档大小".to_owned())
|
||||
})?;
|
||||
let part_len = usize::try_from(remaining.min(PART_SIZE as u64)).map_err(|_| {
|
||||
BackupError::new(code::S3_UPLOAD_FAILED, "分片大小溢出".to_owned())
|
||||
})?;
|
||||
let bytes = read_part(&source, offset, part_len).await?;
|
||||
let part_number = i32::try_from(part).map_err(|_| {
|
||||
BackupError::new(code::S3_UPLOAD_FAILED, "分片号溢出".to_owned())
|
||||
})?;
|
||||
@@ -450,20 +456,19 @@ async fn upload_parts(
|
||||
.map_err(map_s3_error)
|
||||
}
|
||||
|
||||
/// 从文件中读取一个分片(最多 `PART_SIZE` 字节)。
|
||||
async fn read_part(source: &Path, offset: u64, max_len: usize) -> Result<Vec<u8>, BackupError> {
|
||||
/// 从文件中读取指定长度的完整分片。
|
||||
async fn read_part(source: &Path, offset: u64, len: usize) -> Result<Vec<u8>, BackupError> {
|
||||
let mut file = tokio::fs::File::open(source)
|
||||
.await
|
||||
.map_err(|_| BackupError::new(code::S3_UPLOAD_FAILED, "打开暂存归档失败".to_owned()))?;
|
||||
file.seek(SeekFrom::Start(offset))
|
||||
.await
|
||||
.map_err(|_| BackupError::new(code::S3_UPLOAD_FAILED, "定位暂存归档失败".to_owned()))?;
|
||||
let mut buffer = vec![0_u8; max_len];
|
||||
let n = file
|
||||
.read(&mut buffer)
|
||||
let mut buffer = vec![0_u8; len];
|
||||
// Tokio 默认单次文件读取最多 2 MiB;分片必须读取完整,不能把短读当作末尾分片上传。
|
||||
file.read_exact(&mut buffer)
|
||||
.await
|
||||
.map_err(|_| BackupError::new(code::S3_UPLOAD_FAILED, "读取暂存归档失败".to_owned()))?;
|
||||
buffer.truncate(n);
|
||||
Ok(buffer)
|
||||
}
|
||||
|
||||
|
||||
@@ -1,8 +1,11 @@
|
||||
//! 完成连接、迁移与 hydration 的 Store 能力集合与启动屏障。
|
||||
|
||||
use gateway_core::engine::credential::ProviderAccountStore;
|
||||
use gateway_core::provider_ports::ProviderCooldownPort;
|
||||
|
||||
use super::*;
|
||||
|
||||
/// 已完成连接、迁移与 cooldown hydration 的 Store 能力集合。
|
||||
/// 已完成连接、迁移与 hydration 的 Store 能力集合。
|
||||
pub struct StoreBundle {
|
||||
admin_ports: AdminStorePorts,
|
||||
core_ports: CoreStorePorts,
|
||||
@@ -61,12 +64,7 @@ pub async fn initialize(mut config: StoreConfig) -> StoreResult<StoreBundle> {
|
||||
redis_connection.clone(),
|
||||
REDIS_NAMESPACE,
|
||||
)?);
|
||||
let account_store = Arc::new(redis::CooldownCachingProviderAccountStore::new(
|
||||
provider_accounts,
|
||||
cooldowns.clone(),
|
||||
));
|
||||
let hydrated = account_store.hydrate(std::time::SystemTime::now()).await;
|
||||
tracing::info!(hydrated, "账号 cooldown Redis 热缓存重建完成");
|
||||
let account_store: Arc<dyn ProviderAccountStore> = provider_accounts;
|
||||
|
||||
let credential_leases =
|
||||
redis::RedisCredentialLeaseRepository::new(redis_connection.clone(), REDIS_NAMESPACE)?;
|
||||
@@ -88,7 +86,10 @@ pub async fn initialize(mut config: StoreConfig) -> StoreResult<StoreBundle> {
|
||||
)?);
|
||||
|
||||
let admin_ports = AdminStorePorts::new(
|
||||
Arc::new(postgres::PgAdminAccountStore::new(pool.clone())),
|
||||
Arc::new(postgres::PgAdminAccountStore::new(
|
||||
pool.clone(),
|
||||
Some(Arc::clone(&cooldowns) as Arc<dyn ProviderCooldownPort>),
|
||||
)),
|
||||
Arc::new(AdminAuthStoreAdapter {
|
||||
security: postgres::PgAdminSecurityAuditRepository::new(pool.clone()),
|
||||
settings: postgres::PgRuntimeSettingsRepository::new(pool.clone()),
|
||||
@@ -101,6 +102,7 @@ pub async fn initialize(mut config: StoreConfig) -> StoreResult<StoreBundle> {
|
||||
Arc::new(postgres::PgAdminObservabilityStore::new(
|
||||
pool.clone(),
|
||||
Some(credential_leases.clone()),
|
||||
Some(Arc::clone(&cooldowns) as Arc<dyn ProviderCooldownPort>),
|
||||
)),
|
||||
Arc::new(AdminSettingsStoreAdapter {
|
||||
control_plane: postgres::PgControlPlaneRepository::new(pool.clone()),
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
//! 八张终态业务表的 PostgreSQL adapters。
|
||||
//! 九张终态业务表的 PostgreSQL adapters。
|
||||
|
||||
use async_trait::async_trait;
|
||||
use sqlx::{PgPool, postgres::PgPoolOptions};
|
||||
|
||||
@@ -1,5 +1,10 @@
|
||||
//! Pg 观测 adapter:实现 `ObservabilityRepository` 与 `AdminObservabilityStore`。
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::SystemTime;
|
||||
|
||||
use gateway_core::provider_ports::ProviderCooldownPort;
|
||||
|
||||
use super::*;
|
||||
|
||||
use crate::redis::{CredentialLeaseRepository as _, RedisCredentialLeaseRepository};
|
||||
@@ -7,12 +12,82 @@ use crate::redis::{CredentialLeaseRepository as _, RedisCredentialLeaseRepositor
|
||||
#[derive(Clone)]
|
||||
pub struct PgObservabilityRepository {
|
||||
pool: PgPool,
|
||||
cooldowns: Option<Arc<dyn ProviderCooldownPort>>,
|
||||
}
|
||||
|
||||
impl PgObservabilityRepository {
|
||||
#[must_use]
|
||||
pub const fn new(pool: PgPool) -> Self {
|
||||
Self { pool }
|
||||
pub const fn new(pool: PgPool, cooldowns: Option<Arc<dyn ProviderCooldownPort>>) -> Self {
|
||||
Self { pool, cooldowns }
|
||||
}
|
||||
|
||||
/// 账号池指标:SQL 聚合后,仅用 Redis 429 冷却重分类仍可调度的 ready 账号。
|
||||
/// 冷却中从 `active` 或 `quota_exhausted` 移到 `rate_limited`;冷却不可用时保留 SQL 分类。
|
||||
async fn provider_account_metrics_with_cooldowns(
|
||||
&self,
|
||||
observed_at: DateTime<Utc>,
|
||||
) -> StoreResult<ProviderAccountMetrics> {
|
||||
let base = provider_account_metrics(&self.pool, observed_at).await?;
|
||||
let Some(cooldowns) = &self.cooldowns else {
|
||||
return Ok(base);
|
||||
};
|
||||
let candidates = schedulable_metric_candidates(&self.pool, observed_at).await?;
|
||||
if candidates.is_empty() {
|
||||
return Ok(base);
|
||||
}
|
||||
// Redis 只能把 SQL 已归入 active/quota_exhausted 的候选移到冷却桶;
|
||||
// 其余状态不重建,避免 Redis 失败或脏 ID 时丢失持久化计数。
|
||||
let mut active = base.active;
|
||||
let mut rate_limited = base.rate_limited;
|
||||
let mut quota_exhausted = base.quota_exhausted;
|
||||
let now = SystemTime::now();
|
||||
for (account_id, revision, quota_reached) in candidates {
|
||||
let Ok(account_id) =
|
||||
gateway_core::engine::credential::ProviderAccountId::new(account_id)
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
let cooling = cooldowns
|
||||
.read(&account_id)
|
||||
.await
|
||||
.ok()
|
||||
.flatten()
|
||||
.filter(|cooldown| cooldown.credential_revision().get() as i64 == revision)
|
||||
.filter(|cooldown| cooldown.until() > now)
|
||||
.is_some();
|
||||
if !cooling {
|
||||
continue;
|
||||
}
|
||||
if quota_reached {
|
||||
if quota_exhausted == 0 {
|
||||
continue;
|
||||
}
|
||||
quota_exhausted -= 1;
|
||||
} else {
|
||||
if active == 0 {
|
||||
continue;
|
||||
}
|
||||
active -= 1;
|
||||
}
|
||||
rate_limited += 1;
|
||||
}
|
||||
Ok(ProviderAccountMetrics {
|
||||
total: base.total,
|
||||
enabled: base.enabled,
|
||||
unavailable: rate_limited
|
||||
+ quota_exhausted
|
||||
+ base.expired
|
||||
+ base.invalid
|
||||
+ base.disabled
|
||||
+ base.banned,
|
||||
active,
|
||||
rate_limited,
|
||||
expired: base.expired,
|
||||
invalid: base.invalid,
|
||||
quota_exhausted,
|
||||
disabled: base.disabled,
|
||||
banned: base.banned,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -28,9 +103,13 @@ pub struct PgAdminObservabilityStore {
|
||||
|
||||
impl PgAdminObservabilityStore {
|
||||
#[must_use]
|
||||
pub fn new(pool: PgPool, runtime_signals: Option<RedisCredentialLeaseRepository>) -> Self {
|
||||
pub fn new(
|
||||
pool: PgPool,
|
||||
runtime_signals: Option<RedisCredentialLeaseRepository>,
|
||||
cooldowns: Option<Arc<dyn ProviderCooldownPort>>,
|
||||
) -> Self {
|
||||
Self {
|
||||
repository: PgObservabilityRepository::new(pool),
|
||||
repository: PgObservabilityRepository::new(pool, cooldowns),
|
||||
runtime_signals,
|
||||
}
|
||||
}
|
||||
@@ -62,11 +141,13 @@ impl ObservabilityRepository for PgObservabilityRepository {
|
||||
};
|
||||
// 分两批并发:页面延迟从六路之和降为两批各自最慢者之和,同时单次
|
||||
// 渲染最多占用 3 个池连接,不与数据面写路径争抢整个连接池。
|
||||
let (requests, attempts, provider_accounts) = futures::try_join!(
|
||||
let (requests, attempts) = futures::try_join!(
|
||||
request_metrics(&self.pool, range, &filter),
|
||||
attempt_metrics(&self.pool, range, &filter),
|
||||
provider_account_metrics(&self.pool, range.end),
|
||||
)?;
|
||||
let provider_accounts = self
|
||||
.provider_account_metrics_with_cooldowns(range.end)
|
||||
.await?;
|
||||
let (trend, account_usage, recent_requests) = futures::try_join!(
|
||||
request_metric_series(&self.pool, range, &filter),
|
||||
provider_account_usage(&self.pool, account_usage_query),
|
||||
|
||||
@@ -263,7 +263,9 @@ pub(crate) const fn admin_account_pool_metrics(
|
||||
enabled: metrics.enabled,
|
||||
unavailable: metrics.unavailable,
|
||||
active: metrics.active,
|
||||
rate_limited: metrics.rate_limited,
|
||||
expired: metrics.expired,
|
||||
invalid: metrics.invalid,
|
||||
quota_exhausted: metrics.quota_exhausted,
|
||||
refreshing: None,
|
||||
disabled: metrics.disabled,
|
||||
|
||||
@@ -366,7 +366,9 @@ pub struct ProviderAccountMetrics {
|
||||
pub enabled: u64,
|
||||
pub unavailable: u64,
|
||||
pub active: u64,
|
||||
pub rate_limited: u64,
|
||||
pub expired: u64,
|
||||
pub invalid: u64,
|
||||
pub quota_exhausted: u64,
|
||||
pub disabled: u64,
|
||||
pub banned: u64,
|
||||
|
||||
@@ -348,29 +348,28 @@ pub(crate) async fn provider_account_metrics(
|
||||
) -> StoreResult<ProviderAccountMetrics> {
|
||||
let row = sqlx::query(
|
||||
"with normalized as (
|
||||
select enabled,
|
||||
select enabled, availability, quota_limit_reached,
|
||||
case
|
||||
when not enabled then 'disabled'
|
||||
when availability = 'banned' then 'banned'
|
||||
when availability = 'quota_exhausted'
|
||||
or (
|
||||
availability = 'cooldown'
|
||||
and availability_reason = 'usage_limit_exhausted'
|
||||
and (cooldown_until is null or cooldown_until > $1)
|
||||
)
|
||||
then 'quota_exhausted'
|
||||
when availability in ('expired', 'invalid')
|
||||
when availability = 'quota_exhausted' then 'quota_exhausted'
|
||||
when availability = 'invalid' then 'invalid'
|
||||
when availability = 'expired'
|
||||
or (access_token_expires_at is not null and access_token_expires_at <= $1)
|
||||
then 'expired'
|
||||
when availability = 'unknown' then 'invalid'
|
||||
when quota_limit_reached then 'quota_exhausted'
|
||||
else 'active'
|
||||
end as status
|
||||
from provider_accounts
|
||||
)
|
||||
select count(*)::bigint as total,
|
||||
count(*) filter (where enabled)::bigint as enabled,
|
||||
count(*) filter (where status <> 'active')::bigint as unavailable,
|
||||
count(*) filter (where status in ('rate_limited','quota_exhausted','expired','invalid','disabled','banned'))::bigint as unavailable,
|
||||
count(*) filter (where status = 'active')::bigint as active,
|
||||
count(*) filter (where status = 'rate_limited')::bigint as rate_limited,
|
||||
count(*) filter (where status = 'expired')::bigint as expired,
|
||||
count(*) filter (where status = 'invalid')::bigint as invalid,
|
||||
count(*) filter (where status = 'quota_exhausted')::bigint as quota_exhausted,
|
||||
count(*) filter (where status = 'disabled')::bigint as disabled,
|
||||
count(*) filter (where status = 'banned')::bigint as banned
|
||||
@@ -385,7 +384,9 @@ pub(crate) async fn provider_account_metrics(
|
||||
enabled: unsigned(&row, "enabled")?,
|
||||
unavailable: unsigned(&row, "unavailable")?,
|
||||
active: unsigned(&row, "active")?,
|
||||
rate_limited: unsigned(&row, "rate_limited")?,
|
||||
expired: unsigned(&row, "expired")?,
|
||||
invalid: unsigned(&row, "invalid")?,
|
||||
quota_exhausted: unsigned(&row, "quota_exhausted")?,
|
||||
disabled: unsigned(&row, "disabled")?,
|
||||
banned: unsigned(&row, "banned")?,
|
||||
@@ -402,13 +403,7 @@ pub(crate) async fn active_provider_account_ids(
|
||||
case
|
||||
when not enabled then 'disabled'
|
||||
when availability = 'banned' then 'banned'
|
||||
when availability = 'quota_exhausted'
|
||||
or (
|
||||
availability = 'cooldown'
|
||||
and availability_reason = 'usage_limit_exhausted'
|
||||
and (cooldown_until is null or cooldown_until > $1)
|
||||
)
|
||||
then 'quota_exhausted'
|
||||
when availability = 'quota_exhausted' then 'quota_exhausted'
|
||||
when availability in ('expired', 'invalid')
|
||||
or (access_token_expires_at is not null and access_token_expires_at <= $1)
|
||||
then 'expired'
|
||||
@@ -430,6 +425,35 @@ pub(crate) async fn active_provider_account_ids(
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// 查询可被 Redis 冷却重新归类的 ready 账号:id、credential revision 与 quota 标记。
|
||||
/// 已过期或非 ready 的账号由 [`provider_account_metrics`] 直接计数。
|
||||
pub(crate) async fn schedulable_metric_candidates(
|
||||
pool: &PgPool,
|
||||
observed_at: DateTime<Utc>,
|
||||
) -> StoreResult<Vec<(String, i64, bool)>> {
|
||||
sqlx::query(
|
||||
"select id, credential_revision, quota_limit_reached
|
||||
from provider_accounts
|
||||
where enabled
|
||||
and availability = 'ready'
|
||||
and (access_token_expires_at is null or access_token_expires_at > $1)
|
||||
order by id",
|
||||
)
|
||||
.bind(observed_at)
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("load schedulable metric candidates"))?
|
||||
.into_iter()
|
||||
.map(|row| {
|
||||
Ok((
|
||||
get(&row, "id")?,
|
||||
get(&row, "credential_revision")?,
|
||||
get(&row, "quota_limit_reached")?,
|
||||
))
|
||||
})
|
||||
.collect::<StoreResult<Vec<_>>>()
|
||||
}
|
||||
|
||||
pub(crate) async fn provider_observations(
|
||||
pool: &PgPool,
|
||||
range: ObservabilityRange,
|
||||
|
||||
@@ -1,5 +1,11 @@
|
||||
//! `gateway-admin` 账号端口的 PostgreSQL adapter。
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::Arc;
|
||||
use std::time::SystemTime;
|
||||
|
||||
use gateway_core::provider_ports::ProviderCooldownPort;
|
||||
|
||||
use super::*;
|
||||
|
||||
/// Admin 账号用例所需的公共账号、留存观测与 revision 事务能力。
|
||||
@@ -11,19 +17,49 @@ pub struct PgAdminAccountStore {
|
||||
accounts: PgProviderAccountRepository,
|
||||
observability: PgObservabilityRepository,
|
||||
control_plane: PgControlPlaneRepository,
|
||||
cooldowns: Option<Arc<dyn ProviderCooldownPort>>,
|
||||
}
|
||||
|
||||
impl PgAdminAccountStore {
|
||||
#[must_use]
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
pub fn new(pool: PgPool, cooldowns: Option<Arc<dyn ProviderCooldownPort>>) -> Self {
|
||||
Self {
|
||||
pool: pool.clone(),
|
||||
accounts: PgProviderAccountRepository::new(pool.clone()),
|
||||
observability: PgObservabilityRepository::new(pool.clone()),
|
||||
observability: PgObservabilityRepository::new(pool.clone(), cooldowns.clone()),
|
||||
control_plane: PgControlPlaneRepository::new(pool),
|
||||
cooldowns,
|
||||
}
|
||||
}
|
||||
|
||||
/// 批量读取 Redis 429 冷却,返回 `account_id → rate_limited_until` 映射。
|
||||
/// 冷却不可用(Redis 挂/无 cooldown 端口)时返回空映射,账号归入配额耗尽(近似)。
|
||||
async fn load_rate_limited_until(
|
||||
&self,
|
||||
accounts: &[ProviderAccountSummary],
|
||||
) -> AdminStoreResult<BTreeMap<String, SystemTime>> {
|
||||
let Some(cooldowns) = &self.cooldowns else {
|
||||
return Ok(BTreeMap::new());
|
||||
};
|
||||
let mut map = BTreeMap::new();
|
||||
for account in accounts {
|
||||
let Ok(account_id) = CoreProviderAccountId::new(account.id.clone()) else {
|
||||
continue;
|
||||
};
|
||||
let Some(cooldown) = cooldowns.read(&account_id).await.ok().flatten() else {
|
||||
continue;
|
||||
};
|
||||
if cooldown.credential_revision().get() != account.credential_revision.get() {
|
||||
continue;
|
||||
}
|
||||
let until = cooldown.until();
|
||||
if until > SystemTime::now() {
|
||||
map.insert(account.id.clone(), until);
|
||||
}
|
||||
}
|
||||
Ok(map)
|
||||
}
|
||||
|
||||
async fn usage_observations(
|
||||
&self,
|
||||
range: ObservabilityRange,
|
||||
@@ -242,11 +278,17 @@ impl AccountStore for PgAdminAccountStore {
|
||||
)
|
||||
.map_err(|error| admin_store_error(ENTITY, error))?;
|
||||
let now = Utc::now();
|
||||
let summary = admin_account_summary(&accounts, now);
|
||||
// 批量读取 Redis 429 冷却,供列表状态派生与筛选使用。
|
||||
let rate_limited_until = self
|
||||
.load_rate_limited_until(&accounts)
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
let summary = admin_account_summary(&accounts, now, &rate_limited_until);
|
||||
let mut items = accounts
|
||||
.into_iter()
|
||||
.filter_map(|account| {
|
||||
let status = admin_account_status(&account, now);
|
||||
let until = rate_limited_until.get(&account.id).copied();
|
||||
let status = admin_account_status(&account, now, until);
|
||||
account_matches_admin_query(&account, status, &query).then_some(
|
||||
AdminAccountListItem {
|
||||
account,
|
||||
|
||||
@@ -32,8 +32,6 @@ impl ProviderAccountStore for PgProviderAccountRepository {
|
||||
next_refresh_at: account.account.next_refresh_at().map(DateTime::<Utc>::from),
|
||||
enabled: account.account.enabled(),
|
||||
availability: account.account.availability(),
|
||||
availability_reason: None,
|
||||
cooldown_until: account.account.cooldown_until().map(DateTime::<Utc>::from),
|
||||
availability_observed_at: Utc::now(),
|
||||
})
|
||||
.await
|
||||
@@ -165,7 +163,8 @@ impl ProviderAccountStore for PgProviderAccountRepository {
|
||||
.map(|account| account.as_str().to_owned())
|
||||
.collect::<Vec<_>>();
|
||||
let rows = sqlx::query(
|
||||
"select id, credential_revision, provider_quota_json, quota_observed_at \
|
||||
"select id, credential_revision, provider_quota_json, quota_observed_at, \
|
||||
quota_limit_reached \
|
||||
from provider_accounts \
|
||||
where id = any($1) and provider_quota_json is not null",
|
||||
)
|
||||
@@ -193,12 +192,16 @@ impl ProviderAccountStore for PgProviderAccountRepository {
|
||||
let observed_at = row
|
||||
.try_get::<DateTime<Utc>, _>("quota_observed_at")
|
||||
.map_err(|_| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?;
|
||||
let limit_reached = row
|
||||
.try_get::<bool, _>("quota_limit_reached")
|
||||
.map_err(|_| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?;
|
||||
Ok(QuotaObservation {
|
||||
account_id: CoreProviderAccountId::new(account_id)
|
||||
.map_err(|_| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?,
|
||||
expected_revision: revision,
|
||||
quota: Some(quota),
|
||||
observed_at: Some(observed_at.into()),
|
||||
limit_reached: Some(limit_reached),
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
@@ -225,6 +228,7 @@ impl ProviderAccountStore for PgProviderAccountRepository {
|
||||
Revision::new(observation.expected_revision.get()).map_err(core_store_error)?,
|
||||
quota,
|
||||
observation.observed_at.map(DateTime::<Utc>::from),
|
||||
observation.limit_reached,
|
||||
)
|
||||
.await
|
||||
.map_err(core_store_error)?;
|
||||
@@ -242,9 +246,8 @@ impl ProviderAccountStore for PgProviderAccountRepository {
|
||||
expected_revision: Revision::new(change.expected_revision.get())
|
||||
.map_err(core_store_error)?,
|
||||
availability: change.availability,
|
||||
availability_reason: change.reason,
|
||||
cooldown_until: change.cooldown_until.map(DateTime::<Utc>::from),
|
||||
availability_observed_at: DateTime::<Utc>::from(change.observed_at),
|
||||
message: change.message,
|
||||
})
|
||||
.await
|
||||
.map_err(core_store_error)?;
|
||||
|
||||
@@ -2,6 +2,8 @@
|
||||
|
||||
use super::*;
|
||||
|
||||
use gateway_admin::model::accounts::AccountStatusSignals;
|
||||
|
||||
pub(crate) struct AdminAccountListItem {
|
||||
pub(crate) account: ProviderAccountSummary,
|
||||
pub(crate) status: AdminAccountStatus,
|
||||
@@ -58,50 +60,48 @@ pub(crate) fn account_matches_admin_query(
|
||||
pub(crate) fn admin_account_status(
|
||||
account: &ProviderAccountSummary,
|
||||
now: DateTime<Utc>,
|
||||
rate_limited_until: Option<SystemTime>,
|
||||
) -> AdminAccountStatus {
|
||||
if !account.enabled {
|
||||
AdminAccountStatus::Disabled
|
||||
} else if account.availability == AccountAvailability::Banned {
|
||||
AdminAccountStatus::Banned
|
||||
} else if account.availability == AccountAvailability::QuotaExhausted
|
||||
|| (account.availability == AccountAvailability::Cooldown
|
||||
&& account.availability_reason.as_deref() == Some("usage_limit_exhausted")
|
||||
&& account.cooldown_until.is_none_or(|until| until > now))
|
||||
{
|
||||
AdminAccountStatus::QuotaExhausted
|
||||
} else if account.availability == AccountAvailability::Expired
|
||||
|| account.availability == AccountAvailability::Invalid
|
||||
|| account
|
||||
// 派生规则唯一归属 gateway-admin 域层,此处只做类型适配。
|
||||
AccountStatusSignals {
|
||||
enabled: account.enabled,
|
||||
availability: account.availability,
|
||||
access_token_expired: account
|
||||
.access_token_expires_at
|
||||
.is_some_and(|expires_at| expires_at <= now)
|
||||
{
|
||||
AdminAccountStatus::Expired
|
||||
} else {
|
||||
AdminAccountStatus::Active
|
||||
.is_some_and(|expires_at| expires_at <= now),
|
||||
quota_limit_reached: account.quota_limit_reached,
|
||||
rate_limited_until,
|
||||
}
|
||||
.derive()
|
||||
}
|
||||
|
||||
pub(crate) fn admin_account_summary(
|
||||
accounts: &[ProviderAccountSummary],
|
||||
now: DateTime<Utc>,
|
||||
rate_limited_until: &BTreeMap<String, SystemTime>,
|
||||
) -> AccountSummary {
|
||||
let mut summary = AccountSummary {
|
||||
total: u64::try_from(accounts.len()).unwrap_or(u64::MAX),
|
||||
active: 0,
|
||||
normal: 0,
|
||||
quota_exhausted: 0,
|
||||
unavailable: 0,
|
||||
rate_limited: 0,
|
||||
disabled: 0,
|
||||
error: 0,
|
||||
};
|
||||
for account in accounts {
|
||||
match admin_account_status(account, now) {
|
||||
AdminAccountStatus::Active => summary.active = summary.active.saturating_add(1),
|
||||
let until = rate_limited_until.get(&account.id).copied();
|
||||
match admin_account_status(account, now, until) {
|
||||
AdminAccountStatus::Normal => summary.normal = summary.normal.saturating_add(1),
|
||||
AdminAccountStatus::QuotaExhausted => {
|
||||
summary.quota_exhausted = summary.quota_exhausted.saturating_add(1);
|
||||
}
|
||||
AdminAccountStatus::Expired
|
||||
| AdminAccountStatus::Disabled
|
||||
| AdminAccountStatus::Banned => {
|
||||
summary.unavailable = summary.unavailable.saturating_add(1);
|
||||
AdminAccountStatus::RateLimited => {
|
||||
summary.rate_limited = summary.rate_limited.saturating_add(1);
|
||||
}
|
||||
AdminAccountStatus::Disabled => {
|
||||
summary.disabled = summary.disabled.saturating_add(1);
|
||||
}
|
||||
AdminAccountStatus::Error => summary.error = summary.error.saturating_add(1),
|
||||
}
|
||||
}
|
||||
summary
|
||||
@@ -118,9 +118,7 @@ pub(crate) fn sort_admin_account_items(
|
||||
items.sort_by(|left, right| {
|
||||
let ordering = match sort.field {
|
||||
AdminAccountSortField::Email => left.account.email.cmp(&right.account.email),
|
||||
AdminAccountSortField::Status => {
|
||||
admin_account_status_name(left.status).cmp(admin_account_status_name(right.status))
|
||||
}
|
||||
AdminAccountSortField::Status => left.status.as_str().cmp(right.status.as_str()),
|
||||
AdminAccountSortField::PlanType => left.account.plan_type.cmp(&right.account.plan_type),
|
||||
AdminAccountSortField::Usage => left
|
||||
.usage
|
||||
@@ -150,16 +148,6 @@ pub(crate) fn sort_admin_account_items(
|
||||
});
|
||||
}
|
||||
|
||||
pub(crate) const fn admin_account_status_name(status: AdminAccountStatus) -> &'static str {
|
||||
match status {
|
||||
AdminAccountStatus::Active => "active",
|
||||
AdminAccountStatus::Expired => "expired",
|
||||
AdminAccountStatus::QuotaExhausted => "quota_exhausted",
|
||||
AdminAccountStatus::Disabled => "disabled",
|
||||
AdminAccountStatus::Banned => "banned",
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn admin_account_record(
|
||||
summary: ProviderAccountSummary,
|
||||
) -> AdminStoreResult<AccountRecord> {
|
||||
@@ -184,9 +172,8 @@ pub(crate) fn admin_account_record(
|
||||
next_refresh_at: summary.next_refresh_at,
|
||||
enabled: summary.enabled,
|
||||
availability: summary.availability,
|
||||
availability_reason: summary.availability_reason,
|
||||
cooldown_until: summary.cooldown_until,
|
||||
availability_observed_at: summary.availability_observed_at,
|
||||
last_error_message: summary.last_error_message,
|
||||
quota_observed_at: summary.quota_observed_at,
|
||||
created_at: summary.created_at,
|
||||
updated_at: summary.updated_at,
|
||||
@@ -211,8 +198,6 @@ pub(crate) fn prepared_account(
|
||||
next_refresh_at: credential.next_refresh_at,
|
||||
enabled: credential.enabled,
|
||||
availability: credential.availability,
|
||||
availability_reason: credential.availability_reason,
|
||||
cooldown_until: credential.cooldown_until,
|
||||
availability_observed_at: credential.availability_observed_at,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -4,6 +4,7 @@ use std::{
|
||||
collections::{BTreeMap, BTreeSet},
|
||||
fmt,
|
||||
str::FromStr,
|
||||
time::SystemTime,
|
||||
};
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
@@ -27,6 +27,7 @@ pub trait ProviderAccountRepository: Send + Sync {
|
||||
expected_revision: Revision,
|
||||
quota: Option<JsonObject>,
|
||||
observed_at: Option<DateTime<Utc>>,
|
||||
limit_reached: Option<bool>,
|
||||
) -> StoreResult<bool>;
|
||||
async fn delete_provider_account(&self, id: &str) -> StoreResult<bool>;
|
||||
}
|
||||
@@ -93,8 +94,8 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
|
||||
"select id, provider_kind, name, email, upstream_user_id,
|
||||
upstream_account_id, plan_type, authentication_kind, credential_revision, has_refresh_token,
|
||||
access_token_expires_at, next_refresh_at, enabled, availability,
|
||||
availability_reason, cooldown_until, availability_observed_at,
|
||||
quota_observed_at, created_at, updated_at
|
||||
availability_observed_at,
|
||||
quota_observed_at, quota_limit_reached, last_error_message, created_at, updated_at
|
||||
from provider_accounts
|
||||
where ($1::text is null or provider_kind = $1) and ($2 or enabled)
|
||||
order by provider_kind, name, id",
|
||||
@@ -114,11 +115,11 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
|
||||
id, provider_kind, name, email, upstream_user_id,
|
||||
upstream_account_id, plan_type, authentication_kind, provider_credentials_json, credential_revision,
|
||||
has_refresh_token, access_token_expires_at, next_refresh_at, enabled,
|
||||
availability, availability_reason, cooldown_until, provider_quota_json,
|
||||
availability, provider_quota_json,
|
||||
availability_observed_at, quota_observed_at, created_at, updated_at
|
||||
) values (
|
||||
$1, $2, $3, $4, $5, $6, $7, $8, $9, 1, $10, $11, $12, $13,
|
||||
$14, null, $15, null, $16, null, now(), greatest(now(), $16)
|
||||
$14, null, $15, null, now(), greatest(now(), $15)
|
||||
)",
|
||||
)
|
||||
.bind(account.id)
|
||||
@@ -135,7 +136,6 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
|
||||
.bind(account.next_refresh_at)
|
||||
.bind(account.enabled)
|
||||
.bind(account.availability.as_str())
|
||||
.bind(account.cooldown_until)
|
||||
.bind(account.availability_observed_at)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
@@ -210,22 +210,24 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
|
||||
let result = sqlx::query(
|
||||
"update provider_accounts
|
||||
set availability = case when enabled then $3 else availability end,
|
||||
availability_reason = case when enabled then $4 else availability_reason end,
|
||||
cooldown_until = case when enabled then $5 else cooldown_until end,
|
||||
availability_observed_at = case
|
||||
when enabled then $6
|
||||
when enabled then $4
|
||||
else availability_observed_at
|
||||
end,
|
||||
updated_at = case when enabled then greatest(now(), $6) else updated_at end
|
||||
last_error_message = case
|
||||
when enabled then
|
||||
case when $3 = 'ready' then null else coalesce($5, last_error_message) end
|
||||
else last_error_message
|
||||
end,
|
||||
updated_at = case when enabled then greatest(now(), $4) else updated_at end
|
||||
where id = $1 and credential_revision = $2
|
||||
and (availability_observed_at is null or availability_observed_at <= $6)",
|
||||
and (availability_observed_at is null or availability_observed_at <= $4)",
|
||||
)
|
||||
.bind(update.account_id)
|
||||
.bind(to_i64(update.expected_revision.get())?)
|
||||
.bind(update.availability.as_str())
|
||||
.bind(update.availability_reason)
|
||||
.bind(update.cooldown_until)
|
||||
.bind(update.availability_observed_at)
|
||||
.bind(update.message.as_deref())
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("apply provider account state"))?;
|
||||
@@ -251,6 +253,7 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
|
||||
expected_revision: Revision,
|
||||
quota: Option<JsonObject>,
|
||||
observed_at: Option<DateTime<Utc>>,
|
||||
limit_reached: Option<bool>,
|
||||
) -> StoreResult<bool> {
|
||||
require_nonempty(ENTITY, "account_id", account_id)?;
|
||||
if quota.is_some() != observed_at.is_some() {
|
||||
@@ -262,6 +265,7 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
|
||||
let result = sqlx::query(
|
||||
"update provider_accounts
|
||||
set provider_quota_json = $3, quota_observed_at = $4,
|
||||
quota_limit_reached = coalesce($5, quota_limit_reached),
|
||||
updated_at = greatest(now(), coalesce($4, now()))
|
||||
where id = $1 and credential_revision = $2
|
||||
and ($4 is null or quota_observed_at is null or quota_observed_at <= $4)",
|
||||
@@ -270,6 +274,7 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
|
||||
.bind(to_i64(expected_revision.get())?)
|
||||
.bind(quota.map(|value| value.as_value()))
|
||||
.bind(observed_at)
|
||||
.bind(limit_reached)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("compare and swap provider quota"))?;
|
||||
@@ -466,11 +471,11 @@ pub(crate) async fn upsert_provider_account_in_transaction(
|
||||
id, provider_kind, name, email, upstream_user_id,
|
||||
upstream_account_id, plan_type, authentication_kind, provider_credentials_json, credential_revision,
|
||||
has_refresh_token, access_token_expires_at, next_refresh_at, enabled,
|
||||
availability, availability_reason, cooldown_until, provider_quota_json,
|
||||
availability, provider_quota_json,
|
||||
availability_observed_at, quota_observed_at, created_at, updated_at
|
||||
) values (
|
||||
$1, $2, $3, $4, $5, $6, $7, $8, $9, 1, $10, $11, $12, $13,
|
||||
$14, $15, $16, null, $17, null, now(), greatest(now(), $17)
|
||||
$14, null, $15, null, now(), greatest(now(), $15)
|
||||
)
|
||||
on conflict (
|
||||
provider_kind,
|
||||
@@ -488,8 +493,6 @@ pub(crate) async fn upsert_provider_account_in_transaction(
|
||||
next_refresh_at = excluded.next_refresh_at,
|
||||
enabled = excluded.enabled,
|
||||
availability = excluded.availability,
|
||||
availability_reason = excluded.availability_reason,
|
||||
cooldown_until = excluded.cooldown_until,
|
||||
provider_quota_json = null,
|
||||
availability_observed_at = excluded.availability_observed_at,
|
||||
quota_observed_at = null,
|
||||
@@ -510,8 +513,6 @@ pub(crate) async fn upsert_provider_account_in_transaction(
|
||||
.bind(account.next_refresh_at)
|
||||
.bind(account.enabled)
|
||||
.bind(account.availability.as_str())
|
||||
.bind(&account.availability_reason)
|
||||
.bind(account.cooldown_until)
|
||||
.bind(account.availability_observed_at)
|
||||
.fetch_optional(&mut **transaction)
|
||||
.await
|
||||
@@ -565,16 +566,6 @@ pub(crate) async fn rotate_provider_account_in_transaction(
|
||||
when availability <> 'quota_exhausted' then 'ready'
|
||||
else availability
|
||||
end,
|
||||
availability_reason = case
|
||||
when not enabled then availability_reason
|
||||
when availability <> 'quota_exhausted' then null
|
||||
else availability_reason
|
||||
end,
|
||||
cooldown_until = case
|
||||
when not enabled then cooldown_until
|
||||
when availability <> 'quota_exhausted' then null
|
||||
else cooldown_until
|
||||
end,
|
||||
availability_observed_at = case
|
||||
when not enabled then availability_observed_at
|
||||
when availability <> 'quota_exhausted' then now()
|
||||
|
||||
@@ -70,26 +70,6 @@ pub(crate) fn parse_availability(value: &str) -> StoreResult<AccountAvailability
|
||||
AccountAvailability::parse(value).ok_or_else(|| invalid("unknown availability value"))
|
||||
}
|
||||
|
||||
pub(crate) fn validate_availability_cooldown(
|
||||
availability: AccountAvailability,
|
||||
has_cooldown_until: bool,
|
||||
) -> StoreResult<()> {
|
||||
if availability == AccountAvailability::Cooldown && !has_cooldown_until {
|
||||
return Err(invalid("cooldown availability requires cooldown_until"));
|
||||
}
|
||||
if has_cooldown_until
|
||||
&& !matches!(
|
||||
availability,
|
||||
AccountAvailability::Cooldown | AccountAvailability::QuotaExhausted
|
||||
)
|
||||
{
|
||||
return Err(invalid(
|
||||
"cooldown_until is only allowed for cooldown or quota_exhausted availability",
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ProviderAccountSummary {
|
||||
pub id: String,
|
||||
@@ -106,10 +86,12 @@ pub struct ProviderAccountSummary {
|
||||
pub next_refresh_at: Option<DateTime<Utc>>,
|
||||
pub enabled: bool,
|
||||
pub availability: AccountAvailability,
|
||||
pub availability_reason: Option<String>,
|
||||
pub cooldown_until: Option<DateTime<Utc>>,
|
||||
pub availability_observed_at: DateTime<Utc>,
|
||||
pub quota_observed_at: Option<DateTime<Utc>>,
|
||||
/// 快照级限流事实(Provider 物化)。
|
||||
pub quota_limit_reached: bool,
|
||||
/// 最近一次非 ready 状态变更的原始原因(上游错误原文 / 原因码);恢复 ready 后清空。
|
||||
pub last_error_message: Option<String>,
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub updated_at: DateTime<Utc>,
|
||||
}
|
||||
@@ -151,8 +133,6 @@ pub struct NewProviderAccount {
|
||||
pub next_refresh_at: Option<DateTime<Utc>>,
|
||||
pub enabled: bool,
|
||||
pub availability: AccountAvailability,
|
||||
pub availability_reason: Option<String>,
|
||||
pub cooldown_until: Option<DateTime<Utc>>,
|
||||
pub availability_observed_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
@@ -173,7 +153,6 @@ impl fmt::Debug for NewProviderAccount {
|
||||
.field("name", &self.name)
|
||||
.field("enabled", &self.enabled)
|
||||
.field("availability", &self.availability)
|
||||
.field("cooldown_until", &self.cooldown_until)
|
||||
.field("provider_credentials_json", &"[REDACTED]")
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
@@ -189,7 +168,6 @@ impl NewProviderAccount {
|
||||
if !self.has_refresh_token && self.next_refresh_at.is_some() {
|
||||
return Err(invalid("next_refresh_at requires a refresh token"));
|
||||
}
|
||||
validate_availability_cooldown(self.availability, self.cooldown_until.is_some())?;
|
||||
validate_object_size(
|
||||
"provider_credentials_json",
|
||||
&self.provider_credentials_json,
|
||||
@@ -321,11 +299,11 @@ impl fmt::Debug for ProviderCredentialUpdate {
|
||||
pub struct ProviderAccountObservation {
|
||||
pub account_id: String,
|
||||
pub availability: AccountAvailability,
|
||||
pub availability_reason: Option<String>,
|
||||
pub cooldown_until: Option<DateTime<Utc>>,
|
||||
pub provider_quota_json: Option<JsonObject>,
|
||||
pub availability_observed_at: DateTime<Utc>,
|
||||
pub quota_observed_at: Option<DateTime<Utc>>,
|
||||
/// 快照级限流事实(Provider 物化)。
|
||||
pub quota_limit_reached: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
@@ -333,15 +311,14 @@ pub struct ProviderAccountStateUpdate {
|
||||
pub account_id: String,
|
||||
pub expected_revision: Revision,
|
||||
pub availability: AccountAvailability,
|
||||
pub availability_reason: Option<String>,
|
||||
pub cooldown_until: Option<DateTime<Utc>>,
|
||||
pub availability_observed_at: DateTime<Utc>,
|
||||
/// 上游原因(错误原文 / 原因码);`None` 不更新错误信息列,恢复 `Ready` 时清空。
|
||||
pub message: Option<String>,
|
||||
}
|
||||
|
||||
impl ProviderAccountStateUpdate {
|
||||
pub fn validate(&self) -> StoreResult<()> {
|
||||
require_nonempty(ENTITY, "account_id", &self.account_id)?;
|
||||
validate_availability_cooldown(self.availability, self.cooldown_until.is_some())?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -349,7 +326,6 @@ impl ProviderAccountStateUpdate {
|
||||
impl ProviderAccountObservation {
|
||||
pub fn validate(&self) -> StoreResult<()> {
|
||||
require_nonempty(ENTITY, "account_id", &self.account_id)?;
|
||||
validate_availability_cooldown(self.availability, self.cooldown_until.is_some())?;
|
||||
if self.provider_quota_json.is_some() != self.quota_observed_at.is_some() {
|
||||
return Err(invalid("quota JSON and quota_observed_at must agree"));
|
||||
}
|
||||
@@ -363,14 +339,14 @@ impl ProviderAccountObservation {
|
||||
pub(crate) const ACCOUNT_SELECT: &str = "select id, provider_kind, name, email, upstream_user_id,
|
||||
upstream_account_id, plan_type, authentication_kind, provider_credentials_json, credential_revision,
|
||||
has_refresh_token, access_token_expires_at, next_refresh_at, enabled, availability,
|
||||
availability_reason, cooldown_until, provider_quota_json,
|
||||
provider_quota_json, quota_limit_reached, last_error_message,
|
||||
availability_observed_at, quota_observed_at, created_at, updated_at
|
||||
from provider_accounts where id = $1";
|
||||
|
||||
pub(crate) const ACCOUNT_SELECT_BY_IDS: &str = "select id, provider_kind, name, email, upstream_user_id,
|
||||
upstream_account_id, plan_type, authentication_kind, provider_credentials_json, credential_revision,
|
||||
has_refresh_token, access_token_expires_at, next_refresh_at, enabled, availability,
|
||||
availability_reason, cooldown_until, provider_quota_json,
|
||||
provider_quota_json, quota_limit_reached, last_error_message,
|
||||
availability_observed_at, quota_observed_at, created_at, updated_at
|
||||
from provider_accounts
|
||||
where id = any($1::text[]) and provider_kind = $2
|
||||
@@ -420,11 +396,7 @@ pub(crate) fn core_account_from_summary(
|
||||
summary.upstream_account_id,
|
||||
summary.plan_type,
|
||||
)
|
||||
.with_runtime_state(
|
||||
summary.enabled,
|
||||
summary.availability,
|
||||
summary.cooldown_until.map(Into::into),
|
||||
)
|
||||
.with_runtime_state(summary.enabled, summary.availability)
|
||||
.with_refresh_schedule(
|
||||
summary.has_refresh_token,
|
||||
summary.next_refresh_at.map(Into::into),
|
||||
@@ -474,10 +446,10 @@ pub(crate) fn account_summary_from_row(
|
||||
next_refresh_at: get(&row, "next_refresh_at")?,
|
||||
enabled: get(&row, "enabled")?,
|
||||
availability: parse_availability(&availability)?,
|
||||
availability_reason: get(&row, "availability_reason")?,
|
||||
cooldown_until: get(&row, "cooldown_until")?,
|
||||
availability_observed_at: get(&row, "availability_observed_at")?,
|
||||
quota_observed_at: get(&row, "quota_observed_at")?,
|
||||
quota_limit_reached: get(&row, "quota_limit_reached")?,
|
||||
last_error_message: get(&row, "last_error_message")?,
|
||||
created_at: get(&row, "created_at")?,
|
||||
updated_at: get(&row, "updated_at")?,
|
||||
})
|
||||
|
||||
@@ -1,22 +1,13 @@
|
||||
//! 请求调度用的可丢失账号 cooldown Redis 存储。
|
||||
|
||||
use std::{sync::Arc, time::SystemTime};
|
||||
//! 请求调度用的可丢失 Provider 级 cooldown Redis 存储。
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, Utc};
|
||||
use gateway_core::{
|
||||
engine::credential::{
|
||||
AccountAvailability, AccountStateChange, CredentialCasOutcome, CredentialCasUpdate,
|
||||
CredentialRevision, LoadedCredential, NewProviderAccount, ProviderAccount,
|
||||
ProviderAccountId, ProviderAccountStore, ProviderAccountUpdate, QuotaObservation,
|
||||
QuotaWriteOutcome,
|
||||
},
|
||||
error::StoreError as CoreStoreError,
|
||||
engine::credential::{CredentialRevision, ProviderAccountId},
|
||||
provider_ports::{
|
||||
ProviderCooldown, ProviderCooldownPort, ProviderCooldownScope, ProviderScopedCooldown,
|
||||
ProviderStoreError, ProviderStoreErrorKind,
|
||||
},
|
||||
routing::ProviderKind,
|
||||
};
|
||||
use redis::{Script, aio::ConnectionManager};
|
||||
|
||||
@@ -85,6 +76,8 @@ pub trait CredentialCooldownRepository: Send + Sync {
|
||||
provider_account_id: &str,
|
||||
through_revision: Revision,
|
||||
) -> StoreResult<bool>;
|
||||
/// 删除账号时清除该账号全部 account/model scope cooldown key。
|
||||
async fn delete_account_cooldowns(&self, provider_account_id: &str) -> StoreResult<bool>;
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
@@ -227,6 +220,50 @@ impl CredentialCooldownRepository for RedisCredentialCooldownRepository {
|
||||
self.invalidate_at_key(self.key(provider_account_id)?, through_revision)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn delete_account_cooldowns(&self, provider_account_id: &str) -> StoreResult<bool> {
|
||||
require_nonempty(
|
||||
"credential cooldown",
|
||||
"provider_account_id",
|
||||
provider_account_id,
|
||||
)?;
|
||||
// 账号删除:清除该账号的 account key 与全部 model-scoped key。
|
||||
// 用 SCAN 精确匹配命名空间内该账号前缀,避免 KEYS 阻塞。
|
||||
let mut connection = self.connection.clone();
|
||||
let account_key = self.key(provider_account_id)?;
|
||||
let mut keys = vec![account_key.clone()];
|
||||
let pattern = format!(
|
||||
"{}:account:{}:cooldown:*",
|
||||
self.namespace,
|
||||
resource_fingerprint("credential cooldown", provider_account_id)?
|
||||
);
|
||||
let mut cursor = 0_i64;
|
||||
loop {
|
||||
let (next, found): (i64, Vec<String>) = redis::cmd("SCAN")
|
||||
.arg(cursor)
|
||||
.arg("MATCH")
|
||||
.arg(&pattern)
|
||||
.arg("COUNT")
|
||||
.arg(100)
|
||||
.query_async(&mut connection)
|
||||
.await
|
||||
.map_err(|_| redis_unavailable("scan account cooldown keys"))?;
|
||||
keys.extend(found);
|
||||
cursor = next;
|
||||
if cursor == 0 {
|
||||
break;
|
||||
}
|
||||
}
|
||||
if keys.is_empty() {
|
||||
return Ok(false);
|
||||
}
|
||||
let removed: i64 = redis::cmd("DEL")
|
||||
.arg(keys)
|
||||
.query_async(&mut connection)
|
||||
.await
|
||||
.map_err(|_| redis_unavailable("delete account cooldown keys"))?;
|
||||
Ok(removed > 0)
|
||||
}
|
||||
}
|
||||
|
||||
impl ProviderCooldownPort for RedisCredentialCooldownRepository {
|
||||
@@ -350,216 +387,17 @@ impl ProviderCooldownPort for RedisCredentialCooldownRepository {
|
||||
.map_err(|_| provider_unavailable("clear scoped credential cooldown"))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// 账号状态写入到可丢失 cooldown cache 的桥接 adapter。
|
||||
pub struct CooldownCachingProviderAccountStore {
|
||||
authoritative: Arc<dyn ProviderAccountStore>,
|
||||
cooldowns: Arc<dyn CredentialCooldownRepository>,
|
||||
}
|
||||
|
||||
impl CooldownCachingProviderAccountStore {
|
||||
#[must_use]
|
||||
pub const fn new(
|
||||
authoritative: Arc<dyn ProviderAccountStore>,
|
||||
cooldowns: Arc<dyn CredentialCooldownRepository>,
|
||||
) -> Self {
|
||||
Self {
|
||||
authoritative,
|
||||
cooldowns,
|
||||
}
|
||||
fn clear_all<'a>(
|
||||
&'a self,
|
||||
account_id: &'a ProviderAccountId,
|
||||
) -> futures::future::BoxFuture<'a, Result<bool, ProviderStoreError>> {
|
||||
Box::pin(async move {
|
||||
self.delete_account_cooldowns(account_id.as_str())
|
||||
.await
|
||||
.map_err(|_| provider_unavailable("clear all credential cooldowns"))
|
||||
})
|
||||
}
|
||||
|
||||
/// 启动时从 PostgreSQL 尽力重建仍有效的 cooldown cache。
|
||||
pub async fn hydrate(&self, now: SystemTime) -> u64 {
|
||||
let Ok(accounts) = self.authoritative.list_accounts().await else {
|
||||
return 0;
|
||||
};
|
||||
let mut hydrated = 0_u64;
|
||||
for account in accounts {
|
||||
let Some(cooldown) = projection(
|
||||
account.id(),
|
||||
account.revision(),
|
||||
account.availability(),
|
||||
account.cooldown_until(),
|
||||
now,
|
||||
) else {
|
||||
continue;
|
||||
};
|
||||
if matches!(
|
||||
self.cooldowns.cache_credential_cooldown(&cooldown).await,
|
||||
Ok(true)
|
||||
) {
|
||||
hydrated = hydrated.saturating_add(1);
|
||||
}
|
||||
}
|
||||
hydrated
|
||||
}
|
||||
|
||||
async fn mirror(
|
||||
&self,
|
||||
account_id: &ProviderAccountId,
|
||||
credential_revision: CredentialRevision,
|
||||
availability: AccountAvailability,
|
||||
cooldown_until: Option<SystemTime>,
|
||||
) {
|
||||
let Some(cooldown) = projection(
|
||||
account_id,
|
||||
credential_revision,
|
||||
availability,
|
||||
cooldown_until,
|
||||
SystemTime::now(),
|
||||
) else {
|
||||
let Ok(revision) = Revision::new(credential_revision.get()) else {
|
||||
return;
|
||||
};
|
||||
let _ = self
|
||||
.cooldowns
|
||||
.invalidate_credential_cooldown(account_id.as_str(), revision)
|
||||
.await;
|
||||
return;
|
||||
};
|
||||
let _ = self.cooldowns.cache_credential_cooldown(&cooldown).await;
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ProviderAccountStore for CooldownCachingProviderAccountStore {
|
||||
async fn create_account(&self, account: NewProviderAccount) -> Result<(), CoreStoreError> {
|
||||
self.authoritative.create_account(account).await
|
||||
}
|
||||
|
||||
async fn get_account(
|
||||
&self,
|
||||
account: &ProviderAccountId,
|
||||
) -> Result<Option<ProviderAccount>, CoreStoreError> {
|
||||
self.authoritative.get_account(account).await
|
||||
}
|
||||
|
||||
async fn list_accounts(&self) -> Result<Vec<ProviderAccount>, CoreStoreError> {
|
||||
self.authoritative.list_accounts().await
|
||||
}
|
||||
|
||||
async fn list_for_provider(
|
||||
&self,
|
||||
provider: &ProviderKind,
|
||||
) -> Result<Vec<ProviderAccount>, CoreStoreError> {
|
||||
self.authoritative.list_for_provider(provider).await
|
||||
}
|
||||
|
||||
async fn load_credential(
|
||||
&self,
|
||||
account: &ProviderAccountId,
|
||||
expected_revision: CredentialRevision,
|
||||
) -> Result<LoadedCredential, CoreStoreError> {
|
||||
self.authoritative
|
||||
.load_credential(account, expected_revision)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn load_current_credential(
|
||||
&self,
|
||||
account: &ProviderAccountId,
|
||||
) -> Result<LoadedCredential, CoreStoreError> {
|
||||
self.authoritative.load_current_credential(account).await
|
||||
}
|
||||
|
||||
async fn compare_and_swap_credential(
|
||||
&self,
|
||||
update: CredentialCasUpdate,
|
||||
) -> Result<CredentialCasOutcome, CoreStoreError> {
|
||||
let account_id = update.account_id().clone();
|
||||
let outcome = self
|
||||
.authoritative
|
||||
.compare_and_swap_credential(update)
|
||||
.await?;
|
||||
if let CredentialCasOutcome::Updated(credential_revision) = outcome
|
||||
&& let Ok(revision) = Revision::new(credential_revision.get())
|
||||
{
|
||||
let _ = self
|
||||
.cooldowns
|
||||
.invalidate_credential_cooldown(account_id.as_str(), revision)
|
||||
.await;
|
||||
}
|
||||
Ok(outcome)
|
||||
}
|
||||
|
||||
async fn get_quotas(
|
||||
&self,
|
||||
accounts: &[ProviderAccountId],
|
||||
) -> Result<Vec<QuotaObservation>, CoreStoreError> {
|
||||
self.authoritative.get_quotas(accounts).await
|
||||
}
|
||||
|
||||
async fn compare_and_swap_quota(
|
||||
&self,
|
||||
observation: QuotaObservation,
|
||||
) -> Result<QuotaWriteOutcome, CoreStoreError> {
|
||||
self.authoritative.compare_and_swap_quota(observation).await
|
||||
}
|
||||
|
||||
async fn apply_state_change(&self, change: AccountStateChange) -> Result<(), CoreStoreError> {
|
||||
let account_id = change.account_id.clone();
|
||||
let credential_revision = change.expected_revision;
|
||||
let availability = change.availability;
|
||||
let cooldown_until = change.cooldown_until;
|
||||
self.authoritative.apply_state_change(change).await?;
|
||||
let Some(account) = self.authoritative.get_account(&account_id).await? else {
|
||||
return Ok(());
|
||||
};
|
||||
if !account.enabled() {
|
||||
return Ok(());
|
||||
}
|
||||
self.mirror(
|
||||
&account_id,
|
||||
credential_revision,
|
||||
availability,
|
||||
cooldown_until,
|
||||
)
|
||||
.await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn update_account(&self, update: ProviderAccountUpdate) -> Result<(), CoreStoreError> {
|
||||
self.authoritative.update_account(update).await
|
||||
}
|
||||
|
||||
async fn set_enabled(
|
||||
&self,
|
||||
account: &ProviderAccountId,
|
||||
enabled: bool,
|
||||
) -> Result<(), CoreStoreError> {
|
||||
self.authoritative.set_enabled(account, enabled).await
|
||||
}
|
||||
|
||||
async fn delete_account(&self, account: &ProviderAccountId) -> Result<(), CoreStoreError> {
|
||||
self.authoritative.delete_account(account).await?;
|
||||
if let Ok(through_revision) = Revision::new(u64::MAX) {
|
||||
let _ = self
|
||||
.cooldowns
|
||||
.invalidate_credential_cooldown(account.as_str(), through_revision)
|
||||
.await;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn projection(
|
||||
account_id: &ProviderAccountId,
|
||||
credential_revision: CredentialRevision,
|
||||
availability: AccountAvailability,
|
||||
cooldown_until: Option<SystemTime>,
|
||||
now: SystemTime,
|
||||
) -> Option<CredentialCooldown> {
|
||||
let cooldown_until = cooldown_until
|
||||
.filter(|_| availability == AccountAvailability::Cooldown)
|
||||
.filter(|until| *until > now)
|
||||
.map(DateTime::<Utc>::from)?;
|
||||
Some(CredentialCooldown {
|
||||
provider_account_id: account_id.as_str().to_owned(),
|
||||
credential_revision: Revision::new(credential_revision.get()).ok()?,
|
||||
cooldown_until,
|
||||
})
|
||||
}
|
||||
|
||||
fn invalid(message: &str) -> StoreError {
|
||||
|
||||
@@ -512,6 +512,7 @@ impl RedisProviderLeaseCoordinator {
|
||||
last_started_at: signal.last_started_at.map(Into::into),
|
||||
quota_reset_at: None,
|
||||
quota_remaining_rank: None,
|
||||
quota_limit_reached: false,
|
||||
failure_rate_basis_points: None,
|
||||
first_output_latency_ms: None,
|
||||
},
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
mod postgres;
|
||||
mod redis;
|
||||
mod s3;
|
||||
mod support;
|
||||
|
||||
@@ -21,6 +21,8 @@ mod runtime_settings;
|
||||
mod snapshot;
|
||||
mod snapshots;
|
||||
|
||||
static TEST_MIGRATOR: sqlx::migrate::Migrator = sqlx::migrate!("../../migrations");
|
||||
|
||||
pub(super) struct TestDatabase {
|
||||
admin: PgPool,
|
||||
pub(super) pool: PgPool,
|
||||
@@ -56,28 +58,10 @@ impl TestDatabase {
|
||||
.connect(&database_url)
|
||||
.await
|
||||
.expect("connect isolated test schema");
|
||||
let mut migration = pool.begin().await.expect("begin terminal migration");
|
||||
sqlx::raw_sql(include_str!("../../../../migrations/0001_initial.sql"))
|
||||
.execute(&mut *migration)
|
||||
TEST_MIGRATOR
|
||||
.run(&pool)
|
||||
.await
|
||||
.expect("apply terminal migration");
|
||||
sqlx::raw_sql(include_str!(
|
||||
"../../../../migrations/0002_snapshot_provider_account_identity.sql"
|
||||
))
|
||||
.execute(&mut *migration)
|
||||
.await
|
||||
.expect("apply snapshot migration");
|
||||
sqlx::raw_sql(include_str!("../../../../migrations/0003_s3_backup.sql"))
|
||||
.execute(&mut *migration)
|
||||
.await
|
||||
.expect("apply backup migration");
|
||||
sqlx::raw_sql(include_str!(
|
||||
"../../../../migrations/0004_drop_provider_accounts_id_kind_uq.sql"
|
||||
))
|
||||
.execute(&mut *migration)
|
||||
.await
|
||||
.expect("apply index drop migration");
|
||||
migration.commit().await.expect("commit terminal migration");
|
||||
.expect("apply test migrations");
|
||||
Some(Self {
|
||||
admin,
|
||||
pool,
|
||||
@@ -172,24 +156,19 @@ async fn connect_and_migrate_should_apply_all_migrations_once_and_reopen_cleanly
|
||||
.expect("drop migration test database");
|
||||
admin.close().await;
|
||||
|
||||
assert_eq!((first_table_count, migration_count), (10, 4));
|
||||
assert_eq!((first_table_count, migration_count), (10, 1));
|
||||
assert_eq!(response_id_types, ["bytea", "bytea"]);
|
||||
assert!(!raw_response_id_index_exists);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn migrations_should_leave_transaction_ownership_to_sqlx() {
|
||||
let transaction_statements = [
|
||||
include_str!("../../../../migrations/0001_initial.sql"),
|
||||
include_str!("../../../../migrations/0002_snapshot_provider_account_identity.sql"),
|
||||
include_str!("../../../../migrations/0003_s3_backup.sql"),
|
||||
include_str!("../../../../migrations/0004_drop_provider_accounts_id_kind_uq.sql"),
|
||||
]
|
||||
.into_iter()
|
||||
.flat_map(str::lines)
|
||||
.map(str::trim)
|
||||
.filter(|line| matches!(*line, "begin;" | "commit;"))
|
||||
.count();
|
||||
let transaction_statements = TEST_MIGRATOR
|
||||
.iter()
|
||||
.flat_map(|migration| migration.sql.as_str().lines())
|
||||
.map(str::trim)
|
||||
.filter(|line| matches!(*line, "begin;" | "commit;"))
|
||||
.count();
|
||||
|
||||
assert_eq!(transaction_statements, 0);
|
||||
}
|
||||
|
||||
@@ -1,10 +1,23 @@
|
||||
use std::num::NonZeroU32;
|
||||
use std::{
|
||||
collections::BTreeMap,
|
||||
num::NonZeroU32,
|
||||
sync::Arc,
|
||||
time::{Duration, SystemTime},
|
||||
};
|
||||
|
||||
use chrono::{TimeDelta, Utc};
|
||||
use futures::future::BoxFuture;
|
||||
use gateway_admin::{
|
||||
model::{PageSize, observability as admin_observability},
|
||||
ports::store::ObservabilityStore as AdminObservabilityStore,
|
||||
};
|
||||
use gateway_core::{
|
||||
engine::credential::{CredentialRevision, ProviderAccountId},
|
||||
provider_ports::{
|
||||
ProviderCooldown, ProviderCooldownPort, ProviderCooldownScope, ProviderScopedCooldown,
|
||||
ProviderStoreError,
|
||||
},
|
||||
};
|
||||
use gateway_store::postgres::{
|
||||
DiagnosticDimension, ObservabilityPageNumber, ObservabilityPageSize, ObservabilityRange,
|
||||
ObservabilityRepository, OpsErrorFilter, OpsErrorQuery, PgAdminObservabilityStore,
|
||||
@@ -83,7 +96,7 @@ async fn observability_preserves_and_filters_opaque_response_ids() {
|
||||
.expect("persist opaque response IDs");
|
||||
let range = ObservabilityRange::new(now - TimeDelta::hours(1), now + TimeDelta::hours(1))
|
||||
.expect("observability range");
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone());
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone(), None);
|
||||
|
||||
let records = repository
|
||||
.list_usage_records(UsageRecordQuery {
|
||||
@@ -122,33 +135,32 @@ async fn dashboard_account_metrics_should_partition_account_statuses() {
|
||||
id, provider_kind, name, upstream_user_id, authentication_kind,
|
||||
provider_credentials_json, credential_revision, has_refresh_token,
|
||||
access_token_expires_at, enabled, availability,
|
||||
availability_reason, cooldown_until,
|
||||
availability_observed_at, created_at, updated_at
|
||||
) values
|
||||
('acct_available', 'openai', 'available', 'user-available', 'oauth',
|
||||
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'ready',
|
||||
null, null, $1, $1, $1),
|
||||
$1, $1, $1),
|
||||
('acct_expired', 'openai', 'expired', 'user-expired', 'oauth',
|
||||
'{}'::jsonb, 1, false, $1 - interval '1 day', true, 'ready',
|
||||
null, null, $1, $1, $1),
|
||||
$1, $1, $1),
|
||||
('acct_rate_cooldown', 'openai', 'rate-cooldown', 'user-rate-cooldown', 'oauth',
|
||||
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'cooldown',
|
||||
'rate_limited', $1 + interval '2 hours', $1, $1, $1),
|
||||
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'ready',
|
||||
$1, $1, $1),
|
||||
('acct_usage_limit', 'xai', 'usage-limit', 'user-usage-limit', 'oauth',
|
||||
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'cooldown',
|
||||
'usage_limit_exhausted', $1 + interval '2 hours', $1, $1, $1),
|
||||
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'quota_exhausted',
|
||||
$1, $1, $1),
|
||||
('acct_banned', 'xai', 'banned', 'user-banned', 'oauth',
|
||||
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'banned',
|
||||
null, null, $1, $1, $1),
|
||||
$1, $1, $1),
|
||||
('acct_disabled', 'xai', 'disabled', 'user-disabled', 'oauth',
|
||||
'{}'::jsonb, 1, false, $1 + interval '1 day', false, 'ready',
|
||||
null, null, $1, $1, $1)",
|
||||
$1, $1, $1)",
|
||||
)
|
||||
.bind(now)
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.expect("seed account metric states");
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone());
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone(), None);
|
||||
let range = ObservabilityRange::new(now - TimeDelta::hours(1), now + TimeDelta::hours(1))
|
||||
.expect("dashboard range");
|
||||
|
||||
@@ -180,6 +192,184 @@ async fn dashboard_account_metrics_should_partition_account_statuses() {
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dashboard_account_metrics_with_cooldowns_should_only_reclassify_eligible_accounts() {
|
||||
let Some(database) = TestDatabase::create("dashboard_account_metrics_cooldowns").await else {
|
||||
return;
|
||||
};
|
||||
let now = Utc::now();
|
||||
sqlx::query(
|
||||
"insert into provider_accounts (
|
||||
id, provider_kind, name, upstream_user_id, authentication_kind,
|
||||
provider_credentials_json, credential_revision, has_refresh_token,
|
||||
access_token_expires_at, enabled, availability, quota_limit_reached,
|
||||
availability_observed_at, created_at, updated_at
|
||||
) values
|
||||
('acct_metrics_active', 'openai', 'active', 'user-active', 'oauth',
|
||||
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'ready', false,
|
||||
$1, $1, $1),
|
||||
('acct_metrics_active_cooling', 'openai', 'active-cooling', 'user-active-cooling', 'oauth',
|
||||
'{}'::jsonb, 2, false, $1 + interval '1 day', true, 'ready', false,
|
||||
$1, $1, $1),
|
||||
('acct_metrics_ready_quota', 'openai', 'ready-quota', 'user-ready-quota', 'oauth',
|
||||
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'ready', true,
|
||||
$1, $1, $1),
|
||||
('acct_metrics_cooling', 'openai', 'cooling', 'user-cooling', 'oauth',
|
||||
'{}'::jsonb, 3, false, $1 + interval '1 day', true, 'ready', true,
|
||||
$1, $1, $1),
|
||||
('acct_metrics_stale_cooldown', 'openai', 'stale-cooldown', 'user-stale-cooldown', 'oauth',
|
||||
'{}'::jsonb, 4, false, $1 + interval '1 day', true, 'ready', true,
|
||||
$1, $1, $1),
|
||||
('acct_metrics_persistent_quota', 'openai', 'persistent-quota', 'user-persistent-quota', 'oauth',
|
||||
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'quota_exhausted', true,
|
||||
$1, $1, $1),
|
||||
('acct_metrics_expired', 'openai', 'expired', 'user-expired', 'oauth',
|
||||
'{}'::jsonb, 5, false, $1 - interval '1 second', true, 'ready', true,
|
||||
$1, $1, $1),
|
||||
('acct_metrics_unknown', 'openai', 'unknown', 'user-unknown', 'oauth',
|
||||
'{}'::jsonb, 6, false, $1 + interval '1 day', true, 'unknown', false,
|
||||
$1, $1, $1),
|
||||
('acct_metrics_disabled', 'openai', 'disabled', 'user-disabled', 'oauth',
|
||||
'{}'::jsonb, 1, false, $1 + interval '1 day', false, 'ready', false,
|
||||
$1, $1, $1)",
|
||||
)
|
||||
.bind(now)
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.expect("seed cooldown account metric states");
|
||||
let cooldowns = StaticCooldowns::new([
|
||||
test_cooldown("acct_metrics_active_cooling", 2),
|
||||
test_cooldown("acct_metrics_cooling", 3),
|
||||
test_cooldown("acct_metrics_expired", 5),
|
||||
test_cooldown("acct_metrics_unknown", 6),
|
||||
test_cooldown("acct_metrics_stale_cooldown", 3),
|
||||
]);
|
||||
let repository =
|
||||
PgObservabilityRepository::new(database.pool.clone(), Some(Arc::new(cooldowns)));
|
||||
let range = ObservabilityRange::new(now - TimeDelta::hours(1), now + TimeDelta::hours(1))
|
||||
.expect("dashboard range");
|
||||
|
||||
let metrics = repository
|
||||
.dashboard_summary(range)
|
||||
.await
|
||||
.expect("dashboard summary")
|
||||
.provider_accounts;
|
||||
|
||||
assert_eq!(
|
||||
(
|
||||
metrics.total,
|
||||
metrics.enabled,
|
||||
metrics.active,
|
||||
metrics.rate_limited,
|
||||
metrics.quota_exhausted,
|
||||
metrics.expired,
|
||||
metrics.invalid,
|
||||
metrics.disabled,
|
||||
metrics.banned,
|
||||
metrics.unavailable,
|
||||
),
|
||||
(9, 8, 1, 2, 3, 1, 1, 1, 0, 8),
|
||||
);
|
||||
assert_eq!(
|
||||
metrics.total,
|
||||
metrics.active
|
||||
+ metrics.rate_limited
|
||||
+ metrics.quota_exhausted
|
||||
+ metrics.expired
|
||||
+ metrics.invalid
|
||||
+ metrics.disabled
|
||||
+ metrics.banned
|
||||
);
|
||||
assert_eq!(
|
||||
metrics.unavailable,
|
||||
metrics.rate_limited
|
||||
+ metrics.quota_exhausted
|
||||
+ metrics.expired
|
||||
+ metrics.invalid
|
||||
+ metrics.disabled
|
||||
+ metrics.banned
|
||||
);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
struct StaticCooldowns {
|
||||
cooldowns: BTreeMap<ProviderAccountId, ProviderCooldown>,
|
||||
}
|
||||
|
||||
impl StaticCooldowns {
|
||||
fn new(cooldowns: impl IntoIterator<Item = ProviderCooldown>) -> Self {
|
||||
Self {
|
||||
cooldowns: cooldowns
|
||||
.into_iter()
|
||||
.map(|cooldown| (cooldown.account_id().clone(), cooldown))
|
||||
.collect(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ProviderCooldownPort for StaticCooldowns {
|
||||
fn put_if_later(
|
||||
&self,
|
||||
_cooldown: ProviderCooldown,
|
||||
) -> BoxFuture<'_, Result<bool, ProviderStoreError>> {
|
||||
Box::pin(async { Ok(false) })
|
||||
}
|
||||
|
||||
fn read<'a>(
|
||||
&'a self,
|
||||
account_id: &'a ProviderAccountId,
|
||||
) -> BoxFuture<'a, Result<Option<ProviderCooldown>, ProviderStoreError>> {
|
||||
Box::pin(async move { Ok(self.cooldowns.get(account_id).cloned()) })
|
||||
}
|
||||
|
||||
fn clear<'a>(
|
||||
&'a self,
|
||||
_account_id: &'a ProviderAccountId,
|
||||
_through_revision: CredentialRevision,
|
||||
) -> BoxFuture<'a, Result<bool, ProviderStoreError>> {
|
||||
Box::pin(async { Ok(false) })
|
||||
}
|
||||
|
||||
fn put_scoped_if_later(
|
||||
&self,
|
||||
_cooldown: ProviderScopedCooldown,
|
||||
) -> BoxFuture<'_, Result<bool, ProviderStoreError>> {
|
||||
Box::pin(async { Ok(false) })
|
||||
}
|
||||
|
||||
fn read_scoped<'a>(
|
||||
&'a self,
|
||||
_account_id: &'a ProviderAccountId,
|
||||
_scope: &'a ProviderCooldownScope,
|
||||
) -> BoxFuture<'a, Result<Option<ProviderScopedCooldown>, ProviderStoreError>> {
|
||||
Box::pin(async { Ok(None) })
|
||||
}
|
||||
|
||||
fn clear_scoped<'a>(
|
||||
&'a self,
|
||||
_account_id: &'a ProviderAccountId,
|
||||
_scope: &'a ProviderCooldownScope,
|
||||
_through_revision: CredentialRevision,
|
||||
) -> BoxFuture<'a, Result<bool, ProviderStoreError>> {
|
||||
Box::pin(async { Ok(false) })
|
||||
}
|
||||
|
||||
fn clear_all<'a>(
|
||||
&'a self,
|
||||
_account_id: &'a ProviderAccountId,
|
||||
) -> BoxFuture<'a, Result<bool, ProviderStoreError>> {
|
||||
Box::pin(async { Ok(false) })
|
||||
}
|
||||
}
|
||||
|
||||
fn test_cooldown(account_id: &str, revision: u64) -> ProviderCooldown {
|
||||
ProviderCooldown::new(
|
||||
ProviderAccountId::new(account_id).expect("test account ID"),
|
||||
CredentialRevision::new(revision).expect("test revision"),
|
||||
SystemTime::now() + Duration::from_secs(60),
|
||||
)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn calculated_usage_billing_facts_keep_only_completed_calculated_costs() {
|
||||
let Some(database) = TestDatabase::create("calculated_usage_billing_facts").await else {
|
||||
@@ -194,7 +384,7 @@ async fn calculated_usage_billing_facts_keep_only_completed_calculated_costs() {
|
||||
.expect("seed calculated billing facts");
|
||||
let range = ObservabilityRange::new(now - TimeDelta::hours(1), now + TimeDelta::hours(1))
|
||||
.expect("observability range");
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone());
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone(), None);
|
||||
|
||||
let facts = repository
|
||||
.usage_calculated_billing_facts(range, UsageRecordFilter::default())
|
||||
@@ -208,7 +398,7 @@ async fn calculated_usage_billing_facts_keep_only_completed_calculated_costs() {
|
||||
assert_eq!(facts[0].service_tier.as_deref(), Some("priority"));
|
||||
assert_eq!(facts[0].total.amount.as_str(), "1.25");
|
||||
|
||||
let store = PgAdminObservabilityStore::new(database.pool.clone(), None);
|
||||
let store = PgAdminObservabilityStore::new(database.pool.clone(), None, None);
|
||||
let facts = store
|
||||
.usage_calculated_billing_facts(
|
||||
admin_observability::TimeRange::new(range.start, range.end)
|
||||
@@ -236,7 +426,7 @@ async fn admin_observability_adapter_preserves_utc_queries_metrics_costs_and_det
|
||||
let range =
|
||||
admin_observability::TimeRange::new(now - TimeDelta::hours(1), now + TimeDelta::hours(1))
|
||||
.expect("admin observability range");
|
||||
let store = PgAdminObservabilityStore::new(database.pool.clone(), None);
|
||||
let store = PgAdminObservabilityStore::new(database.pool.clone(), None, None);
|
||||
|
||||
let dashboard = store
|
||||
.dashboard_summary(range)
|
||||
@@ -525,7 +715,7 @@ async fn observability_queries_preserve_request_account_cost_and_diagnostic_fact
|
||||
.expect("seed observability facts");
|
||||
let range = ObservabilityRange::new(now - TimeDelta::hours(1), now + TimeDelta::hours(1))
|
||||
.expect("observability range");
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone());
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone(), None);
|
||||
|
||||
let dashboard = repository
|
||||
.dashboard_summary(range)
|
||||
|
||||
@@ -28,9 +28,9 @@ use gateway_store::{
|
||||
postgres::{
|
||||
AdminAuditActorKind, AdminAuditEvent, DeleteProviderAccounts, ImportProviderAccounts,
|
||||
NewProviderAccount, PgAdminAccountStore, PgProviderAccountRepository,
|
||||
ProviderAccountAdminRepository, ProviderAccountAdminScope, ProviderAccountObservation,
|
||||
ProviderAccountRepository, ProviderCredentialUpdate, RotateProviderAccount,
|
||||
SetProviderAccountEnabled, UpdateProviderAccount,
|
||||
ProviderAccountAdminRepository, ProviderAccountAdminScope, ProviderAccountRepository,
|
||||
ProviderCredentialUpdate, RotateProviderAccount, SetProviderAccountEnabled,
|
||||
UpdateProviderAccount,
|
||||
},
|
||||
};
|
||||
use serde_json::json;
|
||||
@@ -74,6 +74,7 @@ async fn core_quota_batch_reads_only_observed_accounts_in_one_contract_call() {
|
||||
.clone(),
|
||||
)),
|
||||
observed_at: Some(SystemTime::now()),
|
||||
limit_reached: None,
|
||||
})
|
||||
.await
|
||||
.expect("persist quota");
|
||||
@@ -140,6 +141,7 @@ async fn delayed_provider_observations_cannot_overwrite_newer_state_or_quota() {
|
||||
.clone(),
|
||||
)),
|
||||
observed_at: Some(observed_at),
|
||||
limit_reached: None,
|
||||
};
|
||||
assert_eq!(
|
||||
repository
|
||||
@@ -158,11 +160,10 @@ async fn delayed_provider_observations_cannot_overwrite_newer_state_or_quota() {
|
||||
|
||||
repository
|
||||
.apply_state_change(AccountStateChange {
|
||||
message: None,
|
||||
account_id: account_id.clone(),
|
||||
expected_revision: revision,
|
||||
availability: AccountAvailability::QuotaExhausted,
|
||||
reason: Some("quota_exhausted".to_owned()),
|
||||
cooldown_until: Some(newer),
|
||||
observed_at: newer,
|
||||
})
|
||||
.await
|
||||
@@ -170,11 +171,10 @@ async fn delayed_provider_observations_cannot_overwrite_newer_state_or_quota() {
|
||||
assert!(
|
||||
repository
|
||||
.apply_state_change(AccountStateChange {
|
||||
message: None,
|
||||
account_id: account_id.clone(),
|
||||
expected_revision: revision,
|
||||
availability: AccountAvailability::Ready,
|
||||
reason: None,
|
||||
cooldown_until: None,
|
||||
observed_at: older,
|
||||
})
|
||||
.await
|
||||
@@ -202,6 +202,61 @@ async fn delayed_provider_observations_cannot_overwrite_newer_state_or_quota() {
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn account_state_message_is_persisted_and_cleared_on_recovery() {
|
||||
let Some(database) = TestDatabase::create("provider_account_state_message").await else {
|
||||
return;
|
||||
};
|
||||
let repository = PgProviderAccountRepository::new(database.pool.clone());
|
||||
let account_id = ProviderAccountId::new("acct_state_message").expect("account ID");
|
||||
repository
|
||||
.insert_provider_account(account(account_id.as_str(), "user-state-message"))
|
||||
.await
|
||||
.expect("insert state message fixture");
|
||||
let revision = CredentialRevision::new(1).expect("revision");
|
||||
|
||||
repository
|
||||
.apply_state_change(AccountStateChange {
|
||||
message: Some("upstream returned 402 payment required".to_owned()),
|
||||
account_id: account_id.clone(),
|
||||
expected_revision: revision,
|
||||
availability: AccountAvailability::QuotaExhausted,
|
||||
observed_at: SystemTime::now(),
|
||||
})
|
||||
.await
|
||||
.expect("persist error message");
|
||||
let stored: Option<String> =
|
||||
sqlx::query_scalar("select last_error_message from provider_accounts where id = $1")
|
||||
.bind(account_id.as_str())
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.expect("read error message");
|
||||
assert_eq!(
|
||||
stored.as_deref(),
|
||||
Some("upstream returned 402 payment required")
|
||||
);
|
||||
|
||||
repository
|
||||
.apply_state_change(AccountStateChange {
|
||||
message: None,
|
||||
account_id: account_id.clone(),
|
||||
expected_revision: revision,
|
||||
availability: AccountAvailability::Ready,
|
||||
observed_at: SystemTime::now(),
|
||||
})
|
||||
.await
|
||||
.expect("recover account");
|
||||
let cleared: Option<String> =
|
||||
sqlx::query_scalar("select last_error_message from provider_accounts where id = $1")
|
||||
.bind(account_id.as_str())
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.expect("read cleared error message");
|
||||
assert_eq!(cleared, None);
|
||||
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_usage() {
|
||||
let Some(database) = TestDatabase::create("provider_account_terminal_list").await else {
|
||||
@@ -233,6 +288,21 @@ async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_u
|
||||
.await
|
||||
.expect("insert account list fixture");
|
||||
}
|
||||
repository
|
||||
.compare_and_swap_quota(QuotaObservation {
|
||||
account_id: ProviderAccountId::new("acct_charlie").expect("account id"),
|
||||
expected_revision: CredentialRevision::new(1).expect("credential revision"),
|
||||
quota: Some(OpaqueProviderData::new(
|
||||
json!({"rate_limit": {"limit_reached": true}})
|
||||
.as_object()
|
||||
.expect("quota object")
|
||||
.clone(),
|
||||
)),
|
||||
observed_at: Some(SystemTime::now()),
|
||||
limit_reached: Some(true),
|
||||
})
|
||||
.await
|
||||
.expect("seed rate-limited account");
|
||||
|
||||
seed_model_request(
|
||||
&database.pool,
|
||||
@@ -291,7 +361,7 @@ async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_u
|
||||
.await
|
||||
.expect("seed charlie usage");
|
||||
|
||||
let store = PgAdminAccountStore::new(database.pool.clone());
|
||||
let store = PgAdminAccountStore::new(database.pool.clone(), None);
|
||||
let usage_page = store
|
||||
.list_accounts(AccountListQuery {
|
||||
page: 1,
|
||||
@@ -309,14 +379,18 @@ async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_u
|
||||
assert_eq!(usage_page.config_revision.get(), 1);
|
||||
assert_eq!(usage_page.total, 6);
|
||||
assert_eq!(usage_page.summary.total, 6);
|
||||
assert_eq!(usage_page.summary.active, 2);
|
||||
assert_eq!(usage_page.summary.quota_exhausted, 1);
|
||||
assert_eq!(usage_page.summary.unavailable, 3);
|
||||
assert_eq!(usage_page.summary.normal, 1);
|
||||
assert_eq!(usage_page.summary.quota_exhausted, 2);
|
||||
assert_eq!(usage_page.summary.rate_limited, 0);
|
||||
assert_eq!(usage_page.summary.disabled, 1);
|
||||
assert_eq!(usage_page.summary.error, 2);
|
||||
assert_eq!(
|
||||
usage_page.summary.total,
|
||||
usage_page.summary.active
|
||||
usage_page.summary.normal
|
||||
+ usage_page.summary.quota_exhausted
|
||||
+ usage_page.summary.unavailable
|
||||
+ usage_page.summary.rate_limited
|
||||
+ usage_page.summary.disabled
|
||||
+ usage_page.summary.error
|
||||
);
|
||||
assert_eq!(
|
||||
usage_page
|
||||
@@ -356,7 +430,7 @@ async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_u
|
||||
page_size: PageSize::new(10).expect("page size"),
|
||||
provider_kind: Some(ProviderKind::new("openai").expect("Provider kind")),
|
||||
search: Some("ALPHA@EXAMPLE".to_owned()),
|
||||
status: Some(AccountStatus::Active),
|
||||
status: Some(AccountStatus::Normal),
|
||||
sort: None,
|
||||
})
|
||||
.await
|
||||
@@ -366,18 +440,20 @@ async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_u
|
||||
assert_eq!(filtered.items[0].id, "acct_alpha");
|
||||
assert_eq!(filtered.items[0].provider_kind.as_str(), "openai");
|
||||
|
||||
let banned = store
|
||||
let error_accounts = store
|
||||
.list_accounts(AccountListQuery {
|
||||
page: 1,
|
||||
page_size: PageSize::new(10).expect("page size"),
|
||||
provider_kind: None,
|
||||
search: None,
|
||||
status: Some(AccountStatus::Banned),
|
||||
status: Some(AccountStatus::Error),
|
||||
sort: None,
|
||||
})
|
||||
.await
|
||||
.expect("filter banned accounts");
|
||||
assert_eq!(banned.items[0].id, "acct_beta");
|
||||
.expect("filter error accounts");
|
||||
assert_eq!(error_accounts.items.len(), 2);
|
||||
assert_eq!(error_accounts.items[0].id, "acct_beta");
|
||||
assert_eq!(error_accounts.items[1].id, "acct_invalid");
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
@@ -405,7 +481,7 @@ async fn terminal_credential_list_preserves_grouped_filters_and_unpaged_collecti
|
||||
.expect("insert xAI credential fixture");
|
||||
}
|
||||
|
||||
let store = PgAdminAccountStore::new(database.pool.clone());
|
||||
let store = PgAdminAccountStore::new(database.pool.clone(), None);
|
||||
let provider = ProviderKind::new("xai").expect("xAI Provider kind");
|
||||
let complete = store
|
||||
.list_credentials(
|
||||
@@ -486,7 +562,7 @@ async fn terminal_admin_usage_chunks_large_selections_and_preserves_exact_costs(
|
||||
let mut account_ids = vec!["acct_usage_exact".to_owned()];
|
||||
account_ids.extend((0..200).map(|index| format!("missing_account_{index}")));
|
||||
|
||||
let usage = PgAdminAccountStore::new(database.pool.clone())
|
||||
let usage = PgAdminAccountStore::new(database.pool.clone(), None)
|
||||
.load_account_usage(
|
||||
TimeRange {
|
||||
start: now - TimeDelta::hours(1),
|
||||
@@ -578,7 +654,7 @@ async fn terminal_admin_quota_window_usage_prefers_provider_total_and_falls_back
|
||||
.await
|
||||
.expect("separate provider total from fallback components");
|
||||
|
||||
let mut usage = PgAdminAccountStore::new(database.pool.clone())
|
||||
let mut usage = PgAdminAccountStore::new(database.pool.clone(), None)
|
||||
.load_account_usage_by_windows(&[
|
||||
AccountUsageWindowQuery {
|
||||
account_id: "acct_quota_window".to_owned(),
|
||||
@@ -618,7 +694,7 @@ async fn terminal_admin_mutations_keep_revision_account_and_audit_atomic() {
|
||||
.insert_provider_account(account("acct_terminal_mutation", "user-terminal-mutation"))
|
||||
.await
|
||||
.expect("insert mutation account");
|
||||
let store = PgAdminAccountStore::new(database.pool.clone());
|
||||
let store = PgAdminAccountStore::new(database.pool.clone(), None);
|
||||
let context = MutationContext {
|
||||
actor: MutationActor::System,
|
||||
request_id: "request_terminal_mutation".to_owned(),
|
||||
@@ -691,7 +767,7 @@ async fn terminal_admin_delete_removes_enabled_accounts_in_one_transaction() {
|
||||
.expect("insert enabled account");
|
||||
}
|
||||
|
||||
let revision = PgAdminAccountStore::new(database.pool.clone())
|
||||
let revision = PgAdminAccountStore::new(database.pool.clone(), None)
|
||||
.delete_accounts(
|
||||
DeleteAccounts {
|
||||
account_ids: vec![
|
||||
@@ -747,7 +823,6 @@ async fn admin_import_updates_the_same_verified_identity_without_rebinding_it()
|
||||
updated.name = "updated import".to_owned();
|
||||
updated.provider_credentials_json = credential_json("updated-import-secret");
|
||||
updated.availability = AccountAvailability::Banned;
|
||||
updated.availability_reason = Some("upstream_account_deactivated".to_owned());
|
||||
let imported = repository
|
||||
.import_provider_accounts(ImportProviderAccounts {
|
||||
scope: scope.clone(),
|
||||
@@ -765,10 +840,9 @@ async fn admin_import_updates_the_same_verified_identity_without_rebinding_it()
|
||||
i64,
|
||||
Option<serde_json::Value>,
|
||||
String,
|
||||
Option<String>,
|
||||
) = sqlx::query_as(
|
||||
"select name, provider_credentials_json, credential_revision, provider_quota_json,
|
||||
availability, availability_reason
|
||||
availability
|
||||
from provider_accounts where id = 'acct_admin_upsert'",
|
||||
)
|
||||
.fetch_one(&database.pool)
|
||||
@@ -782,7 +856,6 @@ async fn admin_import_updates_the_same_verified_identity_without_rebinding_it()
|
||||
"credential replacement must clear stale quota"
|
||||
);
|
||||
assert_eq!(row.4, "banned");
|
||||
assert_eq!(row.5.as_deref(), Some("upstream_account_deactivated"));
|
||||
|
||||
repository
|
||||
.import_provider_accounts(ImportProviderAccounts {
|
||||
@@ -819,8 +892,6 @@ async fn core_refresh_cas_updates_profile_and_credential_under_one_revision() {
|
||||
next_refresh_at: None,
|
||||
enabled: true,
|
||||
availability: AccountAvailability::Ready,
|
||||
availability_reason: None,
|
||||
cooldown_until: None,
|
||||
availability_observed_at: Utc::now(),
|
||||
})
|
||||
.await
|
||||
@@ -942,8 +1013,7 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() {
|
||||
};
|
||||
|
||||
let mut cooldown_account = account("acct_admin_b", "user-admin-b");
|
||||
cooldown_account.availability = AccountAvailability::Cooldown;
|
||||
cooldown_account.cooldown_until = Some(Utc::now() + TimeDelta::minutes(10));
|
||||
cooldown_account.availability = AccountAvailability::Ready;
|
||||
let imported = repository
|
||||
.import_provider_accounts(ImportProviderAccounts {
|
||||
scope: scope.clone(),
|
||||
@@ -953,24 +1023,21 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() {
|
||||
.await
|
||||
.expect("import provider accounts");
|
||||
assert_eq!(imported.config_revision.get(), 2);
|
||||
let ready: (bool, String, Option<chrono::DateTime<Utc>>) = sqlx::query_as(
|
||||
"select enabled, availability, cooldown_until from provider_accounts where id = $1",
|
||||
)
|
||||
.bind("acct_admin_a")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.expect("load ready imported account");
|
||||
assert_eq!(ready, (true, "ready".to_owned(), None));
|
||||
let cooldown: (bool, String, Option<chrono::DateTime<Utc>>) = sqlx::query_as(
|
||||
"select enabled, availability, cooldown_until from provider_accounts where id = $1",
|
||||
)
|
||||
.bind("acct_admin_b")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.expect("load cooldown imported account");
|
||||
let ready: (bool, String) =
|
||||
sqlx::query_as("select enabled, availability from provider_accounts where id = $1")
|
||||
.bind("acct_admin_a")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.expect("load ready imported account");
|
||||
assert_eq!(ready, (true, "ready".to_owned()));
|
||||
let cooldown: (bool, String) =
|
||||
sqlx::query_as("select enabled, availability from provider_accounts where id = $1")
|
||||
.bind("acct_admin_b")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.expect("load cooldown imported account");
|
||||
assert!(cooldown.0);
|
||||
assert_eq!(cooldown.1, "cooldown");
|
||||
assert!(cooldown.2.is_some());
|
||||
assert_eq!(cooldown.1, "ready");
|
||||
|
||||
repository
|
||||
.import_provider_accounts(ImportProviderAccounts {
|
||||
@@ -1017,7 +1084,7 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() {
|
||||
|
||||
sqlx::query(
|
||||
"update provider_accounts
|
||||
set availability = 'expired', availability_reason = 'credential_expired'
|
||||
set availability = 'expired'
|
||||
where id = $1",
|
||||
)
|
||||
.bind("acct_admin_a")
|
||||
@@ -1040,15 +1107,8 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() {
|
||||
.expect("rotate provider account");
|
||||
assert_eq!(rotation.config_revision.get(), 3);
|
||||
assert_eq!(rotation.credential_revision.get(), 2);
|
||||
let restored: (
|
||||
String,
|
||||
Option<String>,
|
||||
Option<chrono::DateTime<Utc>>,
|
||||
String,
|
||||
Option<String>,
|
||||
) = sqlx::query_as(
|
||||
"select availability, availability_reason, cooldown_until,
|
||||
upstream_user_id, upstream_account_id
|
||||
let restored: (String, String, Option<String>) = sqlx::query_as(
|
||||
"select availability, upstream_user_id, upstream_account_id
|
||||
from provider_accounts where id = $1",
|
||||
)
|
||||
.bind("acct_admin_a")
|
||||
@@ -1059,8 +1119,6 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() {
|
||||
restored,
|
||||
(
|
||||
"ready".to_owned(),
|
||||
None,
|
||||
None,
|
||||
"user-admin-rebound".to_owned(),
|
||||
Some("workspace-admin-rebound".to_owned()),
|
||||
)
|
||||
@@ -1158,7 +1216,6 @@ async fn verified_credential_rotation_preserves_quota_exhaustion() {
|
||||
};
|
||||
let mut exhausted = account("acct_rotation_quota", "user-rotation-quota");
|
||||
exhausted.availability = AccountAvailability::QuotaExhausted;
|
||||
exhausted.availability_reason = Some("quota_exhausted".to_owned());
|
||||
repository
|
||||
.import_provider_accounts(ImportProviderAccounts {
|
||||
scope: scope.clone(),
|
||||
@@ -1187,20 +1244,13 @@ async fn verified_credential_rotation_preserves_quota_exhaustion() {
|
||||
.await
|
||||
.expect("rotate exhausted account credential");
|
||||
|
||||
let availability: (String, Option<String>) = sqlx::query_as(
|
||||
"select availability, availability_reason from provider_accounts where id = $1",
|
||||
)
|
||||
.bind("acct_rotation_quota")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.expect("load exhausted account after credential rotation");
|
||||
assert_eq!(
|
||||
availability,
|
||||
(
|
||||
"quota_exhausted".to_owned(),
|
||||
Some("quota_exhausted".to_owned()),
|
||||
),
|
||||
);
|
||||
let availability: String =
|
||||
sqlx::query_scalar("select availability from provider_accounts where id = $1")
|
||||
.bind("acct_rotation_quota")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.expect("load exhausted account after credential rotation");
|
||||
assert_eq!(availability, "quota_exhausted".to_owned());
|
||||
|
||||
database.close().await;
|
||||
}
|
||||
@@ -1221,32 +1271,24 @@ async fn xai_free_quota_exhaustion_state_is_persisted() {
|
||||
|
||||
repository
|
||||
.apply_state_change(AccountStateChange {
|
||||
message: None,
|
||||
account_id: account_id.clone(),
|
||||
expected_revision: CredentialRevision::new(1).expect("credential revision"),
|
||||
availability: AccountAvailability::QuotaExhausted,
|
||||
reason: Some("upstream_free_quota_exhausted".to_owned()),
|
||||
cooldown_until: None,
|
||||
observed_at: SystemTime::now(),
|
||||
})
|
||||
.await
|
||||
.expect("persist xAI free quota exhaustion");
|
||||
|
||||
let current: (String, String, Option<String>) = sqlx::query_as(
|
||||
"select provider_kind, availability, availability_reason
|
||||
let current: (String, String) = sqlx::query_as(
|
||||
"select provider_kind, availability
|
||||
from provider_accounts where id = $1",
|
||||
)
|
||||
.bind(account_id.as_str())
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.expect("load persisted xAI quota state");
|
||||
assert_eq!(
|
||||
current,
|
||||
(
|
||||
"xai".to_owned(),
|
||||
"quota_exhausted".to_owned(),
|
||||
Some("upstream_free_quota_exhausted".to_owned()),
|
||||
),
|
||||
);
|
||||
assert_eq!(current, ("xai".to_owned(), "quota_exhausted".to_owned()),);
|
||||
|
||||
database.close().await;
|
||||
}
|
||||
@@ -1266,38 +1308,26 @@ async fn xai_resettable_usage_limit_state_is_persisted() {
|
||||
.expect("insert xAI account");
|
||||
|
||||
let observed_at = SystemTime::now();
|
||||
let cooldown_until = observed_at
|
||||
.checked_add(Duration::from_secs(24 * 60 * 60))
|
||||
.expect("cooldown deadline");
|
||||
repository
|
||||
.apply_state_change(AccountStateChange {
|
||||
message: None,
|
||||
account_id: account_id.clone(),
|
||||
expected_revision: CredentialRevision::new(1).expect("credential revision"),
|
||||
availability: AccountAvailability::Cooldown,
|
||||
reason: Some("usage_limit_exhausted".to_owned()),
|
||||
cooldown_until: Some(cooldown_until),
|
||||
availability: AccountAvailability::Ready,
|
||||
observed_at,
|
||||
})
|
||||
.await
|
||||
.expect("persist xAI resettable usage limit");
|
||||
|
||||
let current: (String, String, Option<String>, bool) = sqlx::query_as(
|
||||
"select provider_kind, availability, availability_reason, cooldown_until is not null
|
||||
let current: (String, String) = sqlx::query_as(
|
||||
"select provider_kind, availability
|
||||
from provider_accounts where id = $1",
|
||||
)
|
||||
.bind(account_id.as_str())
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.expect("load persisted xAI usage limit state");
|
||||
assert_eq!(
|
||||
current,
|
||||
(
|
||||
"xai".to_owned(),
|
||||
"cooldown".to_owned(),
|
||||
Some("usage_limit_exhausted".to_owned()),
|
||||
true,
|
||||
),
|
||||
);
|
||||
assert_eq!(current, ("xai".to_owned(), "ready".to_owned()));
|
||||
|
||||
database.close().await;
|
||||
}
|
||||
@@ -1315,7 +1345,6 @@ async fn disabled_account_preserves_user_state_during_refresh_writes() {
|
||||
let mut disabled = account(account_id.as_str(), "user-disabled-refresh");
|
||||
disabled.enabled = false;
|
||||
disabled.availability = AccountAvailability::Expired;
|
||||
disabled.availability_reason = Some("credential_expired".to_owned());
|
||||
repository
|
||||
.import_provider_accounts(ImportProviderAccounts {
|
||||
scope: scope.clone(),
|
||||
@@ -1331,11 +1360,10 @@ async fn disabled_account_preserves_user_state_during_refresh_writes() {
|
||||
|
||||
repository
|
||||
.apply_state_change(AccountStateChange {
|
||||
message: None,
|
||||
account_id: account_id.clone(),
|
||||
expected_revision: CredentialRevision::new(1).expect("credential revision"),
|
||||
availability: AccountAvailability::Ready,
|
||||
reason: None,
|
||||
cooldown_until: None,
|
||||
observed_at: SystemTime::now(),
|
||||
})
|
||||
.await
|
||||
@@ -1355,9 +1383,8 @@ async fn disabled_account_preserves_user_state_during_refresh_writes() {
|
||||
.await
|
||||
.expect("refresh disabled account credential");
|
||||
|
||||
let current: (bool, String, Option<String>, serde_json::Value, i64) = sqlx::query_as(
|
||||
"select enabled, availability, availability_reason,
|
||||
provider_credentials_json, credential_revision
|
||||
let current: (bool, String, serde_json::Value, i64) = sqlx::query_as(
|
||||
"select enabled, availability, provider_credentials_json, credential_revision
|
||||
from provider_accounts where id = $1",
|
||||
)
|
||||
.bind(account_id.as_str())
|
||||
@@ -1366,9 +1393,8 @@ async fn disabled_account_preserves_user_state_during_refresh_writes() {
|
||||
.expect("load disabled account after refresh writes");
|
||||
assert!(!current.0);
|
||||
assert_eq!(current.1, "expired");
|
||||
assert_eq!(current.2.as_deref(), Some("credential_expired"));
|
||||
assert_eq!(current.3["access_token"], "disabled-refreshed-secret");
|
||||
assert_eq!(current.4, 2);
|
||||
assert_eq!(current.2["access_token"], "disabled-refreshed-secret");
|
||||
assert_eq!(current.3, 2);
|
||||
|
||||
database.close().await;
|
||||
}
|
||||
@@ -1389,8 +1415,6 @@ fn account(id: &str, upstream_user_id: &str) -> NewProviderAccount {
|
||||
next_refresh_at: None,
|
||||
enabled: true,
|
||||
availability: AccountAvailability::Ready,
|
||||
availability_reason: None,
|
||||
cooldown_until: None,
|
||||
availability_observed_at: Utc::now(),
|
||||
}
|
||||
}
|
||||
@@ -1517,31 +1541,15 @@ fn provider_credentials_are_redacted_from_debug() {
|
||||
assert!(!format!("{object:?}").contains(secret));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cooldown_observation_requires_expiry() {
|
||||
let observation = ProviderAccountObservation {
|
||||
account_id: "account-1".to_owned(),
|
||||
availability: AccountAvailability::Cooldown,
|
||||
availability_reason: None,
|
||||
cooldown_until: None,
|
||||
provider_quota_json: None,
|
||||
availability_observed_at: Utc::now(),
|
||||
quota_observed_at: None,
|
||||
};
|
||||
assert!(observation.validate().is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn quota_exhausted_account_accepts_cooldown_runtime_facts() {
|
||||
let mut imported = account("acct_quota_cooldown", "user-quota-cooldown");
|
||||
imported.availability = AccountAvailability::QuotaExhausted;
|
||||
imported.cooldown_until = Some(Utc::now() + TimeDelta::minutes(5));
|
||||
assert!(imported.validate().is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn imported_account_rejects_cooldown_until_for_ready_runtime_facts() {
|
||||
let mut imported = account("acct_invalid_cooldown", "user-invalid-cooldown");
|
||||
imported.cooldown_until = Some(Utc::now() + TimeDelta::minutes(5));
|
||||
assert!(imported.validate().is_err());
|
||||
fn imported_ready_account_passes_validation() {
|
||||
let imported = account("acct_ready_import", "user-ready-import");
|
||||
assert!(imported.validate().is_ok());
|
||||
}
|
||||
|
||||
@@ -6,19 +6,9 @@ use gateway_store::postgres::{
|
||||
PgObservabilityRepository, PgOpsEventRepository, UsageRecordFilter, UsageRecordQuery,
|
||||
};
|
||||
use sqlx::PgPool;
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::TestDatabase;
|
||||
|
||||
type ProviderAccountSnapshotRow = (
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
);
|
||||
|
||||
#[tokio::test]
|
||||
async fn request_snapshots_should_survive_account_deletion() {
|
||||
let Some(database) = TestDatabase::create("snapshot_request_delete").await else {
|
||||
@@ -49,7 +39,7 @@ async fn request_snapshots_should_survive_account_deletion() {
|
||||
.await
|
||||
.expect("delete provider account");
|
||||
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone());
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone(), None);
|
||||
let page = repository
|
||||
.list_usage_records(usage_query(started_at, UsageRecordFilter::default()))
|
||||
.await
|
||||
@@ -164,7 +154,7 @@ async fn attempts_should_keep_their_own_account_snapshots() {
|
||||
.await
|
||||
.expect("delete provider accounts");
|
||||
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone());
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone(), None);
|
||||
let detail = repository
|
||||
.usage_record_detail("req_snap_ab")
|
||||
.await
|
||||
@@ -286,7 +276,7 @@ async fn ops_errors_should_keep_request_and_event_snapshots_after_account_deleti
|
||||
.await
|
||||
.expect("delete provider accounts");
|
||||
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone());
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone(), None);
|
||||
let errors = repository
|
||||
.list_ops_errors(OpsErrorQuery {
|
||||
range: range_around(started_at),
|
||||
@@ -386,7 +376,7 @@ async fn diagnostics_should_group_same_email_accounts_by_stable_ref() {
|
||||
.await
|
||||
.expect("delete provider accounts");
|
||||
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone());
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone(), None);
|
||||
let diagnostics = repository
|
||||
.usage_diagnostics(
|
||||
range_around(started_at),
|
||||
@@ -455,7 +445,7 @@ async fn diagnostics_should_fallback_to_name_then_ref_for_missing_snapshots() {
|
||||
.expect("insert legacy request without live account");
|
||||
finalize_request(&database.pool, "req_snap_legacy", started_at).await;
|
||||
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone());
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone(), None);
|
||||
let diagnostics = repository
|
||||
.usage_diagnostics(
|
||||
range_around(started_at),
|
||||
@@ -521,7 +511,7 @@ async fn failure_diagnostics_should_only_include_errored_requests() {
|
||||
.await
|
||||
.expect("mark request failed");
|
||||
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone());
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone(), None);
|
||||
let diagnostics = repository
|
||||
.usage_diagnostics(
|
||||
range_around(started_at),
|
||||
@@ -576,7 +566,7 @@ async fn api_key_diagnostics_should_display_key_name_and_fallback_to_ref() {
|
||||
.await
|
||||
.expect("set request b latency");
|
||||
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone());
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone(), None);
|
||||
let diagnostics = repository
|
||||
.usage_diagnostics(
|
||||
range_around(started_at),
|
||||
@@ -643,7 +633,7 @@ async fn renaming_account_should_not_rewrite_historical_snapshots() {
|
||||
.await
|
||||
.expect("rename provider account");
|
||||
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone());
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone(), None);
|
||||
let page = repository
|
||||
.list_usage_records(usage_query(started_at, UsageRecordFilter::default()))
|
||||
.await
|
||||
@@ -664,175 +654,6 @@ async fn renaming_account_should_not_rewrite_historical_snapshots() {
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn migration_should_backfill_snapshots_for_existing_accounts_and_survive_deletion() {
|
||||
let Some(database_url) = crate::support::test_env("CPR_TEST_DATABASE_URL") else {
|
||||
return;
|
||||
};
|
||||
let schema = format!("cpr_store_snapshot_backfill_{}", Uuid::new_v4().simple());
|
||||
let admin = sqlx::postgres::PgPoolOptions::new()
|
||||
.max_connections(1)
|
||||
.connect(&database_url)
|
||||
.await
|
||||
.expect("connect test PostgreSQL");
|
||||
sqlx::raw_sql(sqlx::AssertSqlSafe(format!("create schema \"{schema}\"")))
|
||||
.execute(&admin)
|
||||
.await
|
||||
.expect("create backfill schema");
|
||||
let search_path = schema.clone();
|
||||
let pool = sqlx::postgres::PgPoolOptions::new()
|
||||
.max_connections(2)
|
||||
.after_connect(move |connection, _metadata| {
|
||||
let search_path = search_path.clone();
|
||||
Box::pin(async move {
|
||||
sqlx::query("select set_config('search_path', $1, false)")
|
||||
.bind(search_path)
|
||||
.execute(connection)
|
||||
.await?;
|
||||
Ok(())
|
||||
})
|
||||
})
|
||||
.connect(&database_url)
|
||||
.await
|
||||
.expect("connect backfill pool");
|
||||
let started_at = Utc::now();
|
||||
{
|
||||
let mut migration = pool.begin().await.expect("begin terminal migration");
|
||||
sqlx::raw_sql(include_str!("../../../../migrations/0001_initial.sql"))
|
||||
.execute(&mut *migration)
|
||||
.await
|
||||
.expect("apply initial migration");
|
||||
migration.commit().await.expect("commit initial migration");
|
||||
}
|
||||
seed_account(
|
||||
&pool,
|
||||
"acct_backfill",
|
||||
"Backfill",
|
||||
Some("backfill@example.invalid"),
|
||||
"oauth",
|
||||
started_at,
|
||||
)
|
||||
.await;
|
||||
sqlx::query(
|
||||
"insert into model_requests (
|
||||
id, client_api_key_ref, config_revision, protocol, operation, endpoint,
|
||||
client_transport, requested_model_id,
|
||||
provider_kind, provider_account_id, provider_account_ref,
|
||||
upstream_model_id, upstream_transport, attempt_count, upstream_send_state,
|
||||
started_at, deadline_at
|
||||
) values (
|
||||
'req_backfill', 'key_backfill', 1, 'openai', 'responses', '/v1/responses',
|
||||
'http_sse', 'coding', 'openai', 'acct_backfill', 'acct_backfill',
|
||||
'upstream-backfill', 'http_sse', 1, 'sent', $1, $1 + interval '30 seconds'
|
||||
)",
|
||||
)
|
||||
.bind(started_at)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.expect("insert pre-migration request");
|
||||
sqlx::query(
|
||||
"insert into ops_events (
|
||||
id, level, component, operation,
|
||||
provider_kind, provider_account_id, provider_account_ref,
|
||||
failure_kind, message, occurrence_count, created_at
|
||||
) values (
|
||||
'ops_backfill', 'warning', 'account_probe', 'connection_test',
|
||||
'openai', 'acct_backfill', 'acct_backfill',
|
||||
'auth_failed', 'backfill probe', 1, $1
|
||||
)",
|
||||
)
|
||||
.bind(started_at)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.expect("insert pre-migration ops event");
|
||||
{
|
||||
let mut migration = pool.begin().await.expect("begin snapshot migration");
|
||||
sqlx::raw_sql(include_str!(
|
||||
"../../../../migrations/0002_snapshot_provider_account_identity.sql"
|
||||
))
|
||||
.execute(&mut *migration)
|
||||
.await
|
||||
.expect("apply snapshot migration");
|
||||
migration.commit().await.expect("commit snapshot migration");
|
||||
}
|
||||
let (
|
||||
request_name,
|
||||
request_email,
|
||||
request_authentication_kind,
|
||||
event_name,
|
||||
event_email,
|
||||
event_authentication_kind,
|
||||
): ProviderAccountSnapshotRow = sqlx::query_as(
|
||||
"select
|
||||
(select provider_account_name_snapshot from model_requests where id = 'req_backfill'),
|
||||
(select provider_account_email_snapshot from model_requests where id = 'req_backfill'),
|
||||
(select provider_account_authentication_kind_snapshot from model_requests where id = 'req_backfill'),
|
||||
(select provider_account_name_snapshot from ops_events where id = 'ops_backfill'),
|
||||
(select provider_account_email_snapshot from ops_events where id = 'ops_backfill'),
|
||||
(select provider_account_authentication_kind_snapshot from ops_events where id = 'ops_backfill')",
|
||||
)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.expect("load backfilled snapshots");
|
||||
assert_eq!(
|
||||
(
|
||||
request_name.as_deref(),
|
||||
request_email.as_deref(),
|
||||
request_authentication_kind.as_deref(),
|
||||
),
|
||||
(
|
||||
Some("Backfill"),
|
||||
Some("backfill@example.invalid"),
|
||||
Some("oauth"),
|
||||
)
|
||||
);
|
||||
assert_eq!(
|
||||
(
|
||||
event_name.as_deref(),
|
||||
event_email.as_deref(),
|
||||
event_authentication_kind.as_deref(),
|
||||
),
|
||||
(
|
||||
Some("Backfill"),
|
||||
Some("backfill@example.invalid"),
|
||||
Some("oauth"),
|
||||
)
|
||||
);
|
||||
sqlx::query("delete from provider_accounts where id = 'acct_backfill'")
|
||||
.execute(&pool)
|
||||
.await
|
||||
.expect("delete backfilled account");
|
||||
let request_snapshot: (Option<String>, Option<String>, Option<String>) = sqlx::query_as(
|
||||
"select provider_account_name_snapshot, provider_account_email_snapshot,
|
||||
provider_account_authentication_kind_snapshot
|
||||
from model_requests where id = 'req_backfill'",
|
||||
)
|
||||
.fetch_one(&pool)
|
||||
.await
|
||||
.expect("load request snapshot after deletion");
|
||||
assert_eq!(
|
||||
(
|
||||
request_snapshot.0.as_deref(),
|
||||
request_snapshot.1.as_deref(),
|
||||
request_snapshot.2.as_deref(),
|
||||
),
|
||||
(
|
||||
Some("Backfill"),
|
||||
Some("backfill@example.invalid"),
|
||||
Some("oauth"),
|
||||
)
|
||||
);
|
||||
|
||||
drop(pool);
|
||||
sqlx::raw_sql(sqlx::AssertSqlSafe(format!(
|
||||
"drop schema \"{schema}\" cascade"
|
||||
)))
|
||||
.execute(&admin)
|
||||
.await
|
||||
.expect("drop backfill schema");
|
||||
admin.close().await;
|
||||
}
|
||||
|
||||
fn new_request(id: &str, started_at: DateTime<Utc>) -> NewModelRequest {
|
||||
NewModelRequest {
|
||||
id: id.to_owned(),
|
||||
|
||||
@@ -1,27 +1,13 @@
|
||||
use std::{
|
||||
sync::{Arc, Mutex},
|
||||
time::{Duration as StdDuration, SystemTime},
|
||||
};
|
||||
use std::time::{Duration as StdDuration, SystemTime};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, Duration, Utc};
|
||||
use gateway_core::{
|
||||
engine::credential::{
|
||||
AccountAvailability, AccountStateChange, CredentialCasOutcome, CredentialCasUpdate,
|
||||
CredentialRevision, LoadedCredential, NewProviderAccount, PlaintextCredential,
|
||||
ProviderAccount, ProviderAccountId, ProviderAccountStore, ProviderAccountUpdate,
|
||||
QuotaObservation, QuotaWriteOutcome,
|
||||
},
|
||||
error::{StoreError as CoreStoreError, StoreErrorKind},
|
||||
engine::credential::{CredentialRevision, ProviderAccountId},
|
||||
provider_ports::{ProviderCooldownPort, ProviderCooldownScope, ProviderScopedCooldown},
|
||||
routing::ProviderKind,
|
||||
};
|
||||
use gateway_store::{
|
||||
Revision, StoreBackend, StoreError, StoreResult,
|
||||
redis::{
|
||||
CooldownCachingProviderAccountStore, CredentialCooldown, CredentialCooldownRepository,
|
||||
RedisCredentialCooldownRepository,
|
||||
},
|
||||
Revision,
|
||||
redis::{CredentialCooldown, CredentialCooldownRepository, RedisCredentialCooldownRepository},
|
||||
};
|
||||
use redis::aio::ConnectionManager;
|
||||
use uuid::Uuid;
|
||||
@@ -36,409 +22,6 @@ fn credential_cooldown_is_revision_fenced() {
|
||||
assert_eq!(cooldown.credential_revision.get(), 2);
|
||||
}
|
||||
|
||||
struct CooldownTestAccountStore {
|
||||
accounts: Vec<ProviderAccount>,
|
||||
calls: Arc<Mutex<Vec<&'static str>>>,
|
||||
fail_state_write: bool,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ProviderAccountStore for CooldownTestAccountStore {
|
||||
async fn create_account(&self, _account: NewProviderAccount) -> Result<(), CoreStoreError> {
|
||||
unreachable!("cooldown tests do not create accounts")
|
||||
}
|
||||
|
||||
async fn get_account(
|
||||
&self,
|
||||
account: &ProviderAccountId,
|
||||
) -> Result<Option<ProviderAccount>, CoreStoreError> {
|
||||
Ok(self
|
||||
.accounts
|
||||
.iter()
|
||||
.find(|candidate| candidate.id() == account)
|
||||
.cloned())
|
||||
}
|
||||
|
||||
async fn list_accounts(&self) -> Result<Vec<ProviderAccount>, CoreStoreError> {
|
||||
Ok(self.accounts.clone())
|
||||
}
|
||||
|
||||
async fn list_for_provider(
|
||||
&self,
|
||||
_provider: &ProviderKind,
|
||||
) -> Result<Vec<ProviderAccount>, CoreStoreError> {
|
||||
unreachable!("cooldown tests do not list one provider")
|
||||
}
|
||||
|
||||
async fn load_credential(
|
||||
&self,
|
||||
_account: &ProviderAccountId,
|
||||
_expected_revision: CredentialRevision,
|
||||
) -> Result<LoadedCredential, CoreStoreError> {
|
||||
unreachable!("cooldown tests do not load credentials")
|
||||
}
|
||||
|
||||
async fn load_current_credential(
|
||||
&self,
|
||||
_account: &ProviderAccountId,
|
||||
) -> Result<LoadedCredential, CoreStoreError> {
|
||||
unreachable!("cooldown tests do not load credentials")
|
||||
}
|
||||
|
||||
async fn compare_and_swap_credential(
|
||||
&self,
|
||||
_update: CredentialCasUpdate,
|
||||
) -> Result<CredentialCasOutcome, CoreStoreError> {
|
||||
self.calls.lock().expect("calls lock").push("postgres");
|
||||
Ok(CredentialCasOutcome::Updated(
|
||||
CredentialRevision::new(8).expect("positive revision"),
|
||||
))
|
||||
}
|
||||
|
||||
async fn get_quotas(
|
||||
&self,
|
||||
_accounts: &[ProviderAccountId],
|
||||
) -> Result<Vec<QuotaObservation>, CoreStoreError> {
|
||||
unreachable!("cooldown tests do not read quota")
|
||||
}
|
||||
|
||||
async fn compare_and_swap_quota(
|
||||
&self,
|
||||
_observation: QuotaObservation,
|
||||
) -> Result<QuotaWriteOutcome, CoreStoreError> {
|
||||
unreachable!("cooldown tests do not write quota")
|
||||
}
|
||||
|
||||
async fn apply_state_change(&self, _change: AccountStateChange) -> Result<(), CoreStoreError> {
|
||||
self.calls.lock().expect("calls lock").push("postgres");
|
||||
if self.fail_state_write {
|
||||
Err(CoreStoreError::new(StoreErrorKind::Unavailable))
|
||||
} else {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
async fn update_account(&self, _update: ProviderAccountUpdate) -> Result<(), CoreStoreError> {
|
||||
unreachable!("cooldown tests do not update account metadata")
|
||||
}
|
||||
|
||||
async fn set_enabled(
|
||||
&self,
|
||||
_account: &ProviderAccountId,
|
||||
_enabled: bool,
|
||||
) -> Result<(), CoreStoreError> {
|
||||
unreachable!("cooldown tests do not change account switches")
|
||||
}
|
||||
|
||||
async fn delete_account(&self, _account: &ProviderAccountId) -> Result<(), CoreStoreError> {
|
||||
self.calls.lock().expect("calls lock").push("postgres");
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
struct CooldownTestCache {
|
||||
calls: Arc<Mutex<Vec<&'static str>>>,
|
||||
cached: Mutex<Vec<CredentialCooldown>>,
|
||||
invalidated: Mutex<Vec<(String, Revision)>>,
|
||||
fail: bool,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl CredentialCooldownRepository for CooldownTestCache {
|
||||
async fn cache_credential_cooldown(&self, cooldown: &CredentialCooldown) -> StoreResult<bool> {
|
||||
self.calls.lock().expect("calls lock").push("redis");
|
||||
if self.fail {
|
||||
return Err(StoreError::Unavailable {
|
||||
backend: StoreBackend::Redis,
|
||||
message: "test cache unavailable".to_owned(),
|
||||
});
|
||||
}
|
||||
self.cached
|
||||
.lock()
|
||||
.expect("cached lock")
|
||||
.push(cooldown.clone());
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
async fn read_credential_cooldown(
|
||||
&self,
|
||||
_provider_account_id: &str,
|
||||
) -> StoreResult<Option<CredentialCooldown>> {
|
||||
unreachable!("write-through adapter does not consume cooldown cache")
|
||||
}
|
||||
|
||||
async fn invalidate_credential_cooldown(
|
||||
&self,
|
||||
provider_account_id: &str,
|
||||
through_revision: Revision,
|
||||
) -> StoreResult<bool> {
|
||||
self.calls.lock().expect("calls lock").push("redis");
|
||||
if self.fail {
|
||||
return Err(StoreError::Unavailable {
|
||||
backend: StoreBackend::Redis,
|
||||
message: "test cache unavailable".to_owned(),
|
||||
});
|
||||
}
|
||||
self.invalidated
|
||||
.lock()
|
||||
.expect("invalidated lock")
|
||||
.push((provider_account_id.to_owned(), through_revision));
|
||||
Ok(true)
|
||||
}
|
||||
}
|
||||
|
||||
fn composition_fixture(
|
||||
fail_state_write: bool,
|
||||
fail_cache: bool,
|
||||
) -> (
|
||||
CooldownCachingProviderAccountStore,
|
||||
Arc<CooldownTestCache>,
|
||||
Arc<Mutex<Vec<&'static str>>>,
|
||||
) {
|
||||
composition_fixture_with_accounts(
|
||||
vec![ready_test_account("acct_cooldown_composition")],
|
||||
fail_state_write,
|
||||
fail_cache,
|
||||
)
|
||||
}
|
||||
|
||||
fn composition_fixture_with_accounts(
|
||||
accounts: Vec<ProviderAccount>,
|
||||
fail_state_write: bool,
|
||||
fail_cache: bool,
|
||||
) -> (
|
||||
CooldownCachingProviderAccountStore,
|
||||
Arc<CooldownTestCache>,
|
||||
Arc<Mutex<Vec<&'static str>>>,
|
||||
) {
|
||||
let calls = Arc::new(Mutex::new(Vec::new()));
|
||||
let authoritative = Arc::new(CooldownTestAccountStore {
|
||||
accounts,
|
||||
calls: Arc::clone(&calls),
|
||||
fail_state_write,
|
||||
});
|
||||
let cache = Arc::new(CooldownTestCache {
|
||||
calls: Arc::clone(&calls),
|
||||
cached: Mutex::new(Vec::new()),
|
||||
invalidated: Mutex::new(Vec::new()),
|
||||
fail: fail_cache,
|
||||
});
|
||||
let store = CooldownCachingProviderAccountStore::new(
|
||||
authoritative,
|
||||
Arc::clone(&cache) as Arc<dyn CredentialCooldownRepository>,
|
||||
);
|
||||
(store, cache, calls)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cooldown_composition_writes_cache_only_after_postgres_commit() {
|
||||
let (store, cache, calls) = composition_fixture(false, false);
|
||||
|
||||
store
|
||||
.apply_state_change(cooldown_state_change(60))
|
||||
.await
|
||||
.expect("persist cooldown");
|
||||
|
||||
assert_eq!(*calls.lock().expect("calls lock"), ["postgres", "redis"]);
|
||||
assert_eq!(cache.cached.lock().expect("cached lock").len(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cooldown_composition_does_not_cache_failed_postgres_state() {
|
||||
let (store, cache, calls) = composition_fixture(true, false);
|
||||
|
||||
assert!(
|
||||
store
|
||||
.apply_state_change(cooldown_state_change(60))
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
assert_eq!(*calls.lock().expect("calls lock"), ["postgres"]);
|
||||
assert!(cache.cached.lock().expect("cached lock").is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cooldown_composition_keeps_postgres_terminal_when_redis_fails() {
|
||||
let (store, _cache, calls) = composition_fixture(false, true);
|
||||
|
||||
store
|
||||
.apply_state_change(cooldown_state_change(60))
|
||||
.await
|
||||
.expect("Redis is only a cache");
|
||||
|
||||
assert_eq!(*calls.lock().expect("calls lock"), ["postgres", "redis"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn disabled_account_state_write_does_not_change_runtime_cooldown() {
|
||||
let disabled = test_provider_account("acct_cooldown_composition").with_runtime_state(
|
||||
false,
|
||||
AccountAvailability::Ready,
|
||||
None,
|
||||
);
|
||||
let (store, cache, calls) = composition_fixture_with_accounts(vec![disabled], false, false);
|
||||
|
||||
store
|
||||
.apply_state_change(cooldown_state_change(60))
|
||||
.await
|
||||
.expect("disabled state write is a no-op");
|
||||
|
||||
assert_eq!(*calls.lock().expect("calls lock"), ["postgres"]);
|
||||
assert!(cache.cached.lock().expect("cached lock").is_empty());
|
||||
assert!(
|
||||
cache
|
||||
.invalidated
|
||||
.lock()
|
||||
.expect("invalidated lock")
|
||||
.is_empty()
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ready_state_invalidates_same_revision_cooldown_after_postgres_commit() {
|
||||
let (store, cache, calls) = composition_fixture(false, false);
|
||||
let mut change = cooldown_state_change(60);
|
||||
change.availability = AccountAvailability::Ready;
|
||||
change.reason = None;
|
||||
change.cooldown_until = None;
|
||||
|
||||
store
|
||||
.apply_state_change(change)
|
||||
.await
|
||||
.expect("persist ready state");
|
||||
|
||||
assert_eq!(
|
||||
cache.invalidated.lock().expect("invalidated lock")[0].1,
|
||||
Revision::new(7).expect("positive revision")
|
||||
);
|
||||
assert_eq!(*calls.lock().expect("calls lock"), ["postgres", "redis"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn credential_rotation_invalidates_cooldown_through_new_revision() {
|
||||
let (store, cache, calls) = composition_fixture(false, false);
|
||||
|
||||
let outcome = store
|
||||
.compare_and_swap_credential(credential_update())
|
||||
.await
|
||||
.expect("rotate credential");
|
||||
|
||||
assert_eq!(
|
||||
outcome,
|
||||
CredentialCasOutcome::Updated(CredentialRevision::new(8).expect("positive revision"))
|
||||
);
|
||||
assert_eq!(
|
||||
cache.invalidated.lock().expect("invalidated lock")[0].1,
|
||||
Revision::new(8).expect("positive revision")
|
||||
);
|
||||
assert_eq!(*calls.lock().expect("calls lock"), ["postgres", "redis"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn account_delete_invalidates_all_stale_cooldown_revisions() {
|
||||
let (store, cache, calls) = composition_fixture(false, false);
|
||||
let account_id = ProviderAccountId::new("acct_cooldown_delete").expect("valid account ID");
|
||||
|
||||
store
|
||||
.delete_account(&account_id)
|
||||
.await
|
||||
.expect("delete account");
|
||||
|
||||
assert_eq!(
|
||||
cache.invalidated.lock().expect("invalidated lock")[0],
|
||||
(
|
||||
account_id.to_string(),
|
||||
Revision::new(u64::MAX).expect("positive revision")
|
||||
)
|
||||
);
|
||||
assert_eq!(*calls.lock().expect("calls lock"), ["postgres", "redis"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cooldown_startup_hydrates_only_active_postgres_cooldowns() {
|
||||
let now = SystemTime::now();
|
||||
let calls = Arc::new(Mutex::new(Vec::new()));
|
||||
let authoritative = Arc::new(CooldownTestAccountStore {
|
||||
accounts: vec![
|
||||
cooldown_test_account("acct_hydrate_active", now + StdDuration::from_secs(60)),
|
||||
cooldown_test_account("acct_hydrate_expired", now - StdDuration::from_secs(1)),
|
||||
ready_test_account("acct_hydrate_ready"),
|
||||
],
|
||||
calls: Arc::clone(&calls),
|
||||
fail_state_write: false,
|
||||
});
|
||||
let cache = Arc::new(CooldownTestCache {
|
||||
calls,
|
||||
cached: Mutex::new(Vec::new()),
|
||||
invalidated: Mutex::new(Vec::new()),
|
||||
fail: false,
|
||||
});
|
||||
let store = CooldownCachingProviderAccountStore::new(
|
||||
authoritative,
|
||||
Arc::clone(&cache) as Arc<dyn CredentialCooldownRepository>,
|
||||
);
|
||||
|
||||
assert_eq!(store.hydrate(now).await, 1);
|
||||
assert_eq!(
|
||||
cache.cached.lock().expect("cached lock")[0].provider_account_id,
|
||||
"acct_hydrate_active"
|
||||
);
|
||||
}
|
||||
|
||||
fn cooldown_state_change(seconds: u64) -> AccountStateChange {
|
||||
AccountStateChange {
|
||||
account_id: ProviderAccountId::new("acct_cooldown_composition").expect("valid account ID"),
|
||||
expected_revision: CredentialRevision::new(7).expect("positive revision"),
|
||||
availability: AccountAvailability::Cooldown,
|
||||
reason: Some("rate_limited".to_owned()),
|
||||
cooldown_until: Some(SystemTime::now() + StdDuration::from_secs(seconds)),
|
||||
observed_at: SystemTime::now(),
|
||||
}
|
||||
}
|
||||
|
||||
fn credential_update() -> CredentialCasUpdate {
|
||||
let account_id = ProviderAccountId::new("acct_cooldown_rotation").expect("valid account ID");
|
||||
CredentialCasUpdate::new(
|
||||
account_id.clone(),
|
||||
CredentialRevision::new(7).expect("positive revision"),
|
||||
ProviderAccountUpdate {
|
||||
account_id,
|
||||
name: "Rotated cooldown account".to_owned(),
|
||||
email: None,
|
||||
plan_type: None,
|
||||
},
|
||||
PlaintextCredential::new(serde_json::Map::new()),
|
||||
false,
|
||||
Some(SystemTime::now() + StdDuration::from_secs(3_600)),
|
||||
None,
|
||||
)
|
||||
.expect("valid credential update")
|
||||
}
|
||||
|
||||
fn cooldown_test_account(account_id: &str, cooldown_until: SystemTime) -> ProviderAccount {
|
||||
test_provider_account(account_id).with_runtime_state(
|
||||
true,
|
||||
AccountAvailability::Cooldown,
|
||||
Some(cooldown_until),
|
||||
)
|
||||
}
|
||||
|
||||
fn ready_test_account(account_id: &str) -> ProviderAccount {
|
||||
test_provider_account(account_id).with_runtime_state(true, AccountAvailability::Ready, None)
|
||||
}
|
||||
|
||||
fn test_provider_account(account_id: &str) -> ProviderAccount {
|
||||
ProviderAccount::new(
|
||||
ProviderAccountId::new(account_id).expect("valid account ID"),
|
||||
ProviderKind::new("openai").expect("valid provider kind"),
|
||||
"Cooldown test".to_owned(),
|
||||
"upstream-user".to_owned(),
|
||||
"oauth".to_owned(),
|
||||
CredentialRevision::new(7).expect("positive revision"),
|
||||
Some(SystemTime::now() + StdDuration::from_secs(3_600)),
|
||||
)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn credential_cooldown_round_trips_without_raw_account_id_in_key() {
|
||||
let Some((repository, mut connection, namespace)) = repository().await else {
|
||||
|
||||
@@ -0,0 +1,141 @@
|
||||
use std::io;
|
||||
|
||||
use chrono::Utc;
|
||||
use gateway_admin::model::backup::{BackupObjectMetadata, BackupStorageConfig};
|
||||
use gateway_admin::ports::backup::{BackupObjectStorePort, UploadObjectRequest};
|
||||
use gateway_core::engine::CancellationToken;
|
||||
use gateway_store::backup::s3::S3ObjectStoreAdapter;
|
||||
use secrecy::SecretString;
|
||||
use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
|
||||
use tokio::net::{TcpListener, TcpStream};
|
||||
|
||||
const PART_SIZE_ABOVE_TOKIO_DEFAULT: usize = 2 * 1024 * 1024 + 1;
|
||||
|
||||
#[tokio::test]
|
||||
async fn upload_file_sends_full_part_beyond_tokio_default_buffer_limit() {
|
||||
let directory = tempfile::tempdir().expect("create temporary directory");
|
||||
let source = directory.path().join("archive.dump");
|
||||
let expected = vec![b'x'; PART_SIZE_ABOVE_TOKIO_DEFAULT];
|
||||
tokio::fs::write(&source, &expected)
|
||||
.await
|
||||
.expect("write temporary archive");
|
||||
|
||||
let listener = TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("bind multipart probe server");
|
||||
let endpoint = format!(
|
||||
"http://{}",
|
||||
listener.local_addr().expect("read listener address")
|
||||
);
|
||||
let server = tokio::spawn(serve_multipart_probe(listener));
|
||||
let storage = BackupStorageConfig {
|
||||
storage_revision: 1,
|
||||
endpoint,
|
||||
region: "auto".to_owned(),
|
||||
bucket: "backup-bucket".to_owned(),
|
||||
access_key_id: "test-access-key".to_owned(),
|
||||
secret_access_key: SecretString::from("test-secret-key"),
|
||||
prefix: "backups".to_owned(),
|
||||
force_path_style: true,
|
||||
};
|
||||
let request = UploadObjectRequest {
|
||||
object_key: "backups/archive.dump".to_owned(),
|
||||
source,
|
||||
metadata: BackupObjectMetadata::new(
|
||||
"backup_test".to_owned(),
|
||||
"a".repeat(64),
|
||||
Utc::now(),
|
||||
PART_SIZE_ABOVE_TOKIO_DEFAULT as u64,
|
||||
),
|
||||
cancellation: CancellationToken::new(),
|
||||
};
|
||||
|
||||
S3ObjectStoreAdapter::new()
|
||||
.upload_file(&storage, request)
|
||||
.await
|
||||
.expect("upload multipart archive");
|
||||
let uploaded_len = server
|
||||
.await
|
||||
.expect("multipart probe task did not panic")
|
||||
.expect("serve multipart probe");
|
||||
|
||||
assert_eq!(uploaded_len, PART_SIZE_ABOVE_TOKIO_DEFAULT);
|
||||
}
|
||||
|
||||
async fn serve_multipart_probe(listener: TcpListener) -> io::Result<usize> {
|
||||
let mut uploaded_len = None;
|
||||
for _ in 0..3 {
|
||||
let (mut stream, _) = listener.accept().await?;
|
||||
let (request, body) = read_http_request(&mut stream).await?;
|
||||
let response = if request.starts_with("POST ") && request.contains("uploads") {
|
||||
xml_response(
|
||||
"<InitiateMultipartUploadResult xmlns=\"http://s3.amazonaws.com/doc/2006-03-01/\"><UploadId>probe-upload</UploadId></InitiateMultipartUploadResult>",
|
||||
)
|
||||
} else if request.starts_with("PUT ") && request.contains("partNumber=") {
|
||||
uploaded_len = Some(body.len());
|
||||
"HTTP/1.1 200 OK\r\ncontent-length: 0\r\nconnection: close\r\netag: \"probe-part\"\r\n\r\n".to_owned()
|
||||
} else if request.starts_with("POST ") && request.contains("uploadId=") {
|
||||
xml_response(
|
||||
"<CompleteMultipartUploadResult xmlns=\"http://s3.amazonaws.com/doc/2006-03-01/\"><ETag>\"probe-part\"</ETag></CompleteMultipartUploadResult>",
|
||||
)
|
||||
} else {
|
||||
return Err(io::Error::other("unexpected multipart probe request"));
|
||||
};
|
||||
stream.write_all(response.as_bytes()).await?;
|
||||
stream.shutdown().await?;
|
||||
}
|
||||
uploaded_len.ok_or_else(|| io::Error::other("multipart probe did not receive an upload part"))
|
||||
}
|
||||
|
||||
async fn read_http_request(stream: &mut TcpStream) -> io::Result<(String, Vec<u8>)> {
|
||||
let mut bytes = Vec::new();
|
||||
let header_end = loop {
|
||||
if let Some(position) = bytes.windows(4).position(|window| window == b"\r\n\r\n") {
|
||||
break position + 4;
|
||||
}
|
||||
let mut chunk = [0_u8; 8 * 1024];
|
||||
let read = stream.read(&mut chunk).await?;
|
||||
if read == 0 {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::UnexpectedEof,
|
||||
"multipart probe request ended before headers",
|
||||
));
|
||||
}
|
||||
bytes.extend_from_slice(&chunk[..read]);
|
||||
};
|
||||
let headers = String::from_utf8_lossy(&bytes[..header_end]).into_owned();
|
||||
let content_len = headers
|
||||
.lines()
|
||||
.filter_map(|line| line.split_once(':'))
|
||||
.find(|(name, _)| name.eq_ignore_ascii_case("content-length"))
|
||||
.and_then(|(_, value)| value.trim().parse::<usize>().ok())
|
||||
.unwrap_or_default();
|
||||
if headers
|
||||
.to_ascii_lowercase()
|
||||
.contains("expect: 100-continue")
|
||||
{
|
||||
stream.write_all(b"HTTP/1.1 100 Continue\r\n\r\n").await?;
|
||||
}
|
||||
while bytes.len() < header_end + content_len {
|
||||
let mut chunk = [0_u8; 8 * 1024];
|
||||
let read = stream.read(&mut chunk).await?;
|
||||
if read == 0 {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::UnexpectedEof,
|
||||
"multipart probe request ended before body",
|
||||
));
|
||||
}
|
||||
bytes.extend_from_slice(&chunk[..read]);
|
||||
}
|
||||
Ok((
|
||||
headers,
|
||||
bytes[header_end..header_end + content_len].to_vec(),
|
||||
))
|
||||
}
|
||||
|
||||
fn xml_response(body: &str) -> String {
|
||||
format!(
|
||||
"HTTP/1.1 200 OK\r\ncontent-type: application/xml\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}",
|
||||
body.len()
|
||||
)
|
||||
}
|
||||
@@ -50,7 +50,7 @@ use crate::credential::{
|
||||
CodexCredentialAdminService, CodexCredentialCatalogError, CodexCredentialCatalogService,
|
||||
CodexCredentialQuotaError, CodexCredentialQuotaService, CodexOAuthAdmin, CodexOAuthAdminError,
|
||||
CodexOAuthPendingClaimOutcome, CodexOAuthPendingStore, CodexOAuthPendingStoreError,
|
||||
CodexPendingAuthorization, CodexQuotaWindowKind, CodexQuotaWindowRole,
|
||||
CodexPendingAuthorization, CodexQuotaWindow, CodexQuotaWindowKind, CodexQuotaWindowRole,
|
||||
CompleteCodexOAuthAuthorization, CompletedCodexOAuthCredential, ExportManagedCodexCredential,
|
||||
RotateManagedCodexCredential, StartCodexOAuthAuthorization, StoredCodexPendingAuthorization,
|
||||
refresh_time,
|
||||
@@ -194,22 +194,23 @@ impl ProviderAdmin for OpenAiAdminProvider {
|
||||
let profile = self.profile.snapshot();
|
||||
let release = self.desktop_release.snapshot();
|
||||
let user_agent = profile.user_agent();
|
||||
let client_identity = format!("{}; {}", profile.originator, profile.desktop_version);
|
||||
let release = dashboard_desktop_release(&profile, release);
|
||||
Some(DashboardWireProfile {
|
||||
provider: self.provider_kind.as_str().to_owned(),
|
||||
product: profile.originator,
|
||||
version: profile.desktop_version,
|
||||
build: Some(profile.desktop_build),
|
||||
version: profile.codex_version,
|
||||
build: None,
|
||||
target: DashboardWireTarget {
|
||||
os_type: profile.os_type,
|
||||
os_type: dashboard_os_type(&profile.os_type),
|
||||
os_version: profile.os_version,
|
||||
arch: profile.arch,
|
||||
terminal: profile.terminal,
|
||||
},
|
||||
user_agent,
|
||||
attributes: vec![DashboardWireAttribute {
|
||||
label: "Codex Core".to_owned(),
|
||||
value: profile.codex_version,
|
||||
label: "客户端标识".to_owned(),
|
||||
value: client_identity,
|
||||
}],
|
||||
verified_at: Some(profile.verified_at),
|
||||
release: Some(release),
|
||||
@@ -475,7 +476,14 @@ impl ProviderAdmin for OpenAiAdminProvider {
|
||||
.await
|
||||
.map_err(map_quota_error)?
|
||||
};
|
||||
Ok(project_quota(snapshot, &account))
|
||||
let mut quota = project_quota(snapshot, &account);
|
||||
quota.rate_limited_until = self
|
||||
.quota
|
||||
.rate_limited_until(account.id(), account.revision())
|
||||
.await
|
||||
.map_err(map_quota_error)?
|
||||
.map(DateTime::<Utc>::from);
|
||||
Ok(quota)
|
||||
}
|
||||
|
||||
async fn models(
|
||||
@@ -579,8 +587,6 @@ fn prepared_create(
|
||||
next_refresh_at: account.next_refresh_at().map(DateTime::<Utc>::from),
|
||||
enabled: account.enabled(),
|
||||
availability,
|
||||
availability_reason: None,
|
||||
cooldown_until: account.cooldown_until().map(DateTime::<Utc>::from),
|
||||
availability_observed_at: observed_at,
|
||||
})
|
||||
}
|
||||
@@ -682,17 +688,20 @@ fn account_from_record(account: &AccountRecord) -> Result<ProviderAccount, Provi
|
||||
account.upstream_account_id.clone(),
|
||||
account.plan_type.clone(),
|
||||
)
|
||||
.with_runtime_state(
|
||||
account.enabled,
|
||||
account.availability,
|
||||
account.cooldown_until.map(SystemTime::from),
|
||||
)
|
||||
.with_runtime_state(account.enabled, account.availability)
|
||||
.with_refresh_schedule(
|
||||
account.has_refresh_token,
|
||||
account.next_refresh_at.map(SystemTime::from),
|
||||
))
|
||||
}
|
||||
|
||||
fn dashboard_os_type(os_type: &str) -> String {
|
||||
match os_type {
|
||||
"Mac OS" => "macOS".to_owned(),
|
||||
_ => os_type.to_owned(),
|
||||
}
|
||||
}
|
||||
|
||||
fn dashboard_desktop_release(
|
||||
profile: &CodexWireProfile,
|
||||
snapshot: CodexDesktopReleaseSnapshot,
|
||||
@@ -745,6 +754,8 @@ fn empty_quota() -> ProviderQuota {
|
||||
observed_at: None,
|
||||
refresh_token_expires_at: None,
|
||||
windows: Vec::new(),
|
||||
limit_reached: false,
|
||||
rate_limited_until: None,
|
||||
provider_data: None,
|
||||
}
|
||||
}
|
||||
@@ -753,9 +764,12 @@ fn project_quota(
|
||||
snapshot: Option<CodexAccountQuotaSnapshot>,
|
||||
account: &ProviderAccount,
|
||||
) -> ProviderQuota {
|
||||
let mut quota = snapshot.map_or_else(empty_quota, project_quota_snapshot);
|
||||
// 展示前先滚动已过期窗口,避免 reset 后的旧窗口值悬挂。
|
||||
let mut quota = snapshot
|
||||
.map(|snapshot| project_quota_snapshot(snapshot.roll_expired_windows(SystemTime::now())))
|
||||
.unwrap_or_else(empty_quota);
|
||||
if account.availability() == AccountAvailability::QuotaExhausted {
|
||||
force_confirmed_exhaustion_projection(&mut quota, account.cooldown_until());
|
||||
force_confirmed_exhaustion_projection(&mut quota);
|
||||
}
|
||||
quota
|
||||
}
|
||||
@@ -773,9 +787,10 @@ fn project_quota_snapshot(snapshot: CodexAccountQuotaSnapshot) -> ProviderQuota
|
||||
"exhausted".to_owned(),
|
||||
Value::Bool(snapshot.fact().exhausted()),
|
||||
);
|
||||
let windows = snapshot
|
||||
let windows: Vec<ProviderQuotaWindow> = snapshot
|
||||
.windows()
|
||||
.iter()
|
||||
.filter(|window| should_project_quota_window(window))
|
||||
.map(|window| {
|
||||
let mut data = Map::new();
|
||||
data.insert(
|
||||
@@ -790,7 +805,9 @@ fn project_quota_snapshot(snapshot: CodexAccountQuotaSnapshot) -> ProviderQuota
|
||||
key: window.key().to_owned(),
|
||||
group: quota_group(window.kind()).to_owned(),
|
||||
label: codex_quota_window_label(
|
||||
window.limit_name(),
|
||||
window.kind(),
|
||||
window.role(),
|
||||
window.source(),
|
||||
window.window_seconds(),
|
||||
),
|
||||
@@ -798,27 +815,46 @@ fn project_quota_snapshot(snapshot: CodexAccountQuotaSnapshot) -> ProviderQuota
|
||||
window_seconds: window.window_seconds(),
|
||||
used_percent: window.used_percent(),
|
||||
reset_at: window.reset_at(),
|
||||
limit_reached: window.limit_reached(),
|
||||
local_usage: None,
|
||||
provider_data: Some(ProviderDocument::new(OpaqueProviderData::new(data))),
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
// 快照级 limit_reached 只看滚动后的窗口触顶:顶层标记是观测事实,不能
|
||||
// 在窗口全部过期后继续维持限流。
|
||||
let limit_reached = quota_windows_limit_reached(&windows);
|
||||
ProviderQuota {
|
||||
observed_at: Some(DateTime::<Utc>::from(snapshot.observed_at())),
|
||||
refresh_token_expires_at: None,
|
||||
windows,
|
||||
limit_reached,
|
||||
rate_limited_until: None,
|
||||
provider_data: Some(ProviderDocument::new(OpaqueProviderData::new(
|
||||
provider_data,
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
fn force_confirmed_exhaustion_projection(
|
||||
quota: &mut ProviderQuota,
|
||||
cooldown_until: Option<SystemTime>,
|
||||
) {
|
||||
// 真实 429 对“当前已耗尽”的证明强于可能延迟结算的 usage 文档。这里只覆盖
|
||||
// Admin 展示投影,原始 Provider JSON 仍作为后续判断窗口是否真正恢复的基线。
|
||||
/// `secondary_window` 有时只是上游的空占位:没有时长、重置时间,也没有用量。
|
||||
/// 它不能被可靠翻译成 5 小时或周额度,因此不应占用账号面板;带有实际事实的
|
||||
/// 次级窗口(时长、重置、非零用量或触顶)仍完整保留。
|
||||
fn should_project_quota_window(window: &CodexQuotaWindow) -> bool {
|
||||
window.role() != CodexQuotaWindowRole::Secondary
|
||||
|| window.window_seconds().is_some()
|
||||
|| window.reset_at().is_some()
|
||||
|| window.used_percent().is_some_and(|used| used > 0.0)
|
||||
|| window.limit_reached()
|
||||
}
|
||||
|
||||
fn quota_windows_limit_reached(windows: &[ProviderQuotaWindow]) -> bool {
|
||||
windows.iter().any(|window| window.limit_reached)
|
||||
}
|
||||
|
||||
fn force_confirmed_exhaustion_projection(quota: &mut ProviderQuota) {
|
||||
// 402 已确认耗尽:只覆盖 Admin 展示投影,原始 Provider JSON 仍作为
|
||||
// 后续判断窗口是否真正恢复的基线。
|
||||
quota.limit_reached = true;
|
||||
quota.provider_data = Some(ProviderDocument::new(OpaqueProviderData::new(
|
||||
Map::from_iter([
|
||||
(
|
||||
@@ -828,7 +864,6 @@ fn force_confirmed_exhaustion_projection(
|
||||
("exhausted".to_owned(), Value::Bool(true)),
|
||||
]),
|
||||
)));
|
||||
let reset_at = cooldown_until.map(DateTime::<Utc>::from);
|
||||
if quota.windows.is_empty() {
|
||||
quota.windows.push(ProviderQuotaWindow {
|
||||
key: "confirmed-quota-exhaustion".to_owned(),
|
||||
@@ -837,42 +872,28 @@ fn force_confirmed_exhaustion_projection(
|
||||
source: None,
|
||||
window_seconds: None,
|
||||
used_percent: Some(100.0),
|
||||
reset_at,
|
||||
reset_at: None,
|
||||
limit_reached: true,
|
||||
local_usage: None,
|
||||
provider_data: None,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
let matched_reset = reset_at.and_then(|reset_at| {
|
||||
quota
|
||||
.windows
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter_map(|(index, window)| {
|
||||
let distance = window.reset_at.map(|window_reset| {
|
||||
(window_reset.timestamp() - reset_at.timestamp()).unsigned_abs()
|
||||
})?;
|
||||
Some((index, distance))
|
||||
})
|
||||
.min_by_key(|(_, distance)| *distance)
|
||||
.map(|(index, _)| index)
|
||||
});
|
||||
let index = matched_reset
|
||||
.or_else(|| {
|
||||
quota
|
||||
.windows
|
||||
.iter()
|
||||
.enumerate()
|
||||
.max_by(|(_, left), (_, right)| {
|
||||
left.used_percent
|
||||
.unwrap_or(-1.0)
|
||||
.total_cmp(&right.used_percent.unwrap_or(-1.0))
|
||||
})
|
||||
.map(|(index, _)| index)
|
||||
// 选已用比例最高的窗口强制 100%,保留其自身 reset_at。
|
||||
let index = quota
|
||||
.windows
|
||||
.iter()
|
||||
.enumerate()
|
||||
.max_by(|(_, left), (_, right)| {
|
||||
left.used_percent
|
||||
.unwrap_or(-1.0)
|
||||
.total_cmp(&right.used_percent.unwrap_or(-1.0))
|
||||
})
|
||||
.map(|(index, _)| index)
|
||||
.unwrap_or(0);
|
||||
quota.windows[index].used_percent = Some(100.0);
|
||||
quota.windows[index].limit_reached = true;
|
||||
}
|
||||
|
||||
const fn quota_group(kind: CodexQuotaWindowKind) -> &'static str {
|
||||
@@ -891,54 +912,69 @@ const fn quota_role(role: CodexQuotaWindowRole) -> &'static str {
|
||||
}
|
||||
}
|
||||
|
||||
/// 按窗口时长显示汉化额度名;非核心桶补上上游桶名称作为前缀。
|
||||
fn codex_quota_window_label(
|
||||
limit_name: Option<&str>,
|
||||
kind: CodexQuotaWindowKind,
|
||||
role: CodexQuotaWindowRole,
|
||||
source: &str,
|
||||
window_seconds: Option<u64>,
|
||||
) -> String {
|
||||
let base = match kind {
|
||||
CodexQuotaWindowKind::Monthly => "月限额".to_owned(),
|
||||
CodexQuotaWindowKind::Weekly => "周限额".to_owned(),
|
||||
CodexQuotaWindowKind::Monthly => "月额度".to_owned(),
|
||||
CodexQuotaWindowKind::Weekly => "周额度".to_owned(),
|
||||
CodexQuotaWindowKind::ShortTerm => {
|
||||
if window_seconds.is_some_and(|seconds| seconds > 86_400) {
|
||||
"周限额".to_owned()
|
||||
"周额度".to_owned()
|
||||
} else {
|
||||
"5小时限额".to_owned()
|
||||
"5小时额度".to_owned()
|
||||
}
|
||||
}
|
||||
CodexQuotaWindowKind::Other => custom_quota_window_label(window_seconds),
|
||||
CodexQuotaWindowKind::Other => custom_quota_window_label(window_seconds, role),
|
||||
};
|
||||
if matches!(source, "core" | "codex" | "spend_control" | "monthly_limit") {
|
||||
return base;
|
||||
if let Some(prefix) = codex_quota_window_prefix(limit_name, source) {
|
||||
return format!("{prefix} · {base}");
|
||||
}
|
||||
if is_codex_review_limit(source) {
|
||||
return format!("代码审查 · {base}");
|
||||
}
|
||||
format!("{source} · {base}")
|
||||
base
|
||||
}
|
||||
|
||||
fn custom_quota_window_label(window_seconds: Option<u64>) -> String {
|
||||
fn codex_quota_window_prefix(limit_name: Option<&str>, source: &str) -> Option<String> {
|
||||
let source_identifier = quota_identifier(source);
|
||||
if matches!(source_identifier.as_str(), "core" | "codex") {
|
||||
return None;
|
||||
}
|
||||
let label = limit_name.unwrap_or(source);
|
||||
match quota_identifier(label).as_str() {
|
||||
"premium" => Some("Premium".to_owned()),
|
||||
"code_review" => Some("代码审查".to_owned()),
|
||||
"spend_control" => Some("消费控制".to_owned()),
|
||||
_ => Some(label.replace('_', " ")),
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_identifier(value: &str) -> String {
|
||||
value.trim().to_ascii_lowercase().replace(['-', ' '], "_")
|
||||
}
|
||||
|
||||
fn custom_quota_window_label(window_seconds: Option<u64>, role: CodexQuotaWindowRole) -> String {
|
||||
let Some(seconds) = window_seconds.filter(|seconds| *seconds > 0) else {
|
||||
return "额度".to_owned();
|
||||
// 与官方客户端一致:没有时长时不臆测为 5 小时或周额度;保留
|
||||
// primary/secondary 语义,让用户知道这是上游未标明时长的独立窗口。
|
||||
return match role {
|
||||
CodexQuotaWindowRole::Primary => "主额度".to_owned(),
|
||||
CodexQuotaWindowRole::Secondary => "次级额度".to_owned(),
|
||||
CodexQuotaWindowRole::Monthly => "月额度".to_owned(),
|
||||
};
|
||||
};
|
||||
if seconds % 86_400 == 0 {
|
||||
format!("{}日限额", seconds / 86_400)
|
||||
format!("{}日额度", seconds / 86_400)
|
||||
} else if seconds % 3_600 == 0 {
|
||||
format!("{}小时限额", seconds / 3_600)
|
||||
format!("{}小时额度", seconds / 3_600)
|
||||
} else {
|
||||
format!("{}分钟限额", seconds.div_ceil(60))
|
||||
format!("{}分钟额度", seconds.div_ceil(60))
|
||||
}
|
||||
}
|
||||
|
||||
fn is_codex_review_limit(value: &str) -> bool {
|
||||
let normalized = value.trim().to_ascii_lowercase().replace(['-', ' '], "_");
|
||||
matches!(
|
||||
normalized.as_str(),
|
||||
"review" | "code_review" | "codex_review" | "codex_code_review"
|
||||
) || normalized.contains("code_review")
|
||||
|| normalized.contains("codex_review")
|
||||
}
|
||||
|
||||
/// 将 Provider-owned PKCE/OIDC 状态保存到 Store 提供的 Redis 原子端口。
|
||||
pub(crate) struct OpenAiOAuthPendingStore {
|
||||
port: Arc<dyn OAuthPendingFlowPort>,
|
||||
|
||||
@@ -278,6 +278,7 @@ impl CodexAuthSettings {
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct CodexWireProfileConfig {
|
||||
pub originator: String,
|
||||
/// npm `@openai/codex` 启动版本;运行时发布任务会更新此画像字段。
|
||||
pub codex_version: String,
|
||||
pub desktop_version: String,
|
||||
pub desktop_build: String,
|
||||
@@ -292,9 +293,9 @@ impl Default for CodexWireProfileConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
originator: "Codex Desktop".to_owned(),
|
||||
codex_version: "0.145.0".to_owned(),
|
||||
desktop_version: "26.721.31836".to_owned(),
|
||||
desktop_build: "5828".to_owned(),
|
||||
codex_version: "0.146.0".to_owned(),
|
||||
desktop_version: "26.727.51351".to_owned(),
|
||||
desktop_build: "6119".to_owned(),
|
||||
os_type: "Mac OS".to_owned(),
|
||||
os_version: "15.7.1".to_owned(),
|
||||
arch: "arm64".to_owned(),
|
||||
@@ -355,7 +356,7 @@ impl From<CodexWireProfileConfig> for CodexWireProfile {
|
||||
codex_version: value.codex_version,
|
||||
desktop_version: value.desktop_version,
|
||||
desktop_build: value.desktop_build,
|
||||
os_type: value.os_type,
|
||||
os_type: canonical_codex_os_type(value.os_type),
|
||||
os_version: value.os_version,
|
||||
arch: value.arch,
|
||||
terminal: value.terminal,
|
||||
@@ -364,6 +365,13 @@ impl From<CodexWireProfileConfig> for CodexWireProfile {
|
||||
}
|
||||
}
|
||||
|
||||
fn canonical_codex_os_type(os_type: String) -> String {
|
||||
match os_type.as_str() {
|
||||
"macOS" => "Mac OS".to_owned(),
|
||||
_ => os_type,
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
|
||||
pub enum OpenAiConfigError {
|
||||
#[error("OpenAI configuration field is invalid: {0}")]
|
||||
|
||||
@@ -409,7 +409,7 @@ impl CodexCredentialAdmin {
|
||||
Some(input.verified_account.chatgpt_account_id),
|
||||
input.verified_account.plan_type,
|
||||
)
|
||||
.with_runtime_state(input.enabled, AccountAvailability::Ready, None)
|
||||
.with_runtime_state(input.enabled, AccountAvailability::Ready)
|
||||
.with_refresh_schedule(
|
||||
input.secret.refresh_token.is_some(),
|
||||
optional_time(input.next_refresh_at),
|
||||
@@ -806,10 +806,7 @@ impl CodexCredentialAdminService {
|
||||
)) {
|
||||
return Err(CodexCredentialAdminError::InvalidInput);
|
||||
}
|
||||
account.account =
|
||||
account
|
||||
.account
|
||||
.with_runtime_state(enabled, availability, None);
|
||||
account.account = account.account.with_runtime_state(enabled, availability);
|
||||
account
|
||||
}
|
||||
};
|
||||
@@ -933,7 +930,7 @@ impl CodexCredentialAdminService {
|
||||
Some(upstream_account_id),
|
||||
candidate.plan_type.clone(),
|
||||
)
|
||||
.with_runtime_state(enabled, AccountAvailability::Ready, None)
|
||||
.with_runtime_state(enabled, AccountAvailability::Ready)
|
||||
.with_refresh_schedule(false, None);
|
||||
Ok(NewProviderAccount {
|
||||
account,
|
||||
@@ -1000,9 +997,7 @@ fn cpr_status(account: &ProviderAccount) -> &'static str {
|
||||
AccountAvailability::QuotaExhausted => "quota_exhausted",
|
||||
AccountAvailability::Expired | AccountAvailability::Invalid => "expired",
|
||||
AccountAvailability::Banned => "banned",
|
||||
AccountAvailability::Unknown
|
||||
| AccountAvailability::Ready
|
||||
| AccountAvailability::Cooldown => "active",
|
||||
AccountAvailability::Unknown | AccountAvailability::Ready => "active",
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -10,19 +10,33 @@ use crate::transport::request::derive_conversation_anchor;
|
||||
pub(crate) fn derive_codex_session_affinity_key(
|
||||
request: &CodexResponsesRequest,
|
||||
) -> Option<ProviderSessionAffinityKey> {
|
||||
if let Some(conversation_id) = request
|
||||
let session_key = if let Some(conversation_id) = request
|
||||
.local_conversation_id
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
if let Ok(key) = ProviderSessionAffinityKey::try_new(conversation_id.to_owned()) {
|
||||
return Some(key);
|
||||
Some(key)
|
||||
} else {
|
||||
opaque_affinity_key("local-conversation", conversation_id)
|
||||
}
|
||||
return opaque_affinity_key("local-conversation", conversation_id);
|
||||
}
|
||||
let (domain, value) = derive_conversation_anchor(request)?;
|
||||
opaque_affinity_key(domain, &value)
|
||||
} else {
|
||||
let (domain, value) = derive_conversation_anchor(request)?;
|
||||
opaque_affinity_key(domain, &value)
|
||||
};
|
||||
|
||||
let session_discriminator = request.subagent_kind();
|
||||
session_key.and_then(|session_key| match session_discriminator.as_deref() {
|
||||
None => Some(session_key),
|
||||
Some(discriminator) => {
|
||||
// 仅扩展原会话锚点的键空间;后续查找、绑定和调度仍使用同一套亲和逻辑。
|
||||
opaque_affinity_key(
|
||||
"subagent-session",
|
||||
&format!("{discriminator}\0{}", session_key.expose_to_store()),
|
||||
)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn opaque_affinity_key(domain: &str, value: &str) -> Option<ProviderSessionAffinityKey> {
|
||||
|
||||
@@ -783,7 +783,6 @@ fn eligible_catalog_account(account: &ProviderAccount, now: SystemTime) -> bool
|
||||
&& match account.availability() {
|
||||
AccountAvailability::Unknown
|
||||
| AccountAvailability::Ready
|
||||
| AccountAvailability::Cooldown
|
||||
| AccountAvailability::QuotaExhausted => true,
|
||||
AccountAvailability::Expired
|
||||
| AccountAvailability::Banned
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,59 @@
|
||||
//! 确认额度耗尽后的恢复证据与 availability 转换。
|
||||
//!
|
||||
//! `QuotaExhausted -> Ready` 只能由同一 credential revision 的权威 usage 快照证明:
|
||||
//! 已确认的旧窗口 reset 已到期,且新快照的 reset 已推进到下一窗口。
|
||||
|
||||
use std::time::SystemTime;
|
||||
|
||||
use gateway_core::engine::credential::AccountAvailability;
|
||||
use reqwest::StatusCode;
|
||||
use serde_json::Value;
|
||||
|
||||
use super::snapshot::CodexQuotaFact;
|
||||
use crate::transport::CodexClientError;
|
||||
|
||||
/// 只处理 `QuotaExhausted` 的恢复。`allowed=true`、成功响应和百分比回落都不能单独
|
||||
/// 覆盖一次已确认的额度拒绝;它们在 reset 未推进时可能只是上游滞后的观测。
|
||||
///
|
||||
/// `previous_reset_at` 是写入新快照前读取的已确认耗尽窗口。恢复必须同时满足:
|
||||
/// 旧窗口已经到期,且权威快照的 reset 指向比旧窗口更晚的新周期。
|
||||
pub(crate) fn quota_success_state(
|
||||
current: AccountAvailability,
|
||||
fact: CodexQuotaFact,
|
||||
previous_reset_at: Option<SystemTime>,
|
||||
now: SystemTime,
|
||||
) -> Option<AccountAvailability> {
|
||||
if current != AccountAvailability::QuotaExhausted || fact.exhausted() {
|
||||
return None;
|
||||
}
|
||||
let previous_reset_at = previous_reset_at?;
|
||||
let refreshed_reset_at = fact.resets_at().map(SystemTime::from)?;
|
||||
(previous_reset_at <= now && refreshed_reset_at > previous_reset_at)
|
||||
.then_some(AccountAvailability::Ready)
|
||||
}
|
||||
|
||||
pub(crate) fn quota_state_transition(error: &CodexClientError) -> Option<AccountAvailability> {
|
||||
let CodexClientError::Upstream { status, body, .. } = error else {
|
||||
return None;
|
||||
};
|
||||
if *status == StatusCode::PAYMENT_REQUIRED && is_deactivated_workspace(body) {
|
||||
return Some(AccountAvailability::Banned);
|
||||
}
|
||||
match *status {
|
||||
// 402:真额度耗尽 → QuotaExhausted(worker 按 reset_at 自动恢复)。
|
||||
StatusCode::PAYMENT_REQUIRED => Some(AccountAvailability::QuotaExhausted),
|
||||
// 429/503:临时限流,不写 availability(限流不改变账号可用性),由 quota 数据驱动。
|
||||
StatusCode::TOO_MANY_REQUESTS | StatusCode::SERVICE_UNAVAILABLE => None,
|
||||
// quota/usage 接口只能产生额度域状态;鉴权拒绝不具备判定 RT
|
||||
// 永久失效的证据,终态只能由 OAuth refresh/身份校验状态机写入。
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn is_deactivated_workspace(body: &str) -> bool {
|
||||
serde_json::from_str::<Value>(body)
|
||||
.ok()
|
||||
.is_some_and(|value| {
|
||||
value.pointer("/detail/code").and_then(Value::as_str) == Some("deactivated_workspace")
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,646 @@
|
||||
//! 快照/窗口解析、聚合、O(1) 滚动与调度信号。
|
||||
//!
|
||||
//! Provider raw quota JSON 的唯一解析者:产出 [`CodexAccountQuotaSnapshot`]
|
||||
//! 供服务编排层消费;`limit_reached` 是快照级事实,不能只从子窗口反推。
|
||||
|
||||
use std::time::{Duration, SystemTime};
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use gateway_core::engine::credential::{
|
||||
AccountQuotaSignals, CredentialRevision, ProviderAccountId, QuotaObservation,
|
||||
};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use super::{CodexCredentialQuotaError, QUOTA_SCHEDULING_TTL};
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct CodexQuotaFact {
|
||||
remaining_percent: Option<u8>,
|
||||
resets_at: Option<DateTime<Utc>>,
|
||||
exhausted: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum CodexQuotaWindowKind {
|
||||
ShortTerm,
|
||||
Weekly,
|
||||
Monthly,
|
||||
Other,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum CodexQuotaWindowRole {
|
||||
Primary,
|
||||
Secondary,
|
||||
Monthly,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct CodexQuotaWindow {
|
||||
key: String,
|
||||
source: String,
|
||||
/// 上游提供的限流桶名称(`metered_feature`/`limit_name`/`limitId`);用于非核心桶的展示前缀。
|
||||
limit_name: Option<String>,
|
||||
kind: CodexQuotaWindowKind,
|
||||
role: CodexQuotaWindowRole,
|
||||
window_seconds: Option<u64>,
|
||||
used_percent: Option<f64>,
|
||||
reset_at: Option<DateTime<Utc>>,
|
||||
limit_reached: bool,
|
||||
}
|
||||
|
||||
impl CodexQuotaWindow {
|
||||
#[must_use]
|
||||
pub fn key(&self) -> &str {
|
||||
&self.key
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn source(&self) -> &str {
|
||||
&self.source
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn limit_name(&self) -> Option<&str> {
|
||||
self.limit_name.as_deref()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn kind(&self) -> CodexQuotaWindowKind {
|
||||
self.kind
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn role(&self) -> CodexQuotaWindowRole {
|
||||
self.role
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn window_seconds(&self) -> Option<u64> {
|
||||
self.window_seconds
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn used_percent(&self) -> Option<f64> {
|
||||
self.used_percent
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn reset_at(&self) -> Option<DateTime<Utc>> {
|
||||
self.reset_at
|
||||
}
|
||||
|
||||
/// 该窗口是否已触顶(`limit_reached` 或 `used_percent >= 100`)。
|
||||
#[must_use]
|
||||
pub const fn limit_reached(&self) -> bool {
|
||||
self.limit_reached
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct CodexAccountQuotaSnapshot {
|
||||
account_id: ProviderAccountId,
|
||||
credential_revision: CredentialRevision,
|
||||
observed_at: SystemTime,
|
||||
fact: CodexQuotaFact,
|
||||
windows: Vec<CodexQuotaWindow>,
|
||||
}
|
||||
|
||||
impl CodexAccountQuotaSnapshot {
|
||||
#[must_use]
|
||||
pub const fn account_id(&self) -> &ProviderAccountId {
|
||||
&self.account_id
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn credential_revision(&self) -> CredentialRevision {
|
||||
self.credential_revision
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn observed_at(&self) -> SystemTime {
|
||||
self.observed_at
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn fact(&self) -> CodexQuotaFact {
|
||||
self.fact
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn windows(&self) -> &[CodexQuotaWindow] {
|
||||
&self.windows
|
||||
}
|
||||
|
||||
/// 投影滚动已过期窗口:`reset_at` 已过的窗口 `used_percent=0`、
|
||||
/// `limit_reached=false`、`reset_at` 滚到下一个周期(reset 到期自动归零)。
|
||||
/// 只作用于展示/调度投影,不写回持久化 raw JSON。
|
||||
///
|
||||
/// O(1) 计算:`next_reset = reset + (floor((now - reset) / window) + 1) * window`,
|
||||
/// 不做逐周期线性循环。
|
||||
#[must_use]
|
||||
pub fn roll_expired_windows(mut self, now: SystemTime) -> Self {
|
||||
for window in &mut self.windows {
|
||||
let Some(reset_at) = window.reset_at else {
|
||||
continue;
|
||||
};
|
||||
let reset = SystemTime::from(reset_at);
|
||||
if reset > now {
|
||||
continue;
|
||||
}
|
||||
window.used_percent = Some(0.0);
|
||||
window.limit_reached = false;
|
||||
let Some(window_seconds) = window.window_seconds.filter(|seconds| *seconds > 0) else {
|
||||
// 没有合法窗口时长:清除已过期的 reset,等待下一次权威观测。
|
||||
window.reset_at = None;
|
||||
continue;
|
||||
};
|
||||
let elapsed = now
|
||||
.duration_since(reset)
|
||||
.unwrap_or(Duration::ZERO)
|
||||
.as_secs();
|
||||
let periods = elapsed / window_seconds + 1;
|
||||
let Some(after) =
|
||||
reset.checked_add(Duration::from_secs(periods.saturating_mul(window_seconds)))
|
||||
else {
|
||||
window.reset_at = None;
|
||||
continue;
|
||||
};
|
||||
window.reset_at = Some(DateTime::<Utc>::from(after));
|
||||
}
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
impl CodexQuotaFact {
|
||||
#[must_use]
|
||||
pub const fn remaining_percent(&self) -> Option<u8> {
|
||||
self.remaining_percent
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn resets_at(&self) -> Option<DateTime<Utc>> {
|
||||
self.resets_at
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn exhausted(&self) -> bool {
|
||||
self.exhausted
|
||||
}
|
||||
}
|
||||
|
||||
pub fn parse_codex_quota_usage(usage: &Value) -> Result<CodexQuotaFact, CodexCredentialQuotaError> {
|
||||
let object = usage
|
||||
.as_object()
|
||||
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)?;
|
||||
let mut aggregate = QuotaAggregate::default();
|
||||
for limit in canonical_rate_limits(object)? {
|
||||
aggregate.observe_rate_limit(limit.rate_limit)?;
|
||||
}
|
||||
// spend_control 只提供 exhaustion 信号(`reached`),不生成窗口、
|
||||
// 不参与剩余量聚合(`individual_limit` 等字段无窗口语义)。
|
||||
if let Some(spend_control) = object.get("spend_control") {
|
||||
aggregate.observe_exhaustion_object(spend_control, "reached")?;
|
||||
}
|
||||
if let Some(credits) = object.get("credits") {
|
||||
aggregate.observe_exhaustion_object(credits, "overage_limit_reached")?;
|
||||
}
|
||||
if !aggregate.recognized {
|
||||
return Err(CodexCredentialQuotaError::InvalidCredentialData);
|
||||
}
|
||||
Ok(CodexQuotaFact {
|
||||
remaining_percent: aggregate.remaining_percent,
|
||||
resets_at: aggregate.resets_at,
|
||||
exhausted: aggregate.exhausted,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn quota_snapshot_from_observation(
|
||||
observation: &QuotaObservation,
|
||||
) -> Option<CodexAccountQuotaSnapshot> {
|
||||
let observed_at = observation.observed_at?;
|
||||
let quota = observation.quota.as_ref()?;
|
||||
parse_account_quota_snapshot(
|
||||
observation.account_id.clone(),
|
||||
observation.expected_revision,
|
||||
observed_at,
|
||||
&Value::Object(quota.expose_to_provider().clone()),
|
||||
)
|
||||
.ok()
|
||||
}
|
||||
|
||||
/// 调度投影缓存 TTL:`min(常规 TTL, 最近 active reset - now)`。
|
||||
/// reset 到期时缓存自然失效,不固定多等一个常规 TTL 周期。
|
||||
pub(crate) fn quota_projection_ttl(snapshot: &CodexAccountQuotaSnapshot) -> Option<Duration> {
|
||||
let age = SystemTime::now()
|
||||
.duration_since(snapshot.observed_at())
|
||||
.unwrap_or(Duration::ZERO);
|
||||
let normal_ttl = QUOTA_SCHEDULING_TTL.checked_sub(age)?;
|
||||
let now = SystemTime::now();
|
||||
let reset_ttl = snapshot
|
||||
.windows()
|
||||
.iter()
|
||||
.filter_map(|window| window.reset_at())
|
||||
.map(SystemTime::from)
|
||||
.filter(|reset| *reset > now)
|
||||
.map(|reset| reset.duration_since(now).unwrap_or(Duration::ZERO))
|
||||
.min();
|
||||
let ttl = match reset_ttl {
|
||||
Some(reset_ttl) if reset_ttl < normal_ttl => reset_ttl,
|
||||
_ => normal_ttl,
|
||||
};
|
||||
(!ttl.is_zero()).then_some(ttl)
|
||||
}
|
||||
|
||||
/// 从快照的滚动窗口派生中立调度事实(`limit_reached` = 任一窗口触顶)。
|
||||
pub(crate) fn scheduling_signals_from_snapshot(
|
||||
snapshot: &CodexAccountQuotaSnapshot,
|
||||
) -> Option<AccountQuotaSignals> {
|
||||
let rolled = snapshot.clone().roll_expired_windows(SystemTime::now());
|
||||
quota_scheduling_signals(&rolled)
|
||||
}
|
||||
|
||||
pub(crate) fn quota_scheduling_signals(
|
||||
snapshot: &CodexAccountQuotaSnapshot,
|
||||
) -> Option<AccountQuotaSignals> {
|
||||
let now = SystemTime::now();
|
||||
let mut limit_reached = false;
|
||||
let mut remaining_rank: Option<u64> = None;
|
||||
let mut reset_at: Option<SystemTime> = None;
|
||||
for window in &snapshot.windows {
|
||||
limit_reached |= window.limit_reached();
|
||||
if let Some(used) = window
|
||||
.used_percent()
|
||||
.filter(|used| used.is_finite() && (0.0..=100.0).contains(used))
|
||||
{
|
||||
let rank = ((100.0 - used) * 100.0).round() as u64;
|
||||
remaining_rank = Some(remaining_rank.map_or(rank, |current| current.min(rank)));
|
||||
}
|
||||
if let Some(reset) = window
|
||||
.reset_at()
|
||||
.map(SystemTime::from)
|
||||
.filter(|reset| *reset > now)
|
||||
{
|
||||
reset_at = Some(reset_at.map_or(reset, |current| current.min(reset)));
|
||||
}
|
||||
}
|
||||
// 快照级 limit_reached 是调度排除的充分条件:即使没有任何窗口带
|
||||
// percent/reset(如顶层限流标记),也必须产出信号;不能返回 None 丢失排除。
|
||||
if limit_reached || remaining_rank.is_some() || reset_at.is_some() {
|
||||
Some(AccountQuotaSignals::new(
|
||||
reset_at,
|
||||
remaining_rank,
|
||||
limit_reached,
|
||||
))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn parse_account_quota_snapshot(
|
||||
account_id: ProviderAccountId,
|
||||
credential_revision: CredentialRevision,
|
||||
observed_at: SystemTime,
|
||||
usage: &Value,
|
||||
) -> Result<CodexAccountQuotaSnapshot, CodexCredentialQuotaError> {
|
||||
let fact = parse_codex_quota_usage(usage)?;
|
||||
let object = usage
|
||||
.as_object()
|
||||
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)?;
|
||||
let mut windows = Vec::new();
|
||||
for limit in canonical_rate_limits(object)? {
|
||||
parse_rate_limit_windows(
|
||||
&limit.key,
|
||||
&limit.source,
|
||||
limit.limit_name.as_deref(),
|
||||
limit.rate_limit,
|
||||
&mut windows,
|
||||
)?;
|
||||
}
|
||||
windows.sort_by(|left, right| left.key.cmp(&right.key));
|
||||
Ok(CodexAccountQuotaSnapshot {
|
||||
account_id,
|
||||
credential_revision,
|
||||
observed_at,
|
||||
fact,
|
||||
windows,
|
||||
})
|
||||
}
|
||||
|
||||
/// 一个规范化后的限流桶:`key` 用于稳定区分窗口,`source` 用于展示分组。
|
||||
struct CanonicalRateLimit<'a> {
|
||||
key: String,
|
||||
source: String,
|
||||
limit_name: Option<String>,
|
||||
rate_limit: &'a Value,
|
||||
}
|
||||
|
||||
fn canonical_rate_limits(
|
||||
object: &Map<String, Value>,
|
||||
) -> Result<Vec<CanonicalRateLimit<'_>>, CodexCredentialQuotaError> {
|
||||
let mut limits = Vec::new();
|
||||
let has_primary_codex_rate_limit = object
|
||||
.get("rate_limit")
|
||||
.is_some_and(|value| !value.is_null());
|
||||
if let Some(rate_limit) = object.get("rate_limit") {
|
||||
// 官方 Codex 将顶层 `rate_limit` 固定识别为 `codex`,而不把
|
||||
// `active_limit` 投影成另一个展示桶。后者只是当前活动套餐的元数据。
|
||||
limits.push(CanonicalRateLimit {
|
||||
key: "core".to_owned(),
|
||||
source: "codex".to_owned(),
|
||||
limit_name: None,
|
||||
rate_limit,
|
||||
});
|
||||
}
|
||||
if let Some(rate_limit) = object
|
||||
.get("code_review_rate_limit")
|
||||
.filter(|value| !value.is_null())
|
||||
{
|
||||
limits.push(CanonicalRateLimit {
|
||||
key: "code-review".to_owned(),
|
||||
source: "code_review".to_owned(),
|
||||
limit_name: Some("code_review".to_owned()),
|
||||
rate_limit,
|
||||
});
|
||||
}
|
||||
let Some(additional) = object
|
||||
.get("additional_rate_limits")
|
||||
.filter(|value| !value.is_null())
|
||||
else {
|
||||
return Ok(limits);
|
||||
};
|
||||
for (index, value) in additional
|
||||
.as_array()
|
||||
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)?
|
||||
.iter()
|
||||
.enumerate()
|
||||
{
|
||||
let item = value
|
||||
.as_object()
|
||||
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)?;
|
||||
let (source, limit_name) = canonical_rate_limit_source(item, index);
|
||||
// 部分上游响应会在 `additional_rate_limits` 中重复默认的 `codex`
|
||||
// 桶。它没有独立的 limit id,不能在账号面板再次展示。
|
||||
if has_primary_codex_rate_limit && is_codex_limit_id(&source) {
|
||||
continue;
|
||||
}
|
||||
if let Some(rate_limit) = item.get("rate_limit") {
|
||||
// 非 Codex additional 桶保留为独立快照;额外用索引生成窗口键,
|
||||
// 避免同名 additional 产生重复前端 key。
|
||||
limits.push(CanonicalRateLimit {
|
||||
key: format!("additional-{index}-{source}"),
|
||||
source,
|
||||
limit_name,
|
||||
rate_limit,
|
||||
});
|
||||
}
|
||||
}
|
||||
Ok(limits)
|
||||
}
|
||||
|
||||
fn is_codex_limit_id(value: &str) -> bool {
|
||||
value.trim().eq_ignore_ascii_case("codex")
|
||||
}
|
||||
|
||||
fn canonical_rate_limit_source(
|
||||
item: &Map<String, Value>,
|
||||
index: usize,
|
||||
) -> (String, Option<String>) {
|
||||
let source = item
|
||||
.get("metered_feature")
|
||||
.or_else(|| item.get("limit_name"))
|
||||
.and_then(Value::as_str)
|
||||
.filter(|value| {
|
||||
!value.is_empty() && value.len() <= 128 && !value.chars().any(char::is_control)
|
||||
})
|
||||
.map(str::to_owned)
|
||||
.unwrap_or_else(|| format!("additional-{index}"));
|
||||
let limit_name = item
|
||||
.get("limit_name")
|
||||
.or_else(|| item.get("metered_feature"))
|
||||
.and_then(Value::as_str)
|
||||
.filter(|value| {
|
||||
!value.is_empty() && value.len() <= 128 && !value.chars().any(char::is_control)
|
||||
})
|
||||
.map(str::to_owned);
|
||||
(source, limit_name)
|
||||
}
|
||||
|
||||
fn parse_rate_limit_windows(
|
||||
key_source: &str,
|
||||
source: &str,
|
||||
limit_name: Option<&str>,
|
||||
value: &Value,
|
||||
output: &mut Vec<CodexQuotaWindow>,
|
||||
) -> Result<(), CodexCredentialQuotaError> {
|
||||
let object = value
|
||||
.as_object()
|
||||
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)?;
|
||||
// 顶层限流标记是快照级事实:任一窗口都要继承所属
|
||||
// rate_limit 的 `limit_reached=true` / `allowed=false`。
|
||||
let top_level_reached = optional_bool(object, "limit_reached")?.unwrap_or(false)
|
||||
|| optional_bool(object, "allowed")?.is_some_and(|allowed| !allowed);
|
||||
for (name, role) in [
|
||||
("primary_window", CodexQuotaWindowRole::Primary),
|
||||
("secondary_window", CodexQuotaWindowRole::Secondary),
|
||||
] {
|
||||
let Some(window) = object.get(name) else {
|
||||
continue;
|
||||
};
|
||||
if window.is_null() {
|
||||
continue;
|
||||
}
|
||||
let window = window
|
||||
.as_object()
|
||||
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)?;
|
||||
let window_seconds = optional_positive_u64(window, "limit_window_seconds")?
|
||||
.or(optional_positive_u64(window, "window_seconds")?)
|
||||
.or(optional_positive_u64(window, "window_minutes")?
|
||||
.and_then(|minutes| minutes.checked_mul(60)));
|
||||
let used_percent = window
|
||||
.get("used_percent")
|
||||
.map(|value| {
|
||||
value
|
||||
.as_f64()
|
||||
.filter(|value| value.is_finite() && *value >= 0.0)
|
||||
.map(|value| value.clamp(0.0, 100.0))
|
||||
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)
|
||||
})
|
||||
.transpose()?;
|
||||
let reset_at = window
|
||||
.get("reset_at")
|
||||
.map(|value| {
|
||||
value
|
||||
.as_i64()
|
||||
.filter(|value| *value > 0)
|
||||
.and_then(|value| DateTime::<Utc>::from_timestamp(value, 0))
|
||||
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)
|
||||
})
|
||||
.transpose()?;
|
||||
let limit_reached = top_level_reached
|
||||
|| window
|
||||
.get("limit_reached")
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false)
|
||||
|| used_percent.is_some_and(|used| used >= 100.0);
|
||||
let kind = quota_window_kind(window_seconds);
|
||||
output.push(CodexQuotaWindow {
|
||||
key: format!("{}-{}", quota_key(key_source), quota_role_name(role, kind)),
|
||||
source: source.to_owned(),
|
||||
limit_name: limit_name.map(str::to_owned),
|
||||
kind,
|
||||
role,
|
||||
window_seconds,
|
||||
used_percent,
|
||||
reset_at,
|
||||
limit_reached,
|
||||
});
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn optional_positive_u64(
|
||||
object: &Map<String, Value>,
|
||||
key: &str,
|
||||
) -> Result<Option<u64>, CodexCredentialQuotaError> {
|
||||
object
|
||||
.get(key)
|
||||
.map(|value| {
|
||||
value
|
||||
.as_u64()
|
||||
.filter(|value| *value > 0)
|
||||
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)
|
||||
})
|
||||
.transpose()
|
||||
}
|
||||
|
||||
const fn quota_window_kind(seconds: Option<u64>) -> CodexQuotaWindowKind {
|
||||
match seconds {
|
||||
Some(value) if value >= 17_100 && value <= 18_900 => CodexQuotaWindowKind::ShortTerm,
|
||||
Some(value) if value >= 574_560 && value <= 635_040 => CodexQuotaWindowKind::Weekly,
|
||||
Some(value) if value >= 2_462_400 && value <= 2_721_600 => CodexQuotaWindowKind::Monthly,
|
||||
_ => CodexQuotaWindowKind::Other,
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_key(value: &str) -> String {
|
||||
let mut key = value
|
||||
.chars()
|
||||
.map(|character| {
|
||||
if character.is_ascii_alphanumeric() {
|
||||
character.to_ascii_lowercase()
|
||||
} else {
|
||||
'-'
|
||||
}
|
||||
})
|
||||
.collect::<String>();
|
||||
while key.contains("--") {
|
||||
key = key.replace("--", "-");
|
||||
}
|
||||
let key = key.trim_matches('-');
|
||||
if key.is_empty() {
|
||||
"quota".to_owned()
|
||||
} else {
|
||||
key.to_owned()
|
||||
}
|
||||
}
|
||||
|
||||
const fn quota_role_name(role: CodexQuotaWindowRole, kind: CodexQuotaWindowKind) -> &'static str {
|
||||
match kind {
|
||||
CodexQuotaWindowKind::ShortTerm => "five-hour",
|
||||
CodexQuotaWindowKind::Weekly => "weekly",
|
||||
CodexQuotaWindowKind::Monthly => "monthly",
|
||||
CodexQuotaWindowKind::Other => match role {
|
||||
CodexQuotaWindowRole::Primary => "primary",
|
||||
CodexQuotaWindowRole::Secondary => "secondary",
|
||||
CodexQuotaWindowRole::Monthly => "monthly",
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct QuotaAggregate {
|
||||
recognized: bool,
|
||||
exhausted: bool,
|
||||
remaining_percent: Option<u8>,
|
||||
resets_at: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
impl QuotaAggregate {
|
||||
fn observe_rate_limit(&mut self, value: &Value) -> Result<(), CodexCredentialQuotaError> {
|
||||
let object = value
|
||||
.as_object()
|
||||
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)?;
|
||||
self.recognized = true;
|
||||
self.exhausted |= optional_bool(object, "limit_reached")?.unwrap_or(false);
|
||||
self.exhausted |= optional_bool(object, "allowed")?.is_some_and(|allowed| !allowed);
|
||||
for key in ["primary_window", "secondary_window"] {
|
||||
if let Some(window) = object.get(key) {
|
||||
self.observe_window(window)?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn observe_window(&mut self, value: &Value) -> Result<(), CodexCredentialQuotaError> {
|
||||
if value.is_null() {
|
||||
return Ok(());
|
||||
}
|
||||
let object = value
|
||||
.as_object()
|
||||
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)?;
|
||||
if let Some(used) = object.get("used_percent") {
|
||||
let used = used
|
||||
.as_f64()
|
||||
.filter(|value| value.is_finite() && *value >= 0.0)
|
||||
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)?;
|
||||
// 子窗口触顶(used>=100)是快照级 exhaustion 事实。
|
||||
self.exhausted |= used >= 100.0;
|
||||
let remaining = (100.0 - used).clamp(0.0, 100.0).round() as u8;
|
||||
self.remaining_percent = Some(
|
||||
self.remaining_percent
|
||||
.map_or(remaining, |current| current.min(remaining)),
|
||||
);
|
||||
}
|
||||
if let Some(reset) = object.get("reset_at") {
|
||||
let seconds = reset
|
||||
.as_i64()
|
||||
.filter(|seconds| *seconds > 0)
|
||||
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)?;
|
||||
let reset = DateTime::<Utc>::from_timestamp(seconds, 0)
|
||||
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)?;
|
||||
self.resets_at = Some(self.resets_at.map_or(reset, |current| current.min(reset)));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn observe_exhaustion_object(
|
||||
&mut self,
|
||||
value: &Value,
|
||||
key: &str,
|
||||
) -> Result<(), CodexCredentialQuotaError> {
|
||||
if value.is_null() {
|
||||
return Ok(());
|
||||
}
|
||||
let object = value
|
||||
.as_object()
|
||||
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)?;
|
||||
self.recognized = true;
|
||||
self.exhausted |= optional_bool(object, key)?.unwrap_or(false);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn optional_bool(
|
||||
object: &Map<String, Value>,
|
||||
key: &str,
|
||||
) -> Result<Option<bool>, CodexCredentialQuotaError> {
|
||||
object
|
||||
.get(key)
|
||||
.map(|value| {
|
||||
value
|
||||
.as_bool()
|
||||
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)
|
||||
})
|
||||
.transpose()
|
||||
}
|
||||
@@ -192,7 +192,6 @@ impl CodexCredentialRefreshService {
|
||||
self.persist_terminal(
|
||||
&due.account,
|
||||
AccountAvailability::Expired,
|
||||
"invalid_grant",
|
||||
CodexCredentialRefreshOutcome::Invalidated { account_id },
|
||||
)
|
||||
.await
|
||||
@@ -201,7 +200,6 @@ impl CodexCredentialRefreshService {
|
||||
self.persist_terminal(
|
||||
&due.account,
|
||||
AccountAvailability::Banned,
|
||||
"account_banned",
|
||||
CodexCredentialRefreshOutcome::Banned { account_id },
|
||||
)
|
||||
.await
|
||||
@@ -330,7 +328,6 @@ impl CodexCredentialRefreshService {
|
||||
.persist_terminal(
|
||||
&due.account,
|
||||
AccountAvailability::Invalid,
|
||||
"refreshed_identity_rejected",
|
||||
CodexCredentialRefreshOutcome::Invalidated {
|
||||
account_id: due.account.id().to_string(),
|
||||
},
|
||||
@@ -433,7 +430,6 @@ impl CodexCredentialRefreshService {
|
||||
&self,
|
||||
account: &ProviderAccount,
|
||||
availability: AccountAvailability,
|
||||
reason: &'static str,
|
||||
outcome: CodexCredentialRefreshOutcome,
|
||||
) -> Result<CodexCredentialRefreshOutcome, CodexCredentialRefreshError> {
|
||||
match self.repository.load_runtime_credential(account).await {
|
||||
@@ -447,13 +443,7 @@ impl CodexCredentialRefreshService {
|
||||
}
|
||||
match self
|
||||
.repository
|
||||
.apply_state(
|
||||
account,
|
||||
availability,
|
||||
Some(reason.to_owned()),
|
||||
None,
|
||||
SystemTime::now(),
|
||||
)
|
||||
.apply_state(account, availability, SystemTime::now())
|
||||
.await
|
||||
{
|
||||
Ok(()) => Ok(outcome),
|
||||
|
||||
@@ -233,9 +233,19 @@ impl CodexCredentialRepository {
|
||||
&self,
|
||||
account: &ProviderAccount,
|
||||
availability: AccountAvailability,
|
||||
reason: Option<String>,
|
||||
cooldown_until: Option<SystemTime>,
|
||||
observed_at: SystemTime,
|
||||
) -> Result<(), CredentialRepositoryError> {
|
||||
self.apply_state_with_message(account, availability, observed_at, None)
|
||||
.await
|
||||
}
|
||||
|
||||
/// 带上游错误描述的状态写入;`message` 为 `Some` 时持久化最近错误,恢复 `Ready` 时清空。
|
||||
pub async fn apply_state_with_message(
|
||||
&self,
|
||||
account: &ProviderAccount,
|
||||
availability: AccountAvailability,
|
||||
observed_at: SystemTime,
|
||||
message: Option<String>,
|
||||
) -> Result<(), CredentialRepositoryError> {
|
||||
if !account.enabled() {
|
||||
return Ok(());
|
||||
@@ -245,9 +255,8 @@ impl CodexCredentialRepository {
|
||||
account_id: account.id().clone(),
|
||||
expected_revision: account.revision(),
|
||||
availability,
|
||||
reason,
|
||||
cooldown_until,
|
||||
observed_at,
|
||||
message,
|
||||
})
|
||||
.await?;
|
||||
Ok(())
|
||||
|
||||
@@ -32,8 +32,6 @@ use super::types::{
|
||||
RuntimeCodexCookie,
|
||||
};
|
||||
|
||||
const DEFAULT_RATE_LIMIT_COOLDOWN: Duration = Duration::from_secs(60);
|
||||
const CLOUDFLARE_PATH_BLOCK_COOLDOWN: Duration = Duration::from_secs(30);
|
||||
const CLOUDFLARE_RECOVERY_STALE_AFTER: Duration = Duration::from_secs(60 * 60);
|
||||
const CLOUDFLARE_CHALLENGE_BACKOFF: [Duration; 4] = [
|
||||
Duration::from_secs(10),
|
||||
@@ -84,7 +82,6 @@ struct RiskRecoveryState {
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
enum CookieRecovery {
|
||||
None,
|
||||
ExpireAt(SystemTime),
|
||||
Clear,
|
||||
}
|
||||
@@ -297,6 +294,24 @@ impl CodexCredentialSelector {
|
||||
if !diagnostic {
|
||||
self.quota.prepare_scheduling(&accounts).await;
|
||||
}
|
||||
// 429 临时限流(Redis 冷却)中的账号即使额度正常也不参与调度,
|
||||
// 到期后冷却过期自动恢复。诊断请求(连接测试)绕过冷却,需要真实探测。
|
||||
let accounts = if diagnostic {
|
||||
accounts
|
||||
} else {
|
||||
let mut retained = Vec::with_capacity(accounts.len());
|
||||
for account in accounts {
|
||||
let cooling = self
|
||||
.quota
|
||||
.rate_limited_until(account.id(), account.revision())
|
||||
.await
|
||||
.unwrap_or(None);
|
||||
if cooling.is_none() {
|
||||
retained.push(account);
|
||||
}
|
||||
}
|
||||
retained
|
||||
};
|
||||
let account_ids = accounts
|
||||
.iter()
|
||||
.map(|account| account.id().clone())
|
||||
@@ -321,6 +336,7 @@ impl CodexCredentialSelector {
|
||||
last_started_at: None,
|
||||
quota_reset_at: None,
|
||||
quota_remaining_rank: None,
|
||||
quota_limit_reached: false,
|
||||
failure_rate_basis_points: None,
|
||||
first_output_latency_ms: None,
|
||||
})
|
||||
@@ -329,39 +345,6 @@ impl CodexCredentialSelector {
|
||||
AccountCandidate { account, signals }
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
// 旧版开关只影响“候选资格”。保留原始账号供 revision-fenced 凭据加载,
|
||||
// 用影子候选承载放宽后的资格,避免伪造已持久化账号的运行状态。
|
||||
let quota_exhausted_candidates = (!self.skip_exhausted).then(|| {
|
||||
candidates
|
||||
.iter()
|
||||
.cloned()
|
||||
.map(|mut candidate| {
|
||||
candidate.account = quota_exhausted_scheduling_projection(&candidate.account);
|
||||
candidate
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
});
|
||||
let scheduling_candidates = quota_exhausted_candidates.as_deref().unwrap_or(&candidates);
|
||||
let mut affinity = if diagnostic {
|
||||
AffinitySelection::default()
|
||||
} else {
|
||||
self.resolve_session_affinity(
|
||||
request.session_affinity_key,
|
||||
&candidates,
|
||||
SystemTime::now(),
|
||||
)
|
||||
.await
|
||||
};
|
||||
let cyber_policy_scope = self
|
||||
.prepare_cyber_policy_scope(cyber_policy_session_key)
|
||||
.await;
|
||||
let mut excluded = request.attempt.excluded_accounts().clone();
|
||||
if let Some(state) = cyber_policy_scope
|
||||
.as_ref()
|
||||
.and_then(|scope| scope.state.as_ref())
|
||||
{
|
||||
excluded.extend(state.excluded_accounts().iter().cloned());
|
||||
}
|
||||
let continuation_account = match request.attempt.continuation_attempt() {
|
||||
ContinuationAttempt::Native => request
|
||||
.attempt
|
||||
@@ -383,7 +366,74 @@ impl CodexCredentialSelector {
|
||||
{
|
||||
return Err(CredentialSelectionError::NoEligibleCredential);
|
||||
}
|
||||
let pinned_account = required_account.or(continuation_account);
|
||||
let pinned_account = required_account.or_else(|| continuation_account.clone());
|
||||
let preserve_native_continuation_state = matches!(
|
||||
request.attempt.continuation_attempt(),
|
||||
ContinuationAttempt::Native
|
||||
) && continuation_account.is_some();
|
||||
// skip_exhausted 语义:true 时排除任一窗口触顶(limit_reached)的账号,
|
||||
// 仅由上游 429 兜底;false 时保留全部候选(含 QuotaExhausted 的放宽投影)。
|
||||
// 保留原始账号供 revision-fenced 凭据加载,用影子候选承载资格放宽。
|
||||
let scheduling_candidates = if diagnostic {
|
||||
// 连接测试必须实际探测指定账号;额度快照只影响常规调度。
|
||||
candidates.clone()
|
||||
} else if self.skip_exhausted {
|
||||
candidates
|
||||
.iter()
|
||||
.filter(|candidate| !candidate.signals.quota_limit_reached)
|
||||
.cloned()
|
||||
.collect::<Vec<_>>()
|
||||
} else {
|
||||
candidates
|
||||
.iter()
|
||||
.cloned()
|
||||
.map(|mut candidate| {
|
||||
// 已命中 previous_response_id 的 native continuation 先检查原账号。
|
||||
// 即使常规调度允许探测额度耗尽账号,也必须保留原账号真实状态,
|
||||
// 由 Coordinator 据此切换到可携带重放,而非伪装为 Ready。
|
||||
if !preserve_native_continuation_state
|
||||
|| pinned_account.as_ref() != Some(candidate.account.id())
|
||||
{
|
||||
candidate.account =
|
||||
quota_exhausted_scheduling_projection(&candidate.account);
|
||||
}
|
||||
candidate
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
};
|
||||
let scheduling_candidates = scheduling_candidates.as_slice();
|
||||
let mut affinity = if diagnostic {
|
||||
AffinitySelection::default()
|
||||
} else {
|
||||
self.resolve_session_affinity(
|
||||
request.session_affinity_key,
|
||||
&candidates,
|
||||
SystemTime::now(),
|
||||
)
|
||||
.await
|
||||
};
|
||||
let cyber_policy_scope = self
|
||||
.prepare_cyber_policy_scope(cyber_policy_session_key)
|
||||
.await;
|
||||
let mut excluded = request.attempt.excluded_accounts().clone();
|
||||
if !diagnostic
|
||||
&& preserve_native_continuation_state
|
||||
&& let Some(pinned) = pinned_account.as_ref()
|
||||
&& candidates.iter().any(|candidate| {
|
||||
candidate.account.id() == pinned && candidate.signals.quota_limit_reached
|
||||
})
|
||||
{
|
||||
// `skip_exhausted = false` 只允许普通新请求探测可疑账号;已绑定的
|
||||
// native response handle 先将真实触顶状态交给 Coordinator,随后才可
|
||||
// 通过 ReplayAny 按当前调度策略换号。
|
||||
excluded.insert(pinned.clone());
|
||||
}
|
||||
if let Some(state) = cyber_policy_scope
|
||||
.as_ref()
|
||||
.and_then(|scope| scope.state.as_ref())
|
||||
{
|
||||
excluded.extend(state.excluded_accounts().iter().cloned());
|
||||
}
|
||||
if let Some(required) = pinned_account.as_ref() {
|
||||
for candidate in &candidates {
|
||||
if candidate.account.id() != required {
|
||||
@@ -680,120 +730,78 @@ impl CodexCredentialSelector {
|
||||
&self,
|
||||
account: &ProviderAccount,
|
||||
failure: CodexAccountFailure,
|
||||
message: Option<String>,
|
||||
) -> Result<(), CredentialSelectionError> {
|
||||
let now = SystemTime::now();
|
||||
let (availability, reason, cooldown_until, cookie_recovery) = match failure {
|
||||
CodexAccountFailure::CredentialExpired => (
|
||||
AccountAvailability::Expired,
|
||||
Some("credential_expired".to_owned()),
|
||||
None,
|
||||
CookieRecovery::None,
|
||||
),
|
||||
CodexAccountFailure::IdentityVerificationRequired => (
|
||||
AccountAvailability::Invalid,
|
||||
Some("identity_verification_required".to_owned()),
|
||||
None,
|
||||
CookieRecovery::None,
|
||||
),
|
||||
CodexAccountFailure::Banned => (
|
||||
AccountAvailability::Banned,
|
||||
Some("account_banned".to_owned()),
|
||||
None,
|
||||
CookieRecovery::None,
|
||||
),
|
||||
CodexAccountFailure::QuotaExhausted => (
|
||||
AccountAvailability::QuotaExhausted,
|
||||
Some("quota_exhausted".to_owned()),
|
||||
None,
|
||||
CookieRecovery::None,
|
||||
),
|
||||
CodexAccountFailure::UsageLimitExhausted { retry_after } => {
|
||||
return self
|
||||
.record_resettable_limit_failure(
|
||||
account,
|
||||
retry_after,
|
||||
"usage_limit_exhausted",
|
||||
now,
|
||||
)
|
||||
.await;
|
||||
match failure {
|
||||
CodexAccountFailure::CredentialExpired => {
|
||||
self.apply_state_with_message(account, AccountAvailability::Expired, now, message)
|
||||
.await
|
||||
}
|
||||
CodexAccountFailure::IdentityVerificationRequired => {
|
||||
self.apply_state_with_message(account, AccountAvailability::Invalid, now, message)
|
||||
.await
|
||||
}
|
||||
CodexAccountFailure::Banned => {
|
||||
self.apply_state_with_message(account, AccountAvailability::Banned, now, message)
|
||||
.await
|
||||
}
|
||||
// 402:真额度耗尽,由 quota worker 按窗口 reset 自动恢复。
|
||||
CodexAccountFailure::QuotaExhausted
|
||||
| CodexAccountFailure::UsageLimitExhausted { .. } => {
|
||||
self.apply_state_with_message(
|
||||
account,
|
||||
AccountAvailability::QuotaExhausted,
|
||||
now,
|
||||
message,
|
||||
)
|
||||
.await
|
||||
}
|
||||
// 429:临时限流只写 quota,availability 不改变。
|
||||
CodexAccountFailure::RateLimited { retry_after } => {
|
||||
return self
|
||||
.record_resettable_limit_failure(account, retry_after, "rate_limited", now)
|
||||
.await;
|
||||
self.quota
|
||||
.apply_rate_limit_429(account, retry_after, now)
|
||||
.await
|
||||
.map_err(|_| CredentialSelectionError::Store)?;
|
||||
Ok(())
|
||||
}
|
||||
// Cloudflare 挑战:内存退避表(记录风险计数),不写 availability。
|
||||
CodexAccountFailure::CloudflareChallenge { retry_after } => {
|
||||
let delay = self.cloudflare_challenge_delay(account.id(), now, retry_after);
|
||||
let cooldown_until = now.checked_add(delay);
|
||||
(
|
||||
AccountAvailability::Cooldown,
|
||||
Some("cloudflare_challenge".to_owned()),
|
||||
cooldown_until,
|
||||
cooldown_until.map_or(CookieRecovery::Clear, CookieRecovery::ExpireAt),
|
||||
)
|
||||
let recovery = now
|
||||
.checked_add(delay)
|
||||
.map_or(CookieRecovery::Clear, CookieRecovery::ExpireAt);
|
||||
self.apply_cookie_recovery(account, recovery).await?;
|
||||
Ok(())
|
||||
}
|
||||
// Cloudflare 路径被封:连续超阈值才置 Invalid;否则内存退避。
|
||||
CodexAccountFailure::CloudflarePathBlocked => {
|
||||
let blocked = self.record_cloudflare_path_block(account.id(), now);
|
||||
if blocked >= CLOUDFLARE_PATH_BLOCK_THRESHOLD {
|
||||
(
|
||||
self.apply_state_with_message(
|
||||
account,
|
||||
AccountAvailability::Invalid,
|
||||
Some("cloudflare_path_blocked".to_owned()),
|
||||
None,
|
||||
CookieRecovery::Clear,
|
||||
)
|
||||
} else {
|
||||
(
|
||||
AccountAvailability::Cooldown,
|
||||
Some("cloudflare_path_blocked".to_owned()),
|
||||
now.checked_add(CLOUDFLARE_PATH_BLOCK_COOLDOWN),
|
||||
CookieRecovery::Clear,
|
||||
now,
|
||||
message,
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
self.apply_cookie_recovery(account, CookieRecovery::Clear)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
};
|
||||
self.repository
|
||||
.apply_state(account, availability, reason, cooldown_until, now)
|
||||
.await?;
|
||||
self.apply_cookie_recovery(account, cookie_recovery).await?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
async fn record_resettable_limit_failure(
|
||||
async fn apply_state_with_message(
|
||||
&self,
|
||||
account: &ProviderAccount,
|
||||
retry_after: Option<Duration>,
|
||||
reason: &'static str,
|
||||
availability: AccountAvailability,
|
||||
observed_at: SystemTime,
|
||||
message: Option<String>,
|
||||
) -> Result<(), CredentialSelectionError> {
|
||||
let current = self.current_account(account.id()).await?;
|
||||
if current.revision() != account.revision() {
|
||||
return Err(CredentialRepositoryError::RevisionConflict.into());
|
||||
}
|
||||
let cooldown_until =
|
||||
observed_at.checked_add(retry_after.unwrap_or(DEFAULT_RATE_LIMIT_COOLDOWN));
|
||||
let cooldown_until = match (current.availability(), current.cooldown_until()) {
|
||||
(AccountAvailability::QuotaExhausted, Some(existing)) => {
|
||||
cooldown_until.map_or(Some(existing), |until| Some(existing.max(until)))
|
||||
}
|
||||
_ => cooldown_until,
|
||||
};
|
||||
let reason = if current.availability() == AccountAvailability::QuotaExhausted
|
||||
&& current.cooldown_until().is_none()
|
||||
{
|
||||
Some("quota_exhausted".to_owned())
|
||||
} else {
|
||||
Some(reason.to_owned())
|
||||
};
|
||||
// 一次真实 429 证明上游已经接受当前凭据,可纠正旧的认证/封禁终态。
|
||||
// resettable 用量限制沿用 v2 的 quota lock 语义,并保留/延长复核时间。
|
||||
self.repository
|
||||
.apply_state(
|
||||
¤t,
|
||||
AccountAvailability::QuotaExhausted,
|
||||
reason,
|
||||
cooldown_until,
|
||||
observed_at,
|
||||
)
|
||||
.apply_state_with_message(account, availability, observed_at, message)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -857,7 +865,6 @@ impl CodexCredentialSelector {
|
||||
|| !matches!(
|
||||
current.availability(),
|
||||
AccountAvailability::Unknown
|
||||
| AccountAvailability::Cooldown
|
||||
| AccountAvailability::Expired
|
||||
| AccountAvailability::Invalid
|
||||
| AccountAvailability::Banned
|
||||
@@ -867,13 +874,7 @@ impl CodexCredentialSelector {
|
||||
}
|
||||
if let Err(error) = self
|
||||
.repository
|
||||
.apply_state(
|
||||
¤t,
|
||||
AccountAvailability::Ready,
|
||||
None,
|
||||
None,
|
||||
SystemTime::now(),
|
||||
)
|
||||
.apply_state(¤t, AccountAvailability::Ready, SystemTime::now())
|
||||
.await
|
||||
{
|
||||
tracing::warn!(
|
||||
@@ -985,15 +986,11 @@ impl CodexCredentialSelector {
|
||||
account: &ProviderAccount,
|
||||
recovery: CookieRecovery,
|
||||
) -> Result<(), CredentialSelectionError> {
|
||||
if matches!(recovery, CookieRecovery::None) {
|
||||
return Ok(());
|
||||
}
|
||||
let mut data = self.repository.load_complete_data(account).await?;
|
||||
if data.cookies().is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
match recovery {
|
||||
CookieRecovery::None => return Ok(()),
|
||||
CookieRecovery::ExpireAt(expires_at) => {
|
||||
let expires_at = chrono::DateTime::<chrono::Utc>::from(expires_at);
|
||||
for cookie in data.cookies_mut() {
|
||||
@@ -1013,11 +1010,9 @@ impl CodexCredentialSelector {
|
||||
|
||||
fn quota_exhausted_scheduling_projection(account: &ProviderAccount) -> ProviderAccount {
|
||||
if account.availability() == AccountAvailability::QuotaExhausted {
|
||||
return account.clone().with_runtime_state(
|
||||
account.enabled(),
|
||||
AccountAvailability::Ready,
|
||||
account.cooldown_until(),
|
||||
);
|
||||
return account
|
||||
.clone()
|
||||
.with_runtime_state(account.enabled(), AccountAvailability::Ready);
|
||||
}
|
||||
account.clone()
|
||||
}
|
||||
@@ -1032,7 +1027,6 @@ fn affinity_unavailable_reason(account: &ProviderAccount, now: SystemTime) -> Af
|
||||
}
|
||||
match account.availability() {
|
||||
AccountAvailability::QuotaExhausted => AffinityEscapeReason::QuotaExhausted,
|
||||
AccountAvailability::Cooldown => AffinityEscapeReason::Cooldown,
|
||||
AccountAvailability::Unknown
|
||||
| AccountAvailability::Ready
|
||||
| AccountAvailability::Expired
|
||||
|
||||
@@ -82,7 +82,7 @@ pub async fn initialize(
|
||||
profile.clone(),
|
||||
Arc::new(
|
||||
OfficialCodexCliReleaseTransport::new()
|
||||
.map_err(|_| OpenAiInitializeError::DesktopRelease)?,
|
||||
.map_err(|_| OpenAiInitializeError::CliRelease)?,
|
||||
),
|
||||
));
|
||||
let repository = CodexCredentialRepository::new(Arc::clone(&accounts));
|
||||
@@ -111,6 +111,7 @@ pub async fn initialize(
|
||||
http.clone(),
|
||||
config.base_url().to_owned(),
|
||||
Arc::clone(&agent_identity),
|
||||
ports.cooldowns(),
|
||||
));
|
||||
let quota_skip_exhausted = config.quota_skip_exhausted();
|
||||
let selector = Arc::new(CodexCredentialSelector::new(
|
||||
@@ -267,6 +268,8 @@ pub enum OpenAiInitializeError {
|
||||
Refresh,
|
||||
#[error("OpenAI Desktop release service could not initialize")]
|
||||
DesktopRelease,
|
||||
#[error("OpenAI CLI release service could not initialize")]
|
||||
CliRelease,
|
||||
#[error("OpenAI worker plan is invalid")]
|
||||
Worker,
|
||||
}
|
||||
|
||||
@@ -1697,9 +1697,13 @@ fn map_selection_error(error: CredentialSelectionError) -> ProviderError {
|
||||
struct MappedProviderFailure {
|
||||
error: ProviderError,
|
||||
account_failure: Option<CodexAccountFailure>,
|
||||
/// 原始上游错误描述,随 availability 状态一并持久化。
|
||||
error_message: Option<String>,
|
||||
cyber_policy_failure: bool,
|
||||
set_cookie_headers: Vec<String>,
|
||||
rate_limit_headers: Vec<(String, String)>,
|
||||
/// 结构化 usage-limit 错误给出的窗口 reset(Unix 秒)。
|
||||
confirmed_exhaustion_reset_at: Option<i64>,
|
||||
observation: Option<ProviderResponseObservation>,
|
||||
capture_response_cookies: bool,
|
||||
}
|
||||
@@ -1709,9 +1713,11 @@ impl MappedProviderFailure {
|
||||
Self {
|
||||
error,
|
||||
account_failure: None,
|
||||
error_message: None,
|
||||
cyber_policy_failure: false,
|
||||
set_cookie_headers: Vec::new(),
|
||||
rate_limit_headers: Vec::new(),
|
||||
confirmed_exhaustion_reset_at: None,
|
||||
observation: None,
|
||||
capture_response_cookies: false,
|
||||
}
|
||||
@@ -1869,6 +1875,18 @@ async fn apply_failure(
|
||||
failure.account_failure,
|
||||
Some(CodexAccountFailure::QuotaExhausted | CodexAccountFailure::UsageLimitExhausted { .. })
|
||||
);
|
||||
if needs_authoritative_quota_refresh
|
||||
&& let Err(error) = context
|
||||
.quota
|
||||
.record_confirmed_exhaustion(account, failure.confirmed_exhaustion_reset_at)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(
|
||||
account_id = %account.id(),
|
||||
error = %error,
|
||||
"Failed to project confirmed OpenAI quota exhaustion"
|
||||
);
|
||||
}
|
||||
if failure.cyber_policy_failure {
|
||||
context
|
||||
.selector
|
||||
@@ -1882,7 +1900,7 @@ async fn apply_failure(
|
||||
.await;
|
||||
if let Err(error) = context
|
||||
.selector
|
||||
.record_failure(account, account_failure)
|
||||
.record_failure(account, account_failure, failure.error_message.clone())
|
||||
.await
|
||||
{
|
||||
tracing::warn!(
|
||||
@@ -2162,9 +2180,11 @@ fn map_upstream_failure(
|
||||
MappedProviderFailure {
|
||||
error,
|
||||
account_failure: account_failure(category, failure.retry_after_seconds),
|
||||
error_message: failure.client_message,
|
||||
cyber_policy_failure,
|
||||
set_cookie_headers: failure.set_cookie_headers,
|
||||
rate_limit_headers: failure.rate_limit_headers,
|
||||
confirmed_exhaustion_reset_at: failure.usage_limit_resets_at,
|
||||
observation,
|
||||
capture_response_cookies: !matches!(
|
||||
category,
|
||||
@@ -2533,7 +2553,6 @@ impl ScheduledTask for OpenAiQuotaTask {
|
||||
updated = summary.updated,
|
||||
exhausted = summary.exhausted,
|
||||
banned = summary.banned,
|
||||
cooldown = summary.cooldown,
|
||||
transient = summary.transient,
|
||||
stale = summary.stale,
|
||||
"OpenAI quota cycle contained operational failures"
|
||||
|
||||
@@ -87,8 +87,13 @@ impl CodexBackendClient {
|
||||
let headers = self.request_headers_for_http_response(upstream_request, context)?;
|
||||
let headers_started_at = Instant::now();
|
||||
// 与官方 Codex app-server 一致:/v1/responses 请求体默认 zstd 压缩。
|
||||
// Codex 上游只交付 SSE;即使下游请求 `stream: false`,也要上游流式执行,
|
||||
// 再由 API 层收集 canonical events 并返回完整 JSON。不能把下游的传输偏好
|
||||
// 直接透传给 Codex,否则上游会以 400 拒绝非流式请求。
|
||||
let mut upstream_body = upstream_request.body().clone();
|
||||
upstream_body.insert("stream".to_owned(), serde_json::Value::Bool(true));
|
||||
let body =
|
||||
serde_json::to_vec(&upstream_request).map_err(CodexClientError::RequestBodyEncode)?;
|
||||
serde_json::to_vec(&upstream_body).map_err(CodexClientError::RequestBodyEncode)?;
|
||||
let body = zstd::stream::encode_all(std::io::Cursor::new(body), 3)
|
||||
.map_err(CodexClientError::RequestCompression)?;
|
||||
let response = self
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
use super::client::CodexClientVisibleUpstreamResponse;
|
||||
use super::protocol::responses::ResponsesSseFailure;
|
||||
use base64::{Engine as _, engine::general_purpose::STANDARD};
|
||||
use chrono::{DateTime, Utc};
|
||||
use reqwest::{StatusCode, header::HeaderMap};
|
||||
use serde_json::Value;
|
||||
|
||||
@@ -113,6 +114,8 @@ pub struct CodexUpstreamFailure {
|
||||
pub(crate) client_error_type: Option<String>,
|
||||
pub(crate) client_response: Option<Box<CodexClientVisibleUpstreamResponse>>,
|
||||
pub(crate) identity_error_code: Option<String>,
|
||||
/// 官方 `usage_limit_reached` 错误体携带的窗口重置时间(Unix 秒)。
|
||||
pub(crate) usage_limit_resets_at: Option<i64>,
|
||||
pub(crate) retry_after_seconds: Option<u64>,
|
||||
pub(crate) request_id: Option<String>,
|
||||
pub(crate) set_cookie_headers: Vec<String>,
|
||||
@@ -139,6 +142,7 @@ impl CodexUpstreamFailure {
|
||||
) -> Self {
|
||||
let fields = ParsedUpstreamError::from_body(body);
|
||||
let category = classify_upstream_failure(
|
||||
UpstreamFailureSource::HttpResponse,
|
||||
Some(status),
|
||||
body,
|
||||
&fields,
|
||||
@@ -153,6 +157,9 @@ impl CodexUpstreamFailure {
|
||||
client_error_type: fields.error_type,
|
||||
client_response: client_response.cloned().map(Box::new),
|
||||
identity_error_code: diagnostics.identity_error_code.clone(),
|
||||
usage_limit_resets_at: (category == CodexFailureCategory::UsageLimitExhausted)
|
||||
.then_some(fields.resets_at)
|
||||
.flatten(),
|
||||
retry_after_seconds,
|
||||
request_id: diagnostics.request_id.clone(),
|
||||
set_cookie_headers: set_cookie_headers.to_vec(),
|
||||
@@ -174,11 +181,13 @@ impl CodexUpstreamFailure {
|
||||
error_type: failure.upstream_type.clone(),
|
||||
message: failure.message.clone(),
|
||||
client_message: Some(failure.message.clone()),
|
||||
resets_at: None,
|
||||
};
|
||||
let status = failure
|
||||
.explicit_status_code
|
||||
.and_then(|code| StatusCode::from_u16(code).ok());
|
||||
let category = classify_upstream_failure(
|
||||
UpstreamFailureSource::SseFailure,
|
||||
status,
|
||||
"",
|
||||
&fields,
|
||||
@@ -193,6 +202,7 @@ impl CodexUpstreamFailure {
|
||||
client_error_type: fields.error_type,
|
||||
client_response: None,
|
||||
identity_error_code: diagnostics.identity_error_code.clone(),
|
||||
usage_limit_resets_at: None,
|
||||
retry_after_seconds: failure.retry_after_seconds,
|
||||
request_id: diagnostics.request_id.clone(),
|
||||
set_cookie_headers: set_cookie_headers.to_vec(),
|
||||
@@ -246,6 +256,7 @@ struct ParsedUpstreamError {
|
||||
error_type: Option<String>,
|
||||
message: String,
|
||||
client_message: Option<String>,
|
||||
resets_at: Option<i64>,
|
||||
}
|
||||
|
||||
impl ParsedUpstreamError {
|
||||
@@ -256,6 +267,7 @@ impl ParsedUpstreamError {
|
||||
error_type: None,
|
||||
message: body.to_owned(),
|
||||
client_message: None,
|
||||
resets_at: None,
|
||||
};
|
||||
};
|
||||
let error = value
|
||||
@@ -282,15 +294,26 @@ impl ParsedUpstreamError {
|
||||
.clone()
|
||||
.or_else(|| error.as_str().and_then(non_empty_owned))
|
||||
.unwrap_or_else(|| body.to_owned());
|
||||
let resets_at = error
|
||||
.get("resets_at")
|
||||
.and_then(Value::as_i64)
|
||||
.filter(|seconds| DateTime::<Utc>::from_timestamp(*seconds, 0).is_some());
|
||||
Self {
|
||||
code,
|
||||
error_type,
|
||||
message,
|
||||
client_message,
|
||||
resets_at,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum UpstreamFailureSource {
|
||||
HttpResponse,
|
||||
SseFailure,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct CodexUpstreamDiagnostics {
|
||||
pub status_code: Option<u16>,
|
||||
@@ -363,6 +386,7 @@ fn decode_identity_error_code(encoded: &str) -> Option<String> {
|
||||
}
|
||||
|
||||
fn classify_upstream_failure(
|
||||
source: UpstreamFailureSource,
|
||||
status: Option<StatusCode>,
|
||||
body: &str,
|
||||
fields: &ParsedUpstreamError,
|
||||
@@ -376,6 +400,19 @@ fn classify_upstream_failure(
|
||||
let message = fields.message.to_ascii_lowercase();
|
||||
let body = body.to_ascii_lowercase();
|
||||
|
||||
// 与官方 Codex HTTP/WS 路径一致:429 只有结构化
|
||||
// `error.type=usage_limit_reached` 才能确认额度窗口耗尽;其余 429 都是临时限流。
|
||||
// SSE `response.failed` 的结构化字段形态不同,保留其 code/type 语义单独分类。
|
||||
if matches!(source, UpstreamFailureSource::HttpResponse)
|
||||
&& status == Some(StatusCode::TOO_MANY_REQUESTS)
|
||||
{
|
||||
return if error_type == "usage_limit_reached" {
|
||||
CodexFailureCategory::UsageLimitExhausted
|
||||
} else {
|
||||
CodexFailureCategory::RateLimited
|
||||
};
|
||||
}
|
||||
|
||||
if [code.as_str(), message.as_str(), body.as_str()]
|
||||
.into_iter()
|
||||
.any(is_model_unsupported)
|
||||
|
||||
@@ -14,21 +14,22 @@ use url::Url;
|
||||
pub const CODEX_DESKTOP_APPCAST_URL: &str =
|
||||
"https://persistent.oaistatic.com/codex-app-prod/appcast.xml";
|
||||
pub const CODEX_CLI_RELEASE_URL: &str = "https://registry.npmjs.org/@openai%2Fcodex/latest";
|
||||
/// 官方制品画像检查周期。
|
||||
/// 官方发布元数据检查周期。
|
||||
pub const APPCAST_POLL_INTERVAL: Duration = Duration::from_secs(24 * 60 * 60);
|
||||
|
||||
const APPCAST_TIMEOUT: Duration = Duration::from_secs(30);
|
||||
const MAX_APPCAST_BYTES: usize = 1024 * 1024;
|
||||
const MAX_CLI_RELEASE_BYTES: usize = 64 * 1024;
|
||||
|
||||
/// Codex Desktop 上游请求身份。
|
||||
///
|
||||
/// 启动配置提供经源码审计的 Core、运行环境和初始 Desktop 版本。官方 appcast
|
||||
/// 检查成功后,会仅同步 Desktop 版本和构建号。
|
||||
/// 启动配置提供经源码审计的 Core、运行环境和 Desktop 启动版本。npm CLI metadata
|
||||
/// 与官方 appcast 检查成功后,分别同步 Core 版本和 Desktop 版本及构建号。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct CodexWireProfile {
|
||||
/// `originator` 请求头及 User-Agent 产品名。
|
||||
pub originator: String,
|
||||
/// bundled Codex Core 版本;同时用于 `/codex/models?client_version=`。
|
||||
/// npm `@openai/codex` 启动版本;用于 `/codex/models?client_version=` 与 UA。
|
||||
pub codex_version: String,
|
||||
/// Desktop 应用版本,用于 app-server `clientInfo.version` 对应的 UA 后缀。
|
||||
pub desktop_version: String,
|
||||
@@ -47,16 +48,18 @@ pub struct CodexWireProfile {
|
||||
}
|
||||
|
||||
impl CodexWireProfile {
|
||||
/// 按 Codex Core 的官方格式生成最终 User-Agent。
|
||||
/// 按 bundled Core app-server 的官方格式生成最终 User-Agent。
|
||||
pub fn user_agent(&self) -> String {
|
||||
format!(
|
||||
"{}/{} ({} {}; {}) {}",
|
||||
"{}/{} ({} {}; {}) {} ({}; {})",
|
||||
self.originator,
|
||||
self.codex_version,
|
||||
self.os_type,
|
||||
self.os_version,
|
||||
self.arch,
|
||||
self.terminal,
|
||||
self.originator,
|
||||
self.desktop_version,
|
||||
)
|
||||
}
|
||||
}
|
||||
@@ -96,7 +99,7 @@ impl CodexWireProfileState {
|
||||
profile.desktop_build = desktop_build.to_owned();
|
||||
}
|
||||
|
||||
/// 使用官方 npm metadata 发布的 Codex CLI 版本更新运行时 Core 画像。
|
||||
/// 使用官方 npm metadata 发布的 Codex CLI 版本更新运行时请求画像。
|
||||
pub fn update_cli_release(&self, codex_version: &str) {
|
||||
let mut profile = self
|
||||
.profile
|
||||
@@ -143,12 +146,12 @@ impl CodexCliReleaseTransport for OfficialCodexCliReleaseTransport {
|
||||
}
|
||||
if response
|
||||
.content_length()
|
||||
.is_some_and(|size| size > 64 * 1024)
|
||||
.is_some_and(|size| size > MAX_CLI_RELEASE_BYTES as u64)
|
||||
{
|
||||
return Err(CodexCliReleaseError::ResponseTooLarge);
|
||||
}
|
||||
let bytes = response.bytes().await?;
|
||||
if bytes.len() > 64 * 1024 {
|
||||
if bytes.len() > MAX_CLI_RELEASE_BYTES {
|
||||
return Err(CodexCliReleaseError::ResponseTooLarge);
|
||||
}
|
||||
let document: serde_json::Value = serde_json::from_slice(&bytes)
|
||||
|
||||
@@ -606,6 +606,26 @@ impl CodexResponsesRequest {
|
||||
)
|
||||
}
|
||||
|
||||
/// 返回会话亲和所用的子代理区分值。
|
||||
///
|
||||
/// Codex 原生请求在 turn metadata 中声明 `subagent_kind`;网关的
|
||||
/// `/responses/review` 等入口则通过 `client_metadata.x-openai-subagent`
|
||||
/// 传递给上游。该可选值仅参与会话亲和键派生。
|
||||
pub fn subagent_kind(&self) -> Option<String> {
|
||||
self.semantics()
|
||||
.subagent_kind
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
.or_else(|| {
|
||||
self.client_metadata()
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|metadata| metadata.get("x-openai-subagent"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(str::to_owned)
|
||||
})
|
||||
}
|
||||
|
||||
/// 设置 / 合并 client metadata。
|
||||
pub fn set_client_metadata(&mut self, client_metadata: Option<Value>) {
|
||||
match client_metadata {
|
||||
|
||||
@@ -10,6 +10,8 @@ use sha2::{Digest, Sha256};
|
||||
use crate::transport::protocol::responses::CodexResponsesRequest;
|
||||
|
||||
const PASSTHROUGH_HEADERS_CONTEXT_KEY: &str = "opaque_request_headers";
|
||||
const THREAD_SPAWN_SUBAGENT_KIND: &str = "thread_spawn";
|
||||
const THREAD_SPAWN_CONVERSATION_PREFIX: &str = "thread-spawn:";
|
||||
|
||||
const CROSS_ACCOUNT_IDENTITY_KEYS: &[&str] = &[
|
||||
"authorization",
|
||||
@@ -160,28 +162,52 @@ fn string_value(value: Option<&Value>) -> Option<String> {
|
||||
pub(crate) fn derive_conversation_anchor(
|
||||
request: &CodexResponsesRequest,
|
||||
) -> Option<(&'static str, String)> {
|
||||
thread_spawn_conversation_anchor(request).or_else(|| {
|
||||
request
|
||||
.client_session_id
|
||||
.as_deref()
|
||||
.map(|value| ("session", value.to_owned()))
|
||||
.or_else(|| {
|
||||
request
|
||||
.client_conversation_id
|
||||
.as_deref()
|
||||
.map(|value| ("conversation", value.to_owned()))
|
||||
})
|
||||
.or_else(|| {
|
||||
request
|
||||
.client_thread_id
|
||||
.as_deref()
|
||||
.map(|value| ("thread", value.to_owned()))
|
||||
})
|
||||
.or_else(|| {
|
||||
request
|
||||
.prompt_cache_key()
|
||||
.map(|value| ("prompt-cache", value.to_owned()))
|
||||
})
|
||||
.or_else(|| derive_stable_conversation_key(request).map(|value| ("request", value)))
|
||||
})
|
||||
}
|
||||
|
||||
/// A `thread_spawn` shares its parent session but is a separate child execution.
|
||||
/// Its child thread must therefore own the affinity and WebSocket pool identity.
|
||||
fn thread_spawn_conversation_anchor(
|
||||
request: &CodexResponsesRequest,
|
||||
) -> Option<(&'static str, String)> {
|
||||
if request.subagent_kind().as_deref() != Some(THREAD_SPAWN_SUBAGENT_KIND) {
|
||||
return None;
|
||||
}
|
||||
|
||||
request
|
||||
.client_session_id
|
||||
.client_thread_id
|
||||
.as_deref()
|
||||
.map(|value| ("session", value.to_owned()))
|
||||
.or_else(|| {
|
||||
request
|
||||
.client_conversation_id
|
||||
.as_deref()
|
||||
.map(|value| ("conversation", value.to_owned()))
|
||||
.map(str::trim)
|
||||
.filter(|thread_id| !thread_id.is_empty())
|
||||
.map(|thread_id| {
|
||||
(
|
||||
"subagent-thread",
|
||||
format!("{THREAD_SPAWN_CONVERSATION_PREFIX}{thread_id}"),
|
||||
)
|
||||
})
|
||||
.or_else(|| {
|
||||
request
|
||||
.client_thread_id
|
||||
.as_deref()
|
||||
.map(|value| ("thread", value.to_owned()))
|
||||
})
|
||||
.or_else(|| {
|
||||
request
|
||||
.prompt_cache_key()
|
||||
.map(|value| ("prompt-cache", value.to_owned()))
|
||||
})
|
||||
.or_else(|| derive_stable_conversation_key(request).map(|value| ("request", value)))
|
||||
}
|
||||
|
||||
const LEADING_SYSTEM_REMINDER_OPEN: &str = "<system-reminder>";
|
||||
|
||||
@@ -24,7 +24,7 @@ use super::{
|
||||
WebSocketPoolConnectLease, WebSocketPoolConnectOutcome, WebSocketPoolDecision,
|
||||
WebSocketPoolLease,
|
||||
},
|
||||
pump::{PumpKeepalive, PumpedWebSocket},
|
||||
pump::{PumpKeepalive, PumpLogContext, PumpedWebSocket},
|
||||
};
|
||||
|
||||
pub(crate) const WEBSOCKET_FAST_PATH_BUDGET: Duration = Duration::from_millis(800);
|
||||
@@ -203,10 +203,12 @@ async fn prepare_unpooled_websocket(
|
||||
decision_started_at: Instant,
|
||||
) -> Result<PreparedWebSocket, CodexWebSocketExchangeError> {
|
||||
let permit = acquire_breaker_permit(breaker, origin_key)?;
|
||||
let context = pump_log_context_from_connection(request.connection());
|
||||
let connected = connect_with_budget(
|
||||
request.connection(),
|
||||
PumpKeepalive::disabled(),
|
||||
fast_path_budget,
|
||||
context,
|
||||
)
|
||||
.await;
|
||||
let (connection, connect_elapsed) = finish_breaker_attempt(permit, connected)?;
|
||||
@@ -302,6 +304,10 @@ fn start_pooled_websocket_connect(
|
||||
let cancellation = connect_lease.cancellation_token();
|
||||
let fast_path_reporter = permit.fast_path_reporter();
|
||||
let keepalive = pool.keepalive();
|
||||
let context = PumpLogContext::new(
|
||||
Some(task_key.account_id().to_owned()),
|
||||
Some(task_key.conversation_id_hash()),
|
||||
);
|
||||
let (sender, receiver) = oneshot::channel();
|
||||
pool.spawn_connect_task(async move {
|
||||
let connected = tokio::select! {
|
||||
@@ -319,7 +325,7 @@ fn start_pooled_websocket_connect(
|
||||
);
|
||||
return;
|
||||
}
|
||||
result = connect_with_budget(&connection, keepalive, None) => result,
|
||||
result = connect_with_budget(&connection, keepalive, None, context) => result,
|
||||
};
|
||||
match finish_breaker_attempt(permit, connected) {
|
||||
Ok((connection, connect_elapsed)) => {
|
||||
@@ -392,17 +398,31 @@ fn duration_millis_u64(duration: Duration) -> u64 {
|
||||
.max(1)
|
||||
}
|
||||
|
||||
/// 非池化建连的 pump 日志上下文:从业务头提取账号归属(会话 hash 只有池 key 才有)。
|
||||
fn pump_log_context_from_connection(connection: &CodexWebSocketConnection) -> PumpLogContext {
|
||||
let account_id = connection
|
||||
.headers()
|
||||
.iter()
|
||||
.find(|(name, _)| name.eq_ignore_ascii_case("chatgpt-account-id"))
|
||||
.map(|(_, value)| value.clone());
|
||||
PumpLogContext::new(account_id, None)
|
||||
}
|
||||
|
||||
async fn connect_with_budget(
|
||||
connection: &CodexWebSocketConnection,
|
||||
keepalive: PumpKeepalive,
|
||||
fast_path_budget: Option<Duration>,
|
||||
context: PumpLogContext,
|
||||
) -> Result<(PooledWebSocketConnection, Duration), CodexWebSocketExchangeError> {
|
||||
let started_at = Instant::now();
|
||||
let connected = match fast_path_budget {
|
||||
Some(budget) => timeout(budget, connect_pumped_websocket(connection, keepalive))
|
||||
.await
|
||||
.map_err(|_| CodexWebSocketExchangeError::FastPathTimeout { timeout: budget })?,
|
||||
None => connect_pumped_websocket(connection, keepalive).await,
|
||||
Some(budget) => timeout(
|
||||
budget,
|
||||
connect_pumped_websocket(connection, keepalive, context),
|
||||
)
|
||||
.await
|
||||
.map_err(|_| CodexWebSocketExchangeError::FastPathTimeout { timeout: budget })?,
|
||||
None => connect_pumped_websocket(connection, keepalive, context).await,
|
||||
}?;
|
||||
let (websocket, response) = connected;
|
||||
Ok((
|
||||
|
||||
@@ -30,7 +30,7 @@ use super::{
|
||||
error::CodexWebSocketExchangeError,
|
||||
model::{CodexWebSocketConnection, CodexWebSocketRequest, WebSocketContinuationRequirement},
|
||||
pool::CodexWebSocketConnectionMetadata,
|
||||
pump::{PumpKeepalive, PumpedWebSocket, RawWsStream},
|
||||
pump::{PumpKeepalive, PumpLogContext, PumpedWebSocket, RawWsStream},
|
||||
};
|
||||
|
||||
const WEBSOCKET_EXTENSIONS: &str = "permessage-deflate; client_max_window_bits";
|
||||
@@ -93,9 +93,10 @@ pub fn responses_websocket_endpoint(base_url: &str) -> String {
|
||||
pub(super) async fn connect_pumped_websocket(
|
||||
connection: &CodexWebSocketConnection,
|
||||
keepalive: PumpKeepalive,
|
||||
context: PumpLogContext,
|
||||
) -> Result<(PumpedWebSocket, WsResponse<Option<Vec<u8>>>), CodexWebSocketExchangeError> {
|
||||
let (raw, response) = connect_websocket(connection).await?;
|
||||
Ok((PumpedWebSocket::new(raw, keepalive), response))
|
||||
Ok((PumpedWebSocket::new(raw, keepalive, context), response))
|
||||
}
|
||||
|
||||
pub(super) async fn send_websocket_request(
|
||||
|
||||
@@ -36,6 +36,22 @@ pub(crate) type RawWsStream = WebSocketStream<MaybeTlsStream<TcpStream>>;
|
||||
|
||||
const PUMP_COMMAND_BUFFER: usize = 32;
|
||||
|
||||
/// pump 日志的账号级关联上下文;空闲连接被上游重置等场景用于回溯归属。
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub(crate) struct PumpLogContext {
|
||||
pub(crate) account_id: Option<String>,
|
||||
pub(crate) conversation_id_hash: Option<String>,
|
||||
}
|
||||
|
||||
impl PumpLogContext {
|
||||
pub(crate) fn new(account_id: Option<String>, conversation_id_hash: Option<String>) -> Self {
|
||||
Self {
|
||||
account_id,
|
||||
conversation_id_hash,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// pump 保活策略。
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub(crate) struct PumpKeepalive {
|
||||
@@ -188,7 +204,11 @@ impl std::fmt::Debug for PumpedWebSocket {
|
||||
|
||||
impl PumpedWebSocket {
|
||||
/// 用底层流启动一个 pump 任务并返回句柄。
|
||||
pub(crate) fn new(inner: RawWsStream, keepalive: PumpKeepalive) -> Self {
|
||||
pub(crate) fn new(
|
||||
inner: RawWsStream,
|
||||
keepalive: PumpKeepalive,
|
||||
context: PumpLogContext,
|
||||
) -> Self {
|
||||
let (tx_command, rx_command) = mpsc::channel::<PumpCommand>(PUMP_COMMAND_BUFFER);
|
||||
let (tx_message, rx_message) = mpsc::channel(PUMP_MESSAGE_BUFFER);
|
||||
let connection_id = Uuid::new_v4();
|
||||
@@ -203,6 +223,7 @@ impl PumpedWebSocket {
|
||||
rx_command,
|
||||
tx_message,
|
||||
keepalive,
|
||||
context,
|
||||
closed_for_task,
|
||||
exit_reason_for_task,
|
||||
)
|
||||
@@ -273,12 +294,17 @@ impl Drop for PumpedWebSocket {
|
||||
}
|
||||
}
|
||||
|
||||
#[expect(
|
||||
clippy::too_many_arguments,
|
||||
reason = "pump owns one flat set of channel/state handles and log context"
|
||||
)]
|
||||
async fn pump_loop(
|
||||
connection_id: Uuid,
|
||||
mut inner: RawWsStream,
|
||||
mut rx_command: mpsc::Receiver<PumpCommand>,
|
||||
tx_message: mpsc::Sender<Result<Message, tungstenite::Error>>,
|
||||
keepalive: PumpKeepalive,
|
||||
context: PumpLogContext,
|
||||
closed: Arc<AtomicBool>,
|
||||
exit_reason: Arc<Mutex<Option<PumpExitReason>>>,
|
||||
) {
|
||||
@@ -429,7 +455,7 @@ async fn pump_loop(
|
||||
*exit_reason
|
||||
.lock()
|
||||
.expect("WebSocket pump exit reason lock poisoned") = Some(reason.clone());
|
||||
log_pump_exit(connection_id, &reason, backpressure_events);
|
||||
log_pump_exit(connection_id, &context, &reason, backpressure_events);
|
||||
|
||||
if reason.should_send_close() {
|
||||
let _ = inner.send(Message::Close(None)).await;
|
||||
@@ -510,11 +536,20 @@ async fn send_keepalive_ping(
|
||||
}))
|
||||
}
|
||||
|
||||
fn log_pump_exit(connection_id: Uuid, reason: &PumpExitReason, backpressure_events: u64) {
|
||||
fn log_pump_exit(
|
||||
connection_id: Uuid,
|
||||
context: &PumpLogContext,
|
||||
reason: &PumpExitReason,
|
||||
backpressure_events: u64,
|
||||
) {
|
||||
let detail = reason.detail().unwrap_or_default();
|
||||
let account_id = context.account_id.as_deref().unwrap_or("unknown");
|
||||
let conversation_id_hash = context.conversation_id_hash.as_deref().unwrap_or("unknown");
|
||||
if reason.is_unexpected() {
|
||||
tracing::warn!(
|
||||
websocket_connection_id = %connection_id,
|
||||
account_id = %account_id,
|
||||
conversation_id_hash = %conversation_id_hash,
|
||||
pump_exit_reason = reason.as_str(),
|
||||
pump_exit_detail = detail,
|
||||
backpressure_events,
|
||||
@@ -523,6 +558,8 @@ fn log_pump_exit(connection_id: Uuid, reason: &PumpExitReason, backpressure_even
|
||||
} else if matches!(reason, PumpExitReason::UpstreamCloseFrame { .. }) {
|
||||
tracing::info!(
|
||||
websocket_connection_id = %connection_id,
|
||||
account_id = %account_id,
|
||||
conversation_id_hash = %conversation_id_hash,
|
||||
pump_exit_reason = reason.as_str(),
|
||||
pump_exit_detail = detail,
|
||||
backpressure_events,
|
||||
@@ -531,6 +568,8 @@ fn log_pump_exit(connection_id: Uuid, reason: &PumpExitReason, backpressure_even
|
||||
} else {
|
||||
tracing::debug!(
|
||||
websocket_connection_id = %connection_id,
|
||||
account_id = %account_id,
|
||||
conversation_id_hash = %conversation_id_hash,
|
||||
pump_exit_reason = reason.as_str(),
|
||||
pump_exit_detail = detail,
|
||||
backpressure_events,
|
||||
|
||||
@@ -82,6 +82,22 @@ async fn openai_bundle_exposes_one_core_provider_and_drains_worker_contributions
|
||||
panic!("Desktop release worker must be scheduled");
|
||||
};
|
||||
assert_eq!(schedule.interval(), APPCAST_POLL_INTERVAL);
|
||||
let cli_release_worker = contributions
|
||||
.iter()
|
||||
.find_map(|contribution| match contribution {
|
||||
WorkerContribution::Registration(registration)
|
||||
if registration.id.owner() == "openai-cli-release" =>
|
||||
{
|
||||
Some(registration)
|
||||
}
|
||||
WorkerContribution::Registration(_) | WorkerContribution::Disabled { .. } => None,
|
||||
})
|
||||
.expect("CLI release worker");
|
||||
assert_eq!(cli_release_worker.id.kind(), WorkerKind::QuotaCatalogHealth);
|
||||
let WorkerRunnable::Scheduled { schedule, .. } = &cli_release_worker.runnable else {
|
||||
panic!("CLI release worker must be scheduled");
|
||||
};
|
||||
assert_eq!(schedule.interval(), APPCAST_POLL_INTERVAL);
|
||||
assert!(contributions.iter().any(|contribution| {
|
||||
matches!(
|
||||
contribution,
|
||||
@@ -101,13 +117,21 @@ async fn openai_admin_provider_exposes_live_wire_profile_and_validated_billing()
|
||||
.expect("OpenAI bundle");
|
||||
let admin = bundle.admin_provider();
|
||||
let profile = admin.dashboard_wire_profile().expect("wire profile");
|
||||
assert_eq!(profile.version, "0.102.0");
|
||||
assert_eq!(profile.build, None);
|
||||
assert_eq!(profile.target.os_type, "macOS");
|
||||
assert_eq!(profile.target.os_version, "15.5.0");
|
||||
assert_eq!(
|
||||
profile.user_agent,
|
||||
"Codex Desktop/0.102.0 (Mac OS 15.5.0; arm64) xterm-256color (Codex Desktop; 1.2026.190)"
|
||||
);
|
||||
assert_eq!(
|
||||
profile
|
||||
.attributes
|
||||
.iter()
|
||||
.find(|attribute| attribute.label == "Codex Core")
|
||||
.find(|attribute| attribute.label == "客户端标识")
|
||||
.map(|attribute| attribute.value.as_str()),
|
||||
Some("0.102.0")
|
||||
Some("Codex Desktop; 1.2026.190")
|
||||
);
|
||||
assert_eq!(
|
||||
profile.release.as_ref().map(|release| release.status),
|
||||
@@ -428,7 +452,8 @@ async fn openai_admin_provider_projects_cached_quota_models_and_canonical_export
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn openai_admin_provider_projects_codex_additional_limit_as_the_primary_quota() {
|
||||
async fn openai_admin_provider_projects_official_codex_quota_and_independent_buckets_with_chinese_labels()
|
||||
{
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
store
|
||||
.seed_oauth_credential(ImportCodexOAuthCredential {
|
||||
@@ -444,11 +469,15 @@ async fn openai_admin_provider_projects_codex_additional_limit_as_the_primary_qu
|
||||
.account("acct_admin_canonical_quota")
|
||||
.expect("stored account");
|
||||
let raw = json!({
|
||||
"active_limit": "premium",
|
||||
"rate_limit": {
|
||||
"primary_window": {
|
||||
"used_percent": 91,
|
||||
"reset_at": 1_900_000_000,
|
||||
"limit_window_seconds": 2_592_000
|
||||
},
|
||||
"secondary_window": {
|
||||
"used_percent": 88
|
||||
}
|
||||
},
|
||||
"additional_rate_limits": [{
|
||||
@@ -459,6 +488,19 @@ async fn openai_admin_provider_projects_codex_additional_limit_as_the_primary_qu
|
||||
"used_percent": 2,
|
||||
"reset_at": 1_900_000_000,
|
||||
"limit_window_seconds": 2_592_000
|
||||
},
|
||||
"secondary_window": {
|
||||
"used_percent": 0
|
||||
}
|
||||
}
|
||||
}, {
|
||||
"limit_name": "code_review",
|
||||
"metered_feature": "code_review",
|
||||
"rate_limit": {
|
||||
"primary_window": {
|
||||
"used_percent": 12,
|
||||
"reset_at": 1_900_000_000,
|
||||
"limit_window_seconds": 604_800
|
||||
}
|
||||
}
|
||||
}]
|
||||
@@ -471,6 +513,7 @@ async fn openai_admin_provider_projects_codex_additional_limit_as_the_primary_qu
|
||||
raw.as_object().expect("quota object").clone(),
|
||||
)),
|
||||
observed_at: Some(SystemTime::now()),
|
||||
limit_reached: None,
|
||||
})
|
||||
.await
|
||||
.expect("persist quota");
|
||||
@@ -498,9 +541,31 @@ async fn openai_admin_provider_projects_codex_additional_limit_as_the_primary_qu
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
assert_eq!(monthly.len(), 1);
|
||||
assert_eq!(monthly[0].label, "月限额");
|
||||
assert_eq!(monthly[0].source.as_deref(), Some("core"));
|
||||
assert_eq!(monthly[0].used_percent, Some(2.0));
|
||||
assert!(monthly.iter().any(|window| {
|
||||
window.label == "月额度"
|
||||
&& window.source.as_deref() == Some("codex")
|
||||
&& window.used_percent == Some(91.0)
|
||||
}));
|
||||
assert!(
|
||||
!quota
|
||||
.windows
|
||||
.iter()
|
||||
.any(|window| window.key.starts_with("additional-0-codex")),
|
||||
"the additional codex alias should not become a second display bucket"
|
||||
);
|
||||
let secondary = quota
|
||||
.windows
|
||||
.iter()
|
||||
.find(|window| window.key == "core-secondary")
|
||||
.expect("core secondary quota");
|
||||
assert_eq!(secondary.label, "次级额度");
|
||||
assert_eq!(secondary.used_percent, Some(88.0));
|
||||
let review = quota
|
||||
.windows
|
||||
.iter()
|
||||
.find(|window| window.source.as_deref() == Some("code_review"))
|
||||
.expect("code review quota");
|
||||
assert_eq!(review.label, "代码审查 · 周额度");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -536,16 +601,16 @@ async fn openai_admin_projects_confirmed_quota_exhaustion_as_full_without_mutati
|
||||
.clone(),
|
||||
)),
|
||||
observed_at: Some(SystemTime::now()),
|
||||
limit_reached: None,
|
||||
})
|
||||
.await
|
||||
.expect("persist raw quota");
|
||||
store
|
||||
.apply_state_change(AccountStateChange {
|
||||
message: None,
|
||||
account_id: account.id().clone(),
|
||||
expected_revision: account.revision(),
|
||||
availability: CoreAccountAvailability::QuotaExhausted,
|
||||
reason: Some("usage_limit_exhausted".to_owned()),
|
||||
cooldown_until: Some(SystemTime::UNIX_EPOCH + Duration::from_secs(reset_at)),
|
||||
observed_at: SystemTime::now(),
|
||||
})
|
||||
.await
|
||||
@@ -583,11 +648,10 @@ async fn openai_admin_projects_confirmed_quota_exhaustion_as_full_without_mutati
|
||||
|
||||
store
|
||||
.apply_state_change(AccountStateChange {
|
||||
message: None,
|
||||
account_id: account.id().clone(),
|
||||
expected_revision: account.revision(),
|
||||
availability: CoreAccountAvailability::Ready,
|
||||
reason: None,
|
||||
cooldown_until: None,
|
||||
observed_at: SystemTime::now(),
|
||||
})
|
||||
.await
|
||||
@@ -645,7 +709,6 @@ async fn openai_admin_provider_rejects_unprepared_mutations_before_store_commit(
|
||||
stale_record.next_refresh_at = None;
|
||||
stale_record.enabled = false;
|
||||
stale_record.availability = AdminAccountAvailability::Banned;
|
||||
stale_record.cooldown_until = Some(Utc::now() + chrono::Duration::minutes(5));
|
||||
let rotation_error = admin
|
||||
.prepare_rotation(PrepareCredentialRotation {
|
||||
account: stale_record,
|
||||
@@ -712,9 +775,8 @@ fn account_record(account: &ProviderAccount) -> AccountRecord {
|
||||
next_refresh_at: account.next_refresh_at().map(DateTime::<Utc>::from),
|
||||
enabled: account.enabled(),
|
||||
availability: AdminAccountAvailability::Ready,
|
||||
availability_reason: None,
|
||||
cooldown_until: account.cooldown_until().map(DateTime::<Utc>::from),
|
||||
availability_observed_at: now,
|
||||
last_error_message: None,
|
||||
quota_observed_at: None,
|
||||
created_at: now,
|
||||
updated_at: now,
|
||||
@@ -901,6 +963,13 @@ impl ProviderCooldownPort for TestCooldown {
|
||||
) -> BoxFuture<'a, Result<bool, ProviderStoreError>> {
|
||||
Box::pin(async { Ok(false) })
|
||||
}
|
||||
|
||||
fn clear_all<'a>(
|
||||
&'a self,
|
||||
_account_id: &'a ProviderAccountId,
|
||||
) -> BoxFuture<'a, Result<bool, ProviderStoreError>> {
|
||||
Box::pin(async { Ok(false) })
|
||||
}
|
||||
}
|
||||
|
||||
struct TestRuntimePolicy;
|
||||
|
||||
@@ -13,7 +13,7 @@ fn openai_config_builds_the_audited_wire_profile() {
|
||||
|
||||
assert_eq!(
|
||||
profile.user_agent(),
|
||||
"Codex Desktop/0.102.0 (macOS 15.5.0; arm64) xterm-256color"
|
||||
"Codex Desktop/0.102.0 (Mac OS 15.5.0; arm64) xterm-256color (Codex Desktop; 1.2026.190)"
|
||||
);
|
||||
assert_eq!(profile.desktop_build, "19012345678");
|
||||
}
|
||||
@@ -124,6 +124,10 @@ fn openai_config_defaults_to_the_provider_owned_operating_values() {
|
||||
"https://auth.openai.com/oauth/token",
|
||||
)
|
||||
);
|
||||
assert_eq!(
|
||||
config.wire_profile_state().snapshot().user_agent(),
|
||||
"Codex Desktop/0.146.0 (Mac OS 15.7.1; arm64) unknown (Codex Desktop; 26.727.51351)"
|
||||
);
|
||||
}
|
||||
|
||||
fn valid_config() -> OpenAiConfig {
|
||||
|
||||
@@ -194,7 +194,7 @@ fn agent_export_item(id: &str) -> (ExportManagedCodexCredential, String) {
|
||||
Some("agent-account-export".to_owned()),
|
||||
Some("pro".to_owned()),
|
||||
)
|
||||
.with_runtime_state(true, AccountAvailability::Ready, None)
|
||||
.with_runtime_state(true, AccountAvailability::Ready)
|
||||
.with_refresh_schedule(false, None);
|
||||
(
|
||||
ExportManagedCodexCredential {
|
||||
|
||||
@@ -66,7 +66,7 @@ fn agent_account(
|
||||
Some("agent-account".to_owned()),
|
||||
Some("pro".to_owned()),
|
||||
)
|
||||
.with_runtime_state(true, AccountAvailability::Ready, None)
|
||||
.with_runtime_state(true, AccountAvailability::Ready)
|
||||
.with_refresh_schedule(false, None);
|
||||
NewProviderAccount {
|
||||
account,
|
||||
|
||||
@@ -4,15 +4,21 @@ use std::sync::Arc;
|
||||
use std::time::{Duration, SystemTime};
|
||||
|
||||
use futures::executor::block_on;
|
||||
use gateway_core::engine::continuation::{
|
||||
ContinuationBinding, NativeContinuationPin, PreviousResponseId,
|
||||
};
|
||||
use gateway_core::engine::credential::{
|
||||
AccountAttemptFeedback, AccountAvailability, AccountFeedbackStats, AccountSelectionPolicy,
|
||||
AccountStateChange, ProviderAccountId, ProviderAccountStore as _, RotationStrategy,
|
||||
AccountStateChange, OpaqueProviderData, ProviderAccountId, ProviderAccountStore as _,
|
||||
QuotaObservation, QuotaWriteOutcome, RotationStrategy,
|
||||
};
|
||||
use gateway_core::engine::{
|
||||
AccountAttemptContext, AttemptContext, CancellationToken, ModelRequestId, RequestAttemptContext,
|
||||
};
|
||||
use gateway_core::policy::ClientApiKeyId;
|
||||
use gateway_core::provider_ports::{ProviderSessionAffinityKey, ProviderSessionAffinityPort};
|
||||
use gateway_core::provider_ports::{
|
||||
ProviderCooldownPort, ProviderSessionAffinityKey, ProviderSessionAffinityPort,
|
||||
};
|
||||
use gateway_core::routing::ProviderKind;
|
||||
use provider_openai::OFFICIAL_CODEX_BASE_URL;
|
||||
use provider_openai::credential::{
|
||||
@@ -22,11 +28,12 @@ use provider_openai::credential::{
|
||||
};
|
||||
use provider_openai::transport::profile::{CodexWireProfile, CodexWireProfileState};
|
||||
use secrecy::ExposeSecret;
|
||||
use serde_json::json;
|
||||
use url::Url;
|
||||
|
||||
use crate::support::{
|
||||
MemoryAccountStore, MemorySessionAffinity, MemorySessionExclusions, TestLeaseCoordinator,
|
||||
account_policy, agent_identity_service, catalog_cache, profile, secret,
|
||||
MemoryAccountStore, MemoryCooldownPort, MemorySessionAffinity, MemorySessionExclusions,
|
||||
TestLeaseCoordinator, account_policy, agent_identity_service, catalog_cache, profile, secret,
|
||||
};
|
||||
|
||||
fn create_account(store: &Arc<MemoryAccountStore>, id: &str, token: &str) {
|
||||
@@ -99,6 +106,7 @@ fn selector_with_affinity(
|
||||
session_affinity,
|
||||
Arc::new(AccountFeedbackStats::default()),
|
||||
true,
|
||||
Arc::new(MemoryCooldownPort::new()),
|
||||
)
|
||||
}
|
||||
|
||||
@@ -108,6 +116,7 @@ fn selector_with_runtime(
|
||||
session_affinity: Arc<MemorySessionAffinity>,
|
||||
account_feedback: Arc<AccountFeedbackStats>,
|
||||
skip_exhausted: bool,
|
||||
cooldowns: Arc<dyn ProviderCooldownPort>,
|
||||
) -> CodexCredentialSelector {
|
||||
let profile = CodexWireProfileState::new(CodexWireProfile {
|
||||
originator: "codex_cli_rs".to_owned(),
|
||||
@@ -136,6 +145,7 @@ fn selector_with_runtime(
|
||||
http,
|
||||
OFFICIAL_CODEX_BASE_URL.to_owned(),
|
||||
Arc::clone(&agent_identity),
|
||||
cooldowns,
|
||||
));
|
||||
CodexCredentialSelector::new(
|
||||
ProviderKind::new("openai").expect("provider"),
|
||||
@@ -464,6 +474,7 @@ async fn selector_should_keep_a_schedulable_affinity_account_despite_soft_health
|
||||
affinity,
|
||||
Arc::clone(&account_feedback),
|
||||
true,
|
||||
Arc::new(MemoryCooldownPort::new()),
|
||||
);
|
||||
for _ in 0..4 {
|
||||
account_feedback.report(
|
||||
@@ -507,11 +518,10 @@ async fn selector_should_escape_a_quota_exhausted_affinity_account() {
|
||||
let first = store.account("acct_first").expect("first account");
|
||||
store
|
||||
.apply_state_change(AccountStateChange {
|
||||
message: None,
|
||||
account_id: first.id().clone(),
|
||||
expected_revision: first.revision(),
|
||||
availability: AccountAvailability::QuotaExhausted,
|
||||
reason: Some("quota_exhausted".to_owned()),
|
||||
cooldown_until: None,
|
||||
observed_at: SystemTime::now(),
|
||||
})
|
||||
.await
|
||||
@@ -550,57 +560,6 @@ async fn selector_should_escape_a_quota_exhausted_affinity_account() {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn selector_should_escape_an_affinity_account_during_cooldown() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_first", "at-first");
|
||||
create_account(&store, "acct_second", "at-second");
|
||||
let first = store.account("acct_first").expect("first account");
|
||||
store
|
||||
.apply_state_change(AccountStateChange {
|
||||
account_id: first.id().clone(),
|
||||
expected_revision: first.revision(),
|
||||
availability: AccountAvailability::Cooldown,
|
||||
reason: Some("rate_limited".to_owned()),
|
||||
cooldown_until: Some(SystemTime::now() + Duration::from_secs(60)),
|
||||
observed_at: SystemTime::now(),
|
||||
})
|
||||
.await
|
||||
.expect("mark first account cooling down");
|
||||
let affinity = Arc::new(MemorySessionAffinity::default());
|
||||
let provider = ProviderKind::new("openai").expect("provider");
|
||||
let key = ProviderSessionAffinityKey::try_new("cooldown-session").expect("affinity key");
|
||||
affinity
|
||||
.bind(&provider, &key, first.id(), Duration::from_secs(60))
|
||||
.await
|
||||
.expect("seed affinity");
|
||||
let selector =
|
||||
selector_with_affinity(&store, Arc::new(TestLeaseCoordinator::default()), affinity);
|
||||
let request_url =
|
||||
Url::parse("https://chatgpt.com/backend-api/codex/responses").expect("request URL");
|
||||
let request_attempt = attempt(BTreeSet::new());
|
||||
|
||||
let selected = selector
|
||||
.select(&SelectCodexCredential {
|
||||
upstream_model: "gpt-5.4",
|
||||
request_url: &request_url,
|
||||
attempt: &request_attempt,
|
||||
session_affinity_key: Some(&key),
|
||||
})
|
||||
.await
|
||||
.expect("select fallback account");
|
||||
|
||||
assert_eq!(
|
||||
(
|
||||
selected.account_id().as_str(),
|
||||
selected.affinity_hit(),
|
||||
selected.escape_reason(),
|
||||
selected.account_switch(),
|
||||
),
|
||||
("acct_second", false, Some("cooldown"), true)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selector_honors_attempt_local_account_exclusion() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
@@ -716,8 +675,12 @@ fn credential_expired_failure_marks_unified_account_expired() {
|
||||
}),
|
||||
)
|
||||
.expect("select account");
|
||||
block_on(selector.record_failure(lease.account(), CodexAccountFailure::CredentialExpired))
|
||||
.expect("record credential expiry");
|
||||
block_on(selector.record_failure(
|
||||
lease.account(),
|
||||
CodexAccountFailure::CredentialExpired,
|
||||
None,
|
||||
))
|
||||
.expect("record credential expiry");
|
||||
assert_eq!(
|
||||
store
|
||||
.account("acct_primary")
|
||||
@@ -728,10 +691,18 @@ fn credential_expired_failure_marks_unified_account_expired() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rate_limited_failure_marks_quota_exhausted_with_a_reset_deadline() {
|
||||
fn rate_limited_failure_records_quota_window_without_changing_availability() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_primary", "at-primary");
|
||||
let selector = selector(&store, Arc::new(TestLeaseCoordinator::default()));
|
||||
let cooldowns = Arc::new(MemoryCooldownPort::new());
|
||||
let selector = selector_with_runtime(
|
||||
&store,
|
||||
Arc::new(TestLeaseCoordinator::default()),
|
||||
Arc::new(MemorySessionAffinity::default()),
|
||||
Arc::new(AccountFeedbackStats::default()),
|
||||
true,
|
||||
Arc::clone(&cooldowns) as Arc<dyn ProviderCooldownPort>,
|
||||
);
|
||||
let attempt = attempt(BTreeSet::new());
|
||||
let lease =
|
||||
block_on(
|
||||
@@ -750,15 +721,22 @@ fn rate_limited_failure_marks_quota_exhausted_with_a_reset_deadline() {
|
||||
CodexAccountFailure::RateLimited {
|
||||
retry_after: Some(Duration::from_secs(30)),
|
||||
},
|
||||
None,
|
||||
))
|
||||
.expect("record rate-limit failure");
|
||||
|
||||
// 429 不改变 availability(限流不改变账号可用性),也不污染 quota JSON——
|
||||
// 临时限流写入 Redis 冷却(到期自动解除),额度重置时间保持不变。
|
||||
let account = store.account("acct_primary").expect("account");
|
||||
assert_eq!(account.availability(), AccountAvailability::QuotaExhausted);
|
||||
assert_eq!(account.availability(), AccountAvailability::Ready);
|
||||
assert!(
|
||||
account
|
||||
.cooldown_until()
|
||||
.is_some_and(|until| until > SystemTime::now())
|
||||
store.quota_json("acct_primary").is_none(),
|
||||
"429 must not synthesize quota window"
|
||||
);
|
||||
let cooldown = block_on(cooldowns.read(account.id())).expect("read cooldown");
|
||||
assert!(
|
||||
cooldown.is_some_and(|cooling| cooling.until() > SystemTime::now()),
|
||||
"429 must record a Redis cooldown expiring in the future"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -775,17 +753,13 @@ async fn usage_limit_exhaustion_marks_quota_exhausted_without_usage_probe() {
|
||||
CodexAccountFailure::UsageLimitExhausted {
|
||||
retry_after: Some(Duration::from_secs(30)),
|
||||
},
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.expect("record usage-limit exhaustion");
|
||||
|
||||
let account = store.account("acct_primary").expect("persisted account");
|
||||
assert_eq!(account.availability(), AccountAvailability::QuotaExhausted);
|
||||
assert!(
|
||||
account
|
||||
.cooldown_until()
|
||||
.is_some_and(|until| until > SystemTime::now())
|
||||
);
|
||||
assert_eq!(store.quota_reads(), 0);
|
||||
}
|
||||
|
||||
@@ -795,11 +769,10 @@ fn rate_limited_failure_does_not_downgrade_persisted_quota_exhaustion() {
|
||||
create_account(&store, "acct_primary", "at-primary");
|
||||
let account = store.account("acct_primary").expect("account");
|
||||
block_on(store.apply_state_change(AccountStateChange {
|
||||
message: None,
|
||||
account_id: account.id().clone(),
|
||||
expected_revision: account.revision(),
|
||||
availability: AccountAvailability::QuotaExhausted,
|
||||
reason: Some("quota_exhausted".to_owned()),
|
||||
cooldown_until: None,
|
||||
observed_at: SystemTime::now(),
|
||||
}))
|
||||
.expect("mark exhausted");
|
||||
@@ -810,9 +783,11 @@ fn rate_limited_failure_does_not_downgrade_persisted_quota_exhaustion() {
|
||||
CodexAccountFailure::RateLimited {
|
||||
retry_after: Some(Duration::from_secs(30)),
|
||||
},
|
||||
None,
|
||||
))
|
||||
.expect("record rate-limit failure");
|
||||
|
||||
// 429 不写 availability,402 的 QuotaExhausted 保持不变。
|
||||
assert_eq!(
|
||||
store
|
||||
.account("acct_primary")
|
||||
@@ -820,47 +795,6 @@ fn rate_limited_failure_does_not_downgrade_persisted_quota_exhaustion() {
|
||||
.availability(),
|
||||
AccountAvailability::QuotaExhausted
|
||||
);
|
||||
assert!(
|
||||
store
|
||||
.account("acct_primary")
|
||||
.expect("persisted account")
|
||||
.cooldown_until()
|
||||
.is_some_and(|until| until > SystemTime::now())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rate_limited_failure_keeps_longer_existing_quota_deadline() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_primary", "at-primary");
|
||||
let account = store.account("acct_primary").expect("account");
|
||||
let existing_deadline = SystemTime::now() + Duration::from_secs(300);
|
||||
block_on(store.apply_state_change(AccountStateChange {
|
||||
account_id: account.id().clone(),
|
||||
expected_revision: account.revision(),
|
||||
availability: AccountAvailability::QuotaExhausted,
|
||||
reason: Some("usage_limit_exhausted".to_owned()),
|
||||
cooldown_until: Some(existing_deadline),
|
||||
observed_at: SystemTime::now(),
|
||||
}))
|
||||
.expect("mark exhausted");
|
||||
let selector = selector(&store, Arc::new(TestLeaseCoordinator::default()));
|
||||
|
||||
block_on(selector.record_failure(
|
||||
&account,
|
||||
CodexAccountFailure::RateLimited {
|
||||
retry_after: Some(Duration::from_secs(30)),
|
||||
},
|
||||
))
|
||||
.expect("record shorter rate-limit failure");
|
||||
|
||||
assert_eq!(
|
||||
store
|
||||
.account("acct_primary")
|
||||
.expect("persisted account")
|
||||
.cooldown_until(),
|
||||
Some(existing_deadline)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -868,13 +802,22 @@ fn rate_limited_failure_does_not_consult_stale_quota_snapshot() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_primary", "at-primary");
|
||||
let account = store.account("acct_primary").expect("account");
|
||||
let selector = selector(&store, Arc::new(TestLeaseCoordinator::default()));
|
||||
let cooldowns = Arc::new(MemoryCooldownPort::new());
|
||||
let selector = selector_with_runtime(
|
||||
&store,
|
||||
Arc::new(TestLeaseCoordinator::default()),
|
||||
Arc::new(MemorySessionAffinity::default()),
|
||||
Arc::new(AccountFeedbackStats::default()),
|
||||
true,
|
||||
Arc::clone(&cooldowns) as Arc<dyn ProviderCooldownPort>,
|
||||
);
|
||||
|
||||
block_on(selector.record_failure(
|
||||
&account,
|
||||
CodexAccountFailure::RateLimited {
|
||||
retry_after: Some(Duration::from_secs(30)),
|
||||
},
|
||||
None,
|
||||
))
|
||||
.expect("record rate-limit failure");
|
||||
|
||||
@@ -883,22 +826,27 @@ fn rate_limited_failure_does_not_consult_stale_quota_snapshot() {
|
||||
.account("acct_primary")
|
||||
.expect("persisted account")
|
||||
.availability(),
|
||||
AccountAvailability::QuotaExhausted
|
||||
AccountAvailability::Ready
|
||||
);
|
||||
// 429 临时限流写入 Redis 冷却,不读写 quota JSON。
|
||||
assert_eq!(store.quota_reads(), 0);
|
||||
let cooldown = block_on(cooldowns.read(account.id())).expect("read cooldown");
|
||||
assert!(
|
||||
cooldown.is_some_and(|cooling| cooling.until() > SystemTime::now()),
|
||||
"429 must record a Redis cooldown"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn authenticated_rate_limit_recovers_stale_authentication_state() {
|
||||
fn rate_limited_failure_does_not_overwrite_stale_authentication_state() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_primary", "at-primary");
|
||||
let account = store.account("acct_primary").expect("account");
|
||||
block_on(store.apply_state_change(AccountStateChange {
|
||||
message: None,
|
||||
account_id: account.id().clone(),
|
||||
expected_revision: account.revision(),
|
||||
availability: AccountAvailability::Invalid,
|
||||
reason: Some("identity_verification_required".to_owned()),
|
||||
cooldown_until: None,
|
||||
observed_at: SystemTime::now(),
|
||||
}))
|
||||
.expect("mark account invalid");
|
||||
@@ -910,15 +858,17 @@ fn authenticated_rate_limit_recovers_stale_authentication_state() {
|
||||
CodexAccountFailure::RateLimited {
|
||||
retry_after: Some(Duration::from_secs(30)),
|
||||
},
|
||||
None,
|
||||
))
|
||||
.expect("record rate-limit failure");
|
||||
|
||||
// 429 不写 availability(限流不改变账号可用性),认证终态保持不变。
|
||||
assert_eq!(
|
||||
store
|
||||
.account("acct_primary")
|
||||
.expect("persisted account")
|
||||
.availability(),
|
||||
AccountAvailability::QuotaExhausted
|
||||
AccountAvailability::Invalid
|
||||
);
|
||||
}
|
||||
|
||||
@@ -933,11 +883,10 @@ fn successful_upstream_response_recovers_non_quota_terminal_states() {
|
||||
create_account(&store, "acct_primary", "at-primary");
|
||||
let account = store.account("acct_primary").expect("account");
|
||||
block_on(store.apply_state_change(AccountStateChange {
|
||||
message: None,
|
||||
account_id: account.id().clone(),
|
||||
expected_revision: account.revision(),
|
||||
availability: stale,
|
||||
reason: Some("stale_observation".to_owned()),
|
||||
cooldown_until: None,
|
||||
observed_at: SystemTime::now(),
|
||||
}))
|
||||
.expect("seed stale state");
|
||||
@@ -973,12 +922,14 @@ fn rate_limited_failures_for_distinct_accounts_do_not_conflict() {
|
||||
CodexAccountFailure::RateLimited {
|
||||
retry_after: Some(Duration::from_secs(30)),
|
||||
},
|
||||
None,
|
||||
),
|
||||
selector.record_failure(
|
||||
&second,
|
||||
CodexAccountFailure::RateLimited {
|
||||
retry_after: Some(Duration::from_secs(30)),
|
||||
},
|
||||
None,
|
||||
),
|
||||
);
|
||||
first_result.expect("record first account failure");
|
||||
@@ -996,7 +947,7 @@ fn rate_limited_failures_for_distinct_accounts_do_not_conflict() {
|
||||
.expect("persisted second account")
|
||||
.availability(),
|
||||
],
|
||||
[AccountAvailability::QuotaExhausted; 2]
|
||||
[AccountAvailability::Ready; 2]
|
||||
);
|
||||
}
|
||||
|
||||
@@ -1006,7 +957,7 @@ fn selector_keeps_a_quota_exhausted_account_eligible_when_configured_not_to_skip
|
||||
create_account(&store, "acct_primary", "at-primary");
|
||||
let strict_selector = selector(&store, Arc::new(TestLeaseCoordinator::default()));
|
||||
let account = store.account("acct_primary").expect("account");
|
||||
block_on(strict_selector.record_failure(&account, CodexAccountFailure::QuotaExhausted))
|
||||
block_on(strict_selector.record_failure(&account, CodexAccountFailure::QuotaExhausted, None))
|
||||
.expect("mark account exhausted");
|
||||
|
||||
let selector = selector_with_runtime(
|
||||
@@ -1015,6 +966,7 @@ fn selector_keeps_a_quota_exhausted_account_eligible_when_configured_not_to_skip
|
||||
Arc::new(MemorySessionAffinity::default()),
|
||||
Arc::new(AccountFeedbackStats::default()),
|
||||
false,
|
||||
Arc::new(MemoryCooldownPort::new()),
|
||||
);
|
||||
let request_url =
|
||||
Url::parse("https://chatgpt.com/backend-api/codex/responses").expect("request URL");
|
||||
@@ -1038,6 +990,131 @@ fn selector_keeps_a_quota_exhausted_account_eligible_when_configured_not_to_skip
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn native_continuation_surfaces_the_original_accounts_quota_status_to_the_coordinator() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_original", "at-original");
|
||||
create_account(&store, "acct_fallback", "at-fallback");
|
||||
let strict_selector = selector(&store, Arc::new(TestLeaseCoordinator::default()));
|
||||
let original = store.account("acct_original").expect("original account");
|
||||
block_on(strict_selector.record_failure(&original, CodexAccountFailure::QuotaExhausted, None))
|
||||
.expect("mark original account exhausted");
|
||||
|
||||
let selector = selector_with_runtime(
|
||||
&store,
|
||||
Arc::new(TestLeaseCoordinator::default()),
|
||||
Arc::new(MemorySessionAffinity::default()),
|
||||
Arc::new(AccountFeedbackStats::default()),
|
||||
false,
|
||||
Arc::new(MemoryCooldownPort::new()),
|
||||
);
|
||||
let continuation = NativeContinuationPin::new(
|
||||
PreviousResponseId::new("previous-response"),
|
||||
PreviousResponseId::new("upstream-response"),
|
||||
ProviderKind::new("openai").expect("provider"),
|
||||
original.id().clone(),
|
||||
);
|
||||
let attempt = AttemptContext::new(
|
||||
RequestAttemptContext::new(
|
||||
ModelRequestId::new("req_native_continuation").expect("request id"),
|
||||
ClientApiKeyId::new("key_codex_contract").expect("client key id"),
|
||||
),
|
||||
NonZeroU32::new(1).expect("attempt"),
|
||||
SystemTime::now() + Duration::from_secs(30),
|
||||
account_policy(),
|
||||
AccountAttemptContext::new(BTreeSet::new(), None, None),
|
||||
Some(ContinuationBinding::Pinned(continuation)),
|
||||
CancellationToken::new(),
|
||||
);
|
||||
let error =
|
||||
block_on(
|
||||
selector.select(&SelectCodexCredential {
|
||||
upstream_model: "gpt-5.4",
|
||||
request_url: &Url::parse("https://chatgpt.com/backend-api/codex/responses")
|
||||
.expect("request URL"),
|
||||
attempt: &attempt,
|
||||
session_affinity_key: None,
|
||||
}),
|
||||
)
|
||||
.expect_err("the selector must surface the unavailable native account");
|
||||
|
||||
assert!(matches!(
|
||||
error,
|
||||
CredentialSelectionError::NoEligibleCredential
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn native_continuation_surfaces_the_original_accounts_quota_signal_to_the_coordinator() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_original_signal", "at-original-signal");
|
||||
create_account(&store, "acct_fallback_signal", "at-fallback-signal");
|
||||
let original = store
|
||||
.account("acct_original_signal")
|
||||
.expect("original account");
|
||||
let quota = json!({
|
||||
"rate_limit": {
|
||||
"allowed": false,
|
||||
"limit_reached": true,
|
||||
"primary_window": {"used_percent": 98}
|
||||
}
|
||||
});
|
||||
let outcome = block_on(store.compare_and_swap_quota(QuotaObservation {
|
||||
account_id: original.id().clone(),
|
||||
expected_revision: original.revision(),
|
||||
quota: Some(OpaqueProviderData::new(
|
||||
quota.as_object().expect("quota object").clone(),
|
||||
)),
|
||||
observed_at: Some(SystemTime::now()),
|
||||
limit_reached: None,
|
||||
}))
|
||||
.expect("persist quota signal");
|
||||
assert!(matches!(outcome, QuotaWriteOutcome::Updated));
|
||||
|
||||
let selector = selector_with_runtime(
|
||||
&store,
|
||||
Arc::new(TestLeaseCoordinator::default()),
|
||||
Arc::new(MemorySessionAffinity::default()),
|
||||
Arc::new(AccountFeedbackStats::default()),
|
||||
false,
|
||||
Arc::new(MemoryCooldownPort::new()),
|
||||
);
|
||||
let continuation = NativeContinuationPin::new(
|
||||
PreviousResponseId::new("previous-response"),
|
||||
PreviousResponseId::new("upstream-response"),
|
||||
ProviderKind::new("openai").expect("provider"),
|
||||
original.id().clone(),
|
||||
);
|
||||
let attempt = AttemptContext::new(
|
||||
RequestAttemptContext::new(
|
||||
ModelRequestId::new("req_native_continuation_signal").expect("request id"),
|
||||
ClientApiKeyId::new("key_codex_contract").expect("client key id"),
|
||||
),
|
||||
NonZeroU32::new(1).expect("attempt"),
|
||||
SystemTime::now() + Duration::from_secs(30),
|
||||
account_policy(),
|
||||
AccountAttemptContext::new(BTreeSet::new(), None, None),
|
||||
Some(ContinuationBinding::Pinned(continuation)),
|
||||
CancellationToken::new(),
|
||||
);
|
||||
let error =
|
||||
block_on(
|
||||
selector.select(&SelectCodexCredential {
|
||||
upstream_model: "gpt-5.4",
|
||||
request_url: &Url::parse("https://chatgpt.com/backend-api/codex/responses")
|
||||
.expect("request URL"),
|
||||
attempt: &attempt,
|
||||
session_affinity_key: None,
|
||||
}),
|
||||
)
|
||||
.expect_err("the selector must surface the quota-limited native account");
|
||||
|
||||
assert!(matches!(
|
||||
error,
|
||||
CredentialSelectionError::NoEligibleCredential
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn identity_verification_failure_isolates_only_selected_account() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
@@ -1060,6 +1137,7 @@ fn identity_verification_failure_isolates_only_selected_account() {
|
||||
block_on(selector.record_failure(
|
||||
lease.account(),
|
||||
CodexAccountFailure::IdentityVerificationRequired,
|
||||
None,
|
||||
))
|
||||
.expect("record identity verification failure");
|
||||
|
||||
@@ -1082,7 +1160,7 @@ fn identity_verification_failure_isolates_only_selected_account() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cloudflare_challenge_backoff_escalates_and_success_resets_it() {
|
||||
fn cloudflare_challenge_does_not_change_availability() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_primary", "at-primary");
|
||||
let selector = selector(&store, Arc::new(TestLeaseCoordinator::default()));
|
||||
@@ -1099,48 +1177,22 @@ fn cloudflare_challenge_backoff_escalates_and_success_resets_it() {
|
||||
)
|
||||
.expect("select account");
|
||||
|
||||
// Cloudflare 挑战只走内存退避表,不写 availability(限流不改变账号可用性)。
|
||||
block_on(selector.record_failure(
|
||||
lease.account(),
|
||||
CodexAccountFailure::CloudflareChallenge { retry_after: None },
|
||||
None,
|
||||
))
|
||||
.expect("record first challenge");
|
||||
let first = store
|
||||
.account("acct_primary")
|
||||
.expect("account")
|
||||
.cooldown_until()
|
||||
.expect("first cooldown")
|
||||
.duration_since(SystemTime::now())
|
||||
.expect("future cooldown");
|
||||
assert!(first >= Duration::from_secs(9) && first <= Duration::from_secs(10));
|
||||
|
||||
block_on(selector.record_failure(
|
||||
lease.account(),
|
||||
CodexAccountFailure::CloudflareChallenge { retry_after: None },
|
||||
))
|
||||
.expect("record second challenge");
|
||||
let second = store
|
||||
.account("acct_primary")
|
||||
.expect("account")
|
||||
.cooldown_until()
|
||||
.expect("second cooldown")
|
||||
.duration_since(SystemTime::now())
|
||||
.expect("future cooldown");
|
||||
assert!(second >= Duration::from_secs(29) && second <= Duration::from_secs(30));
|
||||
.expect("record challenge");
|
||||
|
||||
assert_eq!(
|
||||
store
|
||||
.account("acct_primary")
|
||||
.expect("account")
|
||||
.availability(),
|
||||
AccountAvailability::Ready
|
||||
);
|
||||
block_on(selector.record_success(lease.account(), None));
|
||||
block_on(selector.record_failure(
|
||||
lease.account(),
|
||||
CodexAccountFailure::CloudflareChallenge { retry_after: None },
|
||||
))
|
||||
.expect("record reset challenge");
|
||||
let reset = store
|
||||
.account("acct_primary")
|
||||
.expect("account")
|
||||
.cooldown_until()
|
||||
.expect("reset cooldown")
|
||||
.duration_since(SystemTime::now())
|
||||
.expect("future cooldown");
|
||||
assert!(reset >= Duration::from_secs(9) && reset <= Duration::from_secs(10));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -1166,9 +1218,11 @@ fn repeated_cloudflare_path_block_marks_only_the_affected_account_invalid() {
|
||||
.expect("select account");
|
||||
|
||||
for _ in 0..3 {
|
||||
block_on(
|
||||
selector.record_failure(lease.account(), CodexAccountFailure::CloudflarePathBlocked),
|
||||
)
|
||||
block_on(selector.record_failure(
|
||||
lease.account(),
|
||||
CodexAccountFailure::CloudflarePathBlocked,
|
||||
None,
|
||||
))
|
||||
.expect("record path block");
|
||||
}
|
||||
|
||||
@@ -1222,18 +1276,17 @@ fn cloudflare_challenge_expires_provider_owned_cookies_at_cooldown_boundary() {
|
||||
block_on(selector.record_failure(
|
||||
second.account(),
|
||||
CodexAccountFailure::CloudflareChallenge { retry_after: None },
|
||||
None,
|
||||
))
|
||||
.expect("record challenge");
|
||||
|
||||
let account = store.account("acct_primary").expect("account");
|
||||
let cooldown_until = account.cooldown_until().expect("cooldown");
|
||||
let data = block_on(store.repository().load_complete_data(&account)).expect("credential data");
|
||||
assert_eq!(data.cookies().len(), 1);
|
||||
assert!(
|
||||
data.cookies()[0]
|
||||
.expires_at
|
||||
.is_some_and(|expires_at| SystemTime::from(expires_at) <= cooldown_until)
|
||||
);
|
||||
assert!(data.cookies()[0].expires_at.is_some_and(|expires_at| {
|
||||
let expires_at = SystemTime::from(expires_at);
|
||||
expires_at > SystemTime::now() && expires_at <= SystemTime::now() + Duration::from_secs(120)
|
||||
}));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -1267,8 +1320,12 @@ fn cloudflare_path_block_deletes_provider_owned_cookies() {
|
||||
session_affinity_key: None,
|
||||
}))
|
||||
.expect("select revised account");
|
||||
block_on(selector.record_failure(second.account(), CodexAccountFailure::CloudflarePathBlocked))
|
||||
.expect("record path block");
|
||||
block_on(selector.record_failure(
|
||||
second.account(),
|
||||
CodexAccountFailure::CloudflarePathBlocked,
|
||||
None,
|
||||
))
|
||||
.expect("record path block");
|
||||
|
||||
let account = store.account("acct_primary").expect("account");
|
||||
let data = block_on(store.repository().load_complete_data(&account)).expect("credential data");
|
||||
@@ -1304,7 +1361,7 @@ fn response_cookie_rotation_returns_a_current_account_for_later_fenced_writes()
|
||||
|
||||
assert_eq!(outcome.credential_revision, Some(current.revision().get()));
|
||||
assert_ne!(current.revision(), lease.account().revision());
|
||||
block_on(selector.record_failure(¤t, CodexAccountFailure::QuotaExhausted))
|
||||
block_on(selector.record_failure(¤t, CodexAccountFailure::QuotaExhausted, None))
|
||||
.expect("record failure with current revision");
|
||||
assert_eq!(
|
||||
store
|
||||
|
||||
+472
-450
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,213 @@
|
||||
//! 402 恢复回归。
|
||||
//!
|
||||
//! 覆盖 `QuotaExhausted -> Ready` 的证据规则:旧 reset 已到期 + 新快照未耗尽。
|
||||
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn manual_quota_refresh_keeps_confirmed_exhaustion_until_the_recorded_reset() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
let account_id = "acct_manual_quota_state";
|
||||
create_account(&store, account_id).await;
|
||||
let account = store.account(account_id).expect("created account");
|
||||
store
|
||||
.apply_state_change(AccountStateChange {
|
||||
message: None,
|
||||
account_id: account.id().clone(),
|
||||
expected_revision: account.revision(),
|
||||
availability: AccountAvailability::Expired,
|
||||
observed_at: SystemTime::now(),
|
||||
})
|
||||
.await
|
||||
.expect("seed stale expired state");
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/api/codex/usage"))
|
||||
.and(header(
|
||||
"authorization",
|
||||
format!("Bearer token-{account_id}"),
|
||||
))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
|
||||
"rate_limit": {
|
||||
"allowed": false,
|
||||
"limit_reached": true,
|
||||
"primary_window": {"used_percent": 100, "reset_at": 1_900_000_000}
|
||||
}
|
||||
})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let service = quota_service_with_base_url(
|
||||
&store,
|
||||
reqwest::Client::builder()
|
||||
.no_proxy()
|
||||
.build()
|
||||
.expect("client"),
|
||||
server.uri(),
|
||||
);
|
||||
|
||||
let exhausted = service
|
||||
.refresh_account(account.id())
|
||||
.await
|
||||
.expect("refresh exhausted quota");
|
||||
|
||||
assert!(exhausted.fact().exhausted());
|
||||
assert_eq!(
|
||||
store
|
||||
.account(account_id)
|
||||
.expect("exhausted account")
|
||||
.availability(),
|
||||
AccountAvailability::Expired,
|
||||
"manual refresh never changes non-QuotaExhausted availability"
|
||||
);
|
||||
|
||||
server.reset().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/api/codex/usage"))
|
||||
.and(header(
|
||||
"authorization",
|
||||
format!("Bearer token-{account_id}"),
|
||||
))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
|
||||
"rate_limit": {
|
||||
"allowed": true,
|
||||
"limit_reached": false,
|
||||
"primary_window": {"used_percent": 100, "reset_at": 1_900_000_000}
|
||||
}
|
||||
})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let stale = service
|
||||
.refresh_account(account.id())
|
||||
.await
|
||||
.expect("refresh contradictory full quota");
|
||||
|
||||
// used=100 是快照级 exhaustion。
|
||||
assert!(stale.fact().exhausted());
|
||||
assert_eq!(
|
||||
store
|
||||
.account(account_id)
|
||||
.expect("account after contradictory quota")
|
||||
.availability(),
|
||||
AccountAvailability::Expired,
|
||||
"an unchanged usage window cannot override an expired credential state"
|
||||
);
|
||||
|
||||
server.reset().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/api/codex/usage"))
|
||||
.and(header(
|
||||
"authorization",
|
||||
format!("Bearer token-{account_id}"),
|
||||
))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
|
||||
"rate_limit": {
|
||||
"allowed": true,
|
||||
"limit_reached": false,
|
||||
"primary_window": {"used_percent": 99, "reset_at": 1_900_000_000}
|
||||
}
|
||||
})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let unchanged_reset = service
|
||||
.refresh_account(account.id())
|
||||
.await
|
||||
.expect("refresh lower usage quota");
|
||||
|
||||
assert!(!unchanged_reset.fact().exhausted());
|
||||
assert_eq!(
|
||||
store
|
||||
.account(account_id)
|
||||
.expect("account after lower usage")
|
||||
.availability(),
|
||||
AccountAvailability::Expired,
|
||||
"a lower used_percent before the recorded reset is not quota recovery for an expired credential"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn expired_old_reset_with_fresh_allowed_snapshot_recovers_ready() {
|
||||
// 旧快照 reset_at < now,新快照 allowed=true、
|
||||
// limit_reached=false、reset_at > now;最终必须恢复 Ready。
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
let account_id = "acct_cross_window_recovery";
|
||||
create_account(&store, account_id).await;
|
||||
let account = store.account(account_id).expect("created account");
|
||||
let old_reset = Utc::now().timestamp() - 60;
|
||||
store
|
||||
.compare_and_swap_quota(QuotaObservation {
|
||||
account_id: account.id().clone(),
|
||||
expected_revision: account.revision(),
|
||||
quota: Some(OpaqueProviderData::new(
|
||||
json!({
|
||||
"rate_limit": {
|
||||
"allowed": false,
|
||||
"limit_reached": true,
|
||||
"primary_window": {
|
||||
"used_percent": 100,
|
||||
"reset_at": old_reset,
|
||||
"limit_window_seconds": 3600
|
||||
}
|
||||
}
|
||||
})
|
||||
.as_object()
|
||||
.expect("quota object")
|
||||
.clone(),
|
||||
)),
|
||||
observed_at: Some(SystemTime::now()),
|
||||
limit_reached: None,
|
||||
})
|
||||
.await
|
||||
.expect("seed old exhausted snapshot");
|
||||
store
|
||||
.apply_state_change(AccountStateChange {
|
||||
message: None,
|
||||
account_id: account.id().clone(),
|
||||
expected_revision: account.revision(),
|
||||
availability: AccountAvailability::QuotaExhausted,
|
||||
observed_at: SystemTime::now(),
|
||||
})
|
||||
.await
|
||||
.expect("mark account exhausted");
|
||||
let server = MockServer::start().await;
|
||||
let future_reset = Utc::now().timestamp() + 3600;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/api/codex/usage"))
|
||||
.and(header(
|
||||
"authorization",
|
||||
format!("Bearer token-{account_id}"),
|
||||
))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
|
||||
"rate_limit": {
|
||||
"allowed": true,
|
||||
"limit_reached": false,
|
||||
"primary_window": {"used_percent": 5, "reset_at": future_reset}
|
||||
}
|
||||
})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let service = quota_service_with_base_url(
|
||||
&store,
|
||||
reqwest::Client::builder()
|
||||
.no_proxy()
|
||||
.build()
|
||||
.expect("client"),
|
||||
server.uri(),
|
||||
);
|
||||
|
||||
let recovered = service
|
||||
.refresh_account(account.id())
|
||||
.await
|
||||
.expect("refresh new window quota");
|
||||
|
||||
assert!(!recovered.fact().exhausted());
|
||||
assert_eq!(
|
||||
store
|
||||
.account(account_id)
|
||||
.expect("account after cross-window recovery")
|
||||
.availability(),
|
||||
AccountAvailability::Ready,
|
||||
"old reset elapsed + fresh allowed snapshot must recover QuotaExhausted"
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,442 @@
|
||||
//! 快照聚合、窗口滚动和调度信号回归。
|
||||
//!
|
||||
//! 覆盖 raw JSON 解析、`limit_reached` 快照级聚合与窗口投影。
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, SystemTime};
|
||||
|
||||
use futures::future::join_all;
|
||||
use gateway_core::engine::credential::{
|
||||
OpaqueProviderData, ProviderAccountStore as _, QuotaObservation, QuotaWriteOutcome,
|
||||
};
|
||||
use provider_openai::credential::{CodexQuotaWindowKind, parse_codex_quota_usage};
|
||||
use serde_json::json;
|
||||
|
||||
use super::{MemoryAccountStore, create_account, quota_service};
|
||||
|
||||
#[test]
|
||||
fn parser_extracts_dynamic_windows_without_a_fixed_database_shape() {
|
||||
let fact = parse_codex_quota_usage(&json!({
|
||||
"rate_limit": {
|
||||
"allowed": true,
|
||||
"primary_window": {"used_percent": 25.2, "reset_at": 1_800_000_100},
|
||||
"secondary_window": {"used_percent": 80.4, "reset_at": 1_800_000_200}
|
||||
},
|
||||
"additional_rate_limits": [{
|
||||
"limit_name": "future_dynamic_window",
|
||||
"rate_limit": {"primary_window": {"used_percent": 10}}
|
||||
}]
|
||||
}))
|
||||
.expect("valid dynamic quota");
|
||||
|
||||
assert_eq!(fact.remaining_percent(), Some(20));
|
||||
assert_eq!(
|
||||
fact.resets_at().map(|value| value.timestamp()),
|
||||
Some(1_800_000_100)
|
||||
);
|
||||
assert!(!fact.exhausted());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parser_treats_any_confirmed_provider_limit_as_exhausted() {
|
||||
let fact = parse_codex_quota_usage(&json!({
|
||||
"rate_limit": {"primary_window": {"used_percent": 10}},
|
||||
"additional_rate_limits": [{
|
||||
"rate_limit": {"allowed": false, "primary_window": {"used_percent": 100}}
|
||||
}]
|
||||
}))
|
||||
.expect("valid quota");
|
||||
|
||||
assert!(fact.exhausted());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parser_keeps_full_percent_as_display_without_provider_exhaustion_signal() {
|
||||
let fact = parse_codex_quota_usage(&json!({
|
||||
"rate_limit": {
|
||||
"allowed": true,
|
||||
"limit_reached": false,
|
||||
"primary_window": {"used_percent": 100, "reset_at": 1_900_000_000}
|
||||
}
|
||||
}))
|
||||
.expect("valid full-percent quota");
|
||||
|
||||
// 任一子窗口 used>=100 是快照级 exhaustion 事实,
|
||||
// 即使顶层 allowed=true。
|
||||
assert_eq!(fact.remaining_percent(), Some(0));
|
||||
assert!(fact.exhausted());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parser_treats_confirmed_limit_reached_as_exhausted_even_below_full_percent() {
|
||||
let fact = parse_codex_quota_usage(&json!({
|
||||
"rate_limit": {
|
||||
"allowed": false,
|
||||
"limit_reached": true,
|
||||
"primary_window": {"used_percent": 98, "reset_at": 1_900_000_000}
|
||||
}
|
||||
}))
|
||||
.expect("valid provider-confirmed quota");
|
||||
|
||||
assert_eq!(fact.remaining_percent(), Some(2));
|
||||
assert!(fact.exhausted());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parser_does_not_infer_exhaustion_from_unknown_credit_fields() {
|
||||
let fact = parse_codex_quota_usage(&json!({
|
||||
"credits": {
|
||||
"has_credits": false,
|
||||
"balance": 0,
|
||||
"overage_limit_reached": false,
|
||||
"future_provider_field": {"anything": true}
|
||||
}
|
||||
}))
|
||||
.expect("recognized credits object");
|
||||
|
||||
assert!(!fact.exhausted());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parser_accepts_official_null_additional_rate_limits() {
|
||||
let fact = parse_codex_quota_usage(&json!({
|
||||
"rate_limit": {
|
||||
"allowed": true,
|
||||
"limit_reached": false,
|
||||
"primary_window": {
|
||||
"limit_window_seconds": 18_000,
|
||||
"used_percent": 12.5,
|
||||
"reset_at": 1_800_000_100
|
||||
},
|
||||
"secondary_window": {
|
||||
"limit_window_seconds": 604_800,
|
||||
"used_percent": 30.0,
|
||||
"reset_at": 1_800_000_200
|
||||
}
|
||||
},
|
||||
"additional_rate_limits": null,
|
||||
"credits": {
|
||||
"has_credits": false,
|
||||
"balance": null,
|
||||
"overage_limit_reached": false,
|
||||
"unlimited": false
|
||||
},
|
||||
"spend_control": {
|
||||
"individual_limit": null,
|
||||
"reached": false
|
||||
}
|
||||
}))
|
||||
.expect("official null additional quota");
|
||||
|
||||
assert_eq!(fact.remaining_percent(), Some(70));
|
||||
assert!(!fact.exhausted());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parser_rejects_wrong_known_field_type_without_echoing_body() {
|
||||
let marker = "quota-secret-marker";
|
||||
let error = parse_codex_quota_usage(&json!({
|
||||
"rate_limit": {"allowed": marker}
|
||||
}))
|
||||
.expect_err("known field type must be strict");
|
||||
|
||||
assert!(!format!("{error:?} {error}").contains(marker));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parser_rejects_unrecognized_top_level_object() {
|
||||
assert!(parse_codex_quota_usage(&json!({"future_only": {"used": 1}})).is_err());
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn concurrent_cold_scheduling_hydration_reads_quota_once() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_hydration").await;
|
||||
let account = store.account("acct_hydration").expect("created account");
|
||||
let service = quota_service(&store);
|
||||
|
||||
join_all((0..32).map(|_| service.prepare_scheduling(std::slice::from_ref(&account)))).await;
|
||||
|
||||
assert_eq!(store.quota_reads(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn persisted_provider_quota_projects_dynamic_windows_without_network_io() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_quota").await;
|
||||
let account = store.account("acct_quota").expect("created account");
|
||||
let raw = json!({
|
||||
"rate_limit": {
|
||||
"allowed": true,
|
||||
"primary_window": {
|
||||
"used_percent": 37,
|
||||
"reset_at": 1_900_000_000,
|
||||
"limit_window_seconds": 18_000
|
||||
},
|
||||
"secondary_window": {
|
||||
"used_percent": 42,
|
||||
"reset_at": 1_900_604_800,
|
||||
"limit_window_seconds": 604_800
|
||||
}
|
||||
},
|
||||
"additional_rate_limits": null,
|
||||
"spend_control": {
|
||||
"reached": false,
|
||||
"individual_limit": {
|
||||
"used_percent": 12,
|
||||
"reset_at": 1_902_592_000
|
||||
}
|
||||
},
|
||||
"provider_specific_root": {"opaque": [1, 2, 3]}
|
||||
});
|
||||
let outcome = store
|
||||
.compare_and_swap_quota(QuotaObservation {
|
||||
account_id: account.id().clone(),
|
||||
expected_revision: account.revision(),
|
||||
quota: Some(OpaqueProviderData::new(
|
||||
raw.as_object().expect("quota object").clone(),
|
||||
)),
|
||||
observed_at: Some(SystemTime::now()),
|
||||
limit_reached: None,
|
||||
})
|
||||
.await
|
||||
.expect("persist quota");
|
||||
assert_eq!(outcome, QuotaWriteOutcome::Updated);
|
||||
|
||||
let snapshot = quota_service(&store)
|
||||
.read_account(account.id())
|
||||
.await
|
||||
.expect("read quota")
|
||||
.expect("quota snapshot");
|
||||
|
||||
assert_eq!(snapshot.windows().len(), 2);
|
||||
assert_eq!(
|
||||
snapshot.windows()[0].kind(),
|
||||
CodexQuotaWindowKind::ShortTerm
|
||||
);
|
||||
assert_eq!(snapshot.windows()[0].window_seconds(), Some(18_000));
|
||||
assert_eq!(snapshot.windows()[1].kind(), CodexQuotaWindowKind::Weekly);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn roll_expired_windows_is_o1_for_extremely_old_reset() {
|
||||
// 过去数年的 1 秒窗口:O(1) 完成,不能线性循环。
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_o1_roll").await;
|
||||
let account = store.account("acct_o1_roll").expect("account");
|
||||
let old_reset = SystemTime::now() - Duration::from_secs(5 * 365 * 24 * 3600);
|
||||
let raw = json!({
|
||||
"rate_limit": {
|
||||
"primary_window": {
|
||||
"used_percent": 100,
|
||||
"reset_at": old_reset
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.expect("epoch")
|
||||
.as_secs(),
|
||||
"limit_window_seconds": 1
|
||||
}
|
||||
}
|
||||
});
|
||||
store
|
||||
.compare_and_swap_quota(QuotaObservation {
|
||||
account_id: account.id().clone(),
|
||||
expected_revision: account.revision(),
|
||||
quota: Some(OpaqueProviderData::new(
|
||||
raw.as_object().expect("quota object").clone(),
|
||||
)),
|
||||
observed_at: Some(SystemTime::now()),
|
||||
limit_reached: None,
|
||||
})
|
||||
.await
|
||||
.expect("persist quota");
|
||||
let service = quota_service(&store);
|
||||
let snapshot = service
|
||||
.read_account(account.id())
|
||||
.await
|
||||
.expect("read quota")
|
||||
.expect("quota snapshot");
|
||||
let rolled = snapshot.roll_expired_windows(SystemTime::now());
|
||||
let window = &rolled.windows()[0];
|
||||
assert_eq!(window.used_percent(), Some(0.0));
|
||||
assert!(!window.limit_reached());
|
||||
// 下一 reset 必须落在 now 之后(O(1) 公式而非线性累加)。
|
||||
assert!(
|
||||
window
|
||||
.reset_at()
|
||||
.is_some_and(|reset| SystemTime::from(reset) > SystemTime::now())
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn scheduling_signal_survives_limit_reached_without_percent_or_reset() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_signal_only_limit").await;
|
||||
let account = store.account("acct_signal_only_limit").expect("account");
|
||||
let raw = json!({
|
||||
"rate_limit": {
|
||||
"allowed": false,
|
||||
"limit_reached": true,
|
||||
"primary_window": {"used_percent": 0}
|
||||
}
|
||||
});
|
||||
store
|
||||
.compare_and_swap_quota(QuotaObservation {
|
||||
account_id: account.id().clone(),
|
||||
expected_revision: account.revision(),
|
||||
quota: Some(OpaqueProviderData::new(
|
||||
raw.as_object().expect("quota object").clone(),
|
||||
)),
|
||||
observed_at: Some(SystemTime::now()),
|
||||
limit_reached: None,
|
||||
})
|
||||
.await
|
||||
.expect("persist quota");
|
||||
let service = quota_service(&store);
|
||||
let snapshot = service
|
||||
.read_account(account.id())
|
||||
.await
|
||||
.expect("read quota")
|
||||
.expect("quota snapshot");
|
||||
assert!(snapshot.fact().exhausted());
|
||||
let signals = service.scheduling_signals(&account);
|
||||
assert!(
|
||||
signals.is_some_and(|signals| signals.limit_reached()),
|
||||
"limit_reached without percent/reset must still produce scheduling exclusion"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn persisted_codex_alias_keeps_the_top_level_rate_limit_canonical() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_canonical_codex_limit").await;
|
||||
let account = store
|
||||
.account("acct_canonical_codex_limit")
|
||||
.expect("created account");
|
||||
let raw = json!({
|
||||
"rate_limit": {
|
||||
"primary_window": {
|
||||
"used_percent": 91,
|
||||
"reset_at": 1_900_000_000,
|
||||
"limit_window_seconds": 2_592_000
|
||||
}
|
||||
},
|
||||
"additional_rate_limits": [{
|
||||
"metered_feature": "codex",
|
||||
"rate_limit": {
|
||||
"primary_window": {
|
||||
"used_percent": 98,
|
||||
"reset_at": 1_900_000_000,
|
||||
"limit_window_seconds": 2_592_000
|
||||
}
|
||||
}
|
||||
}]
|
||||
});
|
||||
let outcome = store
|
||||
.compare_and_swap_quota(QuotaObservation {
|
||||
account_id: account.id().clone(),
|
||||
expected_revision: account.revision(),
|
||||
quota: Some(OpaqueProviderData::new(
|
||||
raw.as_object().expect("quota object").clone(),
|
||||
)),
|
||||
observed_at: Some(SystemTime::now()),
|
||||
limit_reached: None,
|
||||
})
|
||||
.await
|
||||
.expect("persist quota");
|
||||
assert_eq!(outcome, QuotaWriteOutcome::Updated);
|
||||
|
||||
let snapshot = quota_service(&store)
|
||||
.read_account(account.id())
|
||||
.await
|
||||
.expect("read quota")
|
||||
.expect("quota snapshot");
|
||||
let monthly = snapshot
|
||||
.windows()
|
||||
.iter()
|
||||
.filter(|window| window.kind() == CodexQuotaWindowKind::Monthly)
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
assert_eq!(snapshot.fact().remaining_percent(), Some(9));
|
||||
assert_eq!(monthly.len(), 1);
|
||||
assert!(
|
||||
monthly
|
||||
.iter()
|
||||
.any(|window| { window.source() == "codex" && window.used_percent() == Some(91.0) })
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn code_review_limit_projects_as_independent_window_with_limit_name() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_code_review").await;
|
||||
let account = store.account("acct_code_review").expect("account");
|
||||
let raw = json!({
|
||||
"rate_limit": {
|
||||
"primary_window": {
|
||||
"used_percent": 37,
|
||||
"reset_at": 1_900_000_000,
|
||||
"limit_window_seconds": 18_000
|
||||
}
|
||||
},
|
||||
"code_review_rate_limit": {
|
||||
"primary_window": {
|
||||
"used_percent": 80,
|
||||
"reset_at": 1_900_604_800,
|
||||
"limit_window_seconds": 604_800
|
||||
}
|
||||
},
|
||||
"additional_rate_limits": [{
|
||||
"limit_name": "code review",
|
||||
"metered_feature": "code_review",
|
||||
"rate_limit": {
|
||||
"primary_window": {
|
||||
"used_percent": 55,
|
||||
"reset_at": 1_900_604_800,
|
||||
"limit_window_seconds": 604_800
|
||||
}
|
||||
}
|
||||
}],
|
||||
"spend_control": {
|
||||
"reached": false,
|
||||
"individual_limit": {
|
||||
"used_percent": 12,
|
||||
"reset_at": 1_902_592_000
|
||||
}
|
||||
}
|
||||
});
|
||||
store
|
||||
.compare_and_swap_quota(QuotaObservation {
|
||||
account_id: account.id().clone(),
|
||||
expected_revision: account.revision(),
|
||||
quota: Some(OpaqueProviderData::new(
|
||||
raw.as_object().expect("quota object").clone(),
|
||||
)),
|
||||
observed_at: Some(SystemTime::now()),
|
||||
limit_reached: None,
|
||||
})
|
||||
.await
|
||||
.expect("persist quota");
|
||||
|
||||
let snapshot = quota_service(&store)
|
||||
.read_account(account.id())
|
||||
.await
|
||||
.expect("read quota")
|
||||
.expect("quota snapshot");
|
||||
|
||||
// spend_control 不再生成窗口(只作 exhaustion 信号);顶层和 additional
|
||||
// 的 code_review 桶分别保留,避免一个快照覆盖另一个。
|
||||
let review = snapshot
|
||||
.windows()
|
||||
.iter()
|
||||
.filter(|window| window.source() == "code_review")
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(review.len(), 2);
|
||||
assert!(review.iter().any(|window| {
|
||||
window.limit_name() == Some("code_review") && window.used_percent() == Some(80.0)
|
||||
}));
|
||||
assert!(review.iter().any(|window| {
|
||||
window.limit_name() == Some("code review") && window.used_percent() == Some(55.0)
|
||||
}));
|
||||
assert_ne!(review[0].key(), review[1].key());
|
||||
assert_eq!(snapshot.windows().len(), 3);
|
||||
assert!(!snapshot.fact().exhausted());
|
||||
}
|
||||
@@ -621,8 +621,6 @@ async fn expired_access_token_refreshes_even_when_next_refresh_is_deferred() {
|
||||
.apply_state(
|
||||
&deferred,
|
||||
AccountAvailability::QuotaExhausted,
|
||||
Some("quota_exhausted".to_owned()),
|
||||
None,
|
||||
SystemTime::now(),
|
||||
)
|
||||
.await
|
||||
@@ -685,18 +683,11 @@ async fn expired_access_token_refreshes_during_active_cooldown() {
|
||||
.await
|
||||
.expect("seed expired access token with deferred refresh");
|
||||
let deferred = store.account("acct_refresh").expect("deferred account");
|
||||
let cooldown_until = SystemTime::now() + Duration::from_secs(12 * 60 * 60);
|
||||
store
|
||||
.repository()
|
||||
.apply_state(
|
||||
&deferred,
|
||||
AccountAvailability::Cooldown,
|
||||
Some("rate_limited".to_owned()),
|
||||
Some(cooldown_until),
|
||||
SystemTime::now(),
|
||||
)
|
||||
.apply_state(&deferred, AccountAvailability::Ready, SystemTime::now())
|
||||
.await
|
||||
.expect("mark account cooling down");
|
||||
.expect("mark account ready");
|
||||
let refresher = Arc::new(Refresher::new(Ok(success_tokens())));
|
||||
let service = CodexCredentialRefreshService::new(
|
||||
store.repository(),
|
||||
@@ -724,8 +715,7 @@ async fn expired_access_token_refreshes_during_active_cooldown() {
|
||||
["rt-old-access"]
|
||||
);
|
||||
let refreshed = store.account("acct_refresh").expect("refreshed account");
|
||||
assert_eq!(refreshed.availability(), AccountAvailability::Cooldown);
|
||||
assert_eq!(refreshed.cooldown_until(), Some(cooldown_until));
|
||||
assert_eq!(refreshed.availability(), AccountAvailability::Ready);
|
||||
assert!(
|
||||
refreshed
|
||||
.access_token_expires_at()
|
||||
|
||||
@@ -43,8 +43,9 @@ use wiremock::matchers::{header, method, path};
|
||||
use wiremock::{Mock, MockServer, ResponseTemplate};
|
||||
|
||||
use crate::support::{
|
||||
MemoryAccountStore, MemorySessionAffinity, MemorySessionExclusions, TestLeaseCoordinator,
|
||||
account_policy, agent_identity_service_with_pool, catalog_cache, profile, secret,
|
||||
MemoryAccountStore, MemoryCooldownPort, MemorySessionAffinity, MemorySessionExclusions,
|
||||
TestLeaseCoordinator, account_policy, agent_identity_service_with_pool, catalog_cache, profile,
|
||||
secret,
|
||||
};
|
||||
use crate::transport::accept_codex_test_websocket;
|
||||
|
||||
@@ -115,6 +116,21 @@ fn provider_with_affinity_and_base_url_and_leases(
|
||||
base_url: String,
|
||||
leases: Arc<TestLeaseCoordinator>,
|
||||
) -> CodexProvider {
|
||||
provider_and_quota_with_affinity_and_base_url_and_leases(
|
||||
store,
|
||||
session_affinity,
|
||||
base_url,
|
||||
leases,
|
||||
)
|
||||
.0
|
||||
}
|
||||
|
||||
fn provider_and_quota_with_affinity_and_base_url_and_leases(
|
||||
store: &Arc<MemoryAccountStore>,
|
||||
session_affinity: Arc<MemorySessionAffinity>,
|
||||
base_url: String,
|
||||
leases: Arc<TestLeaseCoordinator>,
|
||||
) -> (CodexProvider, Arc<CodexCredentialQuotaService>) {
|
||||
let profile = wire_profile();
|
||||
let http = reqwest::Client::builder()
|
||||
.no_proxy()
|
||||
@@ -136,6 +152,7 @@ fn provider_with_affinity_and_base_url_and_leases(
|
||||
http.clone(),
|
||||
base_url.clone(),
|
||||
Arc::clone(&agent_identity),
|
||||
Arc::new(MemoryCooldownPort::new()),
|
||||
));
|
||||
let account_feedback = Arc::new(AccountFeedbackStats::default());
|
||||
let selector = Arc::new(CodexCredentialSelector::new(
|
||||
@@ -152,10 +169,10 @@ fn provider_with_affinity_and_base_url_and_leases(
|
||||
true,
|
||||
));
|
||||
|
||||
CodexProvider::new(
|
||||
let provider = CodexProvider::new(
|
||||
selector,
|
||||
catalog,
|
||||
quota,
|
||||
Arc::clone("a),
|
||||
agent_identity,
|
||||
account_feedback,
|
||||
http,
|
||||
@@ -163,7 +180,8 @@ fn provider_with_affinity_and_base_url_and_leases(
|
||||
base_url,
|
||||
websocket_pool,
|
||||
)
|
||||
.expect("official OpenAI provider")
|
||||
.expect("official OpenAI provider");
|
||||
(provider, quota)
|
||||
}
|
||||
|
||||
async fn create_account(store: &Arc<MemoryAccountStore>, id: &str) {
|
||||
@@ -196,6 +214,27 @@ fn generate_operation() -> Operation {
|
||||
))
|
||||
}
|
||||
|
||||
fn generate_with_session_context(
|
||||
session_id: &str,
|
||||
thread_id: Option<&str>,
|
||||
turn_metadata: Option<&str>,
|
||||
) -> GenerateRequest {
|
||||
let mut body = Map::from_iter([
|
||||
("model".to_owned(), json!("gpt-5.4")),
|
||||
("input".to_owned(), json!("hello")),
|
||||
("session_id".to_owned(), json!(session_id)),
|
||||
]);
|
||||
if let Some(thread_id) = thread_id {
|
||||
body.insert("thread_id".to_owned(), json!(thread_id));
|
||||
}
|
||||
if let Some(turn_metadata) = turn_metadata {
|
||||
body.insert("turnMetadata".to_owned(), json!(turn_metadata));
|
||||
}
|
||||
GenerateRequest::from_protocol_payload(
|
||||
ProtocolPayload::json_object("openai", body).expect("OpenAI payload"),
|
||||
)
|
||||
}
|
||||
|
||||
fn http_generate_operation() -> Operation {
|
||||
let payload = ProtocolPayload::json_object(
|
||||
"openai",
|
||||
@@ -1053,6 +1092,66 @@ async fn prompt_cache_key_should_become_an_opaque_session_affinity_lookup_key()
|
||||
assert!(keys[0].bytes().all(|byte| byte.is_ascii_hexdigit()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn subagent_requests_should_use_an_affinity_key_isolated_from_the_root_session() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_subagent").await;
|
||||
let affinity = Arc::new(MemorySessionAffinity::default());
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/codex/responses"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.insert_header("content-type", "text/event-stream")
|
||||
.set_body_string(CAPTURE_COMPLETED_SSE),
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let provider = provider_with_affinity_and_base_url(&store, Arc::clone(&affinity), server.uri());
|
||||
|
||||
for (request_id, subagent_kind) in [
|
||||
("req_root_affinity", None),
|
||||
("req_subagent_affinity_first", Some("review")),
|
||||
("req_subagent_affinity_second", Some("review")),
|
||||
] {
|
||||
let mut body = Map::from_iter([
|
||||
("model".to_owned(), json!("gpt-5.4")),
|
||||
("input".to_owned(), json!("new task")),
|
||||
("prompt_cache_key".to_owned(), json!("root-session-key")),
|
||||
]);
|
||||
if let Some(subagent_kind) = subagent_kind {
|
||||
body.insert(
|
||||
"client_metadata".to_owned(),
|
||||
json!({"x-openai-subagent": subagent_kind}),
|
||||
);
|
||||
}
|
||||
let generation = GenerateRequest::from_protocol_payload(
|
||||
ProtocolPayload::json_object("openai", body)
|
||||
.expect("OpenAI payload")
|
||||
.with_context(Map::from_iter([("use_websocket".to_owned(), json!(false))])),
|
||||
);
|
||||
let mut stream = provider
|
||||
.execute(
|
||||
planned_request("openai", Operation::Generate(generation)),
|
||||
context(request_id, CancellationToken::new()),
|
||||
)
|
||||
.await
|
||||
.expect("prepare subagent provider stream");
|
||||
while let Some(event) = stream.next().await {
|
||||
event.expect("subagent response");
|
||||
}
|
||||
drop(stream);
|
||||
}
|
||||
|
||||
let keys = affinity.lookup_keys();
|
||||
assert_eq!(keys.len(), 3);
|
||||
assert_ne!(keys[0], keys[1], "subagent must not inherit the root key");
|
||||
assert_eq!(
|
||||
keys[1], keys[2],
|
||||
"the same subagent should retain its own key"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn explicit_session_id_should_override_turn_specific_prompt_cache_keys_for_affinity() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
@@ -1095,6 +1194,58 @@ async fn explicit_session_id_should_override_turn_specific_prompt_cache_keys_for
|
||||
assert_eq!(keys[0], keys[1]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn thread_spawn_children_should_have_independent_session_affinity_keys() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_thread_spawn_affinity").await;
|
||||
let affinity = Arc::new(MemorySessionAffinity::default());
|
||||
let provider = provider_with_affinity(&store, Arc::clone(&affinity));
|
||||
let thread_spawn = r#"{"subagent_kind":"thread_spawn"}"#;
|
||||
|
||||
for (request_id, thread_id, turn_metadata) in [
|
||||
(
|
||||
"req_thread_spawn_first",
|
||||
Some("child-one"),
|
||||
Some(thread_spawn),
|
||||
),
|
||||
(
|
||||
"req_thread_spawn_second",
|
||||
Some("child-two"),
|
||||
Some(thread_spawn),
|
||||
),
|
||||
(
|
||||
"req_thread_spawn_repeat",
|
||||
Some("child-one"),
|
||||
Some(thread_spawn),
|
||||
),
|
||||
("req_thread_spawn_fallback", None, Some(thread_spawn)),
|
||||
("req_thread_spawn_parent", None, None),
|
||||
] {
|
||||
let stream = provider
|
||||
.execute(
|
||||
planned_request(
|
||||
"openai",
|
||||
Operation::Generate(generate_with_session_context(
|
||||
"parent-session",
|
||||
thread_id,
|
||||
turn_metadata,
|
||||
)),
|
||||
),
|
||||
context(request_id, CancellationToken::new()),
|
||||
)
|
||||
.await
|
||||
.expect("prepare provider stream");
|
||||
drop(stream);
|
||||
}
|
||||
|
||||
let keys = affinity.lookup_keys();
|
||||
assert_eq!(keys.len(), 5);
|
||||
assert_ne!(keys[0], keys[1]);
|
||||
assert_eq!(keys[0], keys[2]);
|
||||
assert_ne!(keys[3], keys[4]);
|
||||
assert_ne!(keys[0], keys[3]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn invalid_local_conversation_id_should_fall_back_to_an_opaque_affinity_key() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
@@ -1263,21 +1414,18 @@ async fn response_failed_before_semantic_output_is_atomic_and_persists_quota_loc
|
||||
vec!["response.created", "response.failed"]
|
||||
);
|
||||
let account = store.account(account_id).expect("rate-limited account");
|
||||
assert_eq!(account.availability(), AccountAvailability::QuotaExhausted);
|
||||
assert!(
|
||||
account
|
||||
.cooldown_until()
|
||||
.is_some_and(|until| until > SystemTime::now())
|
||||
);
|
||||
// 429 只记录限流窗口,不改变 availability(限流不改变账号可用性)。
|
||||
assert_eq!(account.availability(), AccountAvailability::Ready);
|
||||
let _ = release.send(());
|
||||
server.abort();
|
||||
let _ = server.await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn usage_limit_failure_returns_promptly_and_refreshes_the_authoritative_quota_snapshot() {
|
||||
async fn official_usage_limit_failure_projects_100_percent_and_holds_until_reset_advances() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
let account_id = "acct_usage_limit_request_path";
|
||||
let reset_at = 1_900_000_000;
|
||||
create_account(&store, account_id).await;
|
||||
let account = store.account(account_id).expect("created account");
|
||||
store
|
||||
@@ -1289,7 +1437,8 @@ async fn usage_limit_failure_returns_promptly_and_refreshes_the_authoritative_qu
|
||||
"rate_limit": {
|
||||
"allowed": true,
|
||||
"limit_reached": false,
|
||||
"primary_window": {"used_percent": 99, "reset_at": 1_900_000_000}
|
||||
"primary_window": {"used_percent": 99, "reset_at": reset_at},
|
||||
"secondary_window": {"used_percent": 0}
|
||||
}
|
||||
})
|
||||
.as_object()
|
||||
@@ -1297,6 +1446,7 @@ async fn usage_limit_failure_returns_promptly_and_refreshes_the_authoritative_qu
|
||||
.clone(),
|
||||
)),
|
||||
observed_at: Some(SystemTime::now()),
|
||||
limit_reached: None,
|
||||
})
|
||||
.await
|
||||
.expect("seed stale passive quota");
|
||||
@@ -1304,8 +1454,8 @@ async fn usage_limit_failure_returns_promptly_and_refreshes_the_authoritative_qu
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/api/codex/usage"))
|
||||
.and(header("authorization", format!("Bearer at-{account_id}")))
|
||||
// 生产环境会短暂返回互相矛盾的快照:数字已到 100%,但布尔位仍声称可用。
|
||||
// 后台同步应覆盖展示数据,不能据此撤销真实请求已经确认的额度耗尽。
|
||||
// 生产环境会短暂返回互相矛盾的快照:98/99% 且布尔位仍声称可用。
|
||||
// 后台同步不能覆盖真实请求已确认的额度耗尽。
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
// 验证后台 usage 同步不能把原始的额度错误响应拖到查询完成之后。
|
||||
@@ -1314,7 +1464,8 @@ async fn usage_limit_failure_returns_promptly_and_refreshes_the_authoritative_qu
|
||||
"rate_limit": {
|
||||
"allowed": true,
|
||||
"limit_reached": false,
|
||||
"primary_window": {"used_percent": 100, "reset_at": 1_900_000_000}
|
||||
"primary_window": {"used_percent": 99, "reset_at": reset_at},
|
||||
"secondary_window": {"used_percent": 0}
|
||||
}
|
||||
})),
|
||||
)
|
||||
@@ -1323,16 +1474,13 @@ async fn usage_limit_failure_returns_promptly_and_refreshes_the_authoritative_qu
|
||||
.await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/codex/responses"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.insert_header("content-type", "text/event-stream")
|
||||
.set_body_string(concat!(
|
||||
"event: response.created\n",
|
||||
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_usage_limit_confirm\",\"model\":\"gpt-5.4\",\"status\":\"in_progress\"}}\n\n",
|
||||
"event: response.failed\n",
|
||||
"data: {\"type\":\"response.failed\",\"status_code\":429,\"retry_after_seconds\":86400,\"response\":{\"id\":\"resp_usage_limit_confirm\",\"status\":\"failed\",\"error\":{\"code\":\"usage_limit_reached\",\"message\":\"usage limit reached\"}}}\n\n"
|
||||
)),
|
||||
)
|
||||
.respond_with(ResponseTemplate::new(429).set_body_json(json!({
|
||||
"error": {
|
||||
"type": "usage_limit_reached",
|
||||
"message": "usage limit reached",
|
||||
"resets_at": reset_at
|
||||
}
|
||||
})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
@@ -1358,10 +1506,29 @@ async fn usage_limit_failure_returns_promptly_and_refreshes_the_authoritative_qu
|
||||
assert_eq!(failure.kind(), ProviderErrorKind::QuotaExhausted);
|
||||
let account = store.account(account_id).expect("usage-limit account");
|
||||
assert_eq!(account.availability(), AccountAvailability::QuotaExhausted);
|
||||
let projected = store
|
||||
.get_quotas(&[account.id().clone()])
|
||||
.await
|
||||
.expect("read immediate quota projection")
|
||||
.into_iter()
|
||||
.next()
|
||||
.expect("confirmed quota observation");
|
||||
let projected_observed_at = projected.observed_at.expect("projection observed at");
|
||||
let projected = Value::Object(
|
||||
projected
|
||||
.quota
|
||||
.map(OpaqueProviderData::into_inner)
|
||||
.expect("confirmed quota projection"),
|
||||
);
|
||||
assert_eq!(
|
||||
projected
|
||||
.pointer("/rate_limit/primary_window/used_percent")
|
||||
.and_then(Value::as_u64),
|
||||
Some(100)
|
||||
);
|
||||
assert!(
|
||||
account
|
||||
.cooldown_until()
|
||||
.is_some_and(|until| until > SystemTime::now())
|
||||
projected.pointer("/rate_limit/secondary_window").is_none(),
|
||||
"request failure must not manufacture a secondary quota"
|
||||
);
|
||||
tokio::time::sleep(Duration::from_millis(200)).await;
|
||||
let requests = server.received_requests().await.expect("received requests");
|
||||
@@ -1374,35 +1541,49 @@ async fn usage_limit_failure_returns_promptly_and_refreshes_the_authoritative_qu
|
||||
"usage refresh must wait for the upstream quota settlement delay"
|
||||
);
|
||||
|
||||
let quota = timeout(Duration::from_secs(5), async {
|
||||
timeout(Duration::from_secs(5), async {
|
||||
loop {
|
||||
let quota = store
|
||||
let observed_at = store
|
||||
.get_quotas(&[account.id().clone()])
|
||||
.await
|
||||
.expect("read refreshed quota")
|
||||
.expect("read refreshed quota observation")
|
||||
.into_iter()
|
||||
.next()
|
||||
.and_then(|observation| observation.quota)
|
||||
.expect("authoritative quota snapshot");
|
||||
let quota = Value::Object(quota.into_inner());
|
||||
if quota
|
||||
.pointer("/rate_limit/primary_window/used_percent")
|
||||
.and_then(Value::as_u64)
|
||||
== Some(100)
|
||||
{
|
||||
break quota;
|
||||
.and_then(|observation| observation.observed_at);
|
||||
if observed_at.is_some_and(|observed_at| observed_at > projected_observed_at) {
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(10)).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("background usage refresh must replace the stale quota snapshot");
|
||||
.expect("background usage refresh must run");
|
||||
let quota = store
|
||||
.get_quotas(&[account.id().clone()])
|
||||
.await
|
||||
.expect("read refreshed quota")
|
||||
.into_iter()
|
||||
.next()
|
||||
.and_then(|observation| observation.quota)
|
||||
.map(OpaqueProviderData::into_inner)
|
||||
.map(Value::Object)
|
||||
.expect("authoritative quota projection");
|
||||
assert_eq!(
|
||||
quota
|
||||
.pointer("/rate_limit/primary_window/used_percent")
|
||||
.and_then(Value::as_u64),
|
||||
Some(100)
|
||||
);
|
||||
assert_eq!(
|
||||
quota
|
||||
.pointer("/rate_limit/primary_window/reset_at")
|
||||
.and_then(Value::as_i64),
|
||||
Some(reset_at)
|
||||
);
|
||||
assert!(
|
||||
quota.pointer("/rate_limit/secondary_window").is_none(),
|
||||
"background quota refresh must retain the same window projection"
|
||||
);
|
||||
let requests = server.received_requests().await.expect("received requests");
|
||||
assert_eq!(
|
||||
requests
|
||||
@@ -1799,8 +1980,6 @@ async fn disabled_account_diagnostic_uses_upstream_without_persisting_account_st
|
||||
.apply_state(
|
||||
&account,
|
||||
AccountAvailability::QuotaExhausted,
|
||||
Some("quota_exhausted".to_owned()),
|
||||
None,
|
||||
SystemTime::now(),
|
||||
)
|
||||
.await
|
||||
@@ -1859,6 +2038,85 @@ async fn disabled_account_diagnostic_uses_upstream_without_persisting_account_st
|
||||
assert_eq!(affinity.binding_count(), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn quota_limited_account_diagnostic_uses_upstream() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
let account_id = "acct_quota_limited_diagnostic";
|
||||
create_account(&store, account_id).await;
|
||||
let account = store.account(account_id).expect("test account");
|
||||
let raw_quota = json!({
|
||||
"rate_limit": {
|
||||
"allowed": false,
|
||||
"limit_reached": true,
|
||||
"primary_window": {"used_percent": 100, "reset_at": 1_900_000_000}
|
||||
}
|
||||
});
|
||||
store
|
||||
.compare_and_swap_quota(QuotaObservation {
|
||||
account_id: account.id().clone(),
|
||||
expected_revision: account.revision(),
|
||||
quota: Some(OpaqueProviderData::new(
|
||||
raw_quota
|
||||
.as_object()
|
||||
.expect("quota snapshot object")
|
||||
.clone(),
|
||||
)),
|
||||
observed_at: Some(SystemTime::now()),
|
||||
limit_reached: Some(true),
|
||||
})
|
||||
.await
|
||||
.expect("seed quota-limited snapshot");
|
||||
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/codex/responses"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.insert_header("content-type", "text/event-stream")
|
||||
.set_body_string(concat!(
|
||||
"event: response.created\n",
|
||||
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_quota_limited_diagnostic\",\"model\":\"gpt-5.4\"}}\n\n",
|
||||
"event: response.completed\n",
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_quota_limited_diagnostic\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}\n\n"
|
||||
)),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let (provider, quota) = provider_and_quota_with_affinity_and_base_url_and_leases(
|
||||
&store,
|
||||
Arc::new(MemorySessionAffinity::default()),
|
||||
server.uri(),
|
||||
Arc::new(TestLeaseCoordinator::default()),
|
||||
);
|
||||
quota
|
||||
.prepare_scheduling(std::slice::from_ref(&account))
|
||||
.await;
|
||||
|
||||
let mut stream = provider
|
||||
.execute(
|
||||
planned_request("openai", http_generate_operation()),
|
||||
diagnostic_context("req_quota_limited_diagnostic", account_id),
|
||||
)
|
||||
.await
|
||||
.expect("quota-limited diagnostic should prepare a fixed-account stream");
|
||||
let mut completed = false;
|
||||
while let Some(event) = stream.next().await {
|
||||
let event = event.expect("quota-limited diagnostic upstream response");
|
||||
completed |= event
|
||||
.canonical_facts()
|
||||
.iter()
|
||||
.any(|event| matches!(event, GatewayEvent::Completed(_)));
|
||||
}
|
||||
|
||||
assert!(completed);
|
||||
let requests = server
|
||||
.received_requests()
|
||||
.await
|
||||
.expect("captured quota-limited diagnostic request");
|
||||
assert_eq!(requests.len(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn successful_response_preserves_quota_exhaustion_from_rate_limit_headers() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
@@ -1907,7 +2165,7 @@ async fn successful_response_preserves_quota_exhaustion_from_rate_limit_headers(
|
||||
.account(account_id)
|
||||
.expect("account after successful response")
|
||||
.availability(),
|
||||
AccountAvailability::QuotaExhausted
|
||||
AccountAvailability::Ready
|
||||
);
|
||||
assert!(store.has_quota(account_id));
|
||||
}
|
||||
@@ -1969,16 +2227,8 @@ async fn successful_http_sse_rate_limit_event_persists_structured_exhaustion() {
|
||||
.account(account_id)
|
||||
.expect("account after HTTP SSE response");
|
||||
assert_eq!(
|
||||
(
|
||||
account.availability(),
|
||||
account.cooldown_until(),
|
||||
store.has_quota(account_id),
|
||||
),
|
||||
(
|
||||
AccountAvailability::QuotaExhausted,
|
||||
Some(SystemTime::UNIX_EPOCH + Duration::from_secs(reset_at)),
|
||||
true,
|
||||
)
|
||||
(account.availability(), store.has_quota(account_id),),
|
||||
(AccountAvailability::Ready, true,)
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
@@ -16,9 +16,10 @@ use gateway_core::engine::credential::{
|
||||
};
|
||||
use gateway_core::error::{StoreError, StoreErrorKind};
|
||||
use gateway_core::provider_ports::{
|
||||
ProviderCatalogCacheKey, ProviderCatalogCachePort, ProviderLeaseAcquisition, ProviderLeasePort,
|
||||
ProviderLeaseRequest, ProviderRefreshPolicy, ProviderRuntimePolicyPort,
|
||||
ProviderSchedulingLeaseRequest, ProviderSchedulingState, ProviderSessionAffinityKey,
|
||||
ProviderCatalogCacheKey, ProviderCatalogCachePort, ProviderCooldown, ProviderCooldownPort,
|
||||
ProviderCooldownScope, ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest,
|
||||
ProviderRefreshPolicy, ProviderRuntimePolicyPort, ProviderSchedulingLeaseRequest,
|
||||
ProviderSchedulingState, ProviderScopedCooldown, ProviderSessionAffinityKey,
|
||||
ProviderSessionAffinityPort, ProviderSessionExclusionPort, ProviderSessionExclusions,
|
||||
ProviderStoreError,
|
||||
};
|
||||
@@ -117,6 +118,17 @@ impl MemoryAccountStore {
|
||||
.is_some_and(|stored| stored.quota.is_some())
|
||||
}
|
||||
|
||||
pub(crate) fn quota_json(&self, id: &str) -> Option<serde_json::Value> {
|
||||
let id = ProviderAccountId::new(id).ok()?;
|
||||
self.accounts
|
||||
.lock()
|
||||
.expect("account store lock")
|
||||
.get(&id)
|
||||
.and_then(|stored| stored.quota.as_ref())
|
||||
.and_then(|observation| observation.quota.as_ref())
|
||||
.map(|data| serde_json::Value::Object(data.expose_to_provider().clone()))
|
||||
}
|
||||
|
||||
pub(crate) fn fail_provider_listing(&self) {
|
||||
self.fail_provider_listing.store(true, Ordering::SeqCst);
|
||||
}
|
||||
@@ -235,7 +247,6 @@ impl ProviderAccountStore for MemoryAccountStore {
|
||||
revision: next,
|
||||
enabled: stored.account.enabled(),
|
||||
availability: stored.account.availability(),
|
||||
cooldown_until: stored.account.cooldown_until(),
|
||||
access_token_expires_at,
|
||||
has_refresh_token,
|
||||
next_refresh_at,
|
||||
@@ -303,7 +314,6 @@ impl ProviderAccountStore for MemoryAccountStore {
|
||||
revision: stored.account.revision(),
|
||||
enabled: stored.account.enabled(),
|
||||
availability: change.availability,
|
||||
cooldown_until: change.cooldown_until,
|
||||
access_token_expires_at: stored.account.access_token_expires_at(),
|
||||
has_refresh_token: stored.account.has_refresh_token(),
|
||||
next_refresh_at: stored.account.next_refresh_at(),
|
||||
@@ -325,7 +335,6 @@ impl ProviderAccountStore for MemoryAccountStore {
|
||||
revision: stored.account.revision(),
|
||||
enabled: stored.account.enabled(),
|
||||
availability: stored.account.availability(),
|
||||
cooldown_until: stored.account.cooldown_until(),
|
||||
access_token_expires_at: stored.account.access_token_expires_at(),
|
||||
has_refresh_token: stored.account.has_refresh_token(),
|
||||
next_refresh_at: stored.account.next_refresh_at(),
|
||||
@@ -350,7 +359,6 @@ impl ProviderAccountStore for MemoryAccountStore {
|
||||
revision: stored.account.revision(),
|
||||
enabled,
|
||||
availability: stored.account.availability(),
|
||||
cooldown_until: stored.account.cooldown_until(),
|
||||
access_token_expires_at: stored.account.access_token_expires_at(),
|
||||
has_refresh_token: stored.account.has_refresh_token(),
|
||||
next_refresh_at: stored.account.next_refresh_at(),
|
||||
@@ -378,7 +386,6 @@ struct AccountRebuild {
|
||||
revision: CredentialRevision,
|
||||
enabled: bool,
|
||||
availability: AccountAvailability,
|
||||
cooldown_until: Option<SystemTime>,
|
||||
access_token_expires_at: Option<SystemTime>,
|
||||
has_refresh_token: bool,
|
||||
next_refresh_at: Option<SystemTime>,
|
||||
@@ -407,11 +414,7 @@ fn rebuild_account(current: &ProviderAccount, rebuild: AccountRebuild) -> Provid
|
||||
current.upstream_account_id().map(str::to_owned),
|
||||
plan_type,
|
||||
)
|
||||
.with_runtime_state(
|
||||
rebuild.enabled,
|
||||
rebuild.availability,
|
||||
rebuild.cooldown_until,
|
||||
)
|
||||
.with_runtime_state(rebuild.enabled, rebuild.availability)
|
||||
.with_refresh_schedule(rebuild.has_refresh_token, rebuild.next_refresh_at)
|
||||
}
|
||||
|
||||
@@ -464,6 +467,7 @@ impl ProviderLeasePort for TestLeaseCoordinator {
|
||||
last_started_at: None,
|
||||
quota_reset_at: None,
|
||||
quota_remaining_rank: None,
|
||||
quota_limit_reached: false,
|
||||
failure_rate_basis_points: None,
|
||||
first_output_latency_ms: None,
|
||||
},
|
||||
@@ -765,6 +769,93 @@ pub(crate) fn codex_account(id: &str) -> ProviderAccount {
|
||||
Some(format!("chatgpt-{id}")),
|
||||
Some("pro".to_owned()),
|
||||
)
|
||||
.with_runtime_state(true, AccountAvailability::Ready, None)
|
||||
.with_runtime_state(true, AccountAvailability::Ready)
|
||||
.with_refresh_schedule(true, Some(SystemTime::now() + Duration::from_secs(2_700)))
|
||||
}
|
||||
|
||||
/// 内存 `ProviderCooldownPort`:只实现 `read`/`put_if_later`(openai selector
|
||||
/// 429 冷却路径用),其余 scope 变体测试不涉及,返回占位。
|
||||
#[derive(Clone, Default)]
|
||||
pub(crate) struct MemoryCooldownPort {
|
||||
pub(crate) cooldowns: Arc<Mutex<BTreeMap<ProviderAccountId, ProviderCooldown>>>,
|
||||
}
|
||||
|
||||
impl MemoryCooldownPort {
|
||||
#[must_use]
|
||||
pub(crate) fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
}
|
||||
|
||||
impl ProviderCooldownPort for MemoryCooldownPort {
|
||||
fn put_if_later(
|
||||
&self,
|
||||
cooldown: ProviderCooldown,
|
||||
) -> BoxFuture<'_, Result<bool, ProviderStoreError>> {
|
||||
Box::pin(async move {
|
||||
let mut cooldowns = self.cooldowns.lock().expect("cooldown lock");
|
||||
let should_write = cooldowns.get(cooldown.account_id()).is_none_or(|current| {
|
||||
current.credential_revision() < cooldown.credential_revision()
|
||||
|| (current.credential_revision() == cooldown.credential_revision()
|
||||
&& current.until() < cooldown.until())
|
||||
});
|
||||
if should_write {
|
||||
cooldowns.insert(cooldown.account_id().clone(), cooldown);
|
||||
}
|
||||
Ok(should_write)
|
||||
})
|
||||
}
|
||||
|
||||
fn read<'a>(
|
||||
&'a self,
|
||||
account_id: &'a ProviderAccountId,
|
||||
) -> BoxFuture<'a, Result<Option<ProviderCooldown>, ProviderStoreError>> {
|
||||
Box::pin(async move {
|
||||
Ok(self
|
||||
.cooldowns
|
||||
.lock()
|
||||
.expect("cooldown lock")
|
||||
.get(account_id)
|
||||
.cloned())
|
||||
})
|
||||
}
|
||||
|
||||
fn clear<'a>(
|
||||
&'a self,
|
||||
_account_id: &'a ProviderAccountId,
|
||||
_through_revision: CredentialRevision,
|
||||
) -> BoxFuture<'a, Result<bool, ProviderStoreError>> {
|
||||
Box::pin(async { Ok(false) })
|
||||
}
|
||||
|
||||
fn put_scoped_if_later(
|
||||
&self,
|
||||
_cooldown: ProviderScopedCooldown,
|
||||
) -> BoxFuture<'_, Result<bool, ProviderStoreError>> {
|
||||
Box::pin(async { Ok(false) })
|
||||
}
|
||||
|
||||
fn read_scoped<'a>(
|
||||
&'a self,
|
||||
_account_id: &'a ProviderAccountId,
|
||||
_scope: &'a ProviderCooldownScope,
|
||||
) -> BoxFuture<'a, Result<Option<ProviderScopedCooldown>, ProviderStoreError>> {
|
||||
Box::pin(async { Ok(None) })
|
||||
}
|
||||
|
||||
fn clear_scoped<'a>(
|
||||
&'a self,
|
||||
_account_id: &'a ProviderAccountId,
|
||||
_scope: &'a ProviderCooldownScope,
|
||||
_through_revision: CredentialRevision,
|
||||
) -> BoxFuture<'a, Result<bool, ProviderStoreError>> {
|
||||
Box::pin(async { Ok(false) })
|
||||
}
|
||||
|
||||
fn clear_all<'a>(
|
||||
&'a self,
|
||||
_account_id: &'a ProviderAccountId,
|
||||
) -> BoxFuture<'a, Result<bool, ProviderStoreError>> {
|
||||
Box::pin(async { Ok(false) })
|
||||
}
|
||||
}
|
||||
|
||||
@@ -219,6 +219,10 @@ async fn fetch_should_send_official_catalog_headers_and_query() {
|
||||
.and(header("authorization", "Bearer oauth-access"))
|
||||
.and(header("chatgpt-account-id", "acct_123"))
|
||||
.and(header("originator", "codex_cli_rs"))
|
||||
.and(header(
|
||||
"user-agent",
|
||||
"codex_cli_rs/0.144.0 (linux 6.8; x86_64) xterm (codex_cli_rs; 1.0.0)",
|
||||
))
|
||||
.and(header("accept", "application/json"))
|
||||
.and(header("x-codex-installation-id", "installation-123"))
|
||||
.respond_with(
|
||||
|
||||
@@ -146,13 +146,13 @@ fn structured_expiry_code_classifies_without_an_auth_status() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn structured_quota_signal_takes_precedence_over_http_429() {
|
||||
fn http_429_without_usage_limit_type_remains_a_temporary_rate_limit() {
|
||||
assert_eq!(
|
||||
classify(
|
||||
429,
|
||||
r#"{"error":{"code":"insufficient_quota","message":"rate limit reached"}}"#
|
||||
),
|
||||
CodexFailureCategory::QuotaExhausted
|
||||
CodexFailureCategory::RateLimited
|
||||
);
|
||||
}
|
||||
|
||||
@@ -168,13 +168,24 @@ fn structured_usage_limit_type_is_resettable_quota_exhaustion() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usage_limit_message_without_a_type_is_resettable_quota_exhaustion() {
|
||||
fn http_429_usage_limit_message_without_a_type_remains_a_temporary_rate_limit() {
|
||||
assert_eq!(
|
||||
classify(
|
||||
429,
|
||||
r#"{"error":{"message":"The usage limit has been reached"}}"#
|
||||
),
|
||||
CodexFailureCategory::UsageLimitExhausted
|
||||
CodexFailureCategory::RateLimited
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn http_429_usage_limit_code_without_a_type_remains_a_temporary_rate_limit() {
|
||||
assert_eq!(
|
||||
classify(
|
||||
429,
|
||||
r#"{"error":{"code":"usage_limit_reached","message":"The usage limit has been reached"}}"#
|
||||
),
|
||||
CodexFailureCategory::RateLimited
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
@@ -789,6 +789,70 @@ async fn backend_http_should_zstd_compress_codex_responses_request_body() {
|
||||
assert_eq!(json["instructions"], "compress me");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn backend_http_should_force_upstream_sse_for_non_streaming_client_request() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("bind non-streaming HTTP server");
|
||||
let address = listener
|
||||
.local_addr()
|
||||
.expect("non-streaming HTTP server address");
|
||||
let server = tokio::spawn(async move {
|
||||
let (mut stream, _) = listener.accept().await.expect("accept HTTP client");
|
||||
let request = read_http_request_with_body(&mut stream).await;
|
||||
write_completed_sse_response(&mut stream).await;
|
||||
request
|
||||
});
|
||||
let mut request = CodexResponsesRequest::from_body(
|
||||
json!({
|
||||
"model": "gpt-test",
|
||||
"instructions": "collect this response",
|
||||
"input": [],
|
||||
"stream": false
|
||||
})
|
||||
.as_object()
|
||||
.expect("request body")
|
||||
.clone(),
|
||||
);
|
||||
request.force_http_sse = true;
|
||||
assert!(
|
||||
!request.stream(),
|
||||
"client-facing preference remains non-streaming"
|
||||
);
|
||||
let client = CodexBackendClient::new(
|
||||
reqwest::Client::builder()
|
||||
.no_proxy()
|
||||
.build()
|
||||
.expect("HTTP client"),
|
||||
format!("http://{address}"),
|
||||
test_wire_profile(),
|
||||
);
|
||||
|
||||
client
|
||||
.create_response(
|
||||
&request,
|
||||
request_context("req_non_streaming_client", Some("acct-non-streaming")),
|
||||
)
|
||||
.await
|
||||
.expect("upstream SSE response");
|
||||
|
||||
let raw = server.await.expect("non-streaming HTTP server task");
|
||||
let separator = raw
|
||||
.windows(4)
|
||||
.position(|window| window == b"\r\n\r\n")
|
||||
.expect("HTTP head/body separator");
|
||||
let compressed_body = &raw[separator + 4..];
|
||||
let decompressed = zstd::stream::decode_all(std::io::Cursor::new(compressed_body))
|
||||
.expect("zstd body should decode");
|
||||
let json: serde_json::Value =
|
||||
serde_json::from_slice(&decompressed).expect("decompressed body should be valid JSON");
|
||||
assert_eq!(json["stream"], true);
|
||||
assert!(
|
||||
!request.stream(),
|
||||
"upstream normalization must not mutate local metadata"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn websocket_should_keep_an_exact_chain_while_new_connections_adopt_the_latest_codex_core_wire_profile()
|
||||
{
|
||||
|
||||
@@ -60,7 +60,7 @@ impl CodexCliReleaseTransport for CliReleaseTransport {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn wire_profile_should_generate_codex_core_user_agent() {
|
||||
fn wire_profile_should_generate_bundled_core_app_server_user_agent() {
|
||||
let profile = CodexWireProfile {
|
||||
originator: "Codex Desktop".to_owned(),
|
||||
codex_version: "0.144.2".to_owned(),
|
||||
@@ -78,7 +78,7 @@ fn wire_profile_should_generate_codex_core_user_agent() {
|
||||
|
||||
assert_eq!(
|
||||
profile.user_agent(),
|
||||
"Codex Desktop/0.144.2 (Mac OS 15.7.1; arm64) unknown"
|
||||
"Codex Desktop/0.144.2 (Mac OS 15.7.1; arm64) unknown (Codex Desktop; 26.707.72221)"
|
||||
);
|
||||
}
|
||||
|
||||
@@ -137,6 +137,7 @@ fn desktop_appcast_should_reject_invalid_version_and_build() {
|
||||
#[tokio::test]
|
||||
async fn wire_profile_release_update_should_change_only_desktop_identity() {
|
||||
let original = wire_profile();
|
||||
let original_user_agent = original.user_agent();
|
||||
let state = CodexWireProfileState::new(original.clone());
|
||||
let service = CodexDesktopReleaseService::new(
|
||||
state.clone(),
|
||||
@@ -155,6 +156,11 @@ async fn wire_profile_release_update_should_change_only_desktop_identity() {
|
||||
assert_eq!(updated.arch, original.arch);
|
||||
assert_eq!(updated.terminal, original.terminal);
|
||||
assert_eq!(updated.verified_at, original.verified_at);
|
||||
assert_ne!(updated.user_agent(), original_user_agent);
|
||||
assert_eq!(
|
||||
updated.user_agent(),
|
||||
"Codex Desktop/0.144.2 (Mac OS 15.7.1; arm64) unknown (Codex Desktop; 26.720.1)"
|
||||
);
|
||||
let status = service.status().snapshot();
|
||||
assert_eq!(status.latest, Some(release("26.720.1", "72001")));
|
||||
assert!(status.checked_at.is_some());
|
||||
@@ -164,6 +170,7 @@ async fn wire_profile_release_update_should_change_only_desktop_identity() {
|
||||
#[tokio::test]
|
||||
async fn cli_release_update_should_change_only_codex_core_identity() {
|
||||
let original = wire_profile();
|
||||
let original_user_agent = original.user_agent();
|
||||
let state = CodexWireProfileState::new(original.clone());
|
||||
let service = CodexCliReleaseService::new(
|
||||
state.clone(),
|
||||
@@ -182,6 +189,11 @@ async fn cli_release_update_should_change_only_codex_core_identity() {
|
||||
assert_eq!(updated.arch, original.arch);
|
||||
assert_eq!(updated.terminal, original.terminal);
|
||||
assert_eq!(updated.verified_at, original.verified_at);
|
||||
assert_ne!(updated.user_agent(), original_user_agent);
|
||||
assert_eq!(
|
||||
updated.user_agent(),
|
||||
"Codex Desktop/0.145.0 (Mac OS 15.7.1; arm64) unknown (Codex Desktop; 26.707.72221)"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
@@ -401,6 +401,41 @@ fn observability_semantics_should_reuse_codex_turn_metadata() {
|
||||
assert!(semantics.compact);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn encoder_should_extract_subagent_kind_from_wire_or_turn_metadata() {
|
||||
let metadata_request = request(Map::from_iter([
|
||||
("model".to_owned(), json!("client-model")),
|
||||
("input".to_owned(), json!("prompt")),
|
||||
(
|
||||
"client_metadata".to_owned(),
|
||||
json!({"x-openai-subagent": "review"}),
|
||||
),
|
||||
]));
|
||||
let metadata_encoded =
|
||||
encode_generate_request(&metadata_request, "gpt-test").expect("encode metadata request");
|
||||
assert_eq!(metadata_encoded.subagent_kind().as_deref(), Some("review"));
|
||||
|
||||
let turn_metadata_payload = ProtocolPayload::json_object(
|
||||
"openai",
|
||||
Map::from_iter([
|
||||
("model".to_owned(), json!("client-model")),
|
||||
("input".to_owned(), json!("prompt")),
|
||||
]),
|
||||
)
|
||||
.expect("OpenAI payload")
|
||||
.with_context(Map::from_iter([(
|
||||
"turn_metadata".to_owned(),
|
||||
Value::String(r#"{"subagent_kind":"worker"}"#.to_owned()),
|
||||
)]));
|
||||
let turn_metadata_request = GenerateRequest::from_protocol_payload(turn_metadata_payload);
|
||||
let turn_metadata_encoded = encode_generate_request(&turn_metadata_request, "gpt-test")
|
||||
.expect("encode turn metadata request");
|
||||
assert_eq!(
|
||||
turn_metadata_encoded.subagent_kind().as_deref(),
|
||||
Some("worker")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn observability_semantics_should_use_the_transparent_openai_payload() {
|
||||
let payload = ProtocolPayload::json_object(
|
||||
|
||||
@@ -163,7 +163,12 @@ async fn fetch_should_use_wham_usage_headers_only() {
|
||||
"unexpected {name} header"
|
||||
);
|
||||
}
|
||||
assert!(headers.get("user-agent").is_some());
|
||||
assert_eq!(
|
||||
headers
|
||||
.get("user-agent")
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("codex_cli_rs/0.144.0 (linux 6.8; x86_64) xterm (codex_cli_rs; 1.0.0)")
|
||||
);
|
||||
let quota_header_names = headers
|
||||
.keys()
|
||||
.map(|name| name.as_str())
|
||||
|
||||
@@ -33,7 +33,7 @@ uuid = { version = "1.23.3", features = ["v4", "v7"] }
|
||||
zeroize = { version = "1.8.2", features = ["derive"] }
|
||||
|
||||
[dev-dependencies]
|
||||
p256 = { version = "0.13.2", features = ["ecdsa", "pkcs8"] }
|
||||
p256 = { version = "0.14.0", features = ["ecdsa", "pkcs8"] }
|
||||
tokio = { version = "1.52.3", features = ["macros", "rt"] }
|
||||
wiremock = "0.6.5"
|
||||
|
||||
|
||||
@@ -34,7 +34,7 @@ use gateway_core::error::StoreErrorKind;
|
||||
use gateway_core::operation::{GenerateRequest, Operation, ProtocolPayload};
|
||||
use gateway_core::provider_ports::{
|
||||
NewOAuthPendingFlow, OAuthPendingBinding, OAuthPendingClaimOutcome, OAuthPendingConsumeOutcome,
|
||||
OAuthPendingFlowPort, OAuthPendingPutOutcome, OAuthPendingReleaseOutcome,
|
||||
OAuthPendingFlowPort, OAuthPendingPutOutcome, OAuthPendingReleaseOutcome, ProviderCooldownPort,
|
||||
ProviderRefreshPolicy, ProviderRuntimePolicyPort, ProviderStoreError, ProviderStoreErrorKind,
|
||||
};
|
||||
use gateway_core::routing::{ProviderKind, UpstreamModelId};
|
||||
@@ -143,6 +143,7 @@ pub(crate) struct XaiAdminProvider {
|
||||
quota: Arc<GrokCredentialQuotaService>,
|
||||
catalog: Arc<GrokCredentialCatalogService>,
|
||||
runtime_policy: Arc<dyn ProviderRuntimePolicyPort>,
|
||||
cooldowns: Arc<dyn ProviderCooldownPort>,
|
||||
}
|
||||
|
||||
pub(crate) struct XaiAdminServices {
|
||||
@@ -154,6 +155,7 @@ pub(crate) struct XaiAdminServices {
|
||||
pub(crate) quota: Arc<GrokCredentialQuotaService>,
|
||||
pub(crate) catalog: Arc<GrokCredentialCatalogService>,
|
||||
pub(crate) runtime_policy: Arc<dyn ProviderRuntimePolicyPort>,
|
||||
pub(crate) cooldowns: Arc<dyn ProviderCooldownPort>,
|
||||
}
|
||||
|
||||
impl XaiAdminProvider {
|
||||
@@ -178,6 +180,7 @@ impl XaiAdminProvider {
|
||||
quota: services.quota,
|
||||
catalog: services.catalog,
|
||||
runtime_policy: services.runtime_policy,
|
||||
cooldowns: services.cooldowns,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -325,7 +328,10 @@ impl ProviderAdmin for XaiAdminProvider {
|
||||
&self.provider_kind
|
||||
}
|
||||
|
||||
async fn account_unavailable(&self, _account_id: &ProviderAccountId) {}
|
||||
async fn account_unavailable(&self, account_id: &ProviderAccountId) {
|
||||
// 删除/禁用账号:清除该账号全部 account/model scoped cooldown key。
|
||||
let _ = self.cooldowns.clear_all(account_id).await;
|
||||
}
|
||||
|
||||
fn connection_test_operation(
|
||||
&self,
|
||||
@@ -913,8 +919,6 @@ fn prepared_create(
|
||||
next_refresh_at: account.next_refresh_at().map(Into::into),
|
||||
enabled: account.enabled(),
|
||||
availability: account.availability(),
|
||||
availability_reason: None,
|
||||
cooldown_until: account.cooldown_until().map(Into::into),
|
||||
availability_observed_at: observed_at,
|
||||
})
|
||||
}
|
||||
@@ -1018,11 +1022,7 @@ fn account_from_record(account: &AccountRecord) -> Result<ProviderAccount, Provi
|
||||
account.upstream_account_id.clone(),
|
||||
account.plan_type.clone(),
|
||||
)
|
||||
.with_runtime_state(
|
||||
account.enabled,
|
||||
account.availability,
|
||||
account.cooldown_until.map(Into::into),
|
||||
)
|
||||
.with_runtime_state(account.enabled, account.availability)
|
||||
.with_refresh_schedule(
|
||||
account.has_refresh_token,
|
||||
account.next_refresh_at.map(Into::into),
|
||||
@@ -1047,6 +1047,8 @@ fn project_quota(
|
||||
observed_at: None,
|
||||
refresh_token_expires_at,
|
||||
windows: Vec::new(),
|
||||
limit_reached: false,
|
||||
rate_limited_until: None,
|
||||
provider_data: None,
|
||||
};
|
||||
};
|
||||
@@ -1087,6 +1089,9 @@ fn project_quota(
|
||||
window_seconds,
|
||||
used_percent: billing.used_percent(),
|
||||
reset_at: billing.period_end().and_then(parse_utc),
|
||||
limit_reached: billing
|
||||
.used_percent()
|
||||
.is_some_and(|used| used.is_finite() && used >= 100.0),
|
||||
local_usage: None,
|
||||
provider_data: Some(ProviderDocument::new(OpaqueProviderData::new(data))),
|
||||
}
|
||||
@@ -1099,6 +1104,7 @@ fn project_quota(
|
||||
window_seconds: Some(crate::GROK_FREE_ROLLING_WINDOW_SECONDS),
|
||||
used_percent: None,
|
||||
reset_at: None,
|
||||
limit_reached: false,
|
||||
local_usage: rolling_usage,
|
||||
provider_data: None,
|
||||
}
|
||||
@@ -1107,6 +1113,10 @@ fn project_quota(
|
||||
observed_at: Some(snapshot.observed_at()),
|
||||
refresh_token_expires_at,
|
||||
windows: vec![window],
|
||||
limit_reached: billing
|
||||
.used_percent()
|
||||
.is_some_and(|used| used.is_finite() && used >= 100.0),
|
||||
rate_limited_until: None,
|
||||
provider_data: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -451,6 +451,9 @@ impl GrokCredentialQuotaService {
|
||||
loaded.account.revision(),
|
||||
document,
|
||||
observed_at.into(),
|
||||
presentation
|
||||
.used_percent()
|
||||
.is_some_and(|used| used.is_finite() && used >= 100.0),
|
||||
)
|
||||
.await
|
||||
.map_err(map_quota_repository_error)?;
|
||||
@@ -482,7 +485,6 @@ impl GrokCredentialQuotaService {
|
||||
GrokCredentialAvailability::QuotaExhausted
|
||||
)
|
||||
.then_some("quota_exhausted".to_owned()),
|
||||
cooldown_until: None,
|
||||
observed_at,
|
||||
})
|
||||
.await
|
||||
@@ -1151,7 +1153,6 @@ fn eligible_catalog_candidate(candidate: &LoadedGrokCredential) -> bool {
|
||||
&& match account.availability() {
|
||||
AccountAvailability::Unknown
|
||||
| AccountAvailability::Ready
|
||||
| AccountAvailability::Cooldown
|
||||
| AccountAvailability::QuotaExhausted => true,
|
||||
AccountAvailability::Expired
|
||||
| AccountAvailability::Banned
|
||||
@@ -1294,8 +1295,15 @@ fn quota_scheduling_signals(billing: &GrokBillingPresentation) -> Option<Account
|
||||
.and_then(|value| DateTime::parse_from_rfc3339(value).ok())
|
||||
.map(|value| SystemTime::from(value.to_utc()))
|
||||
.filter(|reset_at| *reset_at > now);
|
||||
(remaining_rank.is_some() || reset_at.is_some())
|
||||
.then(|| AccountQuotaSignals::new(reset_at, remaining_rank))
|
||||
(remaining_rank.is_some() || reset_at.is_some()).then(|| {
|
||||
AccountQuotaSignals::new(
|
||||
reset_at,
|
||||
remaining_rank,
|
||||
billing
|
||||
.used_percent()
|
||||
.is_some_and(|used| used.is_finite() && used >= 100.0),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn quota_is_exhausted(billing: &GrokBillingPresentation) -> bool {
|
||||
@@ -1343,7 +1351,6 @@ fn quota_refresh_may_update_state(before: &ProviderAccount, current: &ProviderAc
|
||||
.is_some_and(|expires_at| expires_at > SystemTime::now())
|
||||
&& current.revision() == before.revision()
|
||||
&& current.availability() == before.availability()
|
||||
&& current.cooldown_until() == before.cooldown_until()
|
||||
}
|
||||
|
||||
fn quota_refresh_availability(
|
||||
@@ -1357,15 +1364,13 @@ fn quota_refresh_availability(
|
||||
AccountAvailability::QuotaExhausted => {
|
||||
matches!(exhausted, Some(false)).then_some(GrokCredentialAvailability::Ready)
|
||||
}
|
||||
AccountAvailability::Cooldown | AccountAvailability::Unknown => {
|
||||
exhausted.map(|exhausted| {
|
||||
if exhausted {
|
||||
GrokCredentialAvailability::QuotaExhausted
|
||||
} else {
|
||||
GrokCredentialAvailability::Ready
|
||||
}
|
||||
})
|
||||
}
|
||||
AccountAvailability::Unknown => exhausted.map(|exhausted| {
|
||||
if exhausted {
|
||||
GrokCredentialAvailability::QuotaExhausted
|
||||
} else {
|
||||
GrokCredentialAvailability::Ready
|
||||
}
|
||||
}),
|
||||
AccountAvailability::Ready => {
|
||||
matches!(exhausted, Some(true)).then_some(GrokCredentialAvailability::QuotaExhausted)
|
||||
}
|
||||
|
||||
@@ -651,18 +651,11 @@ impl GrokCredentialRefreshService {
|
||||
}
|
||||
|
||||
async fn runtime_cooldown_active(&self, credential: &DueGrokCredential) -> bool {
|
||||
// runtime cooldown 是账号/model 级生命周期,与 credential revision 无关:
|
||||
// 轮换不清除、不因 revision 变化判失效(与 selector 的读取语义一致)。
|
||||
let Ok(Some(cooldown)) = self.cooldowns.read(credential.account_id()).await else {
|
||||
return false;
|
||||
};
|
||||
if cooldown.credential_revision() != credential.credential_revision() {
|
||||
if cooldown.credential_revision() < credential.credential_revision() {
|
||||
let _ = self
|
||||
.cooldowns
|
||||
.clear(credential.account_id(), credential.credential_revision())
|
||||
.await;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
cooldown.until() > SystemTime::now()
|
||||
}
|
||||
|
||||
@@ -705,7 +698,6 @@ impl GrokCredentialRefreshService {
|
||||
expected_revision: credential.credential_revision,
|
||||
availability,
|
||||
availability_reason: Some(reason.to_owned()),
|
||||
cooldown_until: None,
|
||||
observed_at: Utc::now(),
|
||||
})
|
||||
.await
|
||||
@@ -852,8 +844,6 @@ fn account_due(
|
||||
| AccountAvailability::Banned
|
||||
| AccountAvailability::Invalid
|
||||
)
|
||||
&& (availability != AccountAvailability::Cooldown
|
||||
|| account.cooldown_until().is_some_and(|until| until <= now))
|
||||
&& account.next_refresh_at().is_some_and(|next| next <= now)
|
||||
}
|
||||
|
||||
|
||||
@@ -168,11 +168,7 @@ impl GrokCredentialAdmin {
|
||||
input.account.upstream_account_id.clone(),
|
||||
input.account.plan_type.clone(),
|
||||
)
|
||||
.with_runtime_state(
|
||||
input.enabled,
|
||||
input.initial_availability.into(),
|
||||
input.initial_cooldown_until.map(to_system_time),
|
||||
)
|
||||
.with_runtime_state(input.enabled, input.initial_availability.into())
|
||||
.with_refresh_schedule(true, Some(to_system_time(input.next_refresh_at)));
|
||||
Ok(NewProviderAccount {
|
||||
account,
|
||||
@@ -222,7 +218,6 @@ impl GrokCredentialAdmin {
|
||||
enabled: input.enabled,
|
||||
initial_availability: GrokCredentialAvailability::Ready,
|
||||
initial_availability_reason: None,
|
||||
initial_cooldown_until: None,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -434,16 +429,14 @@ impl GrokCredentialRepository {
|
||||
input: &UpdateGrokCredentialState,
|
||||
) -> Result<(), GrokCredentialRepositoryError> {
|
||||
validate_reason(input.availability_reason.as_deref())?;
|
||||
validate_cooldown(input.availability, input.cooldown_until)?;
|
||||
self.ensure_xai_account(&input.account_id).await?;
|
||||
self.store
|
||||
.apply_state_change(AccountStateChange {
|
||||
account_id: input.account_id.clone(),
|
||||
expected_revision: input.expected_revision,
|
||||
availability: input.availability.into(),
|
||||
reason: input.availability_reason.clone(),
|
||||
cooldown_until: input.cooldown_until.map(to_system_time),
|
||||
observed_at: to_system_time(input.observed_at),
|
||||
message: input.availability_reason.clone(),
|
||||
})
|
||||
.await
|
||||
.map_err(map_store_error)
|
||||
@@ -527,6 +520,7 @@ impl GrokCredentialRepository {
|
||||
expected_revision: CredentialRevision,
|
||||
document: serde_json::Map<String, Value>,
|
||||
observed_at: SystemTime,
|
||||
limit_reached: bool,
|
||||
) -> Result<(), GrokCredentialRepositoryError> {
|
||||
let outcome = self
|
||||
.store
|
||||
@@ -535,6 +529,7 @@ impl GrokCredentialRepository {
|
||||
expected_revision,
|
||||
quota: Some(OpaqueProviderData::new(document)),
|
||||
observed_at: Some(observed_at),
|
||||
limit_reached: Some(limit_reached),
|
||||
})
|
||||
.await
|
||||
.map_err(map_store_error)?;
|
||||
@@ -708,8 +703,7 @@ fn validate_create(input: &CreateGrokCredential) -> Result<(), GrokCredentialRep
|
||||
validate_profile(&input.account)?;
|
||||
validate_refresh_schedule(input.next_refresh_at, input.account.access_token_expires_at)?;
|
||||
validate_secret(&input.secret)?;
|
||||
validate_reason(input.initial_availability_reason.as_deref())?;
|
||||
validate_cooldown(input.initial_availability, input.initial_cooldown_until)
|
||||
validate_reason(input.initial_availability_reason.as_deref())
|
||||
}
|
||||
|
||||
fn validate_profile(profile: &GrokAccountProfile) -> Result<(), GrokCredentialRepositoryError> {
|
||||
@@ -816,18 +810,6 @@ fn validate_reason(reason: Option<&str>) -> Result<(), GrokCredentialRepositoryE
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_cooldown(
|
||||
availability: GrokCredentialAvailability,
|
||||
cooldown_until: Option<DateTime<Utc>>,
|
||||
) -> Result<(), GrokCredentialRepositoryError> {
|
||||
if matches!(availability, GrokCredentialAvailability::Cooldown) != cooldown_until.is_some() {
|
||||
return Err(GrokCredentialRepositoryError::InvalidInput(
|
||||
"cooldown_until",
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn ensure_xai(account: &ProviderAccount) -> Result<(), GrokCredentialRepositoryError> {
|
||||
if account.provider().as_str() == XAI_PROVIDER_KIND
|
||||
&& account.authentication_kind() == XAI_AUTHENTICATION_KIND_OAUTH
|
||||
|
||||
@@ -4,7 +4,7 @@ use std::fmt;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, SystemTime};
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use chrono::Utc;
|
||||
use gateway_core::engine::credential::{
|
||||
AccountAvailability, AccountAvailabilityPolicy, AccountCandidate, AccountFeedbackStats,
|
||||
AccountSelectionContext, AccountSelector, ProviderAccount, ProviderAccountId,
|
||||
@@ -300,36 +300,26 @@ impl GrokAccountSessionSelector {
|
||||
for ((account, account_read), model_read) in
|
||||
accounts.iter().zip(account_reads).zip(model_reads)
|
||||
{
|
||||
if let Ok(Some(cooldown)) = account_read {
|
||||
if cooldown.credential_revision() != account.revision() {
|
||||
if cooldown.credential_revision() < account.revision() {
|
||||
let _ = self.cooldowns.clear(account.id(), account.revision()).await;
|
||||
}
|
||||
} else if cooldown.until() > now {
|
||||
insert_runtime_cooldown(
|
||||
&mut cooled,
|
||||
account.id(),
|
||||
cooldown.until(),
|
||||
RuntimeCooldownScope::Account,
|
||||
);
|
||||
}
|
||||
if let Ok(Some(cooldown)) = account_read
|
||||
&& cooldown.until() > now
|
||||
{
|
||||
// 账号级限流冷却跨凭据轮换保留:上游限流不因本地凭据更换解除。
|
||||
insert_runtime_cooldown(
|
||||
&mut cooled,
|
||||
account.id(),
|
||||
cooldown.until(),
|
||||
RuntimeCooldownScope::Account,
|
||||
);
|
||||
}
|
||||
if let Ok(Some(cooldown)) = model_read {
|
||||
if cooldown.credential_revision() != account.revision() {
|
||||
if cooldown.credential_revision() < account.revision() {
|
||||
let _ = self
|
||||
.cooldowns
|
||||
.clear_scoped(account.id(), &model_scope, account.revision())
|
||||
.await;
|
||||
}
|
||||
} else if cooldown.until() > now {
|
||||
insert_runtime_cooldown(
|
||||
&mut cooled,
|
||||
account.id(),
|
||||
cooldown.until(),
|
||||
RuntimeCooldownScope::Model,
|
||||
);
|
||||
}
|
||||
if let Ok(Some(cooldown)) = model_read
|
||||
&& cooldown.until() > now
|
||||
{
|
||||
insert_runtime_cooldown(
|
||||
&mut cooled,
|
||||
account.id(),
|
||||
cooldown.until(),
|
||||
RuntimeCooldownScope::Model,
|
||||
);
|
||||
}
|
||||
}
|
||||
cooled
|
||||
@@ -374,7 +364,6 @@ impl GrokAccountSessionSelector {
|
||||
session: &SelectedGrokSession,
|
||||
availability: GrokCredentialAvailability,
|
||||
reason: &'static str,
|
||||
cooldown_until: Option<DateTime<Utc>>,
|
||||
observed_at: chrono::DateTime<Utc>,
|
||||
) {
|
||||
let update = |expected_revision| UpdateGrokCredentialState {
|
||||
@@ -382,7 +371,6 @@ impl GrokAccountSessionSelector {
|
||||
expected_revision,
|
||||
availability,
|
||||
availability_reason: Some(reason.to_owned()),
|
||||
cooldown_until,
|
||||
observed_at,
|
||||
};
|
||||
match self
|
||||
@@ -467,28 +455,33 @@ impl GrokSessionSelector for GrokAccountSessionSelector {
|
||||
GrokCredentialFailure::QuotaExhausted => (
|
||||
GrokCredentialAvailability::QuotaExhausted,
|
||||
"upstream_quota_exhausted",
|
||||
None,
|
||||
),
|
||||
GrokCredentialFailure::FreeQuotaExhausted => (
|
||||
GrokCredentialAvailability::QuotaExhausted,
|
||||
"upstream_free_quota_exhausted",
|
||||
None,
|
||||
),
|
||||
GrokCredentialFailure::PaymentRequired { .. } => (
|
||||
GrokCredentialAvailability::QuotaExhausted,
|
||||
"upstream_payment_quota_exhausted",
|
||||
None,
|
||||
),
|
||||
GrokCredentialFailure::ModelQuotaExhausted { retry_after, .. } => {
|
||||
// bare 402 没有结构化 quota code,不能证明账号额度耗尽:
|
||||
// 只写短期账号级 runtime cooldown(TransientBackoff),
|
||||
// 不持久化 QuotaExhausted(结构化 QuotaExhausted/FreeQuotaExhausted 走上面)。
|
||||
GrokCredentialFailure::PaymentRequired { retry_after } => {
|
||||
let retry_after = retry_after
|
||||
.unwrap_or(DEFAULT_RATE_LIMIT_COOLDOWN)
|
||||
.min(MAX_RATE_LIMIT_COOLDOWN);
|
||||
self.record_runtime_cooldown(session, retry_after).await;
|
||||
return;
|
||||
}
|
||||
// xAI 免费模型额度按该模型滚动窗口恢复:限流写 model-scoped
|
||||
// runtime cooldown(Redis 跨重启保留),不进入持久化账号状态,
|
||||
// 不阻止该账号服务其他模型。
|
||||
GrokCredentialFailure::ModelQuotaExhausted {
|
||||
upstream_model,
|
||||
retry_after,
|
||||
} => {
|
||||
let retry_after =
|
||||
bounded_cooldown(retry_after, MODEL_QUOTA_COOLDOWN, MAX_MODEL_COOLDOWN);
|
||||
// xAI 免费模型额度按账号滚动窗口恢复;持久化使页面和各调度节点收敛。
|
||||
self.record_runtime_cooldown(session, retry_after).await;
|
||||
(
|
||||
GrokCredentialAvailability::Cooldown,
|
||||
"usage_limit_exhausted",
|
||||
persistent_cooldown_until(observed_at, retry_after),
|
||||
)
|
||||
self.record_model_runtime_cooldown(session, upstream_model, retry_after)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
GrokCredentialFailure::ModelAccessDenied {
|
||||
upstream_model,
|
||||
@@ -508,31 +501,17 @@ impl GrokSessionSelector for GrokAccountSessionSelector {
|
||||
.unwrap_or(DEFAULT_RATE_LIMIT_COOLDOWN)
|
||||
.min(MAX_RATE_LIMIT_COOLDOWN);
|
||||
self.record_runtime_cooldown(session, retry_after).await;
|
||||
(
|
||||
GrokCredentialAvailability::Cooldown,
|
||||
"upstream_rate_limited",
|
||||
persistent_cooldown_until(observed_at, retry_after),
|
||||
)
|
||||
return;
|
||||
}
|
||||
GrokCredentialFailure::StreamInterrupted => {
|
||||
self.record_runtime_cooldown(session, STREAM_INTERRUPTION_COOLDOWN)
|
||||
.await;
|
||||
(
|
||||
GrokCredentialAvailability::Cooldown,
|
||||
"upstream_stream_interrupted",
|
||||
persistent_cooldown_until(observed_at, STREAM_INTERRUPTION_COOLDOWN),
|
||||
)
|
||||
return;
|
||||
}
|
||||
};
|
||||
let (availability, reason, cooldown_until) = persistent;
|
||||
self.persist_account_failure(
|
||||
session,
|
||||
availability,
|
||||
reason,
|
||||
cooldown_until,
|
||||
observed_at,
|
||||
)
|
||||
.await;
|
||||
let (availability, reason) = persistent;
|
||||
self.persist_account_failure(session, availability, reason, observed_at)
|
||||
.await;
|
||||
})
|
||||
}
|
||||
|
||||
@@ -552,9 +531,7 @@ impl GrokSessionSelector for GrokAccountSessionSelector {
|
||||
|| !account.enabled()
|
||||
|| !matches!(
|
||||
account.availability(),
|
||||
AccountAvailability::Unknown
|
||||
| AccountAvailability::Cooldown
|
||||
| AccountAvailability::Expired
|
||||
AccountAvailability::Unknown | AccountAvailability::Expired
|
||||
)
|
||||
{
|
||||
return;
|
||||
@@ -566,7 +543,6 @@ impl GrokSessionSelector for GrokAccountSessionSelector {
|
||||
expected_revision: account.revision(),
|
||||
availability: GrokCredentialAvailability::Ready,
|
||||
availability_reason: None,
|
||||
cooldown_until: None,
|
||||
observed_at: Utc::now(),
|
||||
})
|
||||
.await
|
||||
@@ -600,15 +576,6 @@ fn bounded_cooldown(
|
||||
.min(maximum)
|
||||
}
|
||||
|
||||
fn persistent_cooldown_until(
|
||||
observed_at: DateTime<Utc>,
|
||||
duration: Duration,
|
||||
) -> Option<DateTime<Utc>> {
|
||||
chrono::Duration::from_std(duration)
|
||||
.ok()
|
||||
.map(|duration| observed_at + duration)
|
||||
}
|
||||
|
||||
fn insert_runtime_cooldown(
|
||||
cooldowns: &mut std::collections::BTreeMap<ProviderAccountId, RuntimeCooldown>,
|
||||
account_id: &ProviderAccountId,
|
||||
|
||||
@@ -64,11 +64,13 @@ impl fmt::Debug for GrokAccountProfile {
|
||||
}
|
||||
|
||||
/// 与 `provider_accounts.availability` 一一对应的 xAI 状态。
|
||||
///
|
||||
/// 限流(429/滚动窗口)不进入持久化状态:由 `ProviderCooldownPort`(Redis 跨重启
|
||||
/// 保留)驱动调度排除,账号状态保持 `Ready`。
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum GrokCredentialAvailability {
|
||||
Unknown,
|
||||
Ready,
|
||||
Cooldown,
|
||||
QuotaExhausted,
|
||||
Expired,
|
||||
Banned,
|
||||
@@ -81,7 +83,6 @@ impl From<GrokCredentialAvailability> for AccountAvailability {
|
||||
match value {
|
||||
GrokCredentialAvailability::Unknown => Self::Unknown,
|
||||
GrokCredentialAvailability::Ready => Self::Ready,
|
||||
GrokCredentialAvailability::Cooldown => Self::Cooldown,
|
||||
GrokCredentialAvailability::QuotaExhausted => Self::QuotaExhausted,
|
||||
GrokCredentialAvailability::Expired => Self::Expired,
|
||||
GrokCredentialAvailability::Banned => Self::Banned,
|
||||
@@ -100,7 +101,6 @@ pub struct CreateGrokCredential {
|
||||
pub enabled: bool,
|
||||
pub initial_availability: GrokCredentialAvailability,
|
||||
pub initial_availability_reason: Option<String>,
|
||||
pub initial_cooldown_until: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
impl fmt::Debug for CreateGrokCredential {
|
||||
@@ -243,7 +243,6 @@ pub struct UpdateGrokCredentialState {
|
||||
pub expected_revision: CredentialRevision,
|
||||
pub availability: GrokCredentialAvailability,
|
||||
pub availability_reason: Option<String>,
|
||||
pub cooldown_until: Option<DateTime<Utc>>,
|
||||
pub observed_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
|
||||
@@ -171,7 +171,7 @@ pub async fn initialize(
|
||||
refresher,
|
||||
Arc::clone(&catalog),
|
||||
Arc::clone(&leases),
|
||||
cooldowns,
|
||||
cooldowns.clone(),
|
||||
ports.credential_state(),
|
||||
Arc::clone(&runtime_policy),
|
||||
));
|
||||
@@ -200,6 +200,7 @@ pub async fn initialize(
|
||||
quota: Arc::clone("a),
|
||||
catalog: Arc::clone(&catalog),
|
||||
runtime_policy,
|
||||
cooldowns: Arc::clone(&cooldowns),
|
||||
},
|
||||
cli_release_status,
|
||||
));
|
||||
|
||||
@@ -472,9 +472,8 @@ fn account_record(account: &ProviderAccount) -> AccountRecord {
|
||||
next_refresh_at: account.next_refresh_at().map(DateTime::<Utc>::from),
|
||||
enabled: account.enabled(),
|
||||
availability: account.availability(),
|
||||
availability_reason: None,
|
||||
cooldown_until: account.cooldown_until().map(DateTime::<Utc>::from),
|
||||
availability_observed_at: now,
|
||||
last_error_message: None,
|
||||
quota_observed_at: None,
|
||||
created_at: now,
|
||||
updated_at: now,
|
||||
@@ -504,6 +503,7 @@ impl ProviderLeasePort for TestLeases {
|
||||
last_started_at: None,
|
||||
quota_reset_at: None,
|
||||
quota_remaining_rank: None,
|
||||
quota_limit_reached: false,
|
||||
failure_rate_basis_points: None,
|
||||
first_output_latency_ms: None,
|
||||
},
|
||||
@@ -673,6 +673,13 @@ impl ProviderCooldownPort for TestCooldown {
|
||||
) -> BoxFuture<'a, Result<bool, ProviderStoreError>> {
|
||||
Box::pin(async { Ok(false) })
|
||||
}
|
||||
|
||||
fn clear_all<'a>(
|
||||
&'a self,
|
||||
_account_id: &'a ProviderAccountId,
|
||||
) -> BoxFuture<'a, Result<bool, ProviderStoreError>> {
|
||||
Box::pin(async { Ok(false) })
|
||||
}
|
||||
}
|
||||
|
||||
struct TestRuntimePolicy;
|
||||
|
||||
@@ -201,15 +201,13 @@ async fn set_account_state(
|
||||
store: &MemoryProviderAccountStore,
|
||||
id: &ProviderAccountId,
|
||||
availability: AccountAvailability,
|
||||
cooldown_until: Option<SystemTime>,
|
||||
) {
|
||||
store
|
||||
.apply_state_change(AccountStateChange {
|
||||
message: None,
|
||||
account_id: id.clone(),
|
||||
expected_revision: CredentialRevision::new(1).expect("revision"),
|
||||
availability,
|
||||
reason: Some("test isolation".to_owned()),
|
||||
cooldown_until,
|
||||
observed_at: SystemTime::now(),
|
||||
})
|
||||
.await
|
||||
@@ -488,11 +486,10 @@ async fn quota_exhausted_account_remains_eligible_for_catalog_discovery() {
|
||||
let id = account_id("quota-exhausted");
|
||||
store
|
||||
.apply_state_change(AccountStateChange {
|
||||
message: None,
|
||||
account_id: id,
|
||||
expected_revision: CredentialRevision::new(1).expect("revision"),
|
||||
availability: AccountAvailability::QuotaExhausted,
|
||||
reason: Some("quota exhausted".to_owned()),
|
||||
cooldown_until: None,
|
||||
observed_at: SystemTime::now(),
|
||||
})
|
||||
.await
|
||||
@@ -686,18 +683,14 @@ async fn quota_refresh_persists_dynamic_provider_document_and_projects_known_fie
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn recovered_quota_refresh_releases_existing_quota_and_cooldown_isolation() {
|
||||
for (suffix, availability, cooldown_until) in [
|
||||
("recover-quota", AccountAvailability::QuotaExhausted, None),
|
||||
(
|
||||
"recover-cooldown",
|
||||
AccountAvailability::Cooldown,
|
||||
SystemTime::now().checked_add(Duration::from_secs(300)),
|
||||
),
|
||||
async fn recovered_quota_refresh_releases_existing_quota_isolation() {
|
||||
for (suffix, availability) in [
|
||||
("recover-quota", AccountAvailability::QuotaExhausted),
|
||||
("recover-ready", AccountAvailability::Ready),
|
||||
] {
|
||||
let (store, repository) = repository_with_accounts(&[(suffix, suffix)]).await;
|
||||
let id = account_id(suffix);
|
||||
set_account_state(&store, &id, availability, cooldown_until).await;
|
||||
set_account_state(&store, &id, availability).await;
|
||||
crate::support::grok_quota_service(
|
||||
repository,
|
||||
QueueBillingTransport::success(br#"{"config":{"creditUsagePercent":25}}"#),
|
||||
@@ -708,7 +701,6 @@ async fn recovered_quota_refresh_releases_existing_quota_and_cooldown_isolation(
|
||||
|
||||
let account = store.account(&id).expect("account");
|
||||
assert_eq!(account.availability(), AccountAvailability::Ready);
|
||||
assert_eq!(account.cooldown_until(), None);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -717,7 +709,7 @@ async fn authoritative_exhaustion_preserves_quota_exhausted_state() {
|
||||
let (store, repository) =
|
||||
repository_with_accounts(&[("still-exhausted", "still-exhausted")]).await;
|
||||
let id = account_id("still-exhausted");
|
||||
set_account_state(&store, &id, AccountAvailability::QuotaExhausted, None).await;
|
||||
set_account_state(&store, &id, AccountAvailability::QuotaExhausted).await;
|
||||
|
||||
crate::support::grok_quota_service(
|
||||
repository,
|
||||
@@ -742,7 +734,7 @@ async fn recovered_quota_does_not_clear_terminal_account_states() {
|
||||
] {
|
||||
let (store, repository) = repository_with_accounts(&[(suffix, suffix)]).await;
|
||||
let id = account_id(suffix);
|
||||
set_account_state(&store, &id, availability, None).await;
|
||||
set_account_state(&store, &id, availability).await;
|
||||
crate::support::grok_quota_service(
|
||||
repository,
|
||||
QueueBillingTransport::success(br#"{"config":{"creditUsagePercent":10}}"#),
|
||||
@@ -759,20 +751,16 @@ async fn recovered_quota_does_not_clear_terminal_account_states() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn quota_refresh_does_not_overwrite_a_newer_cooldown() {
|
||||
async fn quota_refresh_preserves_ready_availability_across_concurrent_state_write() {
|
||||
let (store, repository) = repository_with_accounts(&[("new-cooldown", "new-cooldown")]).await;
|
||||
let id = account_id("new-cooldown");
|
||||
let new_cooldown = SystemTime::now()
|
||||
.checked_add(Duration::from_secs(600))
|
||||
.expect("cooldown time");
|
||||
let transport = Arc::new(MutatingBillingTransport {
|
||||
store: Arc::clone(&store),
|
||||
mutation: Mutex::new(Some(BillingMutation::State(AccountStateChange {
|
||||
message: None,
|
||||
account_id: id.clone(),
|
||||
expected_revision: CredentialRevision::new(1).expect("revision"),
|
||||
availability: AccountAvailability::Cooldown,
|
||||
reason: Some("new rate limit".to_owned()),
|
||||
cooldown_until: Some(new_cooldown),
|
||||
availability: AccountAvailability::Ready,
|
||||
observed_at: SystemTime::now(),
|
||||
}))),
|
||||
body: br#"{"config":{"creditUsagePercent":10}}"#.to_vec(),
|
||||
@@ -781,11 +769,10 @@ async fn quota_refresh_does_not_overwrite_a_newer_cooldown() {
|
||||
crate::support::grok_quota_service(repository, transport)
|
||||
.refresh_account(&id)
|
||||
.await
|
||||
.expect("refresh around concurrent cooldown");
|
||||
.expect("refresh around concurrent state write");
|
||||
|
||||
let account = store.account(&id).expect("account");
|
||||
assert_eq!(account.availability(), AccountAvailability::Cooldown);
|
||||
assert_eq!(account.cooldown_until(), Some(new_cooldown));
|
||||
assert_eq!(account.availability(), AccountAvailability::Ready);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -878,7 +865,7 @@ async fn non_authoritative_quota_refresh_does_not_clear_existing_quota_exhaustio
|
||||
let (store, repository) =
|
||||
repository_with_accounts(&[("free-quota", "subject-free-quota")]).await;
|
||||
let id = account_id("free-quota");
|
||||
set_account_state(&store, &id, AccountAvailability::QuotaExhausted, None).await;
|
||||
set_account_state(&store, &id, AccountAvailability::QuotaExhausted).await;
|
||||
let transport = QueueBillingTransport::success(
|
||||
br#"{"config":{"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY","start":"2026-07-15T00:00:00Z","end":"2026-07-22T00:00:00Z"},"onDemandCap":{"val":0},"onDemandUsed":{"val":0},"prepaidBalance":{"val":0}}}"#,
|
||||
);
|
||||
@@ -935,6 +922,7 @@ async fn quota_read_rejects_corrupt_provider_document() {
|
||||
expected_revision: CredentialRevision::new(1).expect("revision"),
|
||||
quota: Some(OpaqueProviderData::new(document)),
|
||||
observed_at: Some(SystemTime::now()),
|
||||
limit_reached: None,
|
||||
})
|
||||
.await
|
||||
.expect("seed corrupt quota");
|
||||
|
||||
@@ -219,44 +219,18 @@ async fn state_update_uses_credential_revision_fence() {
|
||||
let (store, repository) = repository();
|
||||
let input = create_input("state", "subject-state");
|
||||
seed_input(&store, &input).await.expect("create");
|
||||
let cooldown_until = Utc::now() + chrono::Duration::minutes(1);
|
||||
repository
|
||||
.update_state(&UpdateGrokCredentialState {
|
||||
account_id: input.account_id.clone(),
|
||||
expected_revision: CredentialRevision::new(1).expect("revision"),
|
||||
availability: GrokCredentialAvailability::Cooldown,
|
||||
availability_reason: Some("rate_limited".to_owned()),
|
||||
cooldown_until: Some(cooldown_until),
|
||||
availability: GrokCredentialAvailability::Ready,
|
||||
availability_reason: Some("upstream_rate_limited".to_owned()),
|
||||
observed_at: Utc::now(),
|
||||
})
|
||||
.await
|
||||
.expect("state update");
|
||||
let account = store.account(&input.account_id).expect("account");
|
||||
assert_eq!(account.availability(), AccountAvailability::Cooldown);
|
||||
assert!(account.cooldown_until().is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cooldown_requires_matching_deadline() {
|
||||
let (store, repository) = repository();
|
||||
let input = create_input("bad-cooldown", "subject-cooldown");
|
||||
seed_input(&store, &input).await.expect("create");
|
||||
let result = repository
|
||||
.update_state(&UpdateGrokCredentialState {
|
||||
account_id: input.account_id,
|
||||
expected_revision: CredentialRevision::new(1).expect("revision"),
|
||||
availability: GrokCredentialAvailability::Cooldown,
|
||||
availability_reason: None,
|
||||
cooldown_until: None,
|
||||
observed_at: Utc::now(),
|
||||
})
|
||||
.await;
|
||||
assert_eq!(
|
||||
result,
|
||||
Err(GrokCredentialRepositoryError::InvalidInput(
|
||||
"cooldown_until"
|
||||
))
|
||||
);
|
||||
assert_eq!(account.availability(), AccountAvailability::Ready);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
@@ -420,7 +420,7 @@ struct TestKey {
|
||||
impl TestKey {
|
||||
fn new(seed: u8, kid: &'static str) -> Self {
|
||||
let signing_key = SigningKey::from_slice(&[seed; 32]).expect("valid deterministic scalar");
|
||||
let point = signing_key.verifying_key().to_encoded_point(false);
|
||||
let point = signing_key.verifying_key().to_sec1_point(false);
|
||||
let x = URL_SAFE_NO_PAD.encode(point.x().expect("P-256 x coordinate"));
|
||||
let y = URL_SAFE_NO_PAD.encode(point.y().expect("P-256 y coordinate"));
|
||||
Self {
|
||||
|
||||
@@ -10,8 +10,8 @@ use gateway_core::engine::credential::{
|
||||
ProviderAccountId, ProviderAccountStore, ProviderAccountUpdate,
|
||||
};
|
||||
use gateway_core::provider_ports::{
|
||||
ProviderCredentialState, ProviderCredentialStatePort, ProviderLeaseAcquisition,
|
||||
ProviderLeasePort, ProviderLeaseRequest, ProviderStoreError,
|
||||
ProviderCooldown, ProviderCooldownPort, ProviderCredentialState, ProviderCredentialStatePort,
|
||||
ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest, ProviderStoreError,
|
||||
};
|
||||
use provider_xai::{
|
||||
GrokCredentialCatalogCache, GrokCredentialRecovery, GrokCredentialRecoveryOutcome,
|
||||
@@ -143,6 +143,48 @@ async fn fixture(
|
||||
fixture_many([input], responses, lease_available).await
|
||||
}
|
||||
|
||||
/// 带外部注入 cooldown 端口的 fixture(AUD-05 组合测试用)。
|
||||
#[allow(clippy::type_complexity)]
|
||||
async fn fixture_with_cooldowns(
|
||||
input: provider_xai::CreateGrokCredential,
|
||||
responses: impl IntoIterator<Item = Result<GrokRefreshTokens, GrokRefreshFailure>>,
|
||||
lease_available: bool,
|
||||
cooldowns: Arc<MemoryCooldownPort>,
|
||||
) -> (
|
||||
Arc<MemoryProviderAccountStore>,
|
||||
GrokCredentialRepository,
|
||||
Arc<QueueRefresher>,
|
||||
GrokCredentialRefreshService,
|
||||
Arc<MemoryCooldownPort>,
|
||||
) {
|
||||
let store = MemoryProviderAccountStore::shared();
|
||||
let account_store: Arc<dyn ProviderAccountStore> = store.clone();
|
||||
let repository = GrokCredentialRepository::new(account_store);
|
||||
seed_input(&store, &input).await.expect("create account");
|
||||
let refresher = QueueRefresher::new(responses);
|
||||
let refresher_port: Arc<dyn GrokCredentialRefresher> = refresher.clone();
|
||||
let cache: Arc<dyn GrokCredentialCatalogCache> = MemoryGrokCatalogCache::shared();
|
||||
let catalog = Arc::new(crate::support::grok_catalog_service(
|
||||
repository.clone(),
|
||||
Arc::new(StaticCatalogTransport),
|
||||
cache,
|
||||
));
|
||||
let leases = Arc::new(TestRefreshLeases {
|
||||
available: lease_available,
|
||||
calls: AtomicUsize::new(0),
|
||||
});
|
||||
let service = GrokCredentialRefreshService::new(
|
||||
repository.clone(),
|
||||
refresher_port,
|
||||
catalog,
|
||||
leases,
|
||||
cooldowns.clone(),
|
||||
Arc::new(CountingCredentialState::default()),
|
||||
runtime_policy(),
|
||||
);
|
||||
(store, repository, refresher, service, cooldowns)
|
||||
}
|
||||
|
||||
/// 真实累加连续失败计数的测试 double,用于断言退避的指数增长与成功清零。
|
||||
#[derive(Default)]
|
||||
struct CountingCredentialState {
|
||||
@@ -732,3 +774,54 @@ async fn malformed_account_refresh_does_not_stop_later_accounts() {
|
||||
2
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn runtime_cooldown_survives_credential_rotation_and_blocks_refresh() {
|
||||
// cooldown 写于 revision N,轮换到 N+1 后 refresh worker 不得清除
|
||||
// 或失效它;cooldown 活跃期间账号被跳过刷新。
|
||||
let input = create_input(
|
||||
"cooldown-survives-rotation",
|
||||
"subject-cooldown-survives-rotation",
|
||||
);
|
||||
let id = input.account_id.clone();
|
||||
let cooldowns = Arc::new(MemoryCooldownPort::default());
|
||||
let (store, _, _, service, cooldowns) = fixture_with_cooldowns(
|
||||
input,
|
||||
[Ok(success_tokens(Some("rotated-refresh")))],
|
||||
true,
|
||||
cooldowns,
|
||||
)
|
||||
.await;
|
||||
// 先写 account cooldown(revision 1,未来到期)。
|
||||
let until = SystemTime::now() + Duration::from_secs(3600);
|
||||
cooldowns
|
||||
.put_if_later(ProviderCooldown::new(
|
||||
id.clone(),
|
||||
CredentialRevision::new(1).expect("revision"),
|
||||
until,
|
||||
))
|
||||
.await
|
||||
.expect("write cooldown");
|
||||
|
||||
// 轮换到 revision 2。
|
||||
let outcome = service
|
||||
.recover_unauthorized(&id, CredentialRevision::new(1).expect("revision"))
|
||||
.await;
|
||||
assert_eq!(outcome, GrokCredentialRecoveryOutcome::Recovered);
|
||||
assert_eq!(store.account(&id).expect("account").revision().get(), 2);
|
||||
|
||||
// cooldown 仍有效(轮换不清除)。
|
||||
assert!(
|
||||
cooldowns
|
||||
.cooldown(&id)
|
||||
.is_some_and(|cooldown| cooldown.until() > SystemTime::now()),
|
||||
"rotation must not clear account runtime cooldown"
|
||||
);
|
||||
// cooldown 活跃期间 refresh worker 跳过该账号(不清除、不刷新)。
|
||||
let outcomes = service.refresh_due().await.expect("refresh due");
|
||||
assert!(
|
||||
outcomes.is_empty(),
|
||||
"cooldown-active account must be skipped by refresh worker"
|
||||
);
|
||||
assert_eq!(store.account(&id).expect("account").revision().get(), 2);
|
||||
}
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user