Compare commits

...
30 Commits
Author SHA1 Message Date
ZyyandCodex 35f28e8aab chore: bump version to 3.2.0
Co-authored-by: Codex <noreply@openai.com>
2026-08-13 22:43:48 +08:00
ZyyandCodex 65645c7403 fix(accounts): gate quota recovery and add manual reset
Co-authored-by: Codex <noreply@openai.com>
2026-08-13 22:43:03 +08:00
ZyyandCodex 71159d5633 refactor(account-groups): remove unused members endpoint
Co-authored-by: Codex <noreply@openai.com>
2026-08-13 17:28:43 +08:00
ZyyandCodex f2ebb09e34 docs: describe group routing and account scheduling
Co-authored-by: Codex <noreply@openai.com>
2026-08-13 16:58:36 +08:00
ZyyandCodex 07acc67ce6 feat(frontend): add account group management
Co-authored-by: Codex <noreply@openai.com>
2026-08-13 16:58:36 +08:00
ZyyandCodex 01ad8bba8a feat(routing): add account groups and weighted scheduling
Co-authored-by: Codex <noreply@openai.com>
2026-08-13 16:58:36 +08:00
ZyyandCodex 986a128d15 fix(observability): separate live status from history range
Co-authored-by: Codex <noreply@openai.com>
2026-08-13 00:26:10 +08:00
ZyyandCodex 6cad2bb6a1 fix(openai): align account import with official token claims
Co-authored-by: Codex <noreply@openai.com>
2026-08-12 23:51:26 +08:00
ZyyandCodex 2b8b3de0e4 refactor(accounts): unify five-state account status
Co-authored-by: Codex <noreply@openai.com>
2026-08-12 23:50:31 +08:00
ZyyandCodex 4127886584 chore: bump version to 3.1.21
Co-authored-by: Codex <noreply@openai.com>
2026-08-11 17:42:54 +08:00
ZyyandCodex ee3e5db6c2 fix(accounts): isolate model-specific quota usage
Co-authored-by: Codex <noreply@openai.com>
2026-08-11 17:42:02 +08:00
ZyyandCodex 153cb67b9c refactor(login): refine panel details
Co-authored-by: Codex <noreply@openai.com>
2026-08-11 17:21:31 +08:00
ZyyandCodex f370d355b6 fix(accounts): align quota ordering and summaries
Co-authored-by: Codex <noreply@openai.com>
2026-08-11 17:21:23 +08:00
ZyyandCodex a31dad992c fix(frontend): stabilize button loading indicator
Co-authored-by: Codex <noreply@openai.com>
2026-08-11 17:21:16 +08:00
ZyyandCodex 55a88f5f6b chore: bump version to 3.1.20
Co-authored-by: Codex <noreply@openai.com>
2026-08-11 15:13:40 +08:00
ZyyandCodex 5d90a92582 fix(logging): rename OAuth recovery log file
Co-authored-by: Codex <noreply@openai.com>
2026-08-11 15:06:40 +08:00
ZyyandCodex 93abe8efcd fix(logging): generalize OAuth recovery logs
Co-authored-by: Codex <noreply@openai.com>
2026-08-11 14:55:59 +08:00
ZyyandCodex f4b75c42ed fix(openai): isolate quota and recovery signals
Co-authored-by: Codex <noreply@openai.com>
2026-08-11 13:37:16 +08:00
ZyyandCodex 36ff038c78 chore: bump version to 3.1.19
Co-authored-by: Codex <noreply@openai.com>
2026-08-11 12:17:28 +08:00
ZyyandCodex 545009590c fix(openai): persist authoritative quota exhaustion
Co-authored-by: Codex <noreply@openai.com>
2026-08-11 12:16:34 +08:00
ZyyandCodex 5b7b9b21ee fix(openai): require low usage for quota recovery
Co-authored-by: Codex <noreply@openai.com>
2026-08-10 15:55:58 +08:00
ZyyandCodex 44a2a87b69 fix(frontend): center button loading spinner
Co-authored-by: Codex <noreply@openai.com>
2026-08-10 15:55:58 +08:00
ZyyandCodex 184e88083b chore: bump version to 3.1.18
Co-authored-by: Codex <noreply@openai.com>
2026-08-10 12:25:42 +08:00
ZyyandCodex 53b40aea46 chore(deps): ignore fingerprint-sensitive updates
Co-authored-by: Codex <noreply@openai.com>
2026-08-10 12:23:45 +08:00
ZyyandCodex aa328b6dbf fix(openai): restore official TLS fingerprint pins
Co-authored-by: Codex <noreply@openai.com>
2026-08-10 12:23:36 +08:00
dependabot[bot]andZyy 64231ff403 chore(deps): bump the minor-and-patch group in /backend with 4 updates (#37)
Bumps the minor-and-patch group in /backend with 4 updates: [base64](https://github.com/marshallpierce/rust-base64), [aws-sdk-s3](https://github.com/awslabs/aws-sdk-rust), [hyper-rustls](https://github.com/rustls/hyper-rustls) and [rustls](https://github.com/rustls/rustls).


Updates `base64` from 0.23.0 to 0.23.1
- [Changelog](https://github.com/marshallpierce/rust-base64/blob/master/RELEASE-NOTES.md)
- [Commits](https://github.com/marshallpierce/rust-base64/compare/v0.23.0...v0.23.1)

Updates `aws-sdk-s3` from 1.140.0 to 1.141.0
- [Release notes](https://github.com/awslabs/aws-sdk-rust/releases)
- [Commits](https://github.com/awslabs/aws-sdk-rust/commits)

Updates `hyper-rustls` from 0.27.7 to 0.27.9
- [Release notes](https://github.com/rustls/hyper-rustls/releases)
- [Commits](https://github.com/rustls/hyper-rustls/compare/v/0.27.7...v/0.27.9)

Updates `rustls` from 0.23.36 to 0.23.43
- [Release notes](https://github.com/rustls/rustls/releases)
- [Changelog](https://github.com/rustls/rustls/blob/main/CHANGELOG.md)
- [Commits](https://github.com/rustls/rustls/compare/v/0.23.36...v/0.23.43)

---
updated-dependencies:
- dependency-name: base64
  dependency-version: 0.23.1
  dependency-type: direct:production
  update-type: version-update:semver-patch
  dependency-group: minor-and-patch
- dependency-name: aws-sdk-s3
  dependency-version: 1.141.0
  dependency-type: direct:production
  update-type: version-update:semver-minor
  dependency-group: minor-and-patch
- dependency-name: hyper-rustls
  dependency-version: 0.27.9
  dependency-type: direct:production
  update-type: version-update:semver-patch
  dependency-group: minor-and-patch
- dependency-name: rustls
  dependency-version: 0.23.43
  dependency-type: direct:production
  update-type: version-update:semver-patch
  dependency-group: minor-and-patch
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: Zyy <39424122+zyycn@users.noreply.github.com>
2026-08-10 12:05:00 +08:00
dependabot[bot]andZyy d7b80a2066 chore(deps): bump the minor-and-patch group with 5 updates (#36)
Bumps the minor-and-patch group with 5 updates:

| Package | From | To |
| --- | --- | --- |
| [github/codeql-action/upload-sarif](https://github.com/github/codeql-action) | `4.37.4` | `4.37.6` |
| [Swatinem/rust-cache](https://github.com/swatinem/rust-cache) | `2.9.1` | `2.9.2` |
| [pnpm/action-setup](https://github.com/pnpm/action-setup) | `6.0.9` | `6.0.10` |
| [dorny/paths-filter](https://github.com/dorny/paths-filter) | `4.0.2` | `4.0.3` |
| [actions/attest](https://github.com/actions/attest) | `4.2.1` | `4.2.2` |


Updates `github/codeql-action/upload-sarif` from 4.37.4 to 4.37.6
- [Release notes](https://github.com/github/codeql-action/releases)
- [Changelog](https://github.com/github/codeql-action/blob/main/CHANGELOG.md)
- [Commits](https://github.com/github/codeql-action/compare/f205ea1c3313d32999d8d6a48b4f6530d4437b38...5595ccaf912efad79be6eef63a5619ff05969be3)

Updates `Swatinem/rust-cache` from 2.9.1 to 2.9.2
- [Release notes](https://github.com/swatinem/rust-cache/releases)
- [Changelog](https://github.com/Swatinem/rust-cache/blob/master/CHANGELOG.md)
- [Commits](https://github.com/swatinem/rust-cache/compare/c19371144df3bb44fab255c43d04cbc2ab54d1c4...6323deb102c322ba6fcbdcafc7e3dddab59af2b6)

Updates `pnpm/action-setup` from 6.0.9 to 6.0.10
- [Release notes](https://github.com/pnpm/action-setup/releases)
- [Commits](https://github.com/pnpm/action-setup/compare/0ebf47130e4866e96fce0953f49152a61190b271...0977fd99725f1db4007ccb2928dbb4e90d06cc86)

Updates `dorny/paths-filter` from 4.0.2 to 4.0.3
- [Release notes](https://github.com/dorny/paths-filter/releases)
- [Changelog](https://github.com/dorny/paths-filter/blob/master/CHANGELOG.md)
- [Commits](https://github.com/dorny/paths-filter/compare/7b450fff21473bca461d4b92ce414b9d0420d706...ceb8a2b8f2d89434be7ff52d3de7ec3738c5cc9d)

Updates `actions/attest` from 4.2.1 to 4.2.2
- [Release notes](https://github.com/actions/attest/releases)
- [Changelog](https://github.com/actions/attest/blob/main/RELEASE.md)
- [Commits](https://github.com/actions/attest/compare/508db95dd578ae2727ebd6217d5ba78e4fbda05d...1e69f48acb82d1966a394da916b4c1698aa569d6)

---
updated-dependencies:
- dependency-name: github/codeql-action/upload-sarif
  dependency-version: 4.37.6
  dependency-type: direct:production
  update-type: version-update:semver-patch
  dependency-group: minor-and-patch
- dependency-name: Swatinem/rust-cache
  dependency-version: 2.9.2
  dependency-type: direct:production
  update-type: version-update:semver-patch
  dependency-group: minor-and-patch
- dependency-name: pnpm/action-setup
  dependency-version: 6.0.10
  dependency-type: direct:production
  update-type: version-update:semver-patch
  dependency-group: minor-and-patch
- dependency-name: dorny/paths-filter
  dependency-version: 4.0.3
  dependency-type: direct:production
  update-type: version-update:semver-patch
  dependency-group: minor-and-patch
- dependency-name: actions/attest
  dependency-version: 4.2.2
  dependency-type: direct:production
  update-type: version-update:semver-patch
  dependency-group: minor-and-patch
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: Zyy <39424122+zyycn@users.noreply.github.com>
2026-08-10 12:01:38 +08:00
dependabot[bot] 08d63ea3f3 chore(deps): bump debian from 7b140f3 to abd67ff in /deploy (#39)
Bumps debian from `7b140f3` to `abd67ff`.

---
updated-dependencies:
- dependency-name: debian
  dependency-version: bookworm-slim
  dependency-type: direct:production
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2026-08-10 11:58:47 +08:00
dependabot[bot] ac9c4c85ad chore(deps): bump the minor-and-patch group in /frontend with 8 updates (#38)
Bumps the minor-and-patch group in /frontend with 8 updates:

| Package | From | To |
| --- | --- | --- |
| [@lucide/vue](https://github.com/lucide-icons/lucide/tree/HEAD/packages/vue) | `1.28.0` | `1.30.0` |
| [dompurify](https://github.com/cure53/DOMPurify) | `3.4.12` | `3.4.13` |
| [marked](https://github.com/markedjs/marked) | `18.0.7` | `18.0.9` |
| [vue](https://github.com/vuejs/core) | `3.5.40` | `3.5.41` |
| [@antfu/eslint-config](https://github.com/antfu/eslint-config) | `9.2.0` | `9.3.0` |
| [@types/node](https://github.com/DefinitelyTyped/DefinitelyTyped/tree/HEAD/types/node) | `26.1.2` | `26.2.0` |
| [eslint](https://github.com/eslint/eslint) | `10.8.0` | `10.8.1` |
| [vite](https://github.com/vitejs/vite/tree/HEAD/packages/vite) | `8.2.0` | `8.2.1` |


Updates `@lucide/vue` from 1.28.0 to 1.30.0
- [Release notes](https://github.com/lucide-icons/lucide/releases)
- [Commits](https://github.com/lucide-icons/lucide/commits/1.30.0/packages/vue)

Updates `dompurify` from 3.4.12 to 3.4.13
- [Release notes](https://github.com/cure53/DOMPurify/releases)
- [Commits](https://github.com/cure53/DOMPurify/compare/3.4.12...3.4.13)

Updates `marked` from 18.0.7 to 18.0.9
- [Release notes](https://github.com/markedjs/marked/releases)
- [Commits](https://github.com/markedjs/marked/compare/v18.0.7...v18.0.9)

Updates `vue` from 3.5.40 to 3.5.41
- [Release notes](https://github.com/vuejs/core/releases)
- [Changelog](https://github.com/vuejs/core/blob/main/CHANGELOG.md)
- [Commits](https://github.com/vuejs/core/compare/v3.5.40...v3.5.41)

Updates `@antfu/eslint-config` from 9.2.0 to 9.3.0
- [Release notes](https://github.com/antfu/eslint-config/releases)
- [Commits](https://github.com/antfu/eslint-config/compare/v9.2.0...v9.3.0)

Updates `@types/node` from 26.1.2 to 26.2.0
- [Release notes](https://github.com/DefinitelyTyped/DefinitelyTyped/releases)
- [Commits](https://github.com/DefinitelyTyped/DefinitelyTyped/commits/HEAD/types/node)

Updates `eslint` from 10.8.0 to 10.8.1
- [Release notes](https://github.com/eslint/eslint/releases)
- [Commits](https://github.com/eslint/eslint/compare/v10.8.0...v10.8.1)

Updates `vite` from 8.2.0 to 8.2.1
- [Release notes](https://github.com/vitejs/vite/releases)
- [Changelog](https://github.com/vitejs/vite/blob/main/packages/vite/CHANGELOG.md)
- [Commits](https://github.com/vitejs/vite/commits/v8.2.1/packages/vite)

---
updated-dependencies:
- dependency-name: "@lucide/vue"
  dependency-version: 1.30.0
  dependency-type: direct:production
  update-type: version-update:semver-minor
  dependency-group: minor-and-patch
- dependency-name: dompurify
  dependency-version: 3.4.13
  dependency-type: direct:production
  update-type: version-update:semver-patch
  dependency-group: minor-and-patch
- dependency-name: marked
  dependency-version: 18.0.9
  dependency-type: direct:production
  update-type: version-update:semver-patch
  dependency-group: minor-and-patch
- dependency-name: vue
  dependency-version: 3.5.41
  dependency-type: direct:production
  update-type: version-update:semver-patch
  dependency-group: minor-and-patch
- dependency-name: "@antfu/eslint-config"
  dependency-version: 9.3.0
  dependency-type: direct:development
  update-type: version-update:semver-minor
  dependency-group: minor-and-patch
- dependency-name: "@types/node"
  dependency-version: 26.2.0
  dependency-type: direct:development
  update-type: version-update:semver-minor
  dependency-group: minor-and-patch
- dependency-name: eslint
  dependency-version: 10.8.1
  dependency-type: direct:development
  update-type: version-update:semver-patch
  dependency-group: minor-and-patch
- dependency-name: vite
  dependency-version: 8.2.1
  dependency-type: direct:development
  update-type: version-update:semver-patch
  dependency-group: minor-and-patch
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2026-08-10 11:54:36 +08:00
ZyyandCodex 492c401599 fix(openai): preserve state on quota auth failures
Co-authored-by: Codex <noreply@openai.com>
Signed-off-by: Zyy <zyycn@users.noreply.github.com>
2026-08-10 11:46:19 +08:00
230 changed files with 15700 additions and 6603 deletions
+5
View File
@@ -21,6 +21,11 @@ updates:
day: monday
time: '09:10'
timezone: Asia/Shanghai
# 固定官方 Codex Desktop 的 TLS 指纹;升级前必须抓取握手并比对扩展顺序。
ignore:
- dependency-name: hyper-rustls
- dependency-name: reqwest
- dependency-name: rustls
groups:
minor-and-patch:
update-types:
+1 -1
View File
@@ -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@f205ea1c3313d32999d8d6a48b4f6530d4437b38 # v4.37.4
uses: github/codeql-action/upload-sarif@5595ccaf912efad79be6eef63a5619ff05969be3 # v4.37.6
with:
sarif_file: trivy-results.sarif
category: runtime-container
+2 -2
View File
@@ -106,7 +106,7 @@ jobs:
components: rustfmt,clippy
- name: Cache Rust
uses: Swatinem/rust-cache@c19371144df3bb44fab255c43d04cbc2ab54d1c4 # v2.9.1
uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2
with:
workspaces: backend -> target
@@ -157,7 +157,7 @@ jobs:
ref: ${{ inputs.ref }}
- name: Setup pnpm
uses: pnpm/action-setup@0ebf47130e4866e96fce0953f49152a61190b271 # v6.0.9
uses: pnpm/action-setup@0977fd99725f1db4007ccb2928dbb4e90d06cc86 # v6.0.10
with:
version: 11.7.0
+1 -1
View File
@@ -32,7 +32,7 @@ jobs:
- name: Detect affected checks
id: filter
uses: dorny/paths-filter@7b450fff21473bca461d4b92ce414b9d0420d706 # v4.0.2
uses: dorny/paths-filter@ceb8a2b8f2d89434be7ff52d3de7ec3738c5cc9d # v4.0.3
with:
filters: |
backend:
+2 -2
View File
@@ -375,7 +375,7 @@ jobs:
PY
- name: Generate release asset provenance
uses: actions/attest@508db95dd578ae2727ebd6217d5ba78e4fbda05d # v4.2.1
uses: actions/attest@1e69f48acb82d1966a394da916b4c1698aa569d6 # v4.2.2
with:
subject-path: dist/release/*
@@ -581,7 +581,7 @@ jobs:
echo "digest=${digest}" >> "$GITHUB_OUTPUT"
- name: Generate container provenance attestation
uses: actions/attest@508db95dd578ae2727ebd6217d5ba78e4fbda05d # v4.2.1
uses: actions/attest@1e69f48acb82d1966a394da916b4c1698aa569d6 # v4.2.2
with:
subject-name: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
subject-digest: ${{ steps.manifest.outputs.digest }}
+2 -2
View File
@@ -31,7 +31,7 @@ jobs:
uses: dtolnay/rust-toolchain@1a3a6d54512beeaffd394f8d516ca16f2c506a20 # 1.97
- name: Cache Rust
uses: Swatinem/rust-cache@c19371144df3bb44fab255c43d04cbc2ab54d1c4 # v2.9.1
uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2
with:
workspaces: backend -> target
@@ -52,7 +52,7 @@ jobs:
ref: ${{ github.event.pull_request.head.sha || github.sha }}
- name: Setup pnpm
uses: pnpm/action-setup@0ebf47130e4866e96fce0953f49152a61190b271 # v6.0.9
uses: pnpm/action-setup@0977fd99725f1db4007ccb2928dbb4e90d06cc86 # v6.0.10
with:
version: 11.7.0
+9 -9
View File
@@ -216,9 +216,9 @@ dependencies = [
[[package]]
name = "aws-sdk-s3"
version = "1.140.0"
version = "1.141.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e9660cf991e512fbe6094f1041ff3d3282bc5aeeeb6184e8de437ec2de024a10"
checksum = "d9f9420d3a2467eed22ed3635ca653653162c386a0b0f65c78189f9bd3c1379e"
dependencies = [
"arc-swap",
"aws-credential-types",
@@ -578,9 +578,9 @@ checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6"
[[package]]
name = "base64"
version = "0.23.0"
version = "0.23.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b25655df2c3cdd83c5e5b293b88acd880332b2ddadd7c30ac43144fdc0033da9"
checksum = "ac07cdecf99051d9a5238b80f35af32cdeba5b336e55d957b318b50137e18da5"
[[package]]
name = "base64-simd"
@@ -1493,7 +1493,7 @@ version = "0.0.0"
dependencies = [
"argon2",
"async-trait",
"base64 0.23.0",
"base64 0.23.1",
"chrono",
"chrono-tz",
"cron",
@@ -1519,7 +1519,7 @@ version = "0.1.0"
dependencies = [
"async-trait",
"axum",
"base64 0.23.0",
"base64 0.23.1",
"bytes",
"chrono",
"futures",
@@ -2575,7 +2575,7 @@ version = "0.0.0"
dependencies = [
"async-stream",
"async-trait",
"base64 0.23.0",
"base64 0.23.1",
"bytes",
"chrono",
"cookie",
@@ -2619,7 +2619,7 @@ version = "0.0.0"
dependencies = [
"async-stream",
"async-trait",
"base64 0.23.0",
"base64 0.23.1",
"bytes",
"chrono",
"futures",
@@ -3623,7 +3623,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd"
dependencies = [
"fastrand",
"getrandom 0.4.3",
"getrandom 0.3.4",
"once_cell",
"rustix",
"windows-sys 0.61.2",
+17 -8
View File
@@ -21,9 +21,9 @@ pub mod ports;
mod use_case;
pub use use_case::{
accounts::AccountsService, auth::AuthService, backup::BackupService,
client_keys::ClientKeyService, observability::ObservabilityService, openai::OpenAiService,
settings::SettingsService, system::SystemService, xai::XaiService,
account_groups::AccountGroupService, accounts::AccountsService, auth::AuthService,
backup::BackupService, client_keys::ClientKeyService, observability::ObservabilityService,
openai::OpenAiService, settings::SettingsService, system::SystemService, xai::XaiService,
};
use model::{AdminError, AdminErrorKind};
@@ -33,10 +33,10 @@ use ports::{
system::SystemOperations,
};
use use_case::{
accounts::DefaultAccountsService, auth::DefaultAuthService, backup::DefaultBackupService,
client_keys::DefaultClientKeyService, observability::DefaultObservabilityService,
openai::DefaultOpenAiService, settings::DefaultSettingsService, system::DefaultSystemService,
xai::DefaultXaiService,
account_groups::DefaultAccountGroupService, accounts::DefaultAccountsService,
auth::DefaultAuthService, backup::DefaultBackupService, client_keys::DefaultClientKeyService,
observability::DefaultObservabilityService, openai::DefaultOpenAiService,
settings::DefaultSettingsService, system::DefaultSystemService, xai::DefaultXaiService,
};
const OPENAI_PROVIDER_KIND: &str = "openai";
@@ -147,6 +147,7 @@ pub enum AdminConfigError {
pub struct AdminServices {
auth: Arc<dyn AuthService>,
accounts: Arc<dyn AccountsService>,
account_groups: Arc<dyn AccountGroupService>,
client_keys: Arc<dyn ClientKeyService>,
observability: Arc<dyn ObservabilityService>,
settings: Arc<dyn SettingsService>,
@@ -167,6 +168,11 @@ impl AdminServices {
self.accounts.as_ref()
}
#[must_use]
pub fn account_groups(&self) -> &dyn AccountGroupService {
self.account_groups.as_ref()
}
#[must_use]
pub fn client_keys(&self) -> &dyn ClientKeyService {
self.client_keys.as_ref()
@@ -255,7 +261,6 @@ pub async fn initialize(
let accounts = Arc::new(DefaultAccountsService::new(
store.accounts(),
store.settings(),
registry.clone(),
snapshot.clone(),
probe.clone(),
@@ -275,6 +280,10 @@ pub async fn initialize(
let services = AdminServices {
auth,
accounts,
account_groups: Arc::new(DefaultAccountGroupService::new(
store.account_groups(),
snapshot.clone(),
)),
client_keys: Arc::new(DefaultClientKeyService::new(
store.client_keys(),
snapshot.clone(),
@@ -0,0 +1,142 @@
//! Provider-neutral account group commands and projections.
use std::collections::BTreeMap;
use chrono::{DateTime, Utc};
use gateway_core::routing::AccountGroupId;
use super::{PageSize, Revision, observability::DecimalAmount};
/// Canonical `#RRGGBBAA` color persisted with an account group.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountGroupColor(String);
impl AccountGroupColor {
#[must_use]
pub fn as_str(&self) -> &str {
&self.0
}
#[must_use]
pub fn parse(value: &str) -> Option<Self> {
(value.len() == 9
&& value.starts_with('#')
&& value[1..].bytes().all(|byte| byte.is_ascii_hexdigit()))
.then(|| Self(value.to_ascii_uppercase()))
}
}
/// Group member availability at the current observation time.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AccountGroupAccountSummary {
pub available: u64,
pub limited: u64,
pub total: u64,
}
/// Group scheduling slots derived from available accounts and runtime leases.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AccountGroupCapacity {
pub used_slots: Option<u64>,
pub total_slots: u64,
}
/// Successful, downstream-committed USD request costs for group accounts.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountGroupUsage {
pub today_usd: DecimalAmount,
pub total_usd: DecimalAmount,
}
/// Lightweight group reference embedded in account and client-key views.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountGroupRef {
pub id: AccountGroupId,
pub name: String,
pub color: AccountGroupColor,
pub enabled: bool,
}
/// Account group list query.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountGroupListQuery {
pub page: u32,
pub page_size: PageSize,
pub search: Option<String>,
pub enabled: Option<bool>,
}
/// Complete account group summary.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountGroupRecord {
pub id: AccountGroupId,
pub name: String,
pub description: Option<String>,
pub color: AccountGroupColor,
pub enabled: bool,
pub member_count: u64,
pub provider_counts: BTreeMap<String, u64>,
pub client_key_count: u64,
pub account_summary: AccountGroupAccountSummary,
pub capacity: AccountGroupCapacity,
pub usage: AccountGroupUsage,
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
}
/// Paginated account groups.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountGroupPage {
pub config_revision: Revision,
pub items: Vec<AccountGroupRecord>,
pub total: u64,
pub page: u32,
pub page_size: u16,
}
/// Create an account group.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CreateAccountGroup {
pub name: String,
pub description: Option<String>,
pub color: AccountGroupColor,
}
/// Store-ready create command with a generated stable ID.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct NewAccountGroup {
pub id: AccountGroupId,
pub name: String,
pub description: Option<String>,
pub color: AccountGroupColor,
}
/// Update an account group's descriptive fields.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UpdateAccountGroup {
pub id: AccountGroupId,
pub name: String,
pub description: Option<String>,
pub color: AccountGroupColor,
}
/// Enable or disable an account group.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SetAccountGroupEnabled {
pub id: AccountGroupId,
pub enabled: bool,
}
/// Delete an account group that is not referenced by a client key.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DeleteAccountGroup {
pub id: AccountGroupId,
}
/// Result of any account-group mutation.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountGroupMutation {
pub config_revision: Revision,
pub id: AccountGroupId,
pub record: Option<AccountGroupRecord>,
}
@@ -1,159 +1,19 @@
//! 多 Provider 账号目录与连接测试的公共事实。
use std::{pin::Pin, time::SystemTime};
use std::pin::Pin;
use chrono::{DateTime, Utc};
use futures::Stream;
use gateway_core::routing::ProviderKind;
use super::{PageSize, Revision, observability::TimeRange};
use super::{PageSize, Revision, account_groups::AccountGroupRef, observability::TimeRange};
/// Provider 账号当前可用性;权威定义与字符串编解码在 gateway-core,此处仅复用。
pub use gateway_core::engine::credential::AccountAvailability;
/// 管理页使用的归一化账号状态。
///
/// 只回答两个问题:是否可调度、是否需要人工干预。具体错误原因不进入此枚举,
/// 由 [`AccountErrorReason`] 承载;新增原因只扩展 reason,不改变分类。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AccountStatus {
/// 正常,可调度。
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),
}
}
pub use gateway_core::engine::credential::{
AccountConcurrencyLimit, AccountErrorReason, AccountStatus, AccountStatusFacts,
AccountStatusProjection, AccountWeight, CredentialState, QuotaAccessState, QuotaEvidence,
QuotaState, resolve_account_status,
};
/// 账号列表排序字段。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -186,16 +46,25 @@ pub struct AccountListQuery {
pub page: u32,
pub page_size: PageSize,
pub provider_kind: Option<ProviderKind>,
pub group_filter: Option<AccountGroupFilter>,
pub search: Option<String>,
pub status: Option<AccountStatus>,
pub sort: Option<AccountSort>,
}
/// Optional account membership filter.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AccountGroupFilter {
Group(gateway_core::routing::AccountGroupId),
Ungrouped,
}
/// 账号公共存储投影;Provider 专属字段不进入此结构。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountRecord {
pub id: String,
pub provider_kind: ProviderKind,
pub groups: Vec<AccountGroupRef>,
pub name: String,
pub email: Option<String>,
pub upstream_user_id: Option<String>,
@@ -207,11 +76,13 @@ pub struct AccountRecord {
pub access_token_expires_at: Option<DateTime<Utc>>,
pub next_refresh_at: Option<DateTime<Utc>>,
pub enabled: bool,
pub availability: AccountAvailability,
pub availability_observed_at: DateTime<Utc>,
/// 最近一次非 ready 状态变更的原始原因(上游错误原文 / 原因码);恢复 ready 后清空。
pub concurrency_limit: Option<AccountConcurrencyLimit>,
pub weight: AccountWeight,
pub credential_state: CredentialState,
pub credential_observed_at: DateTime<Utc>,
pub quota: QuotaState,
pub last_error_reason: Option<AccountErrorReason>,
pub last_error_message: Option<String>,
pub quota_observed_at: Option<DateTime<Utc>>,
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
}
@@ -296,11 +167,18 @@ pub struct AccountUsageWindowResult {
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountPage {
pub config_revision: Revision,
pub items: Vec<AccountRecord>,
pub items: Vec<AccountPageItem>,
pub total: u64,
pub summary: AccountSummary,
}
/// 同一状态快照下的账号事实与唯一状态投影。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountPageItem {
pub account: AccountRecord,
pub projection: AccountStatusProjection,
}
/// 统一账号目录的全局状态计数,不受当前筛选和分页影响。
///
/// 计数与 [`AccountStatus`] 一一对应,由 store 按派生状态聚合。
@@ -314,11 +192,38 @@ pub struct AccountSummary {
pub error: u64,
}
/// 账号启停写入命令。
/// 账号可编辑事实的一次性替换命令。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SetAccountEnabled {
pub struct UpdateAccount {
pub account_id: String,
pub enabled: bool,
pub concurrency_limit: Option<AccountConcurrencyLimit>,
pub weight: AccountWeight,
pub group_ids: Vec<gateway_core::routing::AccountGroupId>,
}
/// 账号更新结果。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountUpdateResult {
pub config_revision: Revision,
pub account_id: gateway_core::engine::credential::ProviderAccountId,
}
/// 一批账号可编辑事实的一次性替换命令。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BatchUpdateAccounts {
pub account_ids: Vec<String>,
pub enabled: bool,
pub concurrency_limit: Option<AccountConcurrencyLimit>,
pub weight: AccountWeight,
pub group_ids: Vec<gateway_core::routing::AccountGroupId>,
}
/// 批量账号更新结果。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountsUpdateResult {
pub config_revision: Revision,
pub account_ids: Vec<gateway_core::engine::credential::ProviderAccountId>,
}
/// 账号批量删除命令。
@@ -342,9 +247,7 @@ pub enum AccountConnectionTestEvent {
Content {
text: String,
},
Completed {
account_status: AccountStatus,
},
Completed,
Failed {
message: String,
provider_error_code: Option<String>,
@@ -352,7 +255,6 @@ pub enum AccountConnectionTestEvent {
upstream_status: Option<u16>,
upstream_content_type: Option<String>,
upstream_body: Option<String>,
account_status: AccountStatus,
},
}
@@ -6,10 +6,10 @@ use chrono::{DateTime, Utc};
use gateway_core::{
policy::{ClientApiKeyId, RateLimits},
routing::ProviderKind,
routing::{AccountGroupId, ProviderKind},
};
use super::{AdminModelError, Revision};
use super::{AdminModelError, Revision, account_groups::AccountGroupRef};
/// Client Key 列表保持旧 HTTP 合同允许的完整非零 `u16` 页大小。
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
@@ -88,7 +88,8 @@ pub struct ClientKeyRecord {
pub id: ClientApiKeyId,
pub name: String,
pub label: Option<String>,
pub provider_kind: ProviderKind,
pub groups: Vec<AccountGroupRef>,
pub provider_kinds: Vec<ProviderKind>,
pub prefix: String,
pub enabled: bool,
pub limits: RateLimits,
@@ -143,7 +144,7 @@ impl fmt::Debug for ClientKeySecret {
pub struct CreateClientKey {
pub name: String,
pub label: Option<String>,
pub provider_kind: ProviderKind,
pub group_ids: Vec<AccountGroupId>,
pub limits: RateLimits,
}
@@ -153,7 +154,7 @@ pub struct NewClientKey {
pub id: ClientApiKeyId,
pub name: String,
pub label: Option<String>,
pub provider_kind: ProviderKind,
pub group_ids: Vec<AccountGroupId>,
pub limits: RateLimits,
pub plaintext: String,
}
@@ -164,7 +165,7 @@ impl fmt::Debug for NewClientKey {
.debug_struct("NewClientKey")
.field("id", &self.id)
.field("name", &self.name)
.field("provider_kind", &self.provider_kind)
.field("group_ids", &self.group_ids)
.field("plaintext", &"[REDACTED]")
.finish_non_exhaustive()
}
@@ -176,7 +177,7 @@ pub struct UpdateClientKey {
pub id: ClientApiKeyId,
pub name: String,
pub label: Option<String>,
pub provider_kind: ProviderKind,
pub group_ids: Vec<AccountGroupId>,
pub limits: RateLimits,
}
@@ -2,6 +2,7 @@
use std::num::{NonZeroU16, NonZeroU64};
pub mod account_groups;
pub mod accounts;
pub mod auth;
pub mod backup;
@@ -439,25 +439,15 @@ pub struct RequestMetricPoint {
pub costs: Vec<CurrencyCost>,
}
/// 账号池容量统计。
///
/// 首页概览的 `unavailable` 表示不可参与调度,包含配额耗尽及过期、禁用和封禁账号。
///
/// `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 normal: u64,
pub quota_exhausted: u64,
pub refreshing: Option<u64>,
pub rate_limited: u64,
pub disabled: u64,
pub banned: u64,
pub error: u64,
}
/// Dashboard 中一个账号的模型级用量事实。
@@ -488,7 +478,7 @@ pub struct DashboardAccountRequestBucket {
pub request_count: u64,
}
/// Dashboard 中一个账号的完整用量与公共状态事实。
/// Dashboard 中一个账号的完整用量事实。
#[derive(Debug, Clone, PartialEq)]
pub struct DashboardAccountUsage {
pub account_id: String,
@@ -497,8 +487,6 @@ pub struct DashboardAccountUsage {
pub name: String,
pub email: Option<String>,
pub plan_type: Option<String>,
pub enabled: bool,
pub availability: String,
pub request_count: u64,
pub success_count: u64,
pub input_tokens: Option<u64>,
@@ -525,8 +513,10 @@ pub struct DashboardAccountUsage {
/// Dashboard 当前账号池的可重建运行时槽位事实。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct DashboardRuntimeSlots {
/// 与持久账号池指标相同谓词下可参与调度的账号数。
pub active_accounts: u64,
/// 继承运行参数默认并发上限的可调度账号数。
pub inherited_accounts: u64,
/// 已配置账号级并发上限的可调度账号槽位总数。
pub overridden_slots: u64,
/// Redis 可用时的实时 in-flight 槽位总数;不可用时为 `None`。
pub used_slots: Option<u64>,
}
@@ -560,6 +550,9 @@ pub struct UsageRecord {
pub id: String,
pub client_api_key_ref: String,
pub config_revision: u64,
pub routing_scope: String,
pub routing_group_refs: Vec<String>,
pub routing_group_names_snapshot: Vec<String>,
pub protocol: String,
pub operation: String,
pub endpoint: String,
@@ -11,10 +11,7 @@ use gateway_core::{
use super::{
AdminError, MutationActor, MutationContext, PageSize, Revision,
accounts::{
AccountAvailability, AccountErrorReason, AccountRecord, AccountStatus, AccountSummary,
AccountUsage,
},
accounts::{AccountRecord, AccountSummary, AccountUsage, CredentialState},
};
/// Provider-owned JSON;公共层只搬运且 Debug 不输出值。
@@ -55,24 +52,24 @@ pub struct CredentialCursor {
/// Provider credential 列表查询。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CredentialListQuery {
pub availability: Option<CredentialAvailabilityFilter>,
pub credential_state: Option<CredentialStateFilter>,
pub enabled: Option<bool>,
pub window: CredentialListWindow,
}
/// Credential 可用性筛选;支持一个 wire 值对应多个持久状态。
/// Credential 状态筛选。
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum CredentialAvailabilityFilter {
Exact(AccountAvailability),
AnyOf(Vec<AccountAvailability>),
pub enum CredentialStateFilter {
Exact(CredentialState),
AnyOf(Vec<CredentialState>),
}
impl CredentialAvailabilityFilter {
impl CredentialStateFilter {
#[must_use]
pub fn matches(&self, availability: AccountAvailability) -> bool {
pub fn matches(&self, credential_state: CredentialState) -> bool {
match self {
Self::Exact(expected) => *expected == availability,
Self::AnyOf(expected) => expected.contains(&availability),
Self::Exact(expected) => *expected == credential_state,
Self::AnyOf(expected) => expected.contains(&credential_state),
}
}
}
@@ -155,8 +152,8 @@ pub struct PreparedCredentialCreate {
pub access_token_expires_at: Option<DateTime<Utc>>,
pub next_refresh_at: Option<DateTime<Utc>>,
pub enabled: bool,
pub availability: AccountAvailability,
pub availability_observed_at: DateTime<Utc>,
pub credential_state: CredentialState,
pub credential_observed_at: DateTime<Utc>,
}
/// Provider 对一份导入文档的完整验证结果。
@@ -564,6 +561,15 @@ impl fmt::Debug for RotateCredential {
}
}
/// 通用账号用量能否可靠归属到一个 Provider quota 窗口。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum QuotaLocalUsageAttribution {
/// 窗口覆盖账号的全部请求,可按账号与时间范围聚合。
AccountWide,
/// 窗口需要 Provider / 模型级归属,通用账号聚合不可用。
Unavailable,
}
/// 一个 Provider quota 窗口的公共投影。
#[derive(Debug, Clone, PartialEq)]
pub struct ProviderQuotaWindow {
@@ -571,6 +577,7 @@ pub struct ProviderQuotaWindow {
pub group: String,
pub label: String,
pub source: Option<String>,
pub local_usage_attribution: QuotaLocalUsageAttribution,
pub window_seconds: Option<u64>,
pub used_percent: Option<f64>,
pub reset_at: Option<DateTime<Utc>>,
@@ -593,15 +600,31 @@ 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>,
}
impl ProviderQuota {
/// 返回账号展开面板使用的当前额度窗口用量。
///
/// 只选择具备完整时间边界且可按账号归属的窗口;优先级与 Dashboard
/// 代表性额度窗口一致,同一优先级保持 Provider 投影顺序。
#[must_use]
pub fn representative_window_usage(&self) -> Option<&AccountUsage> {
self.windows
.iter()
.enumerate()
.filter(|(_, window)| {
window.local_usage_attribution == QuotaLocalUsageAttribution::AccountWide
&& window.window_seconds.is_some()
&& window.reset_at.is_some()
&& window.local_usage.is_some()
})
.min_by_key(|(index, window)| (quota_usage_priority(window), *index))
.and_then(|(_, window)| window.local_usage.as_ref())
}
/// 返回 Dashboard 使用的代表性额度比例。
///
/// 优先使用短周期窗口,再依次使用周、月和其它窗口;同一优先级取较高的已用比例。
@@ -693,9 +716,7 @@ impl fmt::Debug for ProviderExport {
#[derive(Debug, Clone, PartialEq)]
pub struct AccountDirectoryItem {
pub account: AccountRecord,
pub status: AccountStatus,
/// `status == Error` 时的具体原因;其余状态为 `None`。
pub error_reason: Option<AccountErrorReason>,
pub projection: gateway_core::engine::credential::AccountStatusProjection,
pub usage: Option<super::accounts::AccountUsage>,
pub quota: ProviderQuota,
}
@@ -10,9 +10,14 @@ use chrono::{DateTime, Utc};
use super::backup::BackupStorePorts;
use crate::model::{
MutationContext, Revision,
account_groups::{
AccountGroupListQuery, AccountGroupMutation, AccountGroupPage, DeleteAccountGroup,
NewAccountGroup, SetAccountGroupEnabled, UpdateAccountGroup,
},
accounts::{
AccountListQuery, AccountPage, AccountRecord, AccountUsage, AccountUsageWindowQuery,
AccountUsageWindowResult, DeleteAccounts, SetAccountEnabled,
AccountListQuery, AccountPage, AccountPageItem, AccountUpdateResult, AccountUsage,
AccountUsageWindowQuery, AccountUsageWindowResult, AccountsUpdateResult,
BatchUpdateAccounts, DeleteAccounts, UpdateAccount,
},
auth::{AdminAuditEvent, AdminSession},
client_keys::{
@@ -83,7 +88,7 @@ pub type AdminStoreResult<T> = Result<T, AdminStoreError>;
pub trait AccountStore: Send + Sync {
async fn list_accounts(&self, query: AccountListQuery) -> AdminStoreResult<AccountPage>;
async fn load_account(&self, account_id: &str) -> AdminStoreResult<Option<AccountRecord>>;
async fn load_account(&self, account_id: &str) -> AdminStoreResult<Option<AccountPageItem>>;
async fn load_account_usage(
&self,
@@ -138,11 +143,23 @@ pub trait AccountStore: Send + Sync {
context: &MutationContext,
) -> AdminStoreResult<CredentialMutationResult>;
async fn set_account_enabled(
async fn update_account(
&self,
command: SetAccountEnabled,
command: UpdateAccount,
context: &MutationContext,
) -> AdminStoreResult<Revision>;
) -> AdminStoreResult<AccountUpdateResult>;
async fn recover_account(
&self,
account_id: &gateway_core::engine::credential::ProviderAccountId,
context: &MutationContext,
) -> AdminStoreResult<AccountUpdateResult>;
async fn batch_update_accounts(
&self,
command: BatchUpdateAccounts,
context: &MutationContext,
) -> AdminStoreResult<AccountsUpdateResult>;
async fn delete_accounts(
&self,
@@ -215,10 +232,48 @@ pub trait ClientKeyStore: Send + Sync {
) -> AdminStoreResult<Revision>;
}
/// Provider-neutral account group management transactions.
#[async_trait]
pub trait AccountGroupStore: Send + Sync {
async fn list_account_groups(
&self,
query: AccountGroupListQuery,
) -> AdminStoreResult<AccountGroupPage>;
async fn create_account_group(
&self,
command: NewAccountGroup,
context: &MutationContext,
) -> AdminStoreResult<AccountGroupMutation>;
async fn update_account_group(
&self,
command: UpdateAccountGroup,
context: &MutationContext,
) -> AdminStoreResult<AccountGroupMutation>;
async fn set_account_group_enabled(
&self,
command: SetAccountGroupEnabled,
context: &MutationContext,
) -> AdminStoreResult<AccountGroupMutation>;
async fn delete_account_group(
&self,
command: DeleteAccountGroup,
context: &MutationContext,
) -> AdminStoreResult<AccountGroupMutation>;
}
/// 用量、趋势、诊断与运维错误的只读能力。
#[async_trait]
pub trait ObservabilityStore: Send + Sync {
async fn dashboard_summary(&self, range: TimeRange) -> AdminStoreResult<DashboardObservation>;
/// 返回历史统计区间和指定观测时刻下的实时账号状态。
async fn dashboard_summary(
&self,
range: TimeRange,
observed_at: DateTime<Utc>,
) -> AdminStoreResult<DashboardObservation>;
/// 返回 Dashboard 可选的实时槽位事实。
///
@@ -297,6 +352,7 @@ pub trait SettingsStore: Send + Sync {
#[derive(Clone)]
pub struct AdminStorePorts {
accounts: Arc<dyn AccountStore>,
account_groups: Arc<dyn AccountGroupStore>,
auth: Arc<dyn AuthStore>,
client_keys: Arc<dyn ClientKeyStore>,
observability: Arc<dyn ObservabilityStore>,
@@ -308,6 +364,7 @@ impl AdminStorePorts {
#[must_use]
pub fn new(
accounts: Arc<dyn AccountStore>,
account_groups: Arc<dyn AccountGroupStore>,
auth: Arc<dyn AuthStore>,
client_keys: Arc<dyn ClientKeyStore>,
observability: Arc<dyn ObservabilityStore>,
@@ -316,6 +373,7 @@ impl AdminStorePorts {
) -> Self {
Self {
accounts,
account_groups,
auth,
client_keys,
observability,
@@ -329,6 +387,11 @@ impl AdminStorePorts {
self.accounts.clone()
}
#[must_use]
pub fn account_groups(&self) -> Arc<dyn AccountGroupStore> {
self.account_groups.clone()
}
#[must_use]
pub fn auth(&self) -> Arc<dyn AuthStore> {
self.auth.clone()
@@ -0,0 +1,130 @@
//! Account group management use cases.
use std::sync::Arc;
use async_trait::async_trait;
use gateway_core::routing::{AccountGroupId, snapshot::SnapshotControl};
use uuid::Uuid;
use crate::{
model::{
AdminError, MutationContext,
account_groups::{
AccountGroupListQuery, AccountGroupMutation, AccountGroupPage, CreateAccountGroup,
DeleteAccountGroup, NewAccountGroup, SetAccountGroupEnabled, UpdateAccountGroup,
},
},
ports::store::AccountGroupStore,
};
use super::{map_store_error, publish_committed};
/// API-facing account group management service.
#[async_trait]
pub trait AccountGroupService: Send + Sync {
async fn list(&self, query: AccountGroupListQuery) -> Result<AccountGroupPage, AdminError>;
async fn create(
&self,
context: &MutationContext,
command: CreateAccountGroup,
) -> Result<AccountGroupMutation, AdminError>;
async fn update(
&self,
context: &MutationContext,
command: UpdateAccountGroup,
) -> Result<AccountGroupMutation, AdminError>;
async fn set_enabled(
&self,
context: &MutationContext,
command: SetAccountGroupEnabled,
) -> Result<AccountGroupMutation, AdminError>;
async fn delete(
&self,
context: &MutationContext,
command: DeleteAccountGroup,
) -> Result<AccountGroupMutation, AdminError>;
}
pub(crate) struct DefaultAccountGroupService {
store: Arc<dyn AccountGroupStore>,
snapshot: Arc<dyn SnapshotControl>,
}
impl DefaultAccountGroupService {
#[must_use]
pub(crate) fn new(
store: Arc<dyn AccountGroupStore>,
snapshot: Arc<dyn SnapshotControl>,
) -> Self {
Self { store, snapshot }
}
async fn publish(
&self,
result: Result<AccountGroupMutation, crate::ports::store::AdminStoreError>,
) -> Result<AccountGroupMutation, AdminError> {
let result = result.map_err(|error| map_store_error(error, "account group"))?;
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
Ok(result)
}
}
#[async_trait]
impl AccountGroupService for DefaultAccountGroupService {
async fn list(&self, query: AccountGroupListQuery) -> Result<AccountGroupPage, AdminError> {
self.store
.list_account_groups(query)
.await
.map_err(|error| map_store_error(error, "account group"))
}
async fn create(
&self,
context: &MutationContext,
command: CreateAccountGroup,
) -> Result<AccountGroupMutation, AdminError> {
let id = AccountGroupId::new(format!("grp_{}", Uuid::now_v7().simple()))
.map_err(|_| AdminError::internal("Failed to create account group ID"))?;
self.publish(
self.store
.create_account_group(
NewAccountGroup {
id,
name: command.name,
description: command.description,
color: command.color,
},
context,
)
.await,
)
.await
}
async fn update(
&self,
context: &MutationContext,
command: UpdateAccountGroup,
) -> Result<AccountGroupMutation, AdminError> {
self.publish(self.store.update_account_group(command, context).await)
.await
}
async fn set_enabled(
&self,
context: &MutationContext,
command: SetAccountGroupEnabled,
) -> Result<AccountGroupMutation, AdminError> {
self.publish(self.store.set_account_group_enabled(command, context).await)
.await
}
async fn delete(
&self,
context: &MutationContext,
command: DeleteAccountGroup,
) -> Result<AccountGroupMutation, AdminError> {
self.publish(self.store.delete_account_group(command, context).await)
.await
}
}
@@ -1,6 +1,6 @@
//! 统一账号目录与跨 Provider 动态分派。
use std::{collections::BTreeMap, sync::Arc, time::SystemTime};
use std::{collections::BTreeMap, sync::Arc};
use async_trait::async_trait;
use chrono::{Duration, Utc};
@@ -18,20 +18,17 @@ use crate::{
AdminError, MutationContext,
accounts::{
AccountConnectionTestEvent, AccountConnectionTestEventStream, AccountListQuery,
AccountRecord, AccountStatus, AccountStatusSignals, AccountUsageWindowQuery,
derive_error_reason,
AccountPageItem, AccountUpdateResult, AccountUsageWindowQuery, AccountsUpdateResult,
BatchUpdateAccounts, UpdateAccount,
},
observability::TimeRange,
provider_credentials::{
AccountDirectoryItem, AccountDirectoryPage, AccountExportBundle, AccountRefreshResult,
PrepareCredentialRefresh, ProviderModels, ProviderQuota, ProviderQuotaRequest,
ProviderQuotaWindow,
ProviderQuotaWindow, QuotaLocalUsageAttribution,
},
},
ports::{
provider::ProviderAdminRegistry,
store::{AccountStore, SettingsStore},
},
ports::{provider::ProviderAdminRegistry, store::AccountStore},
};
use super::{
@@ -58,6 +55,24 @@ pub trait AccountsService: Send + Sync {
account_id: ProviderAccountId,
) -> Result<AccountRefreshResult, AdminError>;
async fn recover(
&self,
context: &MutationContext,
account_id: ProviderAccountId,
) -> Result<AccountRefreshResult, AdminError>;
async fn update(
&self,
context: &MutationContext,
command: UpdateAccount,
) -> Result<AccountUpdateResult, AdminError>;
async fn batch_update(
&self,
context: &MutationContext,
command: BatchUpdateAccounts,
) -> Result<AccountsUpdateResult, AdminError>;
async fn quota(
&self,
account_id: &ProviderAccountId,
@@ -79,7 +94,6 @@ pub trait AccountsService: Send + Sync {
pub(crate) struct DefaultAccountsService {
accounts: Arc<dyn AccountStore>,
settings: Arc<dyn SettingsStore>,
providers: ProviderAdminRegistry,
snapshot: Arc<dyn SnapshotControl>,
probe: Arc<dyn AccountProbe>,
@@ -89,14 +103,12 @@ impl DefaultAccountsService {
#[must_use]
pub(crate) const fn new(
accounts: Arc<dyn AccountStore>,
settings: Arc<dyn SettingsStore>,
providers: ProviderAdminRegistry,
snapshot: Arc<dyn SnapshotControl>,
probe: Arc<dyn AccountProbe>,
) -> Self {
Self {
accounts,
settings,
providers,
snapshot,
probe,
@@ -106,7 +118,7 @@ impl DefaultAccountsService {
async fn load_account(
&self,
account_id: &ProviderAccountId,
) -> Result<AccountRecord, AdminError> {
) -> Result<AccountPageItem, AdminError> {
self.accounts
.load_account(account_id.as_str())
.await
@@ -119,32 +131,32 @@ impl DefaultAccountsService {
account_id: &ProviderAccountId,
) -> Result<
(
AccountRecord,
AccountPageItem,
Arc<dyn crate::ports::provider::ProviderAdmin>,
),
AdminError,
> {
let account = self.load_account(account_id).await?;
let item = self.load_account(account_id).await?;
let provider = self
.providers
.require(&account.provider_kind)
.require(&item.account.provider_kind)
.map_err(|error| map_provider_error(error, "provider account"))?;
Ok((account, provider))
Ok((item, provider))
}
async fn attach_quota_local_usage(
&self,
accounts: &[AccountRecord],
accounts: &[AccountPageItem],
quotas: &mut [ProviderQuota],
) -> Result<(), AdminError> {
let windows = accounts
.iter()
.zip(quotas.iter())
.flat_map(|(account, quota)| {
.flat_map(|(item, quota)| {
quota
.windows
.iter()
.filter_map(|window| quota_usage_window(&account.id, window))
.filter_map(|window| quota_usage_window(&item.account.id, window))
})
.collect::<Vec<_>>();
if windows.is_empty() {
@@ -158,10 +170,12 @@ impl DefaultAccountsService {
.into_iter()
.map(|result| ((result.account_id, result.key), result.usage))
.collect::<BTreeMap<_, _>>();
for (account, quota) in accounts.iter().zip(quotas) {
for (item, quota) in accounts.iter().zip(quotas) {
for window in &mut quota.windows {
if window.local_usage.is_none() {
let key = (account.id.clone(), window.key.clone());
if window.local_usage.is_none()
&& window.local_usage_attribution == QuotaLocalUsageAttribution::AccountWide
{
let key = (item.account.id.clone(), window.key.clone());
if let Some(usage) = usage_by_window.get(&key) {
window.local_usage = Some(usage.clone());
}
@@ -176,33 +190,19 @@ impl DefaultAccountsService {
account_id: &ProviderAccountId,
refresh_quota: bool,
) -> Result<AccountDirectoryItem, AdminError> {
let (account, provider) = self.provider_for_account(account_id).await?;
let settings = self
.settings
.load_runtime_settings()
.await
.map_err(|error| map_store_error(error, "runtime settings"))?;
let (stored, provider) = self.provider_for_account(account_id).await?;
let account = &stored.account;
let now = Utc::now();
let retained_range = retained_usage_range(now, settings.usage_retention_days);
let rolling_range = TimeRange {
start: now - Duration::hours(24),
end: now,
};
let ids = vec![account.id.clone()];
let (usage, rolling_usage) = futures::try_join!(
async {
self.accounts
.load_account_usage(retained_range, &ids)
.await
.map_err(|error| map_store_error(error, "account usage"))
},
async {
self.accounts
.load_account_usage(rolling_range, &ids)
.await
.map_err(|error| map_store_error(error, "rolling account usage"))
},
)?;
let rolling_usage = self
.accounts
.load_account_usage(rolling_range, &ids)
.await
.map_err(|error| map_store_error(error, "rolling account usage"))?;
let rolling_usage = rolling_usage.into_iter().next();
let mut quota = provider
.quota(ProviderQuotaRequest {
@@ -213,17 +213,15 @@ impl DefaultAccountsService {
.await
.map_err(|error| map_provider_error(error, "provider quota"))?;
self.attach_quota_local_usage(
std::slice::from_ref(&account),
std::slice::from_ref(&stored),
std::slice::from_mut(&mut quota),
)
.await?;
let signals = account_signals(&account, &quota, now);
let error_reason = derive_error_reason(&signals);
let usage = quota.representative_window_usage().cloned();
Ok(AccountDirectoryItem {
status: signals.derive(),
error_reason,
usage: usage.into_iter().next(),
account,
projection: stored.projection,
usage,
account: stored.account,
quota,
})
}
@@ -232,13 +230,12 @@ impl DefaultAccountsService {
#[async_trait]
impl AccountsService for DefaultAccountsService {
async fn list(&self, query: AccountListQuery) -> Result<AccountDirectoryPage, AdminError> {
let (page, settings) = futures::try_join!(
self.accounts.list_accounts(query),
self.settings.load_runtime_settings(),
)
.map_err(|error| map_store_error(error, "account directory"))?;
let page = self
.accounts
.list_accounts(query)
.await
.map_err(|error| map_store_error(error, "account directory"))?;
let now = Utc::now();
let retained_range = retained_usage_range(now, settings.usage_retention_days);
let rolling_range = TimeRange {
start: now - Duration::hours(24),
end: now,
@@ -246,22 +243,19 @@ impl AccountsService for DefaultAccountsService {
let ids = page
.items
.iter()
.map(|account| account.id.clone())
.map(|item| item.account.id.clone())
.collect::<Vec<_>>();
let (usage, rolling_usage) = futures::try_join!(
self.accounts.load_account_usage(retained_range, &ids),
self.accounts.load_account_usage(rolling_range, &ids),
)
.map_err(|error| map_store_error(error, "account usage"))?;
let usage = usage
.into_iter()
.map(|usage| (usage.account_id.clone(), usage))
.collect::<BTreeMap<_, _>>();
let rolling_usage = self
.accounts
.load_account_usage(rolling_range, &ids)
.await
.map_err(|error| map_store_error(error, "rolling account usage"))?;
let rolling_usage = rolling_usage
.into_iter()
.map(|usage| (usage.account_id.clone(), usage))
.collect::<BTreeMap<_, _>>();
let mut quotas = futures::future::join_all(page.items.iter().map(|account| async {
let mut quotas = futures::future::join_all(page.items.iter().map(|item| async {
let account = &item.account;
let account_id = ProviderAccountId::new(account.id.clone())
.map_err(|_| AdminError::invalid("Invalid provider account ID"))?;
// 单个账号的 quota 投影失败(Provider 未注册或 quota 读取失败)不拖垮整页:
@@ -305,14 +299,12 @@ impl AccountsService for DefaultAccountsService {
.items
.into_iter()
.zip(quotas)
.map(|(account, quota)| {
let signals = account_signals(&account, &quota, now);
let error_reason = derive_error_reason(&signals);
.map(|(item, quota)| {
let usage = quota.representative_window_usage().cloned();
AccountDirectoryItem {
status: signals.derive(),
error_reason,
usage: usage.get(&account.id).cloned(),
account,
usage,
account: item.account,
projection: item.projection,
quota,
}
})
@@ -340,7 +332,7 @@ impl AccountsService for DefaultAccountsService {
for account_id in account_ids {
let account = self.load_account(&account_id).await?;
grouped
.entry(account.provider_kind)
.entry(account.account.provider_kind)
.or_default()
.push(account_id);
}
@@ -383,7 +375,8 @@ impl AccountsService for DefaultAccountsService {
context: &MutationContext,
account_id: ProviderAccountId,
) -> Result<AccountRefreshResult, AdminError> {
let (account, provider) = self.provider_for_account(&account_id).await?;
let (stored, provider) = self.provider_for_account(&account_id).await?;
let account = stored.account;
let prepared = provider
.prepare_refresh(PrepareCredentialRefresh {
account: account.clone(),
@@ -409,6 +402,98 @@ impl AccountsService for DefaultAccountsService {
})
}
async fn recover(
&self,
context: &MutationContext,
account_id: ProviderAccountId,
) -> Result<AccountRefreshResult, AdminError> {
let (_, provider) = self.provider_for_account(&account_id).await?;
let result = self
.accounts
.recover_account(&account_id, context)
.await
.map_err(|error| map_store_error(error, "provider account recovery"))?;
provider
.account_facts_changed(std::slice::from_ref(&account_id))
.await;
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
let account = self.load_directory_item(&account_id, false).await?;
Ok(AccountRefreshResult {
config_revision: result.config_revision,
account,
})
}
async fn update(
&self,
context: &MutationContext,
command: UpdateAccount,
) -> Result<AccountUpdateResult, AdminError> {
let account_id = ProviderAccountId::new(command.account_id.clone())
.map_err(|_| AdminError::invalid("Invalid provider account ID"))?;
let (_, provider) = self.provider_for_account(&account_id).await?;
let enabled = command.enabled;
let result = self
.accounts
.update_account(command, context)
.await
.map_err(|error| map_store_error(error, "provider account"))?;
if !enabled {
provider.account_unavailable(&account_id).await;
}
provider
.account_facts_changed(std::slice::from_ref(&account_id))
.await;
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
Ok(result)
}
async fn batch_update(
&self,
context: &MutationContext,
command: BatchUpdateAccounts,
) -> Result<AccountsUpdateResult, AdminError> {
let account_ids = command
.account_ids
.iter()
.map(|id| {
ProviderAccountId::new(id.clone())
.map_err(|_| AdminError::invalid("Invalid provider account ID"))
})
.collect::<Result<Vec<_>, _>>()?;
let mut providers = BTreeMap::<
ProviderKind,
(
Arc<dyn crate::ports::provider::ProviderAdmin>,
Vec<ProviderAccountId>,
),
>::new();
for account_id in &account_ids {
let (item, provider) = self.provider_for_account(account_id).await?;
providers
.entry(item.account.provider_kind)
.or_insert_with(|| (provider, Vec::new()))
.1
.push(account_id.clone());
}
let enabled = command.enabled;
let result = self
.accounts
.batch_update_accounts(command, context)
.await
.map_err(|error| map_store_error(error, "provider accounts"))?;
for (provider, provider_ids) in providers.values() {
if !enabled {
for account_id in provider_ids {
provider.account_unavailable(account_id).await;
}
}
provider.account_facts_changed(provider_ids).await;
}
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
Ok(result)
}
async fn quota(
&self,
account_id: &ProviderAccountId,
@@ -434,8 +519,8 @@ impl AccountsService for DefaultAccountsService {
account_id: ProviderAccountId,
upstream_model: UpstreamModelId,
) -> Result<AccountConnectionTestEventStream, AdminError> {
let (account, provider) = self.provider_for_account(&account_id).await?;
let initial_status = account_status(&account, false, Utc::now());
let (stored, provider) = self.provider_for_account(&account_id).await?;
let account = stored.account;
let model = upstream_model.as_str().to_owned();
let operation = provider
.connection_test_operation(&upstream_model, CONNECTION_TEST_INPUT)
@@ -452,8 +537,6 @@ impl AccountsService for DefaultAccountsService {
},
];
let probe = Arc::clone(&self.probe);
let accounts = Arc::clone(&self.accounts);
let stored_account_id = account.id.clone();
let terminal = futures::stream::once(async move {
let result = probe
.probe(AccountProbeRequest {
@@ -463,22 +546,12 @@ impl AccountsService for DefaultAccountsService {
operation,
})
.await;
let status = accounts
.load_account(&stored_account_id)
.await
.ok()
.flatten()
.map_or(initial_status, |account| {
account_status(&account, false, Utc::now())
});
match result {
Ok(result) => result
.text
.into_iter()
.map(|text| AccountConnectionTestEvent::Content { text })
.chain(std::iter::once(AccountConnectionTestEvent::Completed {
account_status: status,
}))
.chain(std::iter::once(AccountConnectionTestEvent::Completed))
.collect(),
Err(error) => {
let upstream_status = error
@@ -504,7 +577,6 @@ impl AccountsService for DefaultAccountsService {
upstream_status,
upstream_content_type,
upstream_body,
account_status: status,
}]
}
}
@@ -521,54 +593,17 @@ fn empty_quota() -> ProviderQuota {
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, &quota, now).derive()
}
fn retained_usage_range(now: chrono::DateTime<Utc>, retention_days: u32) -> TimeRange {
// 账号留存投影是内部查询,不受外部观测接口 366 天上限约束。
TimeRange {
start: now - Duration::days(i64::from(retention_days)),
end: now,
}
}
fn quota_usage_window(
account_id: &str,
window: &ProviderQuotaWindow,
) -> Option<AccountUsageWindowQuery> {
if window.local_usage.is_some() {
if window.local_usage.is_some()
|| window.local_usage_attribution != QuotaLocalUsageAttribution::AccountWide
{
return None;
}
let reset_at = window.reset_at?;
@@ -97,7 +97,7 @@ impl ClientKeyService for DefaultClientKeyService {
id,
name: command.name,
label: command.label,
provider_kind: command.provider_kind,
group_ids: command.group_ids,
limits: command.limits,
plaintext: plaintext.clone(),
},
@@ -1,5 +1,6 @@
//! 管理控制面的用例实现。
pub mod account_groups;
pub mod accounts;
pub mod auth;
pub mod backup;
@@ -13,11 +14,11 @@ pub mod xai;
use crate::{
model::{
AdminError, AdminErrorKind, MutationContext,
accounts::{DeleteAccounts, SetAccountEnabled},
accounts::DeleteAccounts,
provider_credentials::{
AuthorizationMutationTarget, CredentialDeletion, CredentialDeletionResult,
CredentialDetails, CredentialMutation, CredentialMutationResult,
CredentialRotationCommit, PendingAuthorizationMutation, PreparedAuthorizationCommit,
CredentialDetails, CredentialMutationResult, CredentialRotationCommit,
PendingAuthorizationMutation, PreparedAuthorizationCommit,
PreparedAuthorizationCredential, PreparedCredentialImport, PreparedCredentialRotation,
StartAuthorization,
},
@@ -280,44 +281,6 @@ async fn commit_credential_refresh(
}
}
async fn set_credential_enabled(
accounts: &dyn AccountStore,
provider: &dyn ProviderAdmin,
command: CredentialMutation,
enabled: bool,
resource: &'static str,
) -> Result<CredentialMutationResult, AdminError> {
required_credential(
accounts,
provider.provider_kind(),
&command.account_id,
resource,
)
.await?;
let account_id = command.account_id;
let revision = accounts
.set_account_enabled(
SetAccountEnabled {
account_id: account_id.as_str().to_owned(),
enabled,
},
&command.context,
)
.await
.map_err(|error| map_store_error(error, resource))?;
if !enabled {
provider.account_unavailable(&account_id).await;
}
provider
.account_facts_changed(std::slice::from_ref(&account_id))
.await;
Ok(CredentialMutationResult {
config_revision: revision,
account_id,
credential_revision: None,
})
}
async fn delete_credentials(
accounts: &dyn AccountStore,
provider: &dyn ProviderAdmin,
@@ -98,10 +98,11 @@ impl ObservabilityService for DefaultObservabilityService {
range: TimeRange,
kind: TrendKind,
) -> Result<DashboardResult, AdminError> {
let observed_at = Utc::now();
let (mut observation, settings, runtime_slots) = futures::try_join!(
self.store.dashboard_summary(range),
self.store.dashboard_summary(range, observed_at),
self.settings.load_runtime_settings(),
self.store.dashboard_runtime_slots(range.end),
self.store.dashboard_runtime_slots(observed_at),
)
.map_err(|error| map_store_error(error, "dashboard"))?;
self.enrich_billing(&mut observation.recent_requests);
@@ -125,12 +126,20 @@ impl ObservabilityService for DefaultObservabilityService {
let health_timeline = health_timeline_at(&observation.trend, Utc::now());
let wire_profiles = self.providers.dashboard_wire_profiles();
let max_concurrent_per_account = u64::from(settings.max_concurrent_per_account);
let active_accounts = runtime_slots
.as_ref()
.map_or(observation.provider_accounts.active, |slots| {
slots.active_accounts
});
let total_slots = active_accounts.saturating_mul(max_concurrent_per_account);
let total_slots = runtime_slots.as_ref().map_or_else(
|| {
observation
.provider_accounts
.normal
.saturating_mul(max_concurrent_per_account)
},
|slots| {
slots
.inherited_accounts
.saturating_mul(max_concurrent_per_account)
.saturating_add(slots.overridden_slots)
},
);
let used_slots = runtime_slots.and_then(|slots| slots.used_slots);
Ok(DashboardResult {
capacity: DashboardCapacity {
@@ -484,8 +493,7 @@ impl DefaultObservabilityService {
});
let quota_used = futures::future::join_all(quota_reads).await;
for (account, used_percent) in accounts.iter_mut().zip(quota_used) {
account.quota_used_percent = used_percent
.or_else(|| (account.availability == "quota_exhausted").then_some(100.0));
account.quota_used_percent = used_percent;
}
}
@@ -11,9 +11,9 @@ use crate::{
provider_credentials::{
AuthorizationStarted, CompleteAuthorization, CredentialDeletion,
CredentialDeletionResult, CredentialDetails, CredentialImportCommit,
CredentialImportResult, CredentialListQuery, CredentialMutation,
CredentialMutationResult, CredentialPage, ImportCredentials, PrepareCredentialImport,
PrepareCredentialRotation, ProviderQuotaRequest, RotateCredential, StartAuthorization,
CredentialImportResult, CredentialListQuery, CredentialMutationResult, CredentialPage,
ImportCredentials, PrepareCredentialImport, PrepareCredentialRotation,
ProviderQuotaRequest, RotateCredential, StartAuthorization,
},
},
ports::{provider::ProviderAdmin, store::AccountStore},
@@ -22,8 +22,7 @@ use crate::{
use super::{
commit_authorization, commit_credential_rotation, delete_credentials, map_provider_error,
map_store_error, pending_authorization, publish_committed, required_credential,
set_credential_enabled, validate_authorization_commit, validate_prepared_import,
validate_prepared_rotation,
validate_authorization_commit, validate_prepared_import, validate_prepared_rotation,
};
/// OpenAI 固定管理路由消费的服务。
@@ -50,14 +49,6 @@ pub trait OpenAiService: Send + Sync {
&self,
command: RotateCredential,
) -> Result<CredentialMutationResult, AdminError>;
async fn enable(
&self,
command: CredentialMutation,
) -> Result<CredentialMutationResult, AdminError>;
async fn disable(
&self,
command: CredentialMutation,
) -> Result<CredentialMutationResult, AdminError>;
async fn delete(
&self,
command: CredentialDeletion,
@@ -252,38 +243,6 @@ impl OpenAiService for DefaultOpenAiService {
Ok(result)
}
async fn enable(
&self,
command: CredentialMutation,
) -> Result<CredentialMutationResult, AdminError> {
let result = set_credential_enabled(
self.accounts.as_ref(),
self.provider.as_ref(),
command,
true,
"OpenAI credential",
)
.await?;
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
Ok(result)
}
async fn disable(
&self,
command: CredentialMutation,
) -> Result<CredentialMutationResult, AdminError> {
let result = set_credential_enabled(
self.accounts.as_ref(),
self.provider.as_ref(),
command,
false,
"OpenAI credential",
)
.await?;
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
Ok(result)
}
async fn delete(
&self,
command: CredentialDeletion,
@@ -11,8 +11,8 @@ use crate::{
provider_credentials::{
AuthorizationStarted, CompleteAuthorization, CredentialDeletion,
CredentialDeletionResult, CredentialImportCommit, CredentialImportResult,
CredentialListQuery, CredentialMutation, CredentialMutationResult, CredentialPage,
ImportCredentials, PrepareCredentialImport, ProviderQuotaRequest, StartAuthorization,
CredentialListQuery, CredentialMutationResult, CredentialPage, ImportCredentials,
PrepareCredentialImport, ProviderQuotaRequest, StartAuthorization,
},
},
ports::{provider::ProviderAdmin, store::AccountStore},
@@ -20,8 +20,8 @@ use crate::{
use super::{
commit_authorization, delete_credentials, map_provider_error, map_store_error,
pending_authorization, publish_committed, set_credential_enabled,
validate_authorization_commit, validate_prepared_import,
pending_authorization, publish_committed, validate_authorization_commit,
validate_prepared_import,
};
/// xAI 固定管理路由消费的服务。
@@ -40,14 +40,6 @@ pub trait XaiService: Send + Sync {
&self,
command: CompleteAuthorization,
) -> Result<CredentialMutationResult, AdminError>;
async fn enable(
&self,
command: CredentialMutation,
) -> Result<CredentialMutationResult, AdminError>;
async fn disable(
&self,
command: CredentialMutation,
) -> Result<CredentialMutationResult, AdminError>;
async fn delete(
&self,
command: CredentialDeletion,
@@ -166,38 +158,6 @@ impl XaiService for DefaultXaiService {
Ok(result)
}
async fn enable(
&self,
command: CredentialMutation,
) -> Result<CredentialMutationResult, AdminError> {
let result = set_credential_enabled(
self.accounts.as_ref(),
self.provider.as_ref(),
command,
true,
"xAI credential",
)
.await?;
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
Ok(result)
}
async fn disable(
&self,
command: CredentialMutation,
) -> Result<CredentialMutationResult, AdminError> {
let result = set_credential_enabled(
self.accounts.as_ref(),
self.provider.as_ref(),
command,
false,
"xAI credential",
)
.await?;
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
Ok(result)
}
async fn delete(
&self,
command: CredentialDeletion,
+2 -1
View File
@@ -11,6 +11,7 @@ use gateway_admin::{
observability::{DecimalAmount, RequestOutcome, TimeRange},
provider_credentials::{
CredentialCommitGuard, ProviderDocument, ProviderQuota, ProviderQuotaWindow,
QuotaLocalUsageAttribution,
},
settings::AdminApiKey,
},
@@ -128,7 +129,6 @@ fn representative_quota_should_prefer_short_window_and_highest_usage() {
quota_window("shortTerm", Some(14_400), Some(45.0)),
],
limit_reached: false,
rate_limited_until: None,
provider_data: None,
};
@@ -145,6 +145,7 @@ fn quota_window(
group: group.to_owned(),
label: group.to_owned(),
source: None,
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
window_seconds,
used_percent,
reset_at: None,
@@ -5,7 +5,10 @@ use chrono::{TimeDelta, Utc};
use futures::{StreamExt as _, future::BoxFuture};
use gateway_core::{
engine::{
credential::{OpaqueProviderData, ProviderAccountId},
credential::{
AccountStatusFacts, CredentialState, OpaqueProviderData, ProviderAccountId,
QuotaEvidence, QuotaState, resolve_account_status,
},
probe::{AccountProbe, AccountProbeError, AccountProbeRequest, AccountProbeResult},
},
error::{ClientVisibleUpstreamError, GatewayError, GatewayErrorKind},
@@ -18,9 +21,10 @@ use gateway_admin::{
model::{
AdminError, MutationContext, Revision,
accounts::{
AccountAvailability, AccountConnectionTestEvent, AccountListQuery, AccountPage,
AccountRecord, AccountSummary, AccountUsage, AccountUsageWindowQuery,
AccountUsageWindowResult, DeleteAccounts, SetAccountEnabled,
AccountConnectionTestEvent, AccountListQuery, AccountPage, AccountPageItem,
AccountRecord, AccountSummary, AccountUpdateResult, AccountUsage,
AccountUsageWindowQuery, AccountUsageWindowResult, AccountsUpdateResult,
BatchUpdateAccounts, DeleteAccounts, UpdateAccount,
},
observability::TimeRange,
provider_credentials::{
@@ -33,7 +37,7 @@ use gateway_admin::{
PreparedAuthorizationCredential, PreparedCredentialCreate, PreparedCredentialImport,
PreparedCredentialRotation, PreparedCredentialRotationFacts, ProviderDocument,
ProviderExport, ProviderExportCredentialInput, ProviderModels, ProviderQuota,
ProviderQuotaRequest, ProviderQuotaWindow,
ProviderQuotaRequest, ProviderQuotaWindow, QuotaLocalUsageAttribution,
},
settings::{
AdminApiKey, AdminApiKeyMutation, ReplaceRuntimeSettings, RotationStrategy,
@@ -378,6 +382,7 @@ pub(super) struct FakeAccountStore {
fail_commit: Mutex<bool>,
audit_requests: Mutex<Vec<String>>,
quota_window_usage: Mutex<Vec<AccountUsageWindowResult>>,
quota_window_queries: Mutex<Vec<AccountUsageWindowQuery>>,
}
impl FakeAccountStore {
@@ -393,6 +398,7 @@ impl FakeAccountStore {
fail_commit: Mutex::new(false),
audit_requests: Mutex::new(Vec::new()),
quota_window_usage: Mutex::new(Vec::new()),
quota_window_queries: Mutex::new(Vec::new()),
})
}
@@ -408,6 +414,13 @@ impl FakeAccountStore {
*self.quota_window_usage.lock().expect("quota window usage") = usage;
}
fn quota_window_queries(&self) -> Vec<AccountUsageWindowQuery> {
self.quota_window_queries
.lock()
.expect("quota window queries")
.clone()
}
pub(super) fn set_accounts(&self, accounts: Vec<AccountRecord>) {
*self.accounts.lock().expect("accounts") = accounts;
}
@@ -438,6 +451,22 @@ impl FakeAccountStore {
.expect("audit requests")
.push(context.request_id.clone());
}
fn page_item(account: AccountRecord) -> AccountPageItem {
let facts = AccountStatusFacts {
enabled: account.enabled,
credential_state: account.credential_state,
access_token_expires_at: account.access_token_expires_at.map(Into::into),
quota: account.quota,
rate_limited_until: None,
last_error_reason: account.last_error_reason,
last_error_message: account.last_error_message.clone(),
};
AccountPageItem {
account,
projection: resolve_account_status(&facts, std::time::SystemTime::now()),
}
}
}
#[async_trait]
@@ -445,12 +474,13 @@ 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();
let total = accounts.len() as u64;
Ok(AccountPage {
config_revision: revision(1),
items: accounts.clone(),
total: accounts.len() as u64,
items: accounts.into_iter().map(Self::page_item).collect(),
total,
summary: AccountSummary {
total: accounts.len() as u64,
total,
normal: 1,
quota_exhausted: 0,
rate_limited: 0,
@@ -460,7 +490,7 @@ impl AccountStore for FakeAccountStore {
})
}
async fn load_account(&self, account_id: &str) -> AdminStoreResult<Option<AccountRecord>> {
async fn load_account(&self, account_id: &str) -> AdminStoreResult<Option<AccountPageItem>> {
self.record("store.load_account");
// probe 后的账号状态覆盖只对同一 id 生效;其余按账号列表查询。
let account = self
@@ -477,7 +507,7 @@ impl AccountStore for FakeAccountStore {
.find(|account| account.id == account_id)
.cloned()
});
Ok(account)
Ok(account.map(Self::page_item))
}
async fn load_account_usage(
@@ -490,8 +520,12 @@ impl AccountStore for FakeAccountStore {
async fn load_account_usage_by_windows(
&self,
_: &[AccountUsageWindowQuery],
windows: &[AccountUsageWindowQuery],
) -> AdminStoreResult<Vec<AccountUsageWindowResult>> {
*self
.quota_window_queries
.lock()
.expect("quota window queries") = windows.to_vec();
Ok(self
.quota_window_usage
.lock()
@@ -628,15 +662,61 @@ impl AccountStore for FakeAccountStore {
Ok(rotation_result(command))
}
async fn set_account_enabled(
async fn update_account(
&self,
_: SetAccountEnabled,
command: UpdateAccount,
context: &MutationContext,
) -> AdminStoreResult<Revision> {
self.record("store.set_enabled");
) -> AdminStoreResult<AccountUpdateResult> {
self.record("store.update_account");
self.record_context(context);
self.require_commit()?;
Ok(revision(2))
Ok(AccountUpdateResult {
config_revision: revision(2),
account_id: ProviderAccountId::new(command.account_id).expect("account ID"),
})
}
async fn recover_account(
&self,
account_id: &ProviderAccountId,
context: &MutationContext,
) -> AdminStoreResult<AccountUpdateResult> {
self.record("store.recover_account");
self.record_context(context);
self.require_commit()?;
let mut accounts = self.accounts.lock().expect("accounts");
if let Some(account) = accounts
.iter_mut()
.find(|account| account.id == account_id.as_str())
{
account.enabled = true;
account.credential_state = CredentialState::Ready;
account.quota = QuotaState::allowed(std::time::SystemTime::now());
account.last_error_reason = None;
account.last_error_message = None;
}
Ok(AccountUpdateResult {
config_revision: revision(2),
account_id: account_id.clone(),
})
}
async fn batch_update_accounts(
&self,
command: BatchUpdateAccounts,
context: &MutationContext,
) -> AdminStoreResult<AccountsUpdateResult> {
self.record("store.batch_update_accounts");
self.record_context(context);
self.require_commit()?;
Ok(AccountsUpdateResult {
config_revision: revision(2),
account_ids: command
.account_ids
.into_iter()
.map(|id| ProviderAccountId::new(id).expect("account ID"))
.collect(),
})
}
async fn delete_accounts(
@@ -815,7 +895,11 @@ fn provider_registry_should_reject_duplicate_kind() {
async fn connection_test_should_probe_unavailable_account() {
let provider = FakeProviderAdmin::new("xai", events());
let mut account = account_record("xai");
account.availability = AccountAvailability::QuotaExhausted;
account.quota = QuotaState::exhausted(
QuotaEvidence::UsageLimitReached,
std::time::SystemTime::now(),
None,
);
let store = FakeAccountStore::with_account(account, events());
let services =
accounts_service_with_probe(provider, store, Arc::new(FailingAccountProbe)).await;
@@ -837,7 +921,6 @@ async fn connection_test_should_probe_unavailable_account() {
message,
provider_error_code: Some(code),
provider_error_type: Some(error_type),
account_status: gateway_admin::model::accounts::AccountStatus::QuotaExhausted,
..
}) if message == "included usage exhausted"
&& code == "usage_exhausted"
@@ -846,7 +929,7 @@ async fn connection_test_should_probe_unavailable_account() {
}
#[tokio::test]
async fn connection_test_rate_limited_probe_preserves_normal_account_status() {
async fn connection_test_rate_limited_probe_returns_provider_failure() {
let events = events();
let provider = FakeProviderAdmin::new("xai", events.clone());
let store = FakeAccountStore::new("xai", events);
@@ -869,9 +952,10 @@ async fn connection_test_rate_limited_probe_preserves_normal_account_status() {
assert!(matches!(
events.last(),
Some(AccountConnectionTestEvent::Failed {
account_status: gateway_admin::model::accounts::AccountStatus::Normal,
message,
..
})
if message == "xAI free model quota is exhausted"
));
}
@@ -897,10 +981,7 @@ async fn connection_test_should_preserve_disabled_account_status() {
assert!(matches!(
events.last(),
Some(AccountConnectionTestEvent::Failed {
account_status: gateway_admin::model::accounts::AccountStatus::Disabled,
..
})
Some(AccountConnectionTestEvent::Failed { .. })
));
}
@@ -949,7 +1030,7 @@ async fn accounts_refresh_should_keep_guard_through_store_commit() {
assert_eq!(result.config_revision, revision(2));
assert_eq!(result.account.account.provider_kind.as_str(), "openai");
assert_eq!(
result.account.status,
result.account.projection.status,
gateway_admin::model::accounts::AccountStatus::Normal
);
@@ -967,6 +1048,175 @@ async fn accounts_refresh_should_keep_guard_through_store_commit() {
assert_eq!(store.audit_requests(), ["refresh-request"]);
}
#[tokio::test]
async fn accounts_recover_should_commit_facts_then_return_normal_account() {
let events = events();
let provider = FakeProviderAdmin::new("openai", events.clone());
let mut account = account_record("openai");
account.enabled = false;
account.credential_state = CredentialState::Invalid;
account.quota = QuotaState::exhausted(
QuotaEvidence::UsageLimitReached,
std::time::SystemTime::now(),
None,
);
account.last_error_reason =
Some(gateway_core::engine::credential::AccountErrorReason::CredentialInvalid);
account.last_error_message = Some("invalid credential".to_owned());
let store = FakeAccountStore::with_account(account, events.clone());
let services = accounts_service(provider, store.clone()).await;
let result = services
.accounts()
.recover(
&context("recover-request"),
ProviderAccountId::new("acct_test").expect("account ID"),
)
.await
.expect("recover account");
assert_eq!(result.config_revision, revision(2));
assert!(result.account.account.enabled);
assert_eq!(
result.account.account.credential_state,
CredentialState::Ready
);
assert_eq!(
result.account.account.quota.access(),
gateway_core::engine::credential::QuotaAccessState::Allowed
);
assert_eq!(
result.account.projection.status,
gateway_admin::model::accounts::AccountStatus::Normal
);
assert_eq!(
recorded(&events),
[
"store.load_account",
"store.recover_account",
"provider.account_facts_changed",
"store.load_account",
]
);
assert_eq!(store.audit_requests(), ["recover-request"]);
}
#[tokio::test]
async fn accounts_update_should_commit_then_release_disabled_account_and_publish_facts() {
let events = events();
let provider = FakeProviderAdmin::new("openai", events.clone());
let store = FakeAccountStore::new("openai", events.clone());
let services = accounts_service(provider, store.clone()).await;
let result = services
.accounts()
.update(
&context("update-request"),
UpdateAccount {
account_id: "acct_test".to_owned(),
enabled: false,
concurrency_limit: None,
weight: gateway_core::engine::credential::AccountWeight::DEFAULT,
group_ids: Vec::new(),
},
)
.await
.expect("update account");
assert_eq!(result.config_revision, revision(2));
assert_eq!(
recorded(&events),
[
"store.load_account",
"store.update_account",
"provider.account_unavailable",
"provider.account_facts_changed",
]
);
assert_eq!(store.audit_requests(), ["update-request"]);
}
#[tokio::test]
async fn accounts_update_should_not_notify_provider_when_store_commit_fails() {
let events = events();
let provider = FakeProviderAdmin::new("openai", events.clone());
let store = FakeAccountStore::new("openai", events.clone());
store.fail_next_commit();
let services = accounts_service(provider, store).await;
services
.accounts()
.update(
&context("update-failure"),
UpdateAccount {
account_id: "acct_test".to_owned(),
enabled: false,
concurrency_limit: None,
weight: gateway_core::engine::credential::AccountWeight::DEFAULT,
group_ids: Vec::new(),
},
)
.await
.expect_err("failed account update");
assert_eq!(
recorded(&events),
["store.load_account", "store.update_account"]
);
}
#[tokio::test]
async fn accounts_batch_update_should_commit_once_and_notify_each_provider() {
let events = events();
let openai = FakeProviderAdmin::new("openai", events.clone());
let xai = FakeProviderAdmin::new("xai", events.clone());
let store = FakeAccountStore::new("openai", events.clone());
let mut openai_account = account_record("openai");
openai_account.id = "acct_openai".to_owned();
let mut xai_account = account_record("xai");
xai_account.id = "acct_xai".to_owned();
store.set_accounts(vec![openai_account, xai_account]);
let services = super::AdminHarness::new()
.accounts(store.clone())
.settings(Arc::new(StaticSettingsStore))
.provider(openai)
.provider(xai)
.probe(Arc::new(SuccessfulAccountProbe))
.build()
.await;
let result = services
.accounts()
.batch_update(
&context("batch-update-request"),
BatchUpdateAccounts {
account_ids: vec!["acct_openai".to_owned(), "acct_xai".to_owned()],
enabled: false,
concurrency_limit: None,
weight: gateway_core::engine::credential::AccountWeight::DEFAULT,
group_ids: Vec::new(),
},
)
.await
.expect("batch update accounts");
assert_eq!(result.config_revision, revision(2));
assert_eq!(result.account_ids.len(), 2);
assert_eq!(store.audit_requests(), ["batch-update-request"]);
assert_eq!(
recorded(&events),
[
"store.load_account",
"store.load_account",
"store.batch_update_accounts",
"provider.account_unavailable",
"provider.account_facts_changed",
"provider.account_unavailable",
"provider.account_facts_changed",
]
);
}
#[tokio::test]
async fn accounts_list_should_return_complete_directory_semantics() {
let provider = FakeProviderAdmin::new("openai", events());
@@ -978,6 +1228,7 @@ async fn accounts_list_should_return_complete_directory_semantics() {
page: 1,
page_size: gateway_admin::model::PageSize::new(20).expect("page size"),
provider_kind: None,
group_filter: None,
search: None,
status: None,
sort: None,
@@ -990,7 +1241,7 @@ async fn accounts_list_should_return_complete_directory_semantics() {
let account = page.items.first().expect("account item");
assert_eq!(account.account.provider_kind.as_str(), "openai");
assert_eq!(
account.status,
account.projection.status,
gateway_admin::model::accounts::AccountStatus::Normal
);
}
@@ -1007,6 +1258,7 @@ async fn accounts_list_should_degrade_quota_failure_to_empty_window_without_drop
group: "shortTerm".to_owned(),
label: "5小时限额".to_owned(),
source: None,
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
window_seconds: Some(5 * 60 * 60),
used_percent: Some(97.0),
reset_at: Some(Utc::now() + TimeDelta::hours(1)),
@@ -1015,7 +1267,6 @@ async fn accounts_list_should_degrade_quota_failure_to_empty_window_without_drop
provider_data: None,
}],
limit_reached: false,
rate_limited_until: None,
provider_data: None,
});
let failing = FakeProviderAdmin::new("xai", events.clone());
@@ -1039,6 +1290,7 @@ async fn accounts_list_should_degrade_quota_failure_to_empty_window_without_drop
page: 1,
page_size: gateway_admin::model::PageSize::new(20).expect("page size"),
provider_kind: None,
group_filter: None,
search: None,
status: None,
sort: None,
@@ -1063,10 +1315,14 @@ async fn accounts_list_should_degrade_quota_failure_to_empty_window_without_drop
}
#[tokio::test]
async fn accounts_list_should_prefer_quota_exhaustion_over_stale_token_expiry() {
async fn accounts_list_should_prefer_credential_error_over_quota_exhaustion() {
let provider = FakeProviderAdmin::new("openai", events());
let mut account = account_record("openai");
account.availability = AccountAvailability::QuotaExhausted;
account.quota = QuotaState::exhausted(
QuotaEvidence::ProviderDenied,
std::time::SystemTime::now(),
None,
);
account.access_token_expires_at = Some(Utc::now() - TimeDelta::minutes(5));
let store = FakeAccountStore::with_account(account, events());
@@ -1077,6 +1333,7 @@ async fn accounts_list_should_prefer_quota_exhaustion_over_stale_token_expiry()
page: 1,
page_size: gateway_admin::model::PageSize::new(20).expect("page size"),
provider_kind: None,
group_filter: None,
search: None,
status: None,
sort: None,
@@ -1085,17 +1342,17 @@ async fn accounts_list_should_prefer_quota_exhaustion_over_stale_token_expiry()
.expect("usage-limited account directory");
assert_eq!(
page.items.first().expect("account item").status,
gateway_admin::model::accounts::AccountStatus::QuotaExhausted,
page.items.first().expect("account item").projection.status,
gateway_admin::model::accounts::AccountStatus::Error,
);
}
#[tokio::test]
async fn accounts_list_should_map_unknown_availability_to_error_not_normal() {
async fn accounts_list_should_map_unknown_credential_to_error_not_normal() {
// Unknown 不可调度,Admin 不得显示为 normal。
let provider = FakeProviderAdmin::new("openai", events());
let mut account = account_record("openai");
account.availability = AccountAvailability::Unknown;
account.credential_state = CredentialState::Unknown;
let store = FakeAccountStore::with_account(account, events());
let page = accounts_service(provider, store)
@@ -1105,6 +1362,7 @@ async fn accounts_list_should_map_unknown_availability_to_error_not_normal() {
page: 1,
page_size: gateway_admin::model::PageSize::new(20).expect("page size"),
provider_kind: None,
group_filter: None,
search: None,
status: None,
sort: None,
@@ -1113,15 +1371,13 @@ async fn accounts_list_should_map_unknown_availability_to_error_not_normal() {
.expect("unknown account directory");
assert_eq!(
page.items.first().expect("account item").status,
page.items.first().expect("account item").projection.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 冷却)。
async fn accounts_list_should_not_derive_rate_limited_from_provider_quota_view() {
let provider = FakeProviderAdmin::new("openai", events());
provider.set_quota(ProviderQuota {
observed_at: Some(Utc::now()),
@@ -1131,6 +1387,7 @@ async fn accounts_list_should_map_rate_limited_from_quota_limit_reached() {
group: "shortTerm".to_owned(),
label: "5小时限额".to_owned(),
source: None,
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
window_seconds: Some(5 * 60 * 60),
used_percent: Some(100.0),
reset_at: Some(Utc::now() + TimeDelta::hours(1)),
@@ -1139,7 +1396,6 @@ async fn accounts_list_should_map_rate_limited_from_quota_limit_reached() {
provider_data: None,
}],
limit_reached: true,
rate_limited_until: Some(Utc::now() + TimeDelta::minutes(5)),
provider_data: None,
});
let store = FakeAccountStore::new("openai", events());
@@ -1151,22 +1407,22 @@ async fn accounts_list_should_map_rate_limited_from_quota_limit_reached() {
page: 1,
page_size: gateway_admin::model::PageSize::new(20).expect("page size"),
provider_kind: None,
group_filter: None,
search: None,
status: None,
sort: None,
})
.await
.expect("rate limited account directory");
.expect("account directory with display quota cooldown");
assert_eq!(
page.items.first().expect("account item").status,
gateway_admin::model::accounts::AccountStatus::RateLimited,
page.items.first().expect("account item").projection.status,
gateway_admin::model::accounts::AccountStatus::Normal,
);
}
#[tokio::test]
async fn accounts_list_should_map_quota_exhausted_when_only_limit_reached() {
// ready + 快照级 quota.limitReached、无 429 冷却 → quota_exhausted。
async fn accounts_list_should_not_derive_exhaustion_from_provider_quota_view() {
let provider = FakeProviderAdmin::new("openai", events());
provider.set_quota(ProviderQuota {
observed_at: Some(Utc::now()),
@@ -1176,6 +1432,7 @@ async fn accounts_list_should_map_quota_exhausted_when_only_limit_reached() {
group: "shortTerm".to_owned(),
label: "5小时限额".to_owned(),
source: None,
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
window_seconds: Some(5 * 60 * 60),
used_percent: Some(100.0),
reset_at: Some(Utc::now() + TimeDelta::hours(1)),
@@ -1184,7 +1441,6 @@ async fn accounts_list_should_map_quota_exhausted_when_only_limit_reached() {
provider_data: None,
}],
limit_reached: true,
rate_limited_until: None,
provider_data: None,
});
let store = FakeAccountStore::new("openai", events());
@@ -1196,16 +1452,17 @@ async fn accounts_list_should_map_quota_exhausted_when_only_limit_reached() {
page: 1,
page_size: gateway_admin::model::PageSize::new(20).expect("page size"),
provider_kind: None,
group_filter: None,
search: None,
status: None,
sort: None,
})
.await
.expect("quota exhausted account directory");
.expect("account directory with display quota limit");
assert_eq!(
page.items.first().expect("account item").status,
gateway_admin::model::accounts::AccountStatus::QuotaExhausted,
page.items.first().expect("account item").projection.status,
gateway_admin::model::accounts::AccountStatus::Normal,
);
}
@@ -1221,6 +1478,7 @@ async fn accounts_list_should_attach_local_usage_to_quota_windows() {
group: "shortTerm".to_owned(),
label: "5小时限额".to_owned(),
source: None,
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
window_seconds: Some(5 * 60 * 60),
used_percent: Some(97.0),
reset_at: Some(reset_at),
@@ -1229,7 +1487,6 @@ async fn accounts_list_should_attach_local_usage_to_quota_windows() {
provider_data: None,
}],
limit_reached: false,
rate_limited_until: None,
provider_data: None,
});
let store = FakeAccountStore::new("openai", events());
@@ -1239,13 +1496,14 @@ async fn accounts_list_should_attach_local_usage_to_quota_windows() {
usage: quota_local_usage("acct_test", 4_330_000),
}]);
let page = accounts_service(provider, store)
let page = accounts_service(provider, store.clone())
.await
.accounts()
.list(AccountListQuery {
page: 1,
page_size: gateway_admin::model::PageSize::new(20).expect("page size"),
provider_kind: None,
group_filter: None,
search: None,
status: None,
sort: None,
@@ -1253,11 +1511,120 @@ async fn accounts_list_should_attach_local_usage_to_quota_windows() {
.await
.expect("quota window usage");
let usage = page.items[0].quota.windows[0]
let item = &page.items[0];
let usage = item.quota.windows[0]
.local_usage
.as_ref()
.expect("quota window local usage");
assert_eq!(usage.total_tokens, Some(4_330_000));
assert_eq!(
item.usage.as_ref().and_then(|usage| usage.total_tokens),
Some(4_330_000),
);
let queries = store.quota_window_queries();
assert_eq!(queries.len(), 1);
assert_eq!(queries[0].account_id, "acct_test");
assert_eq!(queries[0].key, "primary");
assert_eq!(queries[0].range.start, reset_at - TimeDelta::hours(5));
assert_eq!(queries[0].range.end, reset_at);
}
#[tokio::test]
async fn accounts_list_should_not_attach_account_usage_to_model_specific_quota_windows() {
let provider = FakeProviderAdmin::new("openai", events());
let reset_at = Utc::now() + TimeDelta::days(7);
provider.set_quota(ProviderQuota {
observed_at: Some(Utc::now()),
refresh_token_expires_at: None,
windows: vec![
ProviderQuotaWindow {
key: "core-primary".to_owned(),
group: "weekly".to_owned(),
label: "周额度".to_owned(),
source: Some("codex".to_owned()),
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
window_seconds: Some(7 * 24 * 60 * 60),
used_percent: Some(1.0),
reset_at: Some(reset_at),
limit_reached: false,
local_usage: None,
provider_data: None,
},
ProviderQuotaWindow {
key: "codex-bengalfox-primary".to_owned(),
group: "weekly".to_owned(),
label: "GPT-5.3-Codex-Spark · 周额度".to_owned(),
source: Some("codex_bengalfox".to_owned()),
local_usage_attribution: QuotaLocalUsageAttribution::Unavailable,
window_seconds: Some(7 * 24 * 60 * 60),
used_percent: Some(0.0),
reset_at: Some(reset_at),
limit_reached: false,
local_usage: None,
provider_data: None,
},
],
limit_reached: false,
provider_data: None,
});
let store = FakeAccountStore::new("openai", events());
store.set_quota_window_usage(vec![
AccountUsageWindowResult {
account_id: "acct_test".to_owned(),
key: "core-primary".to_owned(),
usage: quota_local_usage("acct_test", 12_818_806),
},
AccountUsageWindowResult {
account_id: "acct_test".to_owned(),
key: "codex-bengalfox-primary".to_owned(),
usage: quota_local_usage("acct_test", 127_926),
},
]);
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,
group_filter: None,
search: None,
status: None,
sort: None,
})
.await
.expect("quota window usage");
let local_tokens = page.items[0]
.quota
.windows
.iter()
.map(|window| {
(
window.key.as_str(),
window
.local_usage
.as_ref()
.and_then(|usage| usage.total_tokens),
)
})
.collect::<Vec<_>>();
assert_eq!(
local_tokens,
[
("core-primary", Some(12_818_806)),
("codex-bengalfox-primary", None),
],
);
assert_eq!(
page.items[0]
.usage
.as_ref()
.and_then(|usage| usage.total_tokens),
Some(12_818_806),
);
}
#[tokio::test]
@@ -1338,7 +1705,6 @@ fn empty_quota() -> ProviderQuota {
refresh_token_expires_at: None,
windows: Vec::new(),
limit_reached: false,
rate_limited_until: None,
provider_data: None,
}
}
@@ -1371,6 +1737,7 @@ pub(super) fn account_record(kind: &str) -> AccountRecord {
AccountRecord {
id: "acct_test".to_owned(),
provider_kind: ProviderKind::new(kind).expect("provider kind"),
groups: Vec::new(),
name: "test account".to_owned(),
email: Some("test@example.invalid".to_owned()),
upstream_user_id: Some("upstream-user".to_owned()),
@@ -1382,10 +1749,13 @@ pub(super) fn account_record(kind: &str) -> AccountRecord {
access_token_expires_at: Some(now + TimeDelta::hours(1)),
next_refresh_at: Some(now + TimeDelta::minutes(30)),
enabled: true,
availability: AccountAvailability::Ready,
availability_observed_at: now,
concurrency_limit: None,
weight: gateway_core::engine::credential::AccountWeight::DEFAULT,
credential_state: CredentialState::Ready,
credential_observed_at: now,
quota: QuotaState::allowed(now.into()),
last_error_reason: None,
last_error_message: None,
quota_observed_at: None,
created_at: now,
updated_at: now,
}
@@ -1419,8 +1789,8 @@ fn prepared_create_with_id(
access_token_expires_at: Some(now + TimeDelta::hours(1)),
next_refresh_at: Some(now + TimeDelta::minutes(30)),
enabled: true,
availability: AccountAvailability::Ready,
availability_observed_at: now,
credential_state: CredentialState::Ready,
credential_observed_at: now,
}
}
@@ -1503,9 +1873,7 @@ impl AccountProbe for FreeModelQuotaProbe {
_: AccountProbeRequest,
) -> BoxFuture<'_, Result<AccountProbeResult, AccountProbeError>> {
Box::pin(async move {
let mut account = account_record("xai");
account.availability = AccountAvailability::Ready;
self.store.set_account_after_probe(account);
self.store.set_account_after_probe(account_record("xai"));
Err(GatewayError::new(
GatewayErrorKind::RateLimited,
"xAI free model quota is exhausted",
@@ -14,14 +14,20 @@ use std::{
};
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use futures::future::BoxFuture;
use gateway_admin::{
AdminConfig, AdminServices, InitialAdminPassword,
model::{
MutationContext, Revision,
account_groups::{
AccountGroupListQuery, AccountGroupMutation, AccountGroupPage, DeleteAccountGroup,
NewAccountGroup, SetAccountGroupEnabled, UpdateAccountGroup,
},
accounts::{
AccountListQuery, AccountPage, AccountRecord, AccountUsage, AccountUsageWindowQuery,
AccountUsageWindowResult, DeleteAccounts, SetAccountEnabled,
AccountListQuery, AccountPage, AccountUpdateResult, AccountUsage,
AccountUsageWindowQuery, AccountUsageWindowResult, AccountsUpdateResult,
BatchUpdateAccounts, DeleteAccounts, UpdateAccount,
},
auth::{AdminAuditEvent, AdminSession},
client_keys::{
@@ -50,8 +56,8 @@ use gateway_admin::{
backup::BackupStorePorts,
provider::{ProviderAdmin, ProviderAdminError, ProviderAdminErrorKind},
store::{
AccountStore, AdminStoreError, AdminStoreErrorKind, AdminStorePorts, AdminStoreResult,
AuthStore, ClientKeyStore, ObservabilityStore, SettingsStore,
AccountGroupStore, AccountStore, AdminStoreError, AdminStoreErrorKind, AdminStorePorts,
AdminStoreResult, AuthStore, ClientKeyStore, ObservabilityStore, SettingsStore,
},
system::{
SystemOperationError, SystemOperationErrorKind, SystemOperations,
@@ -170,6 +176,7 @@ impl AdminHarness {
},
AdminStorePorts::new(
self.accounts,
Arc::new(UnavailableAccountGroupStore),
self.auth,
self.client_keys,
self.observability,
@@ -234,13 +241,60 @@ impl AuthStore for BootstrapAuthStore {
struct UnavailableStore;
struct UnavailableAccountGroupStore;
#[async_trait]
impl AccountGroupStore for UnavailableAccountGroupStore {
async fn list_account_groups(
&self,
_: AccountGroupListQuery,
) -> AdminStoreResult<AccountGroupPage> {
Err(unavailable("account groups"))
}
async fn create_account_group(
&self,
_: NewAccountGroup,
_: &MutationContext,
) -> AdminStoreResult<AccountGroupMutation> {
Err(unavailable("account group create"))
}
async fn update_account_group(
&self,
_: UpdateAccountGroup,
_: &MutationContext,
) -> AdminStoreResult<AccountGroupMutation> {
Err(unavailable("account group update"))
}
async fn set_account_group_enabled(
&self,
_: SetAccountGroupEnabled,
_: &MutationContext,
) -> AdminStoreResult<AccountGroupMutation> {
Err(unavailable("account group state"))
}
async fn delete_account_group(
&self,
_: DeleteAccountGroup,
_: &MutationContext,
) -> AdminStoreResult<AccountGroupMutation> {
Err(unavailable("account group delete"))
}
}
#[async_trait]
impl AccountStore for UnavailableStore {
async fn list_accounts(&self, _: AccountListQuery) -> AdminStoreResult<AccountPage> {
Err(unavailable("accounts"))
}
async fn load_account(&self, _: &str) -> AdminStoreResult<Option<AccountRecord>> {
async fn load_account(
&self,
_: &str,
) -> AdminStoreResult<Option<gateway_admin::model::accounts::AccountPageItem>> {
Err(unavailable("account"))
}
@@ -315,14 +369,30 @@ impl AccountStore for UnavailableStore {
Err(unavailable("credential refresh"))
}
async fn set_account_enabled(
async fn update_account(
&self,
_: SetAccountEnabled,
_: UpdateAccount,
_: &MutationContext,
) -> AdminStoreResult<Revision> {
) -> AdminStoreResult<AccountUpdateResult> {
Err(unavailable("account enabled"))
}
async fn recover_account(
&self,
_: &ProviderAccountId,
_: &MutationContext,
) -> AdminStoreResult<AccountUpdateResult> {
Err(unavailable("account recovery"))
}
async fn batch_update_accounts(
&self,
_: BatchUpdateAccounts,
_: &MutationContext,
) -> AdminStoreResult<AccountsUpdateResult> {
Err(unavailable("account batch update"))
}
async fn delete_accounts(
&self,
_: DeleteAccounts,
@@ -388,7 +458,11 @@ impl ClientKeyStore for UnavailableStore {
#[async_trait]
impl ObservabilityStore for UnavailableStore {
async fn dashboard_summary(&self, _: TimeRange) -> AdminStoreResult<DashboardObservation> {
async fn dashboard_summary(
&self,
_: TimeRange,
_: DateTime<Utc>,
) -> AdminStoreResult<DashboardObservation> {
Err(unavailable("dashboard"))
}
@@ -183,7 +183,8 @@ async fn dashboard_summary_should_project_rebuildable_runtime_slots() {
let now = Utc::now();
let store = Arc::new(FixtureObservabilityStore::new(observation_range(now)));
store.replace_runtime_slots(Some(DashboardRuntimeSlots {
active_accounts: 3,
inherited_accounts: 2,
overridden_slots: 5,
used_slots: Some(2),
}));
let services = observability_services(store).await;
@@ -195,9 +196,28 @@ async fn dashboard_summary_should_project_rebuildable_runtime_slots() {
.expect("dashboard summary")
.capacity;
assert_eq!(capacity.total_slots, 3);
assert_eq!(capacity.total_slots, 7);
assert_eq!(capacity.used_slots, Some(2));
assert_eq!(capacity.available_slots, Some(1));
assert_eq!(capacity.available_slots, Some(5));
}
#[tokio::test]
async fn dashboard_summary_should_share_one_current_observation_time_across_runtime_facts() {
let historical_end = Utc::now() - Duration::hours(1);
let store = Arc::new(FixtureObservabilityStore::new(observation_range(
historical_end,
)));
let services = observability_services(store.clone()).await;
services
.observability()
.dashboard_summary(observation_range(historical_end), TrendKind::Usage)
.await
.expect("dashboard summary");
let (summary_observed_at, slots_observed_at) = store.observed_times();
assert_eq!(summary_observed_at, slots_observed_at);
assert!(summary_observed_at.is_some_and(|value| value > historical_end));
}
#[tokio::test]
@@ -394,6 +414,8 @@ struct FixtureObservabilityStore {
calculated_billing_facts: Mutex<Vec<UsageCalculatedBillingFact>>,
diagnostics: Mutex<Vec<DiagnosticObservation>>,
runtime_slots: Mutex<Option<DashboardRuntimeSlots>>,
summary_observed_at: Mutex<Option<DateTime<Utc>>>,
slots_observed_at: Mutex<Option<DateTime<Utc>>>,
usage_records: Mutex<Vec<UsageRecord>>,
}
@@ -410,6 +432,8 @@ impl FixtureObservabilityStore {
calculated_billing_facts: Mutex::new(Vec::new()),
diagnostics: Mutex::new(Vec::new()),
runtime_slots: Mutex::new(None),
summary_observed_at: Mutex::new(None),
slots_observed_at: Mutex::new(None),
usage_records: Mutex::new(Vec::new()),
}
}
@@ -437,6 +461,16 @@ impl FixtureObservabilityStore {
*self.runtime_slots.lock().expect("runtime slots") = runtime_slots;
}
fn observed_times(&self) -> (Option<DateTime<Utc>>, Option<DateTime<Utc>>) {
(
*self
.summary_observed_at
.lock()
.expect("summary observed at"),
*self.slots_observed_at.lock().expect("slots observed at"),
)
}
fn replace_usage_records(&self, records: Vec<UsageRecord>) {
*self.usage_records.lock().expect("usage records") = records;
}
@@ -444,7 +478,15 @@ impl FixtureObservabilityStore {
#[async_trait]
impl ObservabilityStore for FixtureObservabilityStore {
async fn dashboard_summary(&self, range: TimeRange) -> AdminStoreResult<DashboardObservation> {
async fn dashboard_summary(
&self,
range: TimeRange,
observed_at: DateTime<Utc>,
) -> AdminStoreResult<DashboardObservation> {
*self
.summary_observed_at
.lock()
.expect("summary observed at") = Some(observed_at);
Ok(DashboardObservation {
range,
requests: RequestMetrics::default(),
@@ -459,8 +501,9 @@ impl ObservabilityStore for FixtureObservabilityStore {
async fn dashboard_runtime_slots(
&self,
_: DateTime<Utc>,
observed_at: DateTime<Utc>,
) -> AdminStoreResult<Option<DashboardRuntimeSlots>> {
*self.slots_observed_at.lock().expect("slots observed at") = Some(observed_at);
Ok(*self.runtime_slots.lock().expect("runtime slots"))
}
@@ -612,6 +655,9 @@ fn total_record(
id: id.to_owned(),
client_api_key_ref: "key_billing".to_owned(),
config_revision: 1,
routing_scope: "all".to_owned(),
routing_group_refs: Vec::new(),
routing_group_names_snapshot: Vec::new(),
protocol: "openai".to_owned(),
operation: "generate".to_owned(),
endpoint: "/v1/responses".to_owned(),
@@ -5,8 +5,8 @@ use gateway_core::engine::credential::ProviderAccountId;
use gateway_admin::{
AdminServices,
model::provider_credentials::{
AuthorizationMutationTarget, CompleteAuthorization, CredentialDeletion, CredentialMutation,
ImportCredentials, ProviderQuotaRequest, StartAuthorization,
AuthorizationMutationTarget, CompleteAuthorization, CredentialDeletion, ImportCredentials,
ProviderQuotaRequest, StartAuthorization,
},
ports::provider::ProviderAdminErrorKind,
};
@@ -15,55 +15,6 @@ use super::accounts::{
FakeAccountStore, FakeProviderAdmin, context, document, events, recorded, revision,
};
#[tokio::test]
async fn openai_enable_should_use_scoped_store_without_provider_mutation() {
let events = events();
let provider = FakeProviderAdmin::new("openai", events.clone());
let store = FakeAccountStore::new("openai", events.clone());
let services = service(provider, store.clone()).await;
services
.openai()
.enable(mutation("acct_test"))
.await
.expect("enable credential");
assert_eq!(
recorded(&events),
[
"store.credential_details",
"store.set_enabled",
"provider.account_facts_changed",
]
);
assert_eq!(store.audit_requests(), ["request-openai"]);
}
#[tokio::test]
async fn openai_disable_should_commit_then_release_provider_resources() {
let events = events();
let provider = FakeProviderAdmin::new("openai", events.clone());
let store = FakeAccountStore::new("openai", events.clone());
let services = service(provider, store.clone()).await;
services
.openai()
.disable(mutation("acct_test"))
.await
.expect("disable credential");
assert_eq!(
recorded(&events),
[
"store.credential_details",
"store.set_enabled",
"provider.account_unavailable",
"provider.account_facts_changed",
]
);
assert_eq!(store.audit_requests(), ["request-openai"]);
}
#[tokio::test]
async fn openai_delete_should_commit_then_release_provider_resources() {
let events = events();
@@ -89,26 +40,6 @@ async fn openai_delete_should_commit_then_release_provider_resources() {
assert_eq!(store.audit_requests(), ["request-openai"]);
}
#[tokio::test]
async fn openai_disable_store_failure_should_not_release_provider_resources() {
let events = events();
let provider = FakeProviderAdmin::new("openai", events.clone());
let store = FakeAccountStore::new("openai", events.clone());
store.fail_next_commit();
let services = service(provider, store).await;
services
.openai()
.disable(mutation("acct_test"))
.await
.expect_err("failed Store transaction");
assert_eq!(
recorded(&events),
["store.credential_details", "store.set_enabled"]
);
}
#[tokio::test]
async fn openai_import_should_prepare_before_atomic_store_commit() {
let events = events();
@@ -450,13 +381,6 @@ async fn await_quota_requests(provider: &FakeProviderAdmin, expected: usize) {
assert_eq!(provider.quota_requests().len(), expected);
}
fn mutation(account_id: &str) -> CredentialMutation {
CredentialMutation {
context: context("request-openai"),
account_id: ProviderAccountId::new(account_id).expect("account ID"),
}
}
fn deletion(account_id: &str) -> CredentialDeletion {
CredentialDeletion {
context: context("request-openai"),
@@ -4,37 +4,12 @@ use gateway_core::engine::credential::ProviderAccountId;
use gateway_admin::{
AdminServices,
model::provider_credentials::{CredentialMutation, ImportCredentials, ProviderQuotaRequest},
model::provider_credentials::{ImportCredentials, ProviderQuotaRequest},
ports::provider::ProviderAdminErrorKind,
};
use super::accounts::{FakeAccountStore, FakeProviderAdmin, context, document, events, recorded};
#[tokio::test]
async fn xai_disable_should_commit_then_notify_the_stateless_provider() {
let events = events();
let provider = FakeProviderAdmin::new("xai", events.clone());
let store = FakeAccountStore::new("xai", events.clone());
let services = service(provider.clone(), store.clone()).await;
services
.xai()
.disable(mutation("acct_test"))
.await
.expect("disable credential");
assert_eq!(
recorded(&events),
[
"store.credential_details",
"store.set_enabled",
"provider.account_unavailable",
"provider.account_facts_changed",
]
);
assert_eq!(store.audit_requests(), ["request-xai"]);
}
#[tokio::test]
async fn xai_import_should_prepare_before_atomic_store_commit() {
let events = events();
@@ -132,10 +107,3 @@ async fn service(provider: Arc<FakeProviderAdmin>, store: Arc<FakeAccountStore>)
.build()
.await
}
fn mutation(account_id: &str) -> CredentialMutation {
CredentialMutation {
context: context("request-xai"),
account_id: ProviderAccountId::new(account_id).expect("account ID"),
}
}
@@ -0,0 +1,417 @@
//! Account group HTTP wire and fixed routes.
use std::collections::BTreeMap;
use axum::{
Json, Router,
extract::{Query, State},
http::StatusCode,
response::IntoResponse,
routing::{get, post},
};
use chrono::{DateTime, Utc};
use gateway_admin::model::{
PageSize,
account_groups::{
AccountGroupAccountSummary, AccountGroupCapacity, AccountGroupColor, AccountGroupListQuery,
AccountGroupMutation, AccountGroupPage, AccountGroupRecord, AccountGroupUsage,
CreateAccountGroup, DeleteAccountGroup, SetAccountGroupEnabled, UpdateAccountGroup,
},
};
use gateway_core::routing::AccountGroupId;
use serde::{Deserialize, Serialize};
use super::{
AdminAuth, AdminEnvelope, AdminError, AdminResponse, AdminSessionState, PageMeta,
WireValidationError, wire::map_admin_service_error,
};
const DEFAULT_PAGE_SIZE: u32 = 50;
const MAX_PAGE_SIZE: u32 = 200;
#[derive(Debug, Clone, Default, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct ListAccountGroupsQuery {
page: Option<u32>,
page_size: Option<u32>,
search: Option<String>,
enabled: Option<bool>,
}
impl ListAccountGroupsQuery {
fn into_command(self) -> Result<AccountGroupListQuery, WireValidationError> {
let page = self.page.unwrap_or(1);
let page_size = self.page_size.unwrap_or(DEFAULT_PAGE_SIZE);
if page == 0 || page_size == 0 || page_size > MAX_PAGE_SIZE {
return Err(WireValidationError::new("page"));
}
let search = self
.search
.map(|value| value.trim().to_owned())
.filter(|value| !value.is_empty());
if search
.as_deref()
.is_some_and(|value| value.len() > 256 || value.chars().any(char::is_control))
{
return Err(WireValidationError::new("search"));
}
Ok(AccountGroupListQuery {
page,
page_size: PageSize::new(
u16::try_from(page_size).map_err(|_| WireValidationError::new("pageSize"))?,
)
.map_err(|_| WireValidationError::new("pageSize"))?,
search,
enabled: self.enabled,
})
}
}
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct CreateAccountGroupRequest {
name: String,
description: Option<String>,
color: String,
}
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct UpdateAccountGroupRequest {
id: String,
name: String,
description: Option<String>,
color: String,
}
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct AccountGroupIdRequest {
id: String,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct AccountGroupView {
id: String,
name: String,
description: Option<String>,
color: String,
enabled: bool,
member_count: u64,
provider_counts: BTreeMap<String, u64>,
client_key_count: u64,
account_summary: AccountGroupAccountSummaryView,
capacity: AccountGroupCapacityView,
usage: AccountGroupUsageView,
created_at: DateTime<Utc>,
updated_at: DateTime<Utc>,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct AccountGroupAccountSummaryView {
available: u64,
limited: u64,
total: u64,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct AccountGroupCapacityView {
used_slots: Option<u64>,
total_slots: u64,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct AccountGroupUsageView {
today_usd: String,
total_usd: String,
}
impl From<AccountGroupRecord> for AccountGroupView {
fn from(record: AccountGroupRecord) -> Self {
Self {
id: record.id.to_string(),
name: record.name,
description: record.description,
color: record.color.as_str().to_owned(),
enabled: record.enabled,
member_count: record.member_count,
provider_counts: record.provider_counts,
client_key_count: record.client_key_count,
account_summary: account_summary_view(record.account_summary),
capacity: capacity_view(record.capacity),
usage: usage_view(record.usage),
created_at: record.created_at,
updated_at: record.updated_at,
}
}
}
fn account_summary_view(summary: AccountGroupAccountSummary) -> AccountGroupAccountSummaryView {
AccountGroupAccountSummaryView {
available: summary.available,
limited: summary.limited,
total: summary.total,
}
}
fn capacity_view(capacity: AccountGroupCapacity) -> AccountGroupCapacityView {
AccountGroupCapacityView {
used_slots: capacity.used_slots,
total_slots: capacity.total_slots,
}
}
fn usage_view(usage: AccountGroupUsage) -> AccountGroupUsageView {
AccountGroupUsageView {
today_usd: usage.today_usd.to_string(),
total_usd: usage.total_usd.to_string(),
}
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct AccountGroupPageData {
items: Vec<AccountGroupView>,
page: PageMeta,
config_revision: u64,
}
impl From<AccountGroupPage> for AccountGroupPageData {
fn from(page: AccountGroupPage) -> Self {
let total_pages = if page.total == 0 {
0
} else {
page.total.div_ceil(u64::from(page.page_size))
};
Self {
items: page.items.into_iter().map(Into::into).collect(),
page: PageMeta::new(
page.page,
u32::from(page.page_size),
page.total,
u32::try_from(total_pages).unwrap_or(u32::MAX),
),
config_revision: page.config_revision.get(),
}
}
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct AccountGroupMutationData {
id: String,
record: Option<AccountGroupView>,
config_revision: u64,
}
impl From<AccountGroupMutation> for AccountGroupMutationData {
fn from(mutation: AccountGroupMutation) -> Self {
Self {
id: mutation.id.to_string(),
record: mutation.record.map(Into::into),
config_revision: mutation.config_revision.get(),
}
}
}
/// Construct all fixed account-group management routes.
pub fn router<S>() -> Router<S>
where
S: AdminSessionState + Clone + Send + Sync + 'static,
{
Router::new()
.route("/api/admin/account-groups", get(list::<S>))
.route("/api/admin/account-groups/create", post(create::<S>))
.route("/api/admin/account-groups/update", post(update::<S>))
.route("/api/admin/account-groups/enable", post(enable::<S>))
.route("/api/admin/account-groups/disable", post(disable::<S>))
.route("/api/admin/account-groups/delete", post(delete::<S>))
}
async fn list<S>(
_auth: AdminAuth,
State(state): State<S>,
Query(query): Query<ListAccountGroupsQuery>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let result = state
.admin_services()
.account_groups()
.list(query.into_command().map_err(map_wire_error)?)
.await
.map_err(map_service_error)?;
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(AccountGroupPageData::from(result)),
))
}
async fn create<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<CreateAccountGroupRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
validate_group_fields(&request.name, request.description.as_deref())?;
mutation_response(
StatusCode::CREATED,
state
.admin_services()
.account_groups()
.create(
&auth.context().mutation_context(),
CreateAccountGroup {
name: request.name,
description: request.description,
color: group_color(&request.color)?,
},
)
.await,
)
}
async fn update<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<UpdateAccountGroupRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
validate_group_fields(&request.name, request.description.as_deref())?;
mutation_response(
StatusCode::OK,
state
.admin_services()
.account_groups()
.update(
&auth.context().mutation_context(),
UpdateAccountGroup {
id: group_id(request.id)?,
name: request.name,
description: request.description,
color: group_color(&request.color)?,
},
)
.await,
)
}
async fn enable<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<AccountGroupIdRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
set_enabled(auth, state, request, true).await
}
async fn disable<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<AccountGroupIdRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
set_enabled(auth, state, request, false).await
}
async fn set_enabled<S>(
auth: AdminAuth,
state: S,
request: AccountGroupIdRequest,
enabled: bool,
) -> Result<AdminResponse<AdminEnvelope<AccountGroupMutationData>>, AdminError>
where
S: AdminSessionState + Send + Sync,
{
mutation_response(
StatusCode::OK,
state
.admin_services()
.account_groups()
.set_enabled(
&auth.context().mutation_context(),
SetAccountGroupEnabled {
id: group_id(request.id)?,
enabled,
},
)
.await,
)
}
async fn delete<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<AccountGroupIdRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
mutation_response(
StatusCode::OK,
state
.admin_services()
.account_groups()
.delete(
&auth.context().mutation_context(),
DeleteAccountGroup {
id: group_id(request.id)?,
},
)
.await,
)
}
fn mutation_response(
status: StatusCode,
result: Result<AccountGroupMutation, gateway_admin::model::AdminError>,
) -> Result<AdminResponse<AdminEnvelope<AccountGroupMutationData>>, AdminError> {
let result = result.map_err(map_service_error)?;
Ok(AdminResponse::new(
status,
AdminEnvelope::ok(AccountGroupMutationData::from(result)),
))
}
fn group_id(value: String) -> Result<AccountGroupId, AdminError> {
AccountGroupId::new(value).map_err(|_| AdminError::bad_request("Invalid account group ID"))
}
fn group_color(value: &str) -> Result<AccountGroupColor, AdminError> {
AccountGroupColor::parse(value)
.ok_or_else(|| AdminError::bad_request("Invalid account group color"))
}
fn validate_group_fields(name: &str, description: Option<&str>) -> Result<(), AdminError> {
if name.trim() != name
|| name.is_empty()
|| name.chars().count() > 100
|| name.chars().any(char::is_control)
|| description
.is_some_and(|value| value.len() > 4096 || value.chars().any(char::is_control))
{
return Err(AdminError::bad_request("Invalid account group request"));
}
Ok(())
}
fn map_wire_error(_: WireValidationError) -> AdminError {
AdminError::bad_request("Invalid account group query")
}
fn map_service_error(error: gateway_admin::model::AdminError) -> AdminError {
map_admin_service_error(error, "Account group repository unavailable")
}
+273 -99
View File
@@ -18,9 +18,10 @@ use futures::{Stream, StreamExt as _};
use gateway_admin::model::{
AdminError as AdminServiceError, PageSize,
accounts::{
AccountConnectionTestEvent as DomainConnectionTestEvent, AccountCost, AccountListQuery,
AccountModelUsage, AccountSort, AccountSortField, AccountStatus as DomainAccountStatus,
AccountUsage, SortDirection,
AccountConcurrencyLimit, AccountConnectionTestEvent as DomainConnectionTestEvent,
AccountCost, AccountGroupFilter, AccountListQuery, AccountModelUsage, AccountSort,
AccountSortField, AccountStatus as DomainAccountStatus, AccountUpdateResult, AccountUsage,
AccountWeight, AccountsUpdateResult, BatchUpdateAccounts, SortDirection, UpdateAccount,
},
provider_credentials::{
AccountDirectoryItem, AccountDirectoryPage, AccountExportBundle, AccountRefreshResult,
@@ -32,12 +33,12 @@ use gateway_admin::model::{
};
use gateway_core::{
engine::credential::{OpaqueProviderData, ProviderAccountId},
routing::{ProviderKind, UpstreamModelId},
routing::{AccountGroupId, ProviderKind, UpstreamModelId},
};
use serde::{Deserialize, Serialize};
use serde::{Deserialize, Deserializer, Serialize};
use serde_json::{Map, Value};
use super::presenter::{format_compact_number, format_number};
use super::presenter::{format_compact_number, format_decimal_currency, format_number};
use super::{
AdminAuth, AdminEnvelope, AdminError, AdminResponse, AdminSessionState, PageMeta,
WireValidationError,
@@ -51,8 +52,10 @@ const MAX_NAME_BYTES: usize = 512;
const MAX_IMPORT_DATA_BYTES: usize = 64 * 1024 * 1024;
const MAX_ACCESS_TOKEN_BYTES: usize = 16 * 1024;
const MAX_REFRESH_TOKEN_BYTES: usize = 64 * 1024;
const MAX_ID_TOKEN_BYTES: usize = 16 * 1024;
const MAX_CALLBACK_URL_BYTES: usize = 64 * 1024;
const MAX_ACCOUNT_DELETE_BATCH: usize = 200;
const MAX_ACCOUNT_GROUP_BATCH: usize = 1000;
/// 账号列表查询参数。
#[derive(Debug, Clone, Default, Deserialize)]
@@ -61,12 +64,74 @@ pub struct ListQuery {
pub page: Option<u32>,
pub page_size: Option<u32>,
pub provider: Option<String>,
pub group_id: Option<String>,
pub search: Option<String>,
pub status: Option<String>,
pub sort_by: Option<String>,
pub sort_direction: Option<String>,
}
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct BatchUpdateAccountsRequest {
pub account_ids: Vec<String>,
pub enabled: bool,
#[serde(deserialize_with = "deserialize_required_nullable")]
pub concurrency_limit: Option<u64>,
pub weight: u64,
pub group_ids: Vec<String>,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct BatchUpdatedAccountsData {
account_ids: Vec<String>,
config_revision: u64,
}
impl From<AccountsUpdateResult> for BatchUpdatedAccountsData {
fn from(result: AccountsUpdateResult) -> Self {
Self {
account_ids: result
.account_ids
.into_iter()
.map(|id| id.to_string())
.collect(),
config_revision: result.config_revision.get(),
}
}
}
impl BatchUpdateAccountsRequest {
pub fn validate(&self) -> Result<(), WireValidationError> {
if self.account_ids.is_empty()
|| self.account_ids.len() > MAX_ACCOUNT_GROUP_BATCH
|| self
.account_ids
.iter()
.any(|id| require_account_id(id, "accountIds").is_err())
|| self.account_ids.iter().collect::<BTreeSet<_>>().len() != self.account_ids.len()
{
return Err(WireValidationError::new("accountIds"));
}
validate_wire_group_ids(&self.group_ids)?;
parse_concurrency_limit(self.concurrency_limit)?;
parse_account_weight(self.weight)?;
Ok(())
}
fn into_command(self) -> Result<BatchUpdateAccounts, WireValidationError> {
self.validate()?;
Ok(BatchUpdateAccounts {
account_ids: self.account_ids,
enabled: self.enabled,
concurrency_limit: parse_concurrency_limit(self.concurrency_limit)?,
weight: parse_account_weight(self.weight)?,
group_ids: validate_wire_group_ids(&self.group_ids)?,
})
}
}
impl ListQuery {
/// 解析并校验全部 wire 字段,生成 Admin 查询命令。
pub fn validate(self) -> Result<AccountListQuery, WireValidationError> {
@@ -79,6 +144,14 @@ impl ListQuery {
return Err(WireValidationError::new("pageSize"));
}
let provider_kind = parse_provider(self.provider.as_deref().unwrap_or("all"))?;
let group_filter = match self.group_id.as_deref().map(str::trim) {
None | Some("") => None,
Some("ungrouped") => Some(AccountGroupFilter::Ungrouped),
Some(value) => Some(AccountGroupFilter::Group(
AccountGroupId::new(value.to_owned())
.map_err(|_| WireValidationError::new("groupId"))?,
)),
};
let search = self
.search
.map(|value| value.trim().to_owned())
@@ -111,6 +184,7 @@ impl ListQuery {
)
.map_err(|_| WireValidationError::new("pageSize"))?,
provider_kind,
group_filter,
search,
status,
sort,
@@ -176,6 +250,7 @@ pub struct AccountView {
pub id: String,
pub name: String,
pub provider: String,
pub groups: Vec<AccountGroupRefView>,
pub resource_ref: String,
pub email: Option<String>,
pub account_id: Option<String>,
@@ -185,13 +260,13 @@ pub struct AccountView {
pub authentication_kind: String,
pub has_refresh_token: bool,
pub status: String,
pub display_status: String,
/// `status == "error"` 时的具体原因;其余状态为 `null`。
pub error_reason: Option<String>,
/// 最近一次失败的上游错误描述;仅错误状态存在。
pub error_message: Option<String>,
pub availability: String,
pub enabled: bool,
pub concurrency_limit: Option<u32>,
pub weight: u16,
pub access_token_expires_at: Option<String>,
pub access_token_expires_at_display: Option<String>,
pub refresh_token_expires_at: Option<String>,
@@ -204,6 +279,15 @@ pub struct AccountView {
pub usage: AccountUsageView,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AccountGroupRefView {
pub id: String,
pub name: String,
pub color: String,
pub enabled: bool,
}
/// Provider quota 安全视图。
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
@@ -522,10 +606,9 @@ impl From<DomainConnectionTestEvent> for AccountConnectionTestEvent {
DomainConnectionTestEvent::Content { text } => {
serde_json::json!({ "type": "content", "text": text })
}
DomainConnectionTestEvent::Completed { account_status } => serde_json::json!({
DomainConnectionTestEvent::Completed => serde_json::json!({
"type": "test_complete",
"success": true,
"accountStatus": account_status.as_str()
"success": true
}),
DomainConnectionTestEvent::Failed {
message,
@@ -534,7 +617,6 @@ impl From<DomainConnectionTestEvent> for AccountConnectionTestEvent {
upstream_status,
upstream_content_type,
upstream_body,
account_status,
} => serde_json::json!({
"type": "error",
"error": message,
@@ -542,8 +624,7 @@ impl From<DomainConnectionTestEvent> for AccountConnectionTestEvent {
"providerErrorType": provider_error_type,
"upstreamStatus": upstream_status,
"upstreamContentType": upstream_content_type,
"upstreamBody": upstream_body,
"accountStatus": account_status.as_str()
"upstreamBody": upstream_body
}),
};
Self { data }
@@ -686,31 +767,49 @@ impl CompleteAccountAuthorizationRequest {
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct AccountMutationRequest {
pub provider: String,
pub struct UpdateAccountRequest {
pub account_id: String,
pub enabled: bool,
#[serde(deserialize_with = "deserialize_required_nullable")]
pub concurrency_limit: Option<u64>,
pub weight: u64,
pub group_ids: Vec<String>,
}
impl AccountMutationRequest {
impl UpdateAccountRequest {
pub fn validate(&self) -> Result<(), WireValidationError> {
AccountProvider::parse(&self.provider)?;
require_account_id(&self.account_id, "accountId")
require_account_id(&self.account_id, "accountId")?;
parse_concurrency_limit(self.concurrency_limit)?;
parse_account_weight(self.weight)?;
validate_wire_group_ids(&self.group_ids)?;
Ok(())
}
fn into_command(
self,
context: gateway_admin::model::MutationContext,
) -> Result<(AccountProvider, CredentialMutation), WireValidationError> {
fn into_command(self) -> Result<UpdateAccount, WireValidationError> {
self.validate()?;
let provider = AccountProvider::parse(&self.provider)?;
Ok((
provider,
CredentialMutation {
context,
account_id: ProviderAccountId::new(self.account_id)
.map_err(|_| WireValidationError::new("accountId"))?,
},
))
Ok(UpdateAccount {
account_id: self.account_id,
enabled: self.enabled,
concurrency_limit: parse_concurrency_limit(self.concurrency_limit)?,
weight: parse_account_weight(self.weight)?,
group_ids: validate_wire_group_ids(&self.group_ids)?,
})
}
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct UpdatedAccountData {
pub account_id: String,
pub config_revision: u64,
}
impl From<AccountUpdateResult> for UpdatedAccountData {
fn from(result: AccountUpdateResult) -> Self {
Self {
account_id: result.account_id.to_string(),
config_revision: result.config_revision.get(),
}
}
}
@@ -765,6 +864,7 @@ pub struct RotateAccountRequest {
pub account_id: String,
pub access_token: String,
pub refresh_token: Option<String>,
pub id_token: Option<String>,
}
impl RotateAccountRequest {
@@ -773,7 +873,11 @@ impl RotateAccountRequest {
return Err(WireValidationError::new("provider"));
}
require_account_id(&self.account_id, "accountId")?;
validate_oauth_material(&self.access_token, self.refresh_token.as_deref())
validate_oauth_material(
&self.access_token,
self.refresh_token.as_deref(),
self.id_token.as_deref(),
)
}
fn into_command(
@@ -787,6 +891,10 @@ impl RotateAccountRequest {
"refresh_token".to_owned(),
self.refresh_token.map_or(Value::Null, Value::String),
);
material.insert(
"id_token".to_owned(),
self.id_token.map_or(Value::Null, Value::String),
);
Ok(RotateCredential {
mutation: CredentialMutation {
context,
@@ -805,8 +913,8 @@ pub struct AccountImportData {
pub account_ids: Vec<String>,
}
impl From<CredentialImportResult> for AccountImportData {
fn from(result: CredentialImportResult) -> Self {
impl AccountImportData {
pub fn from_result(result: CredentialImportResult) -> Self {
let account_ids = result
.credential_ids
.into_iter()
@@ -883,6 +991,46 @@ fn require_account_id(value: &str, field: &'static str) -> Result<(), WireValida
Ok(())
}
fn parse_concurrency_limit(
value: Option<u64>,
) -> Result<Option<AccountConcurrencyLimit>, WireValidationError> {
value
.map(|value| {
u32::try_from(value)
.ok()
.and_then(AccountConcurrencyLimit::new)
.ok_or_else(|| WireValidationError::new("concurrencyLimit"))
})
.transpose()
}
fn deserialize_required_nullable<'de, D>(deserializer: D) -> Result<Option<u64>, D::Error>
where
D: Deserializer<'de>,
{
Option::<u64>::deserialize(deserializer)
}
fn parse_account_weight(value: u64) -> Result<AccountWeight, WireValidationError> {
u16::try_from(value)
.ok()
.and_then(AccountWeight::new)
.ok_or_else(|| WireValidationError::new("weight"))
}
fn validate_wire_group_ids(values: &[String]) -> Result<Vec<AccountGroupId>, WireValidationError> {
if values.len() > MAX_ACCOUNT_GROUP_BATCH
|| values.iter().collect::<BTreeSet<_>>().len() != values.len()
{
return Err(WireValidationError::new("groupIds"));
}
values
.iter()
.cloned()
.map(|value| AccountGroupId::new(value).map_err(|_| WireValidationError::new("groupIds")))
.collect()
}
/// 构造统一账号管理路由。
pub fn router<S>() -> Router<S>
where
@@ -894,10 +1042,14 @@ where
.route("/api/admin/accounts/export", get(export_accounts::<S>))
.route("/api/admin/accounts/import", post(import_accounts::<S>))
.route("/api/admin/accounts/refresh", post(refresh_account::<S>))
.route("/api/admin/accounts/recover", post(recover_account::<S>))
.route("/api/admin/accounts/rotate", post(rotate_account::<S>))
.route("/api/admin/accounts/enable", post(enable_account::<S>))
.route("/api/admin/accounts/disable", post(disable_account::<S>))
.route("/api/admin/accounts/update", post(update_account::<S>))
.route("/api/admin/accounts/delete", post(delete_accounts::<S>))
.route(
"/api/admin/accounts/batch-update",
post(batch_update_accounts::<S>),
)
.route("/api/admin/accounts/quota", get(account_quota::<S>))
.route(
"/api/admin/accounts/quota/refresh",
@@ -922,6 +1074,27 @@ where
)
}
async fn batch_update_accounts<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<BatchUpdateAccountsRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let command = request.into_command().map_err(map_wire_error)?;
let result = state
.admin_services()
.accounts()
.batch_update(&auth.context().mutation_context(), command)
.await
.map_err(map_service_error)?;
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(BatchUpdatedAccountsData::from(result)),
))
}
async fn list_accounts<S>(
_auth: AdminAuth,
State(state): State<S>,
@@ -1011,7 +1184,7 @@ where
.map_err(map_service_error)?;
Ok(AdminResponse::new(
StatusCode::CREATED,
AdminEnvelope::ok(AccountImportData::from(result)),
AdminEnvelope::ok(AccountImportData::from_result(result)),
))
}
@@ -1106,64 +1279,24 @@ where
))
}
async fn enable_account<S>(
async fn update_account<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<AccountMutationRequest>,
Json(request): Json<UpdateAccountRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
mutate_account(auth, state, request, AccountMutationAction::Enable).await
}
async fn disable_account<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<AccountMutationRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
mutate_account(auth, state, request, AccountMutationAction::Disable).await
}
#[derive(Clone, Copy)]
enum AccountMutationAction {
Enable,
Disable,
}
async fn mutate_account<S>(
auth: AdminAuth,
state: S,
request: AccountMutationRequest,
action: AccountMutationAction,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let (provider, command) = request
.into_command(auth.context().mutation_context())
.map_err(map_wire_error)?;
let result = match (provider, action) {
(AccountProvider::OpenAi, AccountMutationAction::Enable) => {
state.admin_services().openai().enable(command).await
}
(AccountProvider::OpenAi, AccountMutationAction::Disable) => {
state.admin_services().openai().disable(command).await
}
(AccountProvider::Xai, AccountMutationAction::Enable) => {
state.admin_services().xai().enable(command).await
}
(AccountProvider::Xai, AccountMutationAction::Disable) => {
state.admin_services().xai().disable(command).await
}
}
.map_err(map_service_error)?;
let command = request.into_command().map_err(map_wire_error)?;
let result = state
.admin_services()
.accounts()
.update(&auth.context().mutation_context(), command)
.await
.map_err(map_service_error)?;
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(AccountMutationData::from(result)),
AdminEnvelope::ok(UpdatedAccountData::from(result)),
))
}
@@ -1208,6 +1341,25 @@ where
Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data)))
}
async fn recover_account<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<AccountActionRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let account_id = request.into_id().map_err(map_wire_error)?;
let result = state
.admin_services()
.accounts()
.recover(&auth.context().mutation_context(), account_id)
.await
.map_err(map_service_error)?;
let data = account_refresh_data(result, Utc::now());
Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data)))
}
async fn account_quota<S>(
_auth: AdminAuth,
State(state): State<S>,
@@ -1365,20 +1517,30 @@ fn account_models_data(result: ProviderModels) -> AccountModelsData {
fn account_view(item: AccountDirectoryItem, now: DateTime<Utc>) -> AccountView {
let AccountDirectoryItem {
account,
status,
error_reason,
projection,
usage,
quota,
} = item;
let status = status.as_str().to_owned();
let status = projection.status.as_str().to_owned();
let rate_limited_until = projection.rate_limited_until.map(DateTime::<Utc>::from);
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);
let (quota, refresh_token_expires_at) = account_quota_view(quota, now);
let (quota, refresh_token_expires_at) = account_quota_view(quota, rate_limited_until, now);
AccountView {
id: account.id.clone(),
name: account.name,
provider: account.provider_kind.to_string(),
groups: account
.groups
.into_iter()
.map(|group| AccountGroupRefView {
id: group.id.to_string(),
name: group.name,
color: group.color.as_str().to_owned(),
enabled: group.enabled,
})
.collect(),
resource_ref: account.id,
email: account.email,
account_id: account.upstream_account_id,
@@ -1387,12 +1549,14 @@ fn account_view(item: AccountDirectoryItem, now: DateTime<Utc>) -> AccountView {
plan_type: account.plan_type,
authentication_kind: account.authentication_kind,
has_refresh_token: account.has_refresh_token,
status: status.clone(),
display_status: status,
error_reason: error_reason.map(|reason| reason.as_str().to_owned()),
error_message: account.last_error_message,
availability: account.availability.as_str().to_owned(),
status,
error_reason: projection
.error_reason
.map(|reason| reason.as_str().to_owned()),
error_message: projection.error_message,
enabled: account.enabled,
concurrency_limit: account.concurrency_limit.map(|limit| limit.get()),
weight: account.weight.get(),
access_token_expires_at: expires_at,
access_token_expires_at_display: account
.access_token_expires_at
@@ -1411,6 +1575,7 @@ fn account_view(item: AccountDirectoryItem, now: DateTime<Utc>) -> AccountView {
fn account_quota_view(
quota: ProviderQuota,
rate_limited_until: Option<DateTime<Utc>>,
now: DateTime<Utc>,
) -> (AccountQuotaView, Option<String>) {
let refresh_token_expires_at = quota
@@ -1420,7 +1585,7 @@ 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));
let rate_limited_until = rate_limited_until.map(|until| china_datetime(&until));
(
AccountQuotaView {
refreshed_at_display,
@@ -1438,6 +1603,7 @@ fn quota_window_view(window: ProviderQuotaWindow) -> AccountQuotaWindowView {
group,
label,
source: _,
local_usage_attribution: _,
window_seconds,
used_percent,
reset_at,
@@ -1589,7 +1755,7 @@ fn account_model_usage_view(usage: AccountModelUsage, now: DateTime<Utc>) -> Mod
billing_amount_usd: usd.map(|cost| cost.amount.as_str().to_owned()),
billing_amount_usd_display: usd.map_or_else(
|| "—".to_owned(),
|cost| format!("${}", cost.amount.as_str()),
|cost| format_decimal_currency(cost.amount.as_str(), "USD"),
),
cost_estimate_status: cost_estimate_status.to_owned(),
known_cost_count: known_count,
@@ -1684,6 +1850,7 @@ fn require_text(
fn validate_oauth_material(
access_token: &str,
refresh_token: Option<&str>,
id_token: Option<&str>,
) -> Result<(), WireValidationError> {
if access_token.len() > MAX_ACCESS_TOKEN_BYTES
|| !valid_visible_ascii(access_token)
@@ -1698,6 +1865,13 @@ fn validate_oauth_material(
}) {
return Err(WireValidationError::new("refreshToken"));
}
if id_token.is_some_and(|token| {
token.len() > MAX_ID_TOKEN_BYTES
|| !valid_visible_ascii(token)
|| !valid_compact_jwt_shape(token)
}) {
return Err(WireValidationError::new("idToken"));
}
Ok(())
}
@@ -12,7 +12,7 @@ use gateway_admin::model::client_keys::{
};
use gateway_core::{
policy::{ClientApiKeyId, RateLimits},
routing::ProviderKind,
routing::AccountGroupId,
};
use serde::{Deserialize, Serialize};
@@ -137,7 +137,7 @@ impl ClientKeySort {
pub struct CreateClientKeyRequest {
name: String,
label: Option<String>,
provider_kind: String,
group_ids: Vec<String>,
max_concurrency: u64,
requests_per_minute: u64,
}
@@ -147,14 +147,13 @@ impl CreateClientKeyRequest {
pub fn into_command(self) -> Result<CreateClientKey, WireValidationError> {
validate_required_text(&self.name, "name")?;
validate_optional_text(self.label.as_deref(), "label")?;
validate_provider_kind(&self.provider_kind)?;
let group_ids = validate_group_ids(self.group_ids)?;
validate_limit(self.max_concurrency, "maxConcurrency")?;
validate_limit(self.requests_per_minute, "requestsPerMinute")?;
Ok(CreateClientKey {
name: self.name,
label: self.label,
provider_kind: ProviderKind::new(self.provider_kind)
.map_err(|_| WireValidationError::new("providerKind"))?,
group_ids,
limits: RateLimits {
max_concurrency: self.max_concurrency,
requests_per_minute: self.requests_per_minute,
@@ -170,7 +169,7 @@ pub struct UpdateClientKeyRequest {
id: String,
name: String,
label: Option<String>,
provider_kind: String,
group_ids: Vec<String>,
max_concurrency: u64,
requests_per_minute: u64,
}
@@ -181,15 +180,14 @@ impl UpdateClientKeyRequest {
validate_required_text(&self.id, "id")?;
validate_required_text(&self.name, "name")?;
validate_optional_text(self.label.as_deref(), "label")?;
validate_provider_kind(&self.provider_kind)?;
let group_ids = validate_group_ids(self.group_ids)?;
validate_limit(self.max_concurrency, "maxConcurrency")?;
validate_limit(self.requests_per_minute, "requestsPerMinute")?;
Ok(UpdateClientKey {
id: client_key_id(self.id, "clientKeyMutationNotFound")?,
name: self.name,
label: self.label,
provider_kind: ProviderKind::new(self.provider_kind)
.map_err(|_| WireValidationError::new("providerKind"))?,
group_ids,
limits: RateLimits {
max_concurrency: self.max_concurrency,
requests_per_minute: self.requests_per_minute,
@@ -242,7 +240,9 @@ pub struct ClientKeyView {
id: String,
name: String,
label: Option<String>,
provider_kind: String,
routing_scope: &'static str,
groups: Vec<ClientKeyGroupView>,
provider_kinds: Vec<String>,
prefix: String,
enabled: bool,
max_concurrency: u64,
@@ -252,13 +252,42 @@ pub struct ClientKeyView {
last_used_at: Option<DateTime<Utc>>,
}
#[derive(Debug, Clone, PartialEq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct ClientKeyGroupView {
id: String,
name: String,
color: String,
enabled: bool,
}
impl From<ClientKeyRecord> for ClientKeyView {
fn from(record: ClientKeyRecord) -> Self {
let routing_scope = if record.groups.is_empty() {
"all"
} else {
"groups"
};
Self {
id: record.id.to_string(),
name: record.name,
label: record.label,
provider_kind: record.provider_kind.to_string(),
routing_scope,
groups: record
.groups
.into_iter()
.map(|group| ClientKeyGroupView {
id: group.id.to_string(),
name: group.name,
color: group.color.as_str().to_owned(),
enabled: group.enabled,
})
.collect(),
provider_kinds: record
.provider_kinds
.into_iter()
.map(|provider| provider.to_string())
.collect(),
prefix: record.prefix,
enabled: record.enabled,
max_concurrency: record.limits.max_concurrency,
@@ -572,21 +601,20 @@ fn validate_required_text(value: &str, field: &'static str) -> Result<(), WireVa
Ok(())
}
fn validate_provider_kind(value: &str) -> Result<(), WireValidationError> {
let valid = !value.is_empty()
&& value.len() <= 64
&& value.bytes().enumerate().all(|(index, byte)| {
byte.is_ascii_lowercase()
|| byte.is_ascii_digit()
|| (byte == b'-' && index > 0 && index + 1 < value.len())
})
&& !value.starts_with('-')
&& !value.ends_with('-');
if valid {
Ok(())
} else {
Err(WireValidationError::new("providerKind"))
fn validate_group_ids(values: Vec<String>) -> Result<Vec<AccountGroupId>, WireValidationError> {
if values.len() > 1000
|| values
.iter()
.collect::<std::collections::BTreeSet<_>>()
.len()
!= values.len()
{
return Err(WireValidationError::new("groupIds"));
}
values
.into_iter()
.map(|value| AccountGroupId::new(value).map_err(|_| WireValidationError::new("groupIds")))
.collect()
}
fn validate_optional_text(
@@ -7,6 +7,7 @@ use axum::{
response::Response,
};
pub mod account_groups;
pub mod accounts;
pub mod auth;
pub mod backups;
@@ -29,6 +30,7 @@ where
S: AdminSessionState + Clone + Send + Sync + 'static,
{
Router::new()
.merge(account_groups::router::<S>())
.merge(accounts::router::<S>())
.merge(auth::router::<S>())
.merge(backups::router::<S>())
@@ -15,7 +15,7 @@ use gateway_admin::model::{PageSize as DomainPageSize, observability as domain};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use super::presenter::{format_compact_number, format_number};
use super::presenter::{format_compact_number, format_decimal_currency, format_number};
use super::{
AdminAuth, AdminEnvelope, AdminError, AdminResponse, AdminSessionState, PageMeta,
WireValidationError, wire::map_admin_service_error,
@@ -127,20 +127,6 @@ pub(crate) fn token_details(record: &domain::UsageRecord) -> TokenDetailsView {
}
}
pub(crate) fn format_decimal_currency(amount: &str, currency: &str) -> String {
if currency == "USD" {
let value = amount.parse::<f64>().map_or(0.0, |value| value);
let precision = if value != 0.0 && value.abs() < 0.01 {
4
} else {
2
};
format!("${value:.precision$}")
} else {
format!("{currency} {amount}")
}
}
pub(crate) fn format_money(cost: &domain::CurrencyCost) -> String {
format_decimal_currency(cost.amount.as_str(), &cost.currency)
}
@@ -246,6 +232,9 @@ pub(crate) fn usage_record_view(record: domain::UsageRecord) -> UsageRecordView
id: record.id.clone(),
request_id: record.id,
client_api_key_id: Some(record.client_api_key_ref),
routing_scope: record.routing_scope,
routing_group_refs: record.routing_group_refs,
routing_group_names_snapshot: record.routing_group_names_snapshot,
kind: record.operation,
provider: record.provider_kind,
authentication_kind: record.provider_account_authentication_kind,
@@ -747,10 +736,16 @@ pub(crate) fn dashboard_view(
credentials: DashboardCredentialsCardView {
total: format_compact_number(provider_accounts.total),
total_value: provider_accounts.total,
available: format_compact_number(provider_accounts.active),
available_value: provider_accounts.active,
unavailable: format_compact_number(provider_accounts.unavailable),
unavailable_value: provider_accounts.unavailable,
available: format_compact_number(provider_accounts.normal),
available_value: provider_accounts.normal,
unavailable: format_compact_number(
provider_accounts
.total
.saturating_sub(provider_accounts.normal),
),
unavailable_value: provider_accounts
.total
.saturating_sub(provider_accounts.normal),
},
traffic: DashboardTrafficCardView {
today_requests: format_compact_number(today.request_count),
@@ -783,14 +778,11 @@ pub(crate) fn dashboard_view(
usage_records: recent_requests.into_iter().map(usage_record_view).collect(),
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,
normal: provider_accounts.normal,
quota_exhausted: provider_accounts.quota_exhausted,
refreshing: provider_accounts.refreshing,
rate_limited: provider_accounts.rate_limited,
disabled: provider_accounts.disabled,
banned: provider_accounts.banned,
error: provider_accounts.error,
},
capacity_info: DashboardCapacityInfoView {
max_concurrent_per_account: capacity.max_concurrent_per_account,
@@ -84,6 +84,9 @@ pub struct UsageRecordView {
pub id: String,
pub request_id: String,
pub client_api_key_id: Option<String>,
pub routing_scope: String,
pub routing_group_refs: Vec<String>,
pub routing_group_names_snapshot: Vec<String>,
pub kind: String,
pub provider: Option<String>,
pub authentication_kind: Option<String>,
@@ -347,14 +350,11 @@ pub struct DashboardAccountRequestBucketView {
#[serde(rename_all = "camelCase")]
pub struct DashboardPoolSummaryView {
pub total: u64,
pub active: u64,
pub rate_limited: u64,
pub expired: u64,
pub invalid: u64,
pub normal: u64,
pub quota_exhausted: u64,
pub refreshing: Option<u64>,
pub rate_limited: u64,
pub disabled: u64,
pub banned: u64,
pub error: u64,
}
/// 同 target 账号调度容量;Redis 未提供聚合事实时运行中槽位保持空值。
@@ -3,6 +3,8 @@
//! 按 docs/project-redundancy-boundary-audit.md BE-02 的目标,display 字段最终
//! 迁移到前端;过渡期统一由本模块提供,避免 accounts/observability 各自维护。
use gateway_core::accounting::Decimal;
/// 千分位格式化。
#[must_use]
pub fn format_number(value: u64) -> String {
@@ -37,3 +39,33 @@ pub fn format_compact_number(value: u64) -> String {
}
format_number(value)
}
/// 货币展示格式化。USD 常规金额保留两位,小于 1 美元时最多保留四位。
#[must_use]
pub fn format_decimal_currency(amount: &str, currency: &str) -> String {
if currency != "USD" {
return format!("{currency} {amount}");
}
const DECIMAL_SCALE: u128 = 10_000_000_000;
let scaled = amount.parse::<Decimal>().unwrap_or_default().scaled();
let precision = if scaled != 0 && scaled < DECIMAL_SCALE {
4_u32
} else {
2_u32
};
let rounding_unit = 10_u128.pow(10 - precision);
let rounded = (scaled + rounding_unit / 2) / rounding_unit;
let display_scale = 10_u128.pow(precision);
let whole = rounded / display_scale;
let fraction = rounded % display_scale;
let mut display = format!("{whole}.{fraction:0width$}", width = precision as usize);
while display.ends_with('0')
&& display
.split_once('.')
.is_some_and(|(_, fraction)| fraction.len() > 2)
{
display.pop();
}
format!("${display}")
}
@@ -89,6 +89,7 @@ pub fn gateway_error_from_engine(error: &EngineError) -> GatewayError {
| EngineError::ProviderMetadataMismatch
| EngineError::ContinuationPinMismatch
| EngineError::RequiredAccountMismatch
| EngineError::AccountOutsideClientScope
| EngineError::DownstreamCommitRequired
| EngineError::NoActiveAttempt
| EngineError::InvalidDeliveryState => {
@@ -0,0 +1,271 @@
use axum::{
Router,
body::{Body, to_bytes},
http::{Method, Request, StatusCode, header},
};
use gateway_api::admin::account_groups;
use serde_json::{Value, json};
use tower::ServiceExt as _;
use super::{AdminTestFixture, AdminTestState, PRIMARY_GROUP_ID};
const SESSION_COOKIE: &str = "cpr_admin_session=valid-session";
#[tokio::test]
async fn list_route_should_keep_camel_case_group_and_page_wire() {
let fixture = authenticated_fixture().await;
let response = request(
router(&fixture),
Method::GET,
"/api/admin/account-groups?page=1&pageSize=1&search=alpha&enabled=true",
None,
true,
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let value = response_json(response).await;
assert_eq!(value["data"]["page"]["page"], 1);
assert_eq!(value["data"]["page"]["pageSize"], 1);
assert_eq!(value["data"]["page"]["total"], 1);
assert_eq!(value["data"]["page"]["totalPages"], 1);
assert_eq!(value["data"]["configRevision"], 7);
assert_eq!(value["data"]["items"][0]["id"], PRIMARY_GROUP_ID);
assert_eq!(value["data"]["items"][0]["memberCount"], 2);
assert_eq!(value["data"]["items"][0]["providerCounts"]["openai"], 1);
assert_eq!(value["data"]["items"][0]["clientKeyCount"], 2);
assert_eq!(value["data"]["items"][0]["accountSummary"]["available"], 1);
assert_eq!(value["data"]["items"][0]["accountSummary"]["limited"], 1);
assert_eq!(value["data"]["items"][0]["accountSummary"]["total"], 2);
assert_eq!(value["data"]["items"][0]["capacity"]["usedSlots"], 0);
assert_eq!(value["data"]["items"][0]["capacity"]["totalSlots"], 1);
assert_eq!(value["data"]["items"][0]["usage"]["todayUsd"], "1.25");
assert_eq!(value["data"]["items"][0]["usage"]["totalUsd"], "5.5");
assert!(value["data"]["items"][0].get("member_count").is_none());
}
#[tokio::test]
async fn create_route_should_create_an_empty_group() {
let fixture = authenticated_fixture().await;
let response = request(
router(&fixture),
Method::POST,
"/api/admin/account-groups/create",
Some(json!({
"name": "Gamma routing",
"description": "New traffic pool",
"color": "#a855f780"
})),
true,
)
.await;
assert_eq!(response.status(), StatusCode::CREATED);
let value = response_json(response).await;
assert!(
value["data"]["id"]
.as_str()
.is_some_and(|id| id.starts_with("grp_") && id.len() == 36)
);
assert_eq!(value["data"]["record"]["name"], "Gamma routing");
assert_eq!(value["data"]["record"]["memberCount"], 0);
assert_eq!(value["data"]["record"]["color"], "#A855F780");
assert_eq!(value["data"]["configRevision"], 8);
}
#[tokio::test]
async fn update_route_should_replace_group_fields() {
let fixture = authenticated_fixture().await;
let response = request(
router(&fixture),
Method::POST,
"/api/admin/account-groups/update",
Some(json!({
"id": PRIMARY_GROUP_ID,
"name": "Renamed routing",
"description": null,
"color": "#F43F5ECC"
})),
true,
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let value = response_json(response).await;
assert_eq!(value["data"]["record"]["name"], "Renamed routing");
assert!(value["data"]["record"]["description"].is_null());
assert_eq!(value["data"]["record"]["color"], "#F43F5ECC");
}
#[tokio::test]
async fn enable_and_disable_routes_should_set_explicit_group_state() {
let fixture = authenticated_fixture().await;
let disabled = request(
router(&fixture),
Method::POST,
"/api/admin/account-groups/disable",
Some(json!({ "id": PRIMARY_GROUP_ID })),
true,
)
.await;
let enabled = request(
router(&fixture),
Method::POST,
"/api/admin/account-groups/enable",
Some(json!({ "id": PRIMARY_GROUP_ID })),
true,
)
.await;
assert_eq!(disabled.status(), StatusCode::OK);
assert!(
!response_json(disabled).await["data"]["record"]["enabled"]
.as_bool()
.expect("disabled flag")
);
assert_eq!(enabled.status(), StatusCode::OK);
assert!(
response_json(enabled).await["data"]["record"]["enabled"]
.as_bool()
.expect("enabled flag")
);
}
#[tokio::test]
async fn delete_route_should_return_deleted_id_without_a_live_record() {
let fixture = authenticated_fixture().await;
let response = request(
router(&fixture),
Method::POST,
"/api/admin/account-groups/delete",
Some(json!({ "id": PRIMARY_GROUP_ID })),
true,
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let value = response_json(response).await;
assert_eq!(value["data"]["id"], PRIMARY_GROUP_ID);
assert!(value["data"]["record"].is_null());
}
#[tokio::test]
async fn list_route_should_reject_unknown_and_invalid_pagination_fields() {
let fixture = authenticated_fixture().await;
for uri in [
"/api/admin/account-groups?other=true",
"/api/admin/account-groups?page=0",
"/api/admin/account-groups?pageSize=0",
"/api/admin/account-groups?pageSize=201",
] {
let response = request(router(&fixture), Method::GET, uri, None, true).await;
assert_eq!(response.status(), StatusCode::BAD_REQUEST, "{uri}");
}
}
#[tokio::test]
async fn mutation_route_should_reject_unknown_json_fields() {
let fixture = authenticated_fixture().await;
let response = request(
router(&fixture),
Method::POST,
"/api/admin/account-groups/create",
Some(json!({
"name": "Unknown field",
"description": null,
"color": "#2563EBFF",
"providerKind": "openai"
})),
true,
)
.await;
assert_eq!(response.status(), StatusCode::UNPROCESSABLE_ENTITY);
}
#[tokio::test]
async fn create_route_should_reject_six_digit_color_without_alpha() {
let fixture = authenticated_fixture().await;
let response = request(
router(&fixture),
Method::POST,
"/api/admin/account-groups/create",
Some(json!({
"name": "Opaque legacy color",
"description": null,
"color": "#2563EB"
})),
true,
)
.await;
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
}
#[tokio::test]
async fn account_group_routes_should_require_admin_authentication() {
let fixture = AdminTestFixture::new().await;
let get_response = request(
router(&fixture),
Method::GET,
"/api/admin/account-groups",
None,
false,
)
.await;
let post_response = request(
router(&fixture),
Method::POST,
"/api/admin/account-groups/create",
Some(json!({ "name": "Unauthorized", "description": null, "color": "#2563EBFF" })),
false,
)
.await;
assert_eq!(get_response.status(), StatusCode::UNAUTHORIZED);
assert_eq!(post_response.status(), StatusCode::UNAUTHORIZED);
}
async fn authenticated_fixture() -> AdminTestFixture {
let fixture = AdminTestFixture::new().await;
fixture.auth.insert_session("valid-session");
fixture
}
fn router(fixture: &AdminTestFixture) -> Router {
account_groups::router::<AdminTestState>().with_state(fixture.state())
}
async fn request(
router: Router,
method: Method,
uri: &str,
body: Option<Value>,
authenticated: bool,
) -> axum::response::Response {
let mut builder = Request::builder()
.method(method)
.uri(uri)
.header("x-request-id", "req_account_groups");
if authenticated {
builder = builder.header(header::COOKIE, SESSION_COOKIE);
}
let body = match body {
Some(value) => {
builder = builder.header(header::CONTENT_TYPE, "application/json");
Body::from(serde_json::to_vec(&value).expect("serialize request body"))
}
None => Body::empty(),
};
router
.oneshot(builder.body(body).expect("account group request"))
.await
.expect("account group response")
}
async fn response_json(response: axum::response::Response) -> Value {
let body = to_bytes(response.into_body(), 64 * 1024)
.await
.expect("account group response body");
serde_json::from_slice(&body).expect("account group response JSON")
}
@@ -70,8 +70,138 @@ mod query {
}
}
mod batch_update {
use gateway_api::admin::accounts::{BatchUpdateAccountsRequest, UpdateAccountRequest};
use serde_json::json;
#[test]
fn batch_update_should_accept_complete_atomic_payload() {
let request: BatchUpdateAccountsRequest = serde_json::from_value(json!({
"accountIds": ["acct_openai", "acct_xai"],
"enabled": false,
"concurrencyLimit": null,
"weight": 1,
"groupIds": ["grp_00000000000000000000000000000001"]
}))
.expect("deserialize batch update");
request.validate().expect("validate batch update");
}
#[test]
fn batch_update_should_reject_duplicate_or_invalid_account_ids() {
for account_ids in [
json!(["acct_same", "acct_same"]),
json!(["not-an-account"]),
json!([]),
] {
let request: BatchUpdateAccountsRequest = serde_json::from_value(json!({
"accountIds": account_ids,
"enabled": true,
"concurrencyLimit": 8,
"weight": 100,
"groupIds": []
}))
.expect("deserialize invalid batch update");
assert_eq!(
request.validate().expect_err("reject account IDs").field(),
"accountIds"
);
}
}
#[test]
fn batch_update_should_require_enabled_and_reject_unknown_fields() {
assert!(
serde_json::from_value::<BatchUpdateAccountsRequest>(json!({
"accountIds": ["acct_test"],
"concurrencyLimit": null,
"weight": 1,
"groupIds": []
}))
.is_err()
);
assert!(
serde_json::from_value::<BatchUpdateAccountsRequest>(json!({
"accountIds": ["acct_test"],
"enabled": true,
"concurrencyLimit": null,
"weight": 1,
"groupIds": [],
"legacy": true
}))
.is_err()
);
}
#[test]
fn batch_update_should_reject_invalid_scheduling_bounds_and_missing_fields() {
for (concurrency_limit, weight, field) in [
(json!(0), json!(1), "concurrencyLimit"),
(json!(4294967296_u64), json!(1), "concurrencyLimit"),
(json!(null), json!(0), "weight"),
(json!(null), json!(101), "weight"),
] {
let request: BatchUpdateAccountsRequest = serde_json::from_value(json!({
"accountIds": ["acct_test"],
"enabled": true,
"concurrencyLimit": concurrency_limit,
"weight": weight,
"groupIds": []
}))
.expect("deserialize invalid scheduling");
assert_eq!(
request.validate().expect_err("reject scheduling").field(),
field
);
}
assert!(
serde_json::from_value::<BatchUpdateAccountsRequest>(json!({
"accountIds": ["acct_test"],
"enabled": true,
"weight": 1,
"groupIds": []
}))
.is_err()
);
assert!(
serde_json::from_value::<BatchUpdateAccountsRequest>(json!({
"accountIds": ["acct_test"],
"enabled": true,
"concurrencyLimit": null,
"groupIds": []
}))
.is_err()
);
}
#[test]
fn single_update_should_require_the_same_complete_scheduling_contract() {
let request: UpdateAccountRequest = serde_json::from_value(json!({
"accountId": "acct_test",
"enabled": true,
"concurrencyLimit": 4294967295_u64,
"weight": 100,
"groupIds": []
}))
.expect("deserialize single update");
request.validate().expect("validate single update");
assert!(
serde_json::from_value::<UpdateAccountRequest>(json!({
"accountId": "acct_test",
"enabled": true,
"weight": 1,
"groupIds": []
}))
.is_err()
);
}
}
mod response {
use gateway_api::admin::accounts::AccountUsageView;
use gateway_api::admin::presenter::format_decimal_currency;
#[test]
fn account_usage_view_should_keep_unobserved_numbers_null() {
@@ -112,12 +242,22 @@ mod response {
assert!(value["totalTokens"].is_null());
assert_eq!(value["createdTokensDisplay"], "-");
}
#[test]
fn account_usage_currency_should_limit_usd_to_four_fraction_digits() {
assert_eq!(format_decimal_currency("0.1204956", "USD"), "$0.1205");
assert_eq!(format_decimal_currency("0.99996", "USD"), "$1.00");
assert_eq!(format_decimal_currency("0.0300", "USD"), "$0.03");
assert_eq!(format_decimal_currency("12.3456", "USD"), "$12.35");
assert_eq!(format_decimal_currency("0.1204956", "CNY"), "CNY 0.1204956");
}
}
mod actions {
use base64::Engine as _;
use gateway_admin::model::{
Revision,
accounts::{AccountConnectionTestEvent as DomainConnectionTestEvent, AccountStatus},
accounts::AccountConnectionTestEvent as DomainConnectionTestEvent,
provider_credentials::{
CredentialDeletionResult, CredentialImportResult, CredentialMutationResult,
},
@@ -126,7 +266,8 @@ mod actions {
AccountActionRequest, AccountConnectionTestEvent, AccountDeletionData,
AccountDeletionRequest, AccountExportData, AccountExportQuery, AccountIdQuery,
AccountImportData, AccountImportRequest, AccountMutationData, AccountRefreshRequest,
AccountTestQuery, RotateAccountRequest, StartAccountAuthorizationRequest,
AccountTestQuery, CompleteAccountAuthorizationRequest, RotateAccountRequest,
StartAccountAuthorizationRequest,
};
use gateway_core::engine::credential::ProviderAccountId;
use serde_json::json;
@@ -207,7 +348,8 @@ mod actions {
"provider": "openai",
"accountId": "acct_1",
"accessToken": "header.payload.signature",
"refreshToken": "refresh-token"
"refreshToken": "refresh-token",
"idToken": "id-header.id-payload.id-signature"
}))
.expect("decode rotation");
assert!(rotation.validate().is_ok());
@@ -269,6 +411,35 @@ mod actions {
}))
.is_err()
);
assert!(
serde_json::from_value::<AccountImportRequest>(json!({
"provider": "openai",
"data": {},
"groupIds": []
}))
.is_err()
);
}
#[test]
fn oauth_complete_should_not_accept_group_assignment() {
let flow_id = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([7_u8; 32]);
let request: CompleteAccountAuthorizationRequest = serde_json::from_value(json!({
"provider": "openai",
"flowId": flow_id,
"callbackUrl": "http://localhost/callback?code=ok"
}))
.expect("decode OAuth completion");
assert!(request.validate().is_ok());
assert!(
serde_json::from_value::<CompleteAccountAuthorizationRequest>(json!({
"provider": "xai",
"flowId": "flow_test",
"callbackUrl": "http://localhost/callback?code=ok",
"groupIds": []
}))
.is_err()
);
}
#[test]
@@ -305,7 +476,7 @@ mod actions {
#[test]
fn account_import_response_should_emit_account_ids() {
let response = AccountImportData::from(CredentialImportResult {
let response = AccountImportData::from_result(CredentialImportResult {
config_revision: Revision::new(8).expect("revision"),
credential_ids: vec![ProviderAccountId::new("acct_imported").expect("account ID")],
});
@@ -343,9 +514,7 @@ mod actions {
DomainConnectionTestEvent::Content {
text: "OK".to_owned(),
},
DomainConnectionTestEvent::Completed {
account_status: AccountStatus::Normal,
},
DomainConnectionTestEvent::Completed {},
DomainConnectionTestEvent::Failed {
message: "upstream unavailable".to_owned(),
provider_error_code: Some("usage_exhausted".to_owned()),
@@ -353,7 +522,6 @@ mod actions {
upstream_status: Some(429),
upstream_content_type: Some("application/json".to_owned()),
upstream_body: Some(r#"{"error":{"type":"usage_limit_reached"}}"#.to_owned()),
account_status: AccountStatus::QuotaExhausted,
},
]
.map(|event| AccountConnectionTestEvent::from(event).data);
@@ -378,7 +546,7 @@ mod actions {
}
}),
json!({ "type": "content", "text": "OK" }),
json!({ "type": "test_complete", "success": true, "accountStatus": "normal" }),
json!({ "type": "test_complete", "success": true }),
json!({
"type": "error",
"error": "upstream unavailable",
@@ -386,8 +554,7 @@ mod actions {
"providerErrorType": "invalid_request_error",
"upstreamStatus": 429,
"upstreamContentType": "application/json",
"upstreamBody": r#"{"error":{"type":"usage_limit_reached"}}"#,
"accountStatus": "quota_exhausted"
"upstreamBody": r#"{"error":{"type":"usage_limit_reached"}}"#
}),
]
);
@@ -105,7 +105,7 @@ fn client_key_mutations_should_validate_text_limits_and_unknown_fields() {
let valid = serde_json::from_value::<CreateClientKeyRequest>(json!({
"name": "terminal key",
"label": "production",
"providerKind": "openai",
"groupIds": [],
"maxConcurrency": 2,
"requestsPerMinute": 60
}))
@@ -118,7 +118,7 @@ fn client_key_mutations_should_validate_text_limits_and_unknown_fields() {
(
json!({
"name": " ",
"providerKind": "openai",
"groupIds": [],
"maxConcurrency": 0,
"requestsPerMinute": 0
}),
@@ -127,7 +127,7 @@ fn client_key_mutations_should_validate_text_limits_and_unknown_fields() {
(
json!({
"name": "key",
"providerKind": "openai",
"groupIds": [],
"maxConcurrency": u64::MAX,
"requestsPerMinute": 0
}),
@@ -150,7 +150,7 @@ fn client_key_mutations_should_validate_text_limits_and_unknown_fields() {
"id": "key_1",
"expectedConfigRevision": 1,
"name": "key",
"providerKind": "openai",
"groupIds": [],
"maxConcurrency": 0,
"requestsPerMinute": 0,
"tokensPerMinute": 0
@@ -161,7 +161,7 @@ fn client_key_mutations_should_validate_text_limits_and_unknown_fields() {
serde_json::from_value::<CreateClientKeyRequest>(json!({
"expectedConfigRevision": 7,
"name": "terminal key",
"providerKind": "openai",
"groupIds": [],
"maxConcurrency": 2,
"requestsPerMinute": 60
}))
@@ -210,7 +210,10 @@ fn client_key_responses_should_keep_shape_and_redact_creation_debug() {
id: gateway_core::policy::ClientApiKeyId::new("key_visible").expect("Client Key ID"),
name: "visible".to_owned(),
label: None,
provider_kind: gateway_core::routing::ProviderKind::new("openai").expect("Provider kind"),
groups: Vec::new(),
provider_kinds: vec![
gateway_core::routing::ProviderKind::new("openai").expect("Provider kind"),
],
prefix: "sk_visible12".to_owned(),
enabled: true,
limits: gateway_core::policy::RateLimits {
@@ -226,7 +229,11 @@ fn client_key_responses_should_keep_shape_and_redact_creation_debug() {
assert!(list.get("configRevision").is_none());
assert_eq!(list["total"], 1);
assert_eq!(list["items"][0]["id"], "key_visible");
assert_eq!(list["items"][0]["providerKind"], "openai");
assert_eq!(
list["items"][0]["providerKinds"],
serde_json::json!(["openai"])
);
assert_eq!(list["items"][0]["routingScope"], "all");
assert_eq!(list["items"][0]["maxConcurrency"], 2);
assert_eq!(list["items"][0]["requestsPerMinute"], 60);
assert!(list["items"][0].get("tokensPerMinute").is_none());
+273 -14
View File
@@ -7,15 +7,22 @@ use std::{
};
use async_trait::async_trait;
use chrono::{Duration, Utc};
use chrono::{DateTime, Duration, Utc};
use futures::future::BoxFuture;
use gateway_admin::{
AdminConfig, AdminServices, InitialAdminPassword,
model::{
MutationContext, Revision,
account_groups::{
AccountGroupAccountSummary, AccountGroupCapacity, AccountGroupColor,
AccountGroupListQuery, AccountGroupMutation, AccountGroupPage, AccountGroupRecord,
AccountGroupUsage, DeleteAccountGroup, NewAccountGroup, SetAccountGroupEnabled,
UpdateAccountGroup,
},
accounts::{
AccountListQuery, AccountPage, AccountRecord, AccountUsage, AccountUsageWindowQuery,
AccountUsageWindowResult, DeleteAccounts, SetAccountEnabled,
AccountListQuery, AccountPage, AccountPageItem, AccountUpdateResult, AccountUsage,
AccountUsageWindowQuery, AccountUsageWindowResult, AccountsUpdateResult,
BatchUpdateAccounts, DeleteAccounts, UpdateAccount,
},
auth::{AdminAuditEvent, AdminSession},
client_keys::{
@@ -23,9 +30,9 @@ use gateway_admin::{
NewClientKey, SetClientKeyEnabled, UpdateClientKey,
},
observability::{
DashboardObservation, DiagnosticDimension, DiagnosticObservation, OpsError,
OpsErrorPage, OpsErrorQuery, RequestMetricPoint, TimeRange, UsageDetail, UsageFilter,
UsageOverview, UsagePage, UsageQuery, UsageRecord,
DashboardObservation, DecimalAmount, DiagnosticDimension, DiagnosticObservation,
OpsError, OpsErrorPage, OpsErrorQuery, RequestMetricPoint, TimeRange, UsageDetail,
UsageFilter, UsageOverview, UsagePage, UsageQuery, UsageRecord,
},
provider_credentials::{
AuthorizationCommit, AuthorizationStarted, CompleteAuthorization, CredentialDetails,
@@ -45,8 +52,8 @@ use gateway_admin::{
ports::{
provider::{ProviderAdmin, ProviderAdminError, ProviderAdminErrorKind},
store::{
AccountStore, AdminStoreError, AdminStoreErrorKind, AdminStorePorts, AdminStoreResult,
AuthStore, ClientKeyStore, ObservabilityStore, SettingsStore,
AccountGroupStore, AccountStore, AdminStoreError, AdminStoreErrorKind, AdminStorePorts,
AdminStoreResult, AuthStore, ClientKeyStore, ObservabilityStore, SettingsStore,
},
system::{
SystemOperationError, SystemOperationErrorKind, SystemOperations,
@@ -66,6 +73,7 @@ use gateway_core::{
},
};
mod account_groups;
mod accounts;
mod auth;
mod client_keys;
@@ -96,6 +104,7 @@ impl AdminTestFixture {
let auth = Arc::new(MemoryAuthStore::new(api_key.clone()));
let settings = Arc::new(MemorySettingsStore::new(api_key));
let client_keys = Arc::new(MemoryClientKeyStore);
let account_groups = Arc::new(MemoryAccountGroupStore::new());
let usage_records = Arc::new(Mutex::new(Vec::new()));
let usage_detail = Arc::new(Mutex::new(None));
let diagnostics = Arc::new(Mutex::new(Vec::new()));
@@ -112,6 +121,7 @@ impl AdminTestFixture {
});
let stores = AdminStorePorts::new(
unused.clone(),
account_groups.clone(),
auth.clone(),
client_keys.clone(),
unused,
@@ -345,6 +355,204 @@ impl SettingsStore for MemorySettingsStore {
pub(super) struct MemoryClientKeyStore;
pub(super) const PRIMARY_GROUP_ID: &str = "grp_11111111111111111111111111111111";
pub(super) const SECONDARY_GROUP_ID: &str = "grp_22222222222222222222222222222222";
struct MemoryAccountGroupState {
revision: Revision,
groups: BTreeMap<gateway_core::routing::AccountGroupId, AccountGroupRecord>,
}
pub(super) struct MemoryAccountGroupStore {
state: Mutex<MemoryAccountGroupState>,
}
impl MemoryAccountGroupStore {
fn new() -> Self {
let now = Utc::now();
let primary_id = group_id(PRIMARY_GROUP_ID);
let secondary_id = group_id(SECONDARY_GROUP_ID);
let groups = BTreeMap::from([
(
primary_id.clone(),
AccountGroupRecord {
id: primary_id,
name: "Alpha routing".to_owned(),
description: Some("Primary traffic".to_owned()),
color: group_color("#2563ebff"),
enabled: true,
member_count: 2,
provider_counts: BTreeMap::from([
("openai".to_owned(), 1),
("xai".to_owned(), 1),
]),
client_key_count: 2,
account_summary: account_summary(1, 1, 2),
capacity: capacity(Some(0), 1),
usage: usage("1.25", "5.5"),
created_at: now,
updated_at: now,
},
),
(
secondary_id.clone(),
AccountGroupRecord {
id: secondary_id,
name: "Beta routing".to_owned(),
description: None,
color: group_color("#64748B80"),
enabled: false,
member_count: 0,
provider_counts: BTreeMap::new(),
client_key_count: 0,
account_summary: account_summary(0, 0, 0),
capacity: capacity(Some(0), 0),
usage: usage("0", "0"),
created_at: now,
updated_at: now,
},
),
]);
Self {
state: Mutex::new(MemoryAccountGroupState {
revision: Revision::new(7).expect("revision"),
groups,
}),
}
}
}
#[async_trait]
impl AccountGroupStore for MemoryAccountGroupStore {
async fn list_account_groups(
&self,
query: AccountGroupListQuery,
) -> AdminStoreResult<AccountGroupPage> {
let state = self.state.lock().expect("account groups");
let search = query.search.as_ref().map(|value| value.to_lowercase());
let mut matching = state
.groups
.values()
.filter(|record| {
query
.enabled
.is_none_or(|enabled| record.enabled == enabled)
})
.filter(|record| {
search.as_ref().is_none_or(|search| {
record.name.to_lowercase().contains(search)
|| record
.description
.as_ref()
.is_some_and(|description| description.to_lowercase().contains(search))
})
})
.cloned()
.collect::<Vec<_>>();
matching.sort_by(|left, right| left.name.cmp(&right.name));
let total = matching.len() as u64;
let page_size = usize::from(query.page_size.get());
let offset = usize::try_from(query.page.saturating_sub(1))
.unwrap_or(usize::MAX)
.saturating_mul(page_size);
let items = matching.into_iter().skip(offset).take(page_size).collect();
Ok(AccountGroupPage {
config_revision: state.revision,
items,
total,
page: query.page,
page_size: query.page_size.get(),
})
}
async fn create_account_group(
&self,
command: NewAccountGroup,
_: &MutationContext,
) -> AdminStoreResult<AccountGroupMutation> {
let mut state = self.state.lock().expect("account groups");
let now = Utc::now();
let record = AccountGroupRecord {
id: command.id.clone(),
name: command.name,
description: command.description,
color: command.color,
enabled: true,
member_count: 0,
provider_counts: BTreeMap::new(),
client_key_count: 0,
account_summary: account_summary(0, 0, 0),
capacity: capacity(Some(0), 0),
usage: usage("0", "0"),
created_at: now,
updated_at: now,
};
state.groups.insert(command.id.clone(), record);
mutation(&mut state, command.id, true)
}
async fn update_account_group(
&self,
command: UpdateAccountGroup,
_: &MutationContext,
) -> AdminStoreResult<AccountGroupMutation> {
let mut state = self.state.lock().expect("account groups");
let record = state
.groups
.get_mut(&command.id)
.ok_or_else(|| not_found("account group"))?;
record.name = command.name;
record.description = command.description;
record.color = command.color;
record.updated_at = Utc::now();
mutation(&mut state, command.id, true)
}
async fn set_account_group_enabled(
&self,
command: SetAccountGroupEnabled,
_: &MutationContext,
) -> AdminStoreResult<AccountGroupMutation> {
let mut state = self.state.lock().expect("account groups");
let record = state
.groups
.get_mut(&command.id)
.ok_or_else(|| not_found("account group"))?;
record.enabled = command.enabled;
record.updated_at = Utc::now();
mutation(&mut state, command.id, true)
}
async fn delete_account_group(
&self,
command: DeleteAccountGroup,
_: &MutationContext,
) -> AdminStoreResult<AccountGroupMutation> {
let mut state = self.state.lock().expect("account groups");
if state.groups.remove(&command.id).is_none() {
return Err(not_found("account group"));
}
mutation(&mut state, command.id, false)
}
}
fn group_id(value: &str) -> gateway_core::routing::AccountGroupId {
gateway_core::routing::AccountGroupId::new(value).expect("account group ID")
}
fn mutation(
state: &mut MemoryAccountGroupState,
id: gateway_core::routing::AccountGroupId,
include_record: bool,
) -> AdminStoreResult<AccountGroupMutation> {
state.revision = next_revision(state.revision);
Ok(AccountGroupMutation {
config_revision: state.revision,
record: include_record.then(|| state.groups.get(&id).expect("group record").clone()),
id,
})
}
#[async_trait]
impl ClientKeyStore for MemoryClientKeyStore {
async fn list_client_keys(&self, _: ClientKeyListQuery) -> AdminStoreResult<ClientKeyPage> {
@@ -366,7 +574,8 @@ impl ClientKeyStore for MemoryClientKeyStore {
id: id.clone(),
name: "revealed".to_owned(),
label: None,
provider_kind: ProviderKind::new("openai").expect("provider kind"),
groups: Vec::new(),
provider_kinds: vec![ProviderKind::new("openai").expect("provider kind")],
prefix: "sk_aaaaaaaaa".to_owned(),
enabled: true,
limits: RateLimits::unlimited(),
@@ -426,7 +635,7 @@ impl AccountStore for UnusedStore {
Err(unavailable("account list"))
}
async fn load_account(&self, _: &str) -> AdminStoreResult<Option<AccountRecord>> {
async fn load_account(&self, _: &str) -> AdminStoreResult<Option<AccountPageItem>> {
Err(unavailable("account"))
}
@@ -501,14 +710,30 @@ impl AccountStore for UnusedStore {
Err(unavailable("credential refresh"))
}
async fn set_account_enabled(
async fn update_account(
&self,
_: SetAccountEnabled,
_: UpdateAccount,
_: &MutationContext,
) -> AdminStoreResult<Revision> {
) -> AdminStoreResult<AccountUpdateResult> {
Err(unavailable("account enabled"))
}
async fn recover_account(
&self,
_: &ProviderAccountId,
_: &MutationContext,
) -> AdminStoreResult<AccountUpdateResult> {
Err(unavailable("account recovery"))
}
async fn batch_update_accounts(
&self,
_: BatchUpdateAccounts,
_: &MutationContext,
) -> AdminStoreResult<AccountsUpdateResult> {
Err(unavailable("account batch update"))
}
async fn delete_accounts(
&self,
_: DeleteAccounts,
@@ -528,7 +753,11 @@ impl AccountStore for UnusedStore {
#[async_trait]
impl ObservabilityStore for UnusedStore {
async fn dashboard_summary(&self, range: TimeRange) -> AdminStoreResult<DashboardObservation> {
async fn dashboard_summary(
&self,
range: TimeRange,
_: DateTime<Utc>,
) -> AdminStoreResult<DashboardObservation> {
*self
.dashboard_summary_range
.lock()
@@ -788,6 +1017,32 @@ fn next_revision(revision: Revision) -> Revision {
Revision::new(revision.get().saturating_add(1)).expect("next revision")
}
fn account_summary(available: u64, limited: u64, total: u64) -> AccountGroupAccountSummary {
AccountGroupAccountSummary {
available,
limited,
total,
}
}
fn group_color(value: &str) -> AccountGroupColor {
AccountGroupColor::parse(value).expect("group color")
}
fn capacity(used_slots: Option<u64>, total_slots: u64) -> AccountGroupCapacity {
AccountGroupCapacity {
used_slots,
total_slots,
}
}
fn usage(today_usd: &str, total_usd: &str) -> AccountGroupUsage {
AccountGroupUsage {
today_usd: today_usd.parse::<DecimalAmount>().expect("today cost"),
total_usd: total_usd.parse::<DecimalAmount>().expect("total cost"),
}
}
fn unavailable(resource: &'static str) -> AdminStoreError {
AdminStoreError::new(
AdminStoreErrorKind::Unavailable,
@@ -796,6 +1051,10 @@ fn unavailable(resource: &'static str) -> AdminStoreError {
)
}
fn not_found(resource: &'static str) -> AdminStoreError {
AdminStoreError::new(AdminStoreErrorKind::NotFound, resource, "not found")
}
fn unsupported_provider() -> ProviderAdminError {
ProviderAdminError::new(ProviderAdminErrorKind::Unsupported)
}
@@ -222,16 +222,11 @@ async fn dashboard_summary_should_include_quota_exhaustion_in_unavailable_headli
totals: Default::default(),
provider_accounts: AccountPoolMetrics {
total: 807,
enabled: 804,
unavailable: 792,
active: 15,
rate_limited: 0,
expired: 2,
invalid: 0,
normal: 15,
quota_exhausted: 785,
refreshing: None,
rate_limited: 0,
disabled: 3,
banned: 2,
error: 4,
},
trend: Vec::new(),
account_usage: Vec::new(),
@@ -679,6 +674,9 @@ fn usage_record_with_account(
id: id.to_owned(),
client_api_key_ref: "key_detail".to_owned(),
config_revision: 1,
routing_scope: "all".to_owned(),
routing_group_refs: Vec::new(),
routing_group_names_snapshot: Vec::new(),
protocol: "openai".to_owned(),
operation: "responses".to_owned(),
endpoint: "/v1/responses".to_owned(),
@@ -805,6 +803,9 @@ async fn usage_route_should_expose_image_and_websocket_facts() {
id: "request_endpoint".to_owned(),
client_api_key_ref: "key_endpoint".to_owned(),
config_revision: 1,
routing_scope: "groups".to_owned(),
routing_group_refs: vec!["grp_usage".to_owned()],
routing_group_names_snapshot: vec!["Usage group".to_owned()],
protocol: "openai".to_owned(),
operation: "generate".to_owned(),
endpoint: "/v1/responses".to_owned(),
@@ -924,6 +925,9 @@ async fn usage_route_should_expose_image_and_websocket_facts() {
assert_eq!(
serde_json::json!({
"route": value["data"]["items"][0]["route"],
"routingScope": value["data"]["items"][0]["routingScope"],
"routingGroupRefs": value["data"]["items"][0]["routingGroupRefs"],
"routingGroupNamesSnapshot": value["data"]["items"][0]["routingGroupNamesSnapshot"],
"serviceTier": value["data"]["items"][0]["serviceTier"],
"accountId": value["data"]["items"][0]["accountId"],
"accountName": value["data"]["items"][0]["accountName"],
@@ -954,6 +958,9 @@ async fn usage_route_should_expose_image_and_websocket_facts() {
}),
serde_json::json!({
"route": "/v1/responses",
"routingScope": "groups",
"routingGroupRefs": ["grp_usage"],
"routingGroupNamesSnapshot": ["Usage group"],
"serviceTier": "priority",
"accountId": "acct_snapshot",
"accountName": "Snapshot Alpha",
@@ -33,6 +33,7 @@ fn source_tree_should_match_frozen_machine_manifest() {
let mut actual = rust_files(&root.join("src"));
actual.sort();
let mut expected = vec![
"src/admin/account_groups.rs",
"src/admin/accounts.rs",
"src/admin/auth.rs",
"src/admin/backups.rs",
@@ -73,6 +74,7 @@ fn test_tree_should_match_frozen_rust_mirror() {
let mut actual = rust_files(&root.join("tests"));
actual.sort();
let mut expected = vec![
"tests/admin/account_groups.rs",
"tests/admin/accounts.rs",
"tests/admin/auth.rs",
"tests/admin/client_keys.rs",
+16 -5
View File
@@ -4,7 +4,7 @@ mod models;
mod responses;
mod router;
use std::collections::BTreeSet;
use std::collections::{BTreeMap, BTreeSet};
use std::num::NonZeroU32;
use std::sync::Arc;
use std::time::{Duration, SystemTime};
@@ -18,7 +18,9 @@ use gateway_core::engine::admission::{
use gateway_core::engine::continuation::{
NativeContinuationPin, NativeContinuationPort, NativeContinuationStoreError, PreviousResponseId,
};
use gateway_core::engine::credential::{AccountSelectionPolicy, RotationStrategy};
use gateway_core::engine::credential::{
AccountSelectionPolicy, ProviderAccountId, RotationStrategy,
};
use gateway_core::engine::execution::{
AuthenticatedClient, ClientApiKeyUsageSink, DefaultExecutionService, ExecutionService,
ProviderCircuitDecision, ProviderCircuitError, ProviderCircuitPort,
@@ -34,8 +36,8 @@ use gateway_core::lifecycle::{ConnectionDraining, ConnectionGuard, ConnectionLif
use gateway_core::policy::{ClientApiKeyId, ClientPolicy, PlaintextClientApiKey, RateLimits};
use gateway_core::routing::snapshot::RuntimeSnapshotHandle;
use gateway_core::routing::{
ConfigRevision, ModelCapabilities, ProviderKind, ProviderModel, RuntimeSnapshot,
UpstreamModelId,
ClientRoutingScope, ConfigRevision, FrozenAccountScope, ModelCapabilities, ProviderKind,
ProviderModel, RuntimeAccount, RuntimeAccountDirectory, RuntimeSnapshot, UpstreamModelId,
};
pub(super) async fn api_router(execution: Arc<dyn ExecutionService>) -> axum::Router {
@@ -94,6 +96,10 @@ impl ClientApiKeyUsageSink for IgnoredClientApiKeyUsage {
fn snapshot(plaintext: &str, provider_name: &str) -> RuntimeSnapshot {
let provider = ProviderKind::new(provider_name).expect("provider");
let account_directory = Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([(
ProviderAccountId::new("acct_api_test").expect("account ID"),
RuntimeAccount::new(provider.clone(), BTreeSet::new()),
)])));
let capabilities = ModelCapabilities::new(
BTreeSet::from([gateway_core::operation::OperationKind::Generate]),
Some(16_000),
@@ -119,12 +125,16 @@ fn snapshot(plaintext: &str, provider_name: &str) -> RuntimeSnapshot {
vec![ClientPolicy::new(
ClientApiKeyId::new("key_api_test").expect("key ID"),
PlaintextClientApiKey::new(plaintext).expect("plaintext key"),
provider,
Arc::new(FrozenAccountScope::new(
Arc::clone(&account_directory),
ClientRoutingScope::all_accounts(),
)),
true,
RateLimits::unlimited(),
)],
)
.expect("runtime snapshot")
.with_account_directory(account_directory)
}
#[derive(Default)]
@@ -260,6 +270,7 @@ struct UnusedContinuation;
impl NativeContinuationPort for UnusedContinuation {
fn resolve<'a>(
&'a self,
_: &'a ClientApiKeyId,
_: &'a PreviousResponseId,
) -> BoxFuture<'a, Result<Option<NativeContinuationPin>, NativeContinuationStoreError>> {
Box::pin(async { unreachable!("authentication fixture does not execute") })
@@ -451,10 +451,12 @@ fn decoder_should_not_reject_large_bodies_using_catalog_context_limits() {
.to_string();
let decoded = decode_request(body.as_bytes()).expect("large request should decode");
let plan = super::snapshot("sk_context_test", "openai")
let snapshot = super::snapshot("sk_context_test", "openai");
let plan = snapshot
.plan(
&PublicModelId::new(decoded.metadata().public_model()).expect("public model"),
decoded.operation(),
snapshot.all_account_scope(),
&RoutingContext::default(),
)
.expect("large body must not be locally context-gated");
@@ -9,6 +9,7 @@ use futures::future::BoxFuture;
use crate::engine::credential::ProviderAccountId;
use crate::operation::ProviderSessionState;
use crate::policy::ClientApiKeyId;
use crate::routing::ProviderKind;
/// 客户端或 Provider 传递的 opaque response handle。
@@ -48,14 +49,15 @@ pub enum NativeContinuationScope {
/// 从可丢失会话亲和存储恢复的 native previous-response pin。
///
/// 该值不依赖调用方 API Key;它只阻止 native handle 与不透明 Provider state
/// 被发送到错误的 Provider 或账号。
/// 该值同时绑定调用方 API Key、Provider 与账号,阻止 native handle 和不透明
/// Provider state 被跨租户或跨上游目标复用。
#[derive(Clone, PartialEq)]
pub struct NativeContinuationPin {
/// 客户端提交、用于查找可丢失会话亲和记录的 response ID。
previous_response_id: PreviousResponseId,
/// Provider 原生 response handle。
upstream_response_id: PreviousResponseId,
client_api_key_id: ClientApiKeyId,
provider: ProviderKind,
account: ProviderAccountId,
scope: NativeContinuationScope,
@@ -67,12 +69,14 @@ impl NativeContinuationPin {
pub const fn new(
previous_response_id: PreviousResponseId,
upstream_response_id: PreviousResponseId,
client_api_key_id: ClientApiKeyId,
provider: ProviderKind,
account: ProviderAccountId,
) -> Self {
Self {
previous_response_id,
upstream_response_id,
client_api_key_id,
provider,
account,
scope: NativeContinuationScope::ConnectionLocal,
@@ -105,6 +109,16 @@ impl NativeContinuationPin {
&self.upstream_response_id
}
#[must_use]
pub const fn client_api_key_id(&self) -> &ClientApiKeyId {
&self.client_api_key_id
}
#[must_use]
pub fn matches_client(&self, client_api_key_id: &ClientApiKeyId) -> bool {
self.client_api_key_id == *client_api_key_id
}
#[must_use]
pub const fn provider(&self) -> &ProviderKind {
&self.provider
@@ -139,6 +153,7 @@ impl fmt::Debug for NativeContinuationPin {
.debug_struct("NativeContinuationPin")
.field("previous_response_id", &"<redacted>")
.field("upstream_response_id", &"<redacted>")
.field("client_api_key_id", &self.client_api_key_id)
.field("provider", &self.provider)
.field("account", &self.account)
.field("scope", &self.scope)
@@ -190,21 +205,57 @@ impl fmt::Debug for ContinuationBinding {
}
}
/// Native continuation lookup 的稳定失败分类。
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum NativeContinuationStoreErrorKind {
/// Redis 等可丢失协调存储暂不可用;调用方可以按外部 continuation fail-open。
Unavailable,
/// 存储中的记录无法还原为可信 pin;不得将原始 handle 继续透传上游。
InvalidData,
/// 已存在的记录属于另一个客户端 API Key;必须 fail closed。
OwnershipMismatch,
}
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
#[error("native continuation store is unavailable: {detail}")]
#[error("native continuation store failed: {kind:?}: {detail}")]
pub struct NativeContinuationStoreError {
kind: NativeContinuationStoreErrorKind,
detail: String,
}
impl NativeContinuationStoreError {
/// 构造带可排查细节的亲和存储错误;`detail` 不应包含 response handle。
/// 构造可 fail-open 的协调存储不可用错误;`detail` 不应包含 response handle。
#[must_use]
pub fn new(detail: impl Into<String>) -> Self {
pub fn unavailable(detail: impl Into<String>) -> Self {
Self {
kind: NativeContinuationStoreErrorKind::Unavailable,
detail: detail.into(),
}
}
/// 构造必须 fail closed 的无效记录错误;`detail` 不应包含 response handle。
#[must_use]
pub fn invalid_data(detail: impl Into<String>) -> Self {
Self {
kind: NativeContinuationStoreErrorKind::InvalidData,
detail: detail.into(),
}
}
/// 构造不泄露记录归属细节的跨调用方错误。
#[must_use]
pub fn ownership_mismatch() -> Self {
Self {
kind: NativeContinuationStoreErrorKind::OwnershipMismatch,
detail: "continuation belongs to a different client API key".to_owned(),
}
}
#[must_use]
pub const fn kind(&self) -> NativeContinuationStoreErrorKind {
self.kind
}
#[must_use]
pub fn detail(&self) -> &str {
&self.detail
@@ -213,11 +264,12 @@ impl NativeContinuationStoreError {
/// 可丢失的 previous-response 亲和存储端口。
///
/// 亲和记录不按调用方 API Key 隔离:与旧版一样,response ID 是会话级 opaque
/// handle;Redis 不可用或超时必须由调用方 fail-open。
/// 亲和记录按调用方 API Key 隔离;Redis 不可用或超时必须由调用方 fail-open,
/// 但命中其他 Key 的记录必须 fail closed。
pub trait NativeContinuationPort: Send + Sync {
fn resolve<'a>(
&'a self,
client_api_key_id: &'a ClientApiKeyId,
previous_response_id: &'a PreviousResponseId,
) -> BoxFuture<'a, Result<Option<NativeContinuationPin>, NativeContinuationStoreError>>;
@@ -46,6 +46,10 @@ impl AccountSelection {
Self::Diagnostic(account) => Some(account),
}
}
const fn is_diagnostic(&self) -> bool {
matches!(self, Self::Diagnostic(_))
}
}
impl<S: ?Sized> AttemptCoordinator<S>
@@ -126,6 +130,15 @@ where
.map(ProviderAccountStateOwner::from_continuation);
let continuation_attempt =
initial_continuation_attempt(&operation, &plan, continuation.as_ref());
let candidate_index = continuation
.as_ref()
.and_then(ContinuationBinding::pinned)
.map_or(Ok(0), |pin| {
plan.candidates()
.iter()
.position(|candidate| candidate.provider() == pin.provider())
.ok_or(EngineError::ContinuationPinMismatch)
})?;
let image_generation_requested = operation.image_generation_requested();
let mut session = ResponseExecutionSession {
engine: Arc::clone(&self.engine),
@@ -149,6 +162,7 @@ where
account_state_owner,
cancellation,
attempts: 0,
candidate_index,
excluded_accounts: BTreeSet::new(),
credential_recovery_attempted_accounts: BTreeSet::new(),
recovery_account: None,
@@ -225,6 +239,7 @@ pub struct ResponseExecutionSession<S: ?Sized> {
account_state_owner: Option<ProviderAccountStateOwner>,
cancellation: CancellationToken,
attempts: u32,
candidate_index: usize,
excluded_accounts: BTreeSet<crate::engine::credential::ProviderAccountId>,
credential_recovery_attempted_accounts: BTreeSet<crate::engine::credential::ProviderAccountId>,
/// 凭据恢复后的一次性同账号钉选;只约束紧随其后的 replay attempt,
@@ -461,6 +476,7 @@ where
NativeContinuationPin::new(
previous_response_id,
upstream_response_id,
self.client_api_key_ref.clone(),
provider,
account,
)
@@ -565,7 +581,7 @@ where
if self.attempts >= self.plan.max_attempts().get() {
return Err(EngineError::EmptyRoutingPlan);
}
let Some(candidate) = self.plan.candidates().first().cloned() else {
let Some(candidate) = self.plan.candidates().get(self.candidate_index).cloned() else {
let error = GatewayError::new(
GatewayErrorKind::NoAvailableProvider,
"no upstream Provider is available",
@@ -606,7 +622,8 @@ where
self.excluded_accounts.clone(),
pinned_account.clone(),
self.account_state_owner.clone(),
),
)
.with_account_scope(Arc::clone(self.plan.account_scope())),
}
.with_credential_recovery_attempted(pinned_account.as_ref().is_some_and(|account| {
self.credential_recovery_attempted_accounts
@@ -663,6 +680,20 @@ where
if self.prepare_unavailable_native_continuation_replay(&error) {
return Ok(Some(PullOutcome::AttemptDiscarded));
}
if matches!(
error.kind(),
ProviderErrorKind::AccountCapacityUnavailable
| ProviderErrorKind::NoEligibleAccount
| ProviderErrorKind::ProviderInfrastructureUnavailable
) && error.send_state() == UpstreamSendState::NotSent
&& matches!(
self.continuation_attempt,
ContinuationAttempt::None | ContinuationAttempt::ReplayAny
)
&& self.advance_provider_candidate()
{
return Ok(Some(PullOutcome::AttemptDiscarded));
}
if matches!(
error.kind(),
ProviderErrorKind::AccountCapacityUnavailable
@@ -716,6 +747,27 @@ where
}
let metadata = stream.metadata().clone();
if !self.account_selection.is_diagnostic()
&& metadata
.provider_account_id()
.is_none_or(|account| !candidate.account_scope().allows(account))
{
let error = GatewayError::new(
GatewayErrorKind::Internal,
"provider selected an account outside the frozen client scope",
);
self.finish_failure(FailureFinalization {
outcome: ExecutionOutcome::Failed,
send_state: self.current_send_state(),
error,
upstream_status_code: None,
provider_error_code: None,
retry_after_ms: None,
provider_response_id: None,
})
.await?;
return Err(EngineError::AccountOutsideClientScope);
}
if pinned_account
.as_ref()
.is_some_and(|required| metadata.provider_account_id() != Some(required))
@@ -819,6 +871,17 @@ where
Ok(None)
}
fn advance_provider_candidate(&mut self) -> bool {
let Some(next) = self.candidate_index.checked_add(1) else {
return false;
};
if next >= self.plan.candidates().len() {
return false;
}
self.candidate_index = next;
true
}
async fn observe_event(&mut self, event: &GatewayEvent) {
self.observe_event_facts(event);
self.mark_send_observed().await;
@@ -3,7 +3,7 @@
use std::cmp::Reverse;
use std::collections::{BTreeSet, HashMap};
use std::fmt;
use std::num::{NonZeroU32, NonZeroU64};
use std::num::{NonZeroU16, NonZeroU32, NonZeroU64};
use std::sync::RwLock;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, SystemTime};
@@ -38,6 +38,58 @@ impl ProviderAccountId {
}
}
/// Per-account scheduling concurrency override.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct AccountConcurrencyLimit(NonZeroU32);
impl AccountConcurrencyLimit {
#[must_use]
pub const fn new(value: u32) -> Option<Self> {
match NonZeroU32::new(value) {
Some(value) => Some(Self(value)),
None => None,
}
}
#[must_use]
pub const fn get(self) -> u32 {
self.0.get()
}
#[must_use]
pub const fn into_non_zero(self) -> NonZeroU32 {
self.0
}
}
/// Relative scheduling priority for an account.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct AccountWeight(NonZeroU16);
impl AccountWeight {
pub const DEFAULT: Self = Self(NonZeroU16::MIN);
pub const MAX: u16 = 100;
#[must_use]
pub const fn new(value: u16) -> Option<Self> {
match NonZeroU16::new(value) {
Some(value) if value.get() <= Self::MAX => Some(Self(value)),
_ => None,
}
}
#[must_use]
pub const fn get(self) -> u16 {
self.0.get()
}
}
impl Default for AccountWeight {
fn default() -> Self {
Self::DEFAULT
}
}
impl fmt::Display for ProviderAccountId {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(&self.0)
@@ -187,45 +239,41 @@ impl fmt::Debug for OpaqueProviderData {
}
}
/// 数据库中固定的账号运行状态。
/// 已持久化的凭据/账号身份状态。
///
/// 只承载终态或需要外部干预的状态;临时限流(429)不进入此状态,由
/// Provider quota 数据驱动(`AccountQuotaSignals::limit_reached`)。
/// 这里只保存凭据与账号身份事实。额度耗尽属于 [`QuotaState`],临时 429 属于
/// 运行时冷却;二者都不得写入此枚举。
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum AccountAvailability {
pub enum CredentialState {
Unknown,
Ready,
QuotaExhausted,
Expired,
Banned,
Invalid,
}
impl AccountAvailability {
impl CredentialState {
/// 返回数据库稳定值。
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::Unknown => "unknown",
Self::Ready => "ready",
Self::QuotaExhausted => "quota_exhausted",
Self::Expired => "expired",
Self::Banned => "banned",
Self::Invalid => "invalid",
}
}
/// 返回账号进入不可恢复错误态、但 Provider 未给出更具体上游原因时的稳定兜底。
///
/// Provider 应优先持久化已脱敏的上游 message 或原因码;该值只避免错误态沿用
/// 上一次无关的错误信息。
/// 返回该凭据状态对应的稳定错误原因。
#[must_use]
pub const fn fallback_error_reason(self) -> Option<&'static str> {
pub const fn error_reason(self) -> Option<AccountErrorReason> {
match self {
Self::Expired => Some("credential_expired"),
Self::Banned => Some("account_banned"),
Self::Invalid => Some("credential_invalid"),
Self::Unknown | Self::Ready | Self::QuotaExhausted => None,
Self::Unknown => Some(AccountErrorReason::AccountUnverified),
Self::Expired => Some(AccountErrorReason::CredentialExpired),
Self::Banned => Some(AccountErrorReason::AccountBanned),
Self::Invalid => Some(AccountErrorReason::CredentialInvalid),
Self::Ready => None,
}
}
@@ -235,7 +283,6 @@ impl AccountAvailability {
match value {
"unknown" => Some(Self::Unknown),
"ready" => Some(Self::Ready),
"quota_exhausted" => Some(Self::QuotaExhausted),
"expired" => Some(Self::Expired),
"banned" => Some(Self::Banned),
"invalid" => Some(Self::Invalid),
@@ -244,6 +291,373 @@ impl AccountAvailability {
}
}
/// 额度访问结论;百分比和余额等展示值不进入此枚举。
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum QuotaAccessState {
Unknown,
Allowed,
Exhausted,
}
impl QuotaAccessState {
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::Unknown => "unknown",
Self::Allowed => "allowed",
Self::Exhausted => "exhausted",
}
}
#[must_use]
pub fn parse(value: &str) -> Option<Self> {
match value {
"unknown" => Some(Self::Unknown),
"allowed" => Some(Self::Allowed),
"exhausted" => Some(Self::Exhausted),
_ => None,
}
}
}
/// 权威额度耗尽证据。
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum QuotaEvidence {
ProviderDenied,
AccountLimitReached,
UsageLimitReached,
PaymentRequired,
}
impl QuotaEvidence {
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::ProviderDenied => "provider_denied",
Self::AccountLimitReached => "account_limit_reached",
Self::UsageLimitReached => "usage_limit_reached",
Self::PaymentRequired => "payment_required",
}
}
#[must_use]
pub fn parse(value: &str) -> Option<Self> {
match value {
"provider_denied" => Some(Self::ProviderDenied),
"account_limit_reached" => Some(Self::AccountLimitReached),
"usage_limit_reached" => Some(Self::UsageLimitReached),
"payment_required" => Some(Self::PaymentRequired),
_ => None,
}
}
}
/// 可审计的额度访问事实。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct QuotaState {
access: QuotaAccessState,
evidence: Option<QuotaEvidence>,
observed_at: Option<SystemTime>,
reset_at: Option<SystemTime>,
}
impl Default for QuotaState {
fn default() -> Self {
Self::unknown()
}
}
impl QuotaState {
#[must_use]
pub const fn unknown() -> Self {
Self {
access: QuotaAccessState::Unknown,
evidence: None,
observed_at: None,
reset_at: None,
}
}
#[must_use]
pub const fn observed_unknown(observed_at: SystemTime) -> Self {
Self {
access: QuotaAccessState::Unknown,
evidence: None,
observed_at: Some(observed_at),
reset_at: None,
}
}
#[must_use]
pub const fn allowed(observed_at: SystemTime) -> Self {
Self {
access: QuotaAccessState::Allowed,
evidence: None,
observed_at: Some(observed_at),
reset_at: None,
}
}
#[must_use]
pub const fn exhausted(
evidence: QuotaEvidence,
observed_at: SystemTime,
reset_at: Option<SystemTime>,
) -> Self {
Self {
access: QuotaAccessState::Exhausted,
evidence: Some(evidence),
observed_at: Some(observed_at),
reset_at,
}
}
/// 从持久化列恢复额度状态;非法组合返回 `None`。
#[must_use]
pub const fn from_persisted(
access: QuotaAccessState,
evidence: Option<QuotaEvidence>,
observed_at: Option<SystemTime>,
reset_at: Option<SystemTime>,
) -> Option<Self> {
match access {
QuotaAccessState::Unknown if evidence.is_none() && reset_at.is_none() => Some(Self {
access,
evidence,
observed_at,
reset_at,
}),
QuotaAccessState::Allowed
if evidence.is_none() && observed_at.is_some() && reset_at.is_none() =>
{
Some(Self {
access,
evidence,
observed_at,
reset_at,
})
}
QuotaAccessState::Exhausted if evidence.is_some() && observed_at.is_some() => {
Some(Self {
access,
evidence,
observed_at,
reset_at,
})
}
_ => None,
}
}
#[must_use]
pub const fn access(self) -> QuotaAccessState {
self.access
}
#[must_use]
pub const fn evidence(self) -> Option<QuotaEvidence> {
self.evidence
}
#[must_use]
pub const fn observed_at(self) -> Option<SystemTime> {
self.observed_at
}
#[must_use]
pub const fn reset_at(self) -> Option<SystemTime> {
self.reset_at
}
/// 返回当前已确认的额度耗尽结论。
///
/// `reset_at` 只表示何时应重新向 Provider 求证;时间到期本身不是额度恢复证据。
#[must_use]
pub fn is_exhausted(self) -> bool {
self.access == QuotaAccessState::Exhausted
}
/// 判断已耗尽额度是否到达 Provider 的下一次复核时间。
///
/// 有明确 `reset_at` 时严格等待该时刻;没有时由 Provider 传入自己的保守复核周期。
#[must_use]
pub fn exhaustion_refresh_due(self, now: SystemTime, fallback_interval: Duration) -> bool {
if !self.is_exhausted() {
return false;
}
let due_at = self.reset_at.or_else(|| {
self.observed_at
.and_then(|observed_at| observed_at.checked_add(fallback_interval))
});
due_at.is_none_or(|due_at| due_at <= now)
}
/// 合并一次额度访问观察;`Unknown` 不能擦除已经确认的访问结论。
#[must_use]
pub fn merge_observation(self, observation: Self) -> Self {
if observation.access == QuotaAccessState::Unknown
&& self.access != QuotaAccessState::Unknown
{
self
} else {
observation
}
}
}
/// 对外唯一的五态账号状态。
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum AccountStatus {
Normal,
QuotaExhausted,
RateLimited,
Disabled,
Error,
}
impl AccountStatus {
#[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",
}
}
#[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,
}
}
/// 管理页稳定业务排序:正常、限流、耗尽、错误、停用。
#[must_use]
pub const fn sort_rank(self) -> u8 {
match self {
Self::Normal => 0,
Self::RateLimited => 1,
Self::QuotaExhausted => 2,
Self::Error => 3,
Self::Disabled => 4,
}
}
}
/// `error` 状态下的稳定原因码。
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum AccountErrorReason {
AccountUnverified,
AccessTokenExpired,
CredentialExpired,
CredentialInvalid,
AccountBanned,
}
impl AccountErrorReason {
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::AccountUnverified => "account_unverified",
Self::AccessTokenExpired => "access_token_expired",
Self::CredentialExpired => "credential_expired",
Self::CredentialInvalid => "credential_invalid",
Self::AccountBanned => "account_banned",
}
}
#[must_use]
pub fn parse(value: &str) -> Option<Self> {
match value {
"account_unverified" => Some(Self::AccountUnverified),
"access_token_expired" => Some(Self::AccessTokenExpired),
"credential_expired" => Some(Self::CredentialExpired),
"credential_invalid" => Some(Self::CredentialInvalid),
"account_banned" => Some(Self::AccountBanned),
_ => None,
}
}
}
/// 唯一状态解析器的完整输入事实。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountStatusFacts {
pub enabled: bool,
pub credential_state: CredentialState,
pub access_token_expires_at: Option<SystemTime>,
pub quota: QuotaState,
pub rate_limited_until: Option<SystemTime>,
pub last_error_reason: Option<AccountErrorReason>,
pub last_error_message: Option<String>,
}
/// 唯一状态解析器的输出。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountStatusProjection {
pub status: AccountStatus,
pub error_reason: Option<AccountErrorReason>,
pub error_message: Option<String>,
/// 仅 `rate_limited` 状态携带仍有效的运行时冷却截止时间。
pub rate_limited_until: Option<SystemTime>,
}
/// 从独立事实派生唯一、互斥的对外状态。
#[must_use]
pub fn resolve_account_status(
facts: &AccountStatusFacts,
now: SystemTime,
) -> AccountStatusProjection {
if !facts.enabled {
return status_projection(AccountStatus::Disabled);
}
let credential_error = facts.credential_state.error_reason().or_else(|| {
facts
.access_token_expires_at
.is_some_and(|expires_at| expires_at <= now)
.then_some(AccountErrorReason::AccessTokenExpired)
});
if let Some(default_reason) = credential_error {
return AccountStatusProjection {
status: AccountStatus::Error,
error_reason: Some(facts.last_error_reason.unwrap_or(default_reason)),
error_message: facts.last_error_message.clone(),
rate_limited_until: None,
};
}
if facts.quota.is_exhausted() {
return status_projection(AccountStatus::QuotaExhausted);
}
if facts
.rate_limited_until
.is_some_and(|rate_limited_until| rate_limited_until > now)
{
return AccountStatusProjection {
status: AccountStatus::RateLimited,
error_reason: None,
error_message: None,
rate_limited_until: facts.rate_limited_until,
};
}
status_projection(AccountStatus::Normal)
}
const fn status_projection(status: AccountStatus) -> AccountStatusProjection {
AccountStatusProjection {
status,
error_reason: None,
error_message: None,
rate_limited_until: None,
}
}
/// 不含 secret 的账号持久事实。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderAccount {
@@ -257,7 +671,12 @@ pub struct ProviderAccount {
authentication_kind: String,
revision: CredentialRevision,
enabled: bool,
availability: AccountAvailability,
concurrency_limit: Option<AccountConcurrencyLimit>,
weight: AccountWeight,
credential_state: CredentialState,
quota: QuotaState,
last_error_reason: Option<AccountErrorReason>,
last_error_message: Option<String>,
access_token_expires_at: Option<SystemTime>,
next_refresh_at: Option<SystemTime>,
has_refresh_token: bool,
@@ -286,7 +705,12 @@ impl ProviderAccount {
authentication_kind,
revision,
enabled: true,
availability: AccountAvailability::Unknown,
concurrency_limit: None,
weight: AccountWeight::DEFAULT,
credential_state: CredentialState::Unknown,
quota: QuotaState::unknown(),
last_error_reason: None,
last_error_message: None,
access_token_expires_at,
next_refresh_at: None,
has_refresh_token: false,
@@ -307,17 +731,34 @@ impl ProviderAccount {
}
#[must_use]
pub const fn with_runtime_state(
pub fn with_account_facts(
mut self,
enabled: bool,
availability: AccountAvailability,
credential_state: CredentialState,
quota: QuotaState,
last_error_reason: Option<AccountErrorReason>,
last_error_message: Option<String>,
) -> Self {
self.enabled = enabled;
self.availability = if self.upstream_user_id.is_some() {
availability
self.credential_state = if self.upstream_user_id.is_some() {
credential_state
} else {
AccountAvailability::Unknown
CredentialState::Unknown
};
self.quota = quota;
self.last_error_reason = last_error_reason;
self.last_error_message = last_error_message;
self
}
#[must_use]
pub const fn with_scheduling(
mut self,
concurrency_limit: Option<AccountConcurrencyLimit>,
weight: AccountWeight,
) -> Self {
self.concurrency_limit = concurrency_limit;
self.weight = weight;
self
}
@@ -382,8 +823,23 @@ impl ProviderAccount {
}
#[must_use]
pub const fn availability(&self) -> AccountAvailability {
self.availability
pub const fn credential_state(&self) -> CredentialState {
self.credential_state
}
#[must_use]
pub const fn quota(&self) -> QuotaState {
self.quota
}
#[must_use]
pub const fn last_error_reason(&self) -> Option<AccountErrorReason> {
self.last_error_reason
}
#[must_use]
pub fn last_error_message(&self) -> Option<&str> {
self.last_error_message.as_deref()
}
#[must_use]
@@ -391,6 +847,24 @@ impl ProviderAccount {
self.enabled
}
#[must_use]
pub const fn concurrency_limit(&self) -> Option<AccountConcurrencyLimit> {
self.concurrency_limit
}
#[must_use]
pub const fn weight(&self) -> AccountWeight {
self.weight
}
#[must_use]
pub const fn effective_concurrency(&self, default: NonZeroU32) -> NonZeroU32 {
match self.concurrency_limit {
Some(limit) => limit.into_non_zero(),
None => default,
}
}
#[must_use]
pub const fn access_token_expires_at(&self) -> Option<SystemTime> {
self.access_token_expires_at
@@ -407,23 +881,25 @@ impl ProviderAccount {
self.has_refresh_token
}
/// 判断账号当前能否进入同 target 的候选池。
/// 组合持久事实与请求级冷却,交给唯一解析器派生状态。
#[must_use]
pub fn is_schedulable(&self, now: SystemTime) -> bool {
let available = match self.availability {
AccountAvailability::Ready => true,
AccountAvailability::Unknown
| AccountAvailability::QuotaExhausted
| AccountAvailability::Expired
| AccountAvailability::Banned
| AccountAvailability::Invalid => false,
};
self.upstream_user_id.is_some()
&& self.enabled
&& available
&& self
.access_token_expires_at
.is_none_or(|expires_at| expires_at > now)
pub fn status_projection(
&self,
now: SystemTime,
rate_limited_until: Option<SystemTime>,
) -> AccountStatusProjection {
resolve_account_status(
&AccountStatusFacts {
enabled: self.enabled,
credential_state: self.credential_state,
access_token_expires_at: self.access_token_expires_at,
quota: self.quota,
rate_limited_until,
last_error_reason: self.last_error_reason,
last_error_message: self.last_error_message.clone(),
},
now,
)
}
}
@@ -600,10 +1076,11 @@ pub enum CredentialCasOutcome {
pub struct QuotaObservation {
pub account_id: ProviderAccountId,
pub expected_revision: CredentialRevision,
pub quota: Option<OpaqueProviderData>,
pub observed_at: Option<SystemTime>,
/// 快照级限流事实(Provider 解析后物化);None 表示不更新物化列。
pub limit_reached: Option<bool>,
pub quota: OpaqueProviderData,
/// Provider 原始 quota document 的观察时间;不代表访问结论发生变化。
pub observed_at: SystemTime,
/// Provider 已从私有 JSON 归一化出的额度访问事实。
pub state: QuotaState,
}
impl fmt::Debug for QuotaObservation {
@@ -614,6 +1091,7 @@ impl fmt::Debug for QuotaObservation {
.field("expected_revision", &self.expected_revision)
.field("quota", &self.quota)
.field("observed_at", &self.observed_at)
.field("state", &self.state)
.finish()
}
}
@@ -624,14 +1102,24 @@ pub enum QuotaWriteOutcome {
Conflict,
}
/// 不改 Provider 原始 quota JSON 的额度访问事实写入。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct QuotaAccessChange {
pub account_id: ProviderAccountId,
pub expected_revision: CredentialRevision,
pub state: QuotaState,
}
/// 账号状态的 revision-fenced 写入。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountStateChange {
pub account_id: ProviderAccountId,
pub expected_revision: CredentialRevision,
pub availability: AccountAvailability,
pub credential_state: CredentialState,
pub observed_at: SystemTime,
/// 上游原因(错误原文 / 原因码);`Some` 写入持久化错误信息,`None` 不更新,恢复 `Ready` 时清空。
/// 受控错误原因;恢复 Ready 时必须清空。
pub error_reason: Option<AccountErrorReason>,
/// 已脱敏的上游诊断信息;恢复 Ready 时必须清空。
pub message: Option<String>,
}
@@ -682,6 +1170,11 @@ pub trait ProviderAccountStore: Send + Sync {
observation: QuotaObservation,
) -> Result<QuotaWriteOutcome, StoreError>;
async fn apply_quota_access(
&self,
change: QuotaAccessChange,
) -> Result<QuotaWriteOutcome, StoreError>;
async fn apply_state_change(&self, change: AccountStateChange) -> Result<(), StoreError>;
async fn update_account(&self, update: ProviderAccountUpdate) -> Result<(), StoreError>;
@@ -772,7 +1265,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 rate_limited_until: Option<SystemTime>,
pub failure_rate_basis_points: Option<u16>,
pub first_output_latency_ms: Option<u64>,
}
@@ -920,11 +1413,16 @@ 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
}
#[must_use]
pub const fn with_rate_limit(mut self, rate_limited_until: Option<SystemTime>) -> Self {
self.rate_limited_until = rate_limited_until;
self
}
#[must_use]
pub const fn with_runtime_health(
mut self,
@@ -942,20 +1440,14 @@ 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>,
limit_reached: bool,
) -> Self {
pub const fn new(reset_at: Option<SystemTime>, remaining_rank: Option<u64>) -> Self {
Self {
reset_at,
remaining_rank,
limit_reached,
}
}
@@ -968,12 +1460,6 @@ 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
}
}
/// 账号持久事实与可重建运行信号的请求级组合。
@@ -991,24 +1477,25 @@ pub struct AccountSelectionContext {
pub excluded_accounts: BTreeSet<ProviderAccountId>,
pub preferred_account: Option<ProviderAccountId>,
pub round_robin_cursor: u64,
pub availability: AccountAvailabilityPolicy,
pub eligibility: AccountEligibilityPolicy,
pub account_scope: Option<std::sync::Arc<crate::routing::FrozenAccountScope>>,
}
/// 选择账号时是否信任本地可用性投影。
/// 选择账号时是否执行本地调度资格投影。
///
/// 管理端对指定账号执行诊断时,会直接向上游确认实际状态;该模式跳过
/// `enabled`、可用性、冷却和 token 到期等本地投影,仍保留租约、并发和
/// 请求间隔约束,且只能与固定账号约束组合使用。
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum AccountAvailabilityPolicy {
pub enum AccountEligibilityPolicy {
#[default]
Enforce,
BypassForDiagnostic,
}
impl AccountAvailabilityPolicy {
impl AccountEligibilityPolicy {
#[must_use]
pub const fn bypasses_local_availability(self) -> bool {
pub const fn bypasses_local_eligibility(self) -> bool {
matches!(self, Self::BypassForDiagnostic)
}
}
@@ -1016,10 +1503,12 @@ impl AccountAvailabilityPolicy {
/// 候选账号未进入调度池的约束。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AccountSchedulingBlocker {
OutsideClientScope,
LocalAvailability,
Excluded,
ConcurrencyLimit,
RequestInterval,
LowerWeight,
}
/// 优先账号在本次选择中的处理结果。
@@ -1062,26 +1551,6 @@ impl AccountSelector {
candidates: &'a [AccountCandidate],
context: &AccountSelectionContext,
) -> Option<AccountSelection<'a>> {
let preferred = if let Some(preferred) = context.preferred_account.as_ref() {
match candidates
.iter()
.find(|candidate| candidate.account.id() == preferred)
{
Some(candidate) => match self.scheduling_blocker(candidate, context) {
None => {
return Some(AccountSelection {
candidate,
preferred: PreferredAccountSelection::Hit,
});
}
Some(blocker) => PreferredAccountSelection::Blocked(blocker),
},
None => PreferredAccountSelection::Missing,
}
} else {
PreferredAccountSelection::NotRequested
};
let mut eligible = candidates
.iter()
.filter(|candidate| self.scheduling_blocker(candidate, context).is_none())
@@ -1089,6 +1558,33 @@ impl AccountSelector {
if eligible.is_empty() {
return None;
}
let highest_weight = eligible
.iter()
.map(|candidate| candidate.account.weight())
.max()?;
let preferred = if let Some(preferred) = context.preferred_account.as_ref() {
match candidates
.iter()
.find(|candidate| candidate.account.id() == preferred)
{
Some(candidate) => match self.scheduling_blocker(candidate, context) {
Some(blocker) => PreferredAccountSelection::Blocked(blocker),
None if candidate.account.weight() < highest_weight => {
PreferredAccountSelection::Blocked(AccountSchedulingBlocker::LowerWeight)
}
None => {
return Some(AccountSelection {
candidate,
preferred: PreferredAccountSelection::Hit,
});
}
},
None => PreferredAccountSelection::Missing,
}
} else {
PreferredAccountSelection::NotRequested
};
eligible.retain(|candidate| candidate.account.weight() == highest_weight);
let candidate = match context.policy.strategy() {
RotationStrategy::QuotaResetPriority => {
@@ -1132,15 +1628,31 @@ impl AccountSelector {
candidate: &AccountCandidate,
context: &AccountSelectionContext,
) -> Option<AccountSchedulingBlocker> {
if !context.availability.bypasses_local_availability()
&& !candidate.account.is_schedulable(context.now)
if context
.account_scope
.as_ref()
.is_some_and(|scope| !scope.allows(candidate.account.id()))
{
return Some(AccountSchedulingBlocker::LocalAvailability);
return Some(AccountSchedulingBlocker::OutsideClientScope);
}
if !context.eligibility.bypasses_local_eligibility() {
let status = candidate
.account
.status_projection(context.now, candidate.signals.rate_limited_until)
.status;
if status != AccountStatus::Normal {
return Some(AccountSchedulingBlocker::LocalAvailability);
}
}
if context.excluded_accounts.contains(candidate.account.id()) {
return Some(AccountSchedulingBlocker::Excluded);
}
if candidate.signals.in_flight >= context.policy.max_concurrent_per_account().get() {
if candidate.signals.in_flight
>= candidate
.account
.effective_concurrency(context.policy.max_concurrent_per_account())
.get()
{
return Some(AccountSchedulingBlocker::ConcurrencyLimit);
}
if candidate
@@ -15,7 +15,8 @@ use crate::engine::admission::{
ClientAdmissionDecision, ClientAdmissionPort, ClientAdmissionRejection, ClientAdmissionRequest,
};
use crate::engine::continuation::{
ContinuationBinding, NativeContinuationPin, NativeContinuationPort, PreviousResponseId,
ContinuationBinding, NativeContinuationPin, NativeContinuationPort,
NativeContinuationStoreErrorKind, PreviousResponseId,
};
use crate::engine::coordinator::ResponseExecutionSession;
use crate::engine::probe::{
@@ -261,22 +262,32 @@ impl DefaultExecutionService {
})?;
let request_id = new_request_id()?;
let routing_context = self
.route_context(request.client.policy.provider_kind())
.route_context(request.client.policy.account_scope().provider_kinds())
.await?;
let plan = request
.client
.snapshot
.plan(&request.public_model, &request.operation, &routing_context)
.plan(
&request.public_model,
&request.operation,
Arc::clone(request.client.policy.account_scope()),
&routing_context,
)
.map_err(map_routing_error)?;
let continuation = match request.metadata.previous_response_id.as_ref() {
Some(previous) => {
let resolve = self.continuation.resolve(previous).fuse();
let resolve = self
.continuation
.resolve(request.client.policy.key_id(), previous)
.fuse();
let timeout = Delay::new(COORDINATION_TIMEOUT).fuse();
pin_mut!(resolve, timeout);
let pin = select_biased! {
result = resolve => match result {
Ok(pin) => pin,
Err(error) => {
Err(error)
if error.kind() == NativeContinuationStoreErrorKind::Unavailable =>
{
tracing::debug!(
request_id = request_id.as_str(),
%error,
@@ -284,10 +295,53 @@ impl DefaultExecutionService {
);
None
}
Err(error)
if error.kind()
== NativeContinuationStoreErrorKind::OwnershipMismatch =>
{
return Err(GatewayError::new(
GatewayErrorKind::PolicyDenied,
"continuation does not belong to this client API key",
));
}
Err(error) => {
tracing::warn!(
request_id = request_id.as_str(),
%error,
"Continuation affinity 记录无效,已拒绝续接"
);
return Err(GatewayError::new(
GatewayErrorKind::Internal,
"continuation state is invalid",
));
}
},
_ = timeout => None,
};
match pin {
Some(pin) if !pin.matches_client(request.client.policy.key_id()) => {
return Err(GatewayError::new(
GatewayErrorKind::PolicyDenied,
"continuation does not belong to this client API key",
));
}
Some(pin) if !request.client.policy.account_scope().allows(pin.account()) => {
return Err(GatewayError::new(
GatewayErrorKind::PolicyDenied,
"continuation account is outside the client account scope",
));
}
Some(pin)
if !plan
.candidates()
.iter()
.any(|candidate| candidate.provider() == pin.provider()) =>
{
return Err(GatewayError::new(
GatewayErrorKind::NoAvailableProvider,
"continuation provider is not available",
));
}
Some(pin) => {
attach_continuation_session_state(&mut request.operation, &pin);
ContinuationBinding::Pinned(pin)
@@ -360,14 +414,19 @@ impl DefaultExecutionService {
client_api_key_id: request.client.policy.key_id().clone(),
model_request_id: request_id.clone(),
};
let observation = self
.providers
.request_observation(request.client.policy.provider_kind(), &request.operation);
let observation = plan
.candidates()
.first()
.map_or_else(Default::default, |candidate| {
self.providers
.request_observation(candidate.provider(), &request.operation)
});
let new_request = NewModelRequest {
id: request_id.clone(),
client_api_key_id: Some(request.client.policy.key_id().clone()),
client_api_key_ref: request.client.policy.key_id().clone(),
config_revision: plan.config_revision(),
routing: request.client.policy.account_scope().routing_snapshot(),
protocol: request.metadata.protocol,
operation: request.operation.kind(),
endpoint: request.metadata.endpoint,
@@ -417,38 +476,42 @@ impl DefaultExecutionService {
async fn route_context(
&self,
provider_kind: &ProviderKind,
provider_kinds: &BTreeSet<ProviderKind>,
) -> Result<RoutingContext, GatewayError> {
let decision = self.circuits.decision(provider_kind).fuse();
let timeout = Delay::new(COORDINATION_TIMEOUT).fuse();
pin_mut!(decision, timeout);
let decision = select_biased! {
result = decision => Some(result),
_ = timeout => None,
};
let blocked_providers = match decision {
Some(Ok(ProviderCircuitDecision::Allow)) => BTreeSet::new(),
Some(Ok(ProviderCircuitDecision::BlockedUntil(_))) => {
BTreeSet::from([provider_kind.clone()])
let decisions = futures::future::join_all(provider_kinds.iter().map(|provider_kind| {
let circuits = Arc::clone(&self.circuits);
async move {
let decision = circuits.decision(provider_kind).fuse();
let timeout = Delay::new(COORDINATION_TIMEOUT).fuse();
pin_mut!(decision, timeout);
let decision = select_biased! {
result = decision => Some(result),
_ = timeout => None,
};
(provider_kind, decision)
}
Some(Err(error)) => {
tracing::warn!(
}))
.await;
let mut blocked_providers = BTreeSet::new();
for (provider_kind, decision) in decisions {
match decision {
Some(Ok(ProviderCircuitDecision::BlockedUntil(_))) => {
blocked_providers.insert(provider_kind.clone());
}
Some(Err(error)) => tracing::warn!(
provider = provider_kind.as_str(),
%error,
"Provider circuit 读取失败,按可重建协调状态 fail-open"
);
BTreeSet::new()
}
None => {
tracing::warn!(
),
None => tracing::warn!(
provider = provider_kind.as_str(),
"Provider circuit 读取超时,按可重建协调状态 fail-open"
);
BTreeSet::new()
),
Some(Ok(ProviderCircuitDecision::Allow)) => {}
}
};
}
Ok(RoutingContext {
provider_kind: Some(provider_kind.clone()),
required_provider: None,
blocked_providers,
})
}
@@ -479,11 +542,16 @@ impl DefaultExecutionService {
GatewayError::new(GatewayErrorKind::Unsupported, "requested model is invalid")
})?;
let routing_context = RoutingContext {
provider_kind: Some(provider_kind),
required_provider: Some(provider_kind),
..RoutingContext::default()
};
let plan = snapshot
.plan(&public_model, &operation, &routing_context)
.plan(
&public_model,
&operation,
snapshot.all_account_scope(),
&routing_context,
)
.map_err(map_routing_error)?;
let started_at = SystemTime::now();
let deadline_at = started_at
@@ -499,6 +567,7 @@ impl DefaultExecutionService {
client_api_key_id: None,
client_api_key_ref: actor,
config_revision: plan.config_revision(),
routing: crate::routing::AccountRoutingSnapshot::all(),
protocol: "admin_connection_test".to_owned(),
operation: operation.kind(),
endpoint: "/api/admin/accounts/connection-test".to_owned(),
@@ -693,19 +762,19 @@ impl ExecutionService for DefaultExecutionService {
fn public_models(&self, client: &AuthenticatedClient) -> Vec<PublicModelId> {
client
.snapshot
.public_models_for_provider(client.policy.provider_kind())
.public_models_for_scope(client.policy.account_scope())
}
fn public_model_profiles(&self, client: &AuthenticatedClient) -> Vec<PublicModelProfile> {
client
.snapshot
.public_model_profiles_for_provider(client.policy.provider_kind())
.public_model_profiles_for_scope(client.policy.account_scope())
}
fn contains_public_model(&self, client: &AuthenticatedClient, model: &PublicModelId) -> bool {
client
.snapshot
.contains_public_model_for_provider(model, client.policy.provider_kind())
.contains_public_model_for_scope(model, client.policy.account_scope())
}
fn start(
@@ -1006,7 +1075,8 @@ fn new_request_id() -> Result<ModelRequestId, GatewayError> {
fn map_routing_error(error: crate::error::RoutingError) -> GatewayError {
match error {
crate::error::RoutingError::NoCapableProvider { .. } => GatewayError::new(
crate::error::RoutingError::NoCapableProvider { .. }
| crate::error::RoutingError::EmptyAccountScope => GatewayError::new(
GatewayErrorKind::NoAvailableProvider,
"no provider can execute this request",
),
@@ -233,6 +233,7 @@ pub struct AccountAttemptContext {
state_owner: Option<ProviderAccountStateOwner>,
credential_recovery_attempted: bool,
diagnostic_required_account: bool,
account_scope: Option<Arc<crate::routing::FrozenAccountScope>>,
}
impl AccountAttemptContext {
@@ -248,6 +249,7 @@ impl AccountAttemptContext {
state_owner,
credential_recovery_attempted: false,
diagnostic_required_account: false,
account_scope: None,
}
}
@@ -266,9 +268,17 @@ impl AccountAttemptContext {
state_owner,
credential_recovery_attempted: false,
diagnostic_required_account: true,
account_scope: None,
}
}
/// 附着普通请求认证时冻结的账号范围。
#[must_use]
pub fn with_account_scope(mut self, scope: Arc<crate::routing::FrozenAccountScope>) -> Self {
self.account_scope = Some(scope);
self
}
/// 标记本请求已对即将选择的固定账号执行过一次凭据恢复。
#[must_use]
pub const fn with_credential_recovery_attempted(mut self, attempted: bool) -> Self {
@@ -301,6 +311,11 @@ impl AccountAttemptContext {
pub const fn is_diagnostic_required_account(&self) -> bool {
self.diagnostic_required_account
}
#[must_use]
pub const fn account_scope(&self) -> Option<&Arc<crate::routing::FrozenAccountScope>> {
self.account_scope.as_ref()
}
}
/// 请求中 Provider 账号绑定状态的唯一归属。
@@ -506,6 +521,12 @@ impl AttemptContext {
self.account.is_diagnostic_required_account()
}
/// 普通请求认证时冻结的账号范围;管理端诊断为 `None`。
#[must_use]
pub const fn account_scope(&self) -> Option<&Arc<crate::routing::FrozenAccountScope>> {
self.account.account_scope()
}
#[must_use]
pub const fn continuation(&self) -> Option<&ContinuationBinding> {
self.continuation.as_ref()
@@ -529,6 +550,7 @@ pub struct NewModelRequest {
pub client_api_key_id: Option<ClientApiKeyId>,
pub client_api_key_ref: ClientApiKeyId,
pub config_revision: ConfigRevision,
pub routing: crate::routing::AccountRoutingSnapshot,
pub protocol: String,
pub operation: OperationKind,
pub endpoint: String,
@@ -749,6 +771,8 @@ pub enum EngineError {
ContinuationPinMismatch,
#[error("provider did not use the account required by this execution")]
RequiredAccountMismatch,
#[error("provider selected an account outside the frozen client scope")]
AccountOutsideClientScope,
#[error("provider execution failed")]
Provider(ProviderError),
#[error("request was cancelled")]
+9
View File
@@ -33,6 +33,9 @@ pub enum IdentifierError {
/// 规定前缀。
expected: &'static str,
},
/// 标识不满足该领域值对象的固定格式。
#[error("identifier has an invalid format")]
InvalidFormat,
}
/// Operation 构造或校验失败。
@@ -67,6 +70,12 @@ pub enum RoutingError {
/// 配置 revision 必须为正数。
#[error("config revision must be greater than zero")]
InvalidRevision,
/// Restricted scope 必须保留至少一个持久化分组 binding。
#[error("account routing scope is invalid")]
InvalidAccountScope,
/// Key 的账号范围中当前没有任何可参与路由的账号。
#[error("account routing scope is empty")]
EmptyAccountScope,
/// 快照中存在重复实体。
#[error("duplicate {entity} `{id}`")]
DuplicateEntity {
@@ -1,11 +1,12 @@
//! 下游 Client API Key 的准入策略。
//!
//! Client API Key 绑定一个 Provider 平台;模型名称不参与改写或 allowlist 判断。
//! Client API Key 冻结账号分组权限;模型名称不参与权限判断。
use std::fmt;
use std::sync::Arc;
use crate::error::{IdentifierError, PolicyError, validate_text};
use crate::routing::ProviderKind;
use crate::routing::FrozenAccountScope;
/// `client_api_keys.id` 的核心值对象。
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
@@ -88,7 +89,7 @@ impl RateLimits {
pub struct ClientPolicy {
key_id: ClientApiKeyId,
plaintext_key: PlaintextClientApiKey,
provider_kind: ProviderKind,
account_scope: Arc<FrozenAccountScope>,
enabled: bool,
limits: RateLimits,
}
@@ -98,14 +99,14 @@ impl ClientPolicy {
pub const fn new(
key_id: ClientApiKeyId,
plaintext_key: PlaintextClientApiKey,
provider_kind: ProviderKind,
account_scope: Arc<FrozenAccountScope>,
enabled: bool,
limits: RateLimits,
) -> Self {
Self {
key_id,
plaintext_key,
provider_kind,
account_scope,
enabled,
limits,
}
@@ -122,8 +123,8 @@ impl ClientPolicy {
}
#[must_use]
pub const fn provider_kind(&self) -> &ProviderKind {
&self.provider_kind
pub const fn account_scope(&self) -> &Arc<FrozenAccountScope> {
&self.account_scope
}
#[must_use]
@@ -9,10 +9,11 @@ use std::time::{Duration, SystemTime};
use futures::future::BoxFuture;
use crate::engine::credential::{
AccountAvailability, AccountFeedbackStats, AccountRuntimeSignals, CredentialRevision,
AccountFeedbackStats, AccountRuntimeSignals, CredentialRevision, CredentialState,
OpaqueProviderData, ProviderAccountId, ProviderAccountStore,
};
use crate::error::{IdentifierError, validate_text};
use crate::policy::ClientApiKeyId;
use crate::routing::{ProviderKind, UpstreamModelId};
const MAX_PENDING_FLOW_TTL: Duration = Duration::from_secs(30 * 60);
@@ -170,6 +171,7 @@ pub enum ProviderLeaseRequest {
pub trait ProviderLeasePort: Send + Sync {
fn load_state<'a>(
&'a self,
client_api_key_id: &'a ClientApiKeyId,
provider_kind: &'a ProviderKind,
accounts: &'a [ProviderAccountId],
) -> BoxFuture<'a, Result<ProviderSchedulingState, ProviderStoreError>>;
@@ -411,7 +413,7 @@ pub struct ProviderCredentialState {
account_id: ProviderAccountId,
credential_revision: CredentialRevision,
enabled: bool,
availability: AccountAvailability,
credential_state: CredentialState,
observed_at: SystemTime,
}
@@ -421,14 +423,14 @@ impl ProviderCredentialState {
account_id: ProviderAccountId,
credential_revision: CredentialRevision,
enabled: bool,
availability: AccountAvailability,
credential_state: CredentialState,
observed_at: SystemTime,
) -> Self {
Self {
account_id,
credential_revision,
enabled,
availability,
credential_state,
observed_at,
}
}
@@ -449,8 +451,8 @@ impl ProviderCredentialState {
}
#[must_use]
pub const fn availability(&self) -> AccountAvailability {
self.availability
pub const fn credential_state(&self) -> CredentialState {
self.credential_state
}
#[must_use]
+296 -3
View File
@@ -9,12 +9,293 @@ use std::fmt;
use std::num::{NonZeroU32, NonZeroU64};
use std::sync::Arc;
use crate::engine::credential::AccountSelectionPolicy;
use crate::engine::credential::{AccountSelectionPolicy, ProviderAccountId};
use crate::error::{IdentifierError, RoutingError, validate_text};
use crate::operation::{CapabilityRequirements, Feature, OperationKind};
const MAX_REQUEST_ATTEMPTS: u32 = 32;
/// `account_groups.id` 的核心值对象。
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct AccountGroupId(String);
impl AccountGroupId {
/// 校验并创建账号分组 ID。
pub fn new(value: impl Into<String>) -> Result<Self, IdentifierError> {
let value = value.into();
let Some(suffix) = value.strip_prefix("grp_") else {
return Err(IdentifierError::MissingPrefix { expected: "grp_" });
};
if suffix.len() != 32
|| !suffix
.bytes()
.all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
{
return Err(IdentifierError::InvalidFormat);
}
Ok(Self(value))
}
#[must_use]
pub fn as_str(&self) -> &str {
&self.0
}
}
impl fmt::Display for AccountGroupId {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(&self.0)
}
}
/// 快照中一个账号的 Provider 与分组归属。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RuntimeAccount {
provider_kind: ProviderKind,
group_ids: Arc<BTreeSet<AccountGroupId>>,
}
impl RuntimeAccount {
#[must_use]
pub fn new(provider_kind: ProviderKind, group_ids: BTreeSet<AccountGroupId>) -> Self {
Self {
provider_kind,
group_ids: Arc::new(group_ids),
}
}
#[must_use]
pub const fn provider_kind(&self) -> &ProviderKind {
&self.provider_kind
}
#[must_use]
pub fn group_ids(&self) -> &BTreeSet<AccountGroupId> {
&self.group_ids
}
}
/// 全快照共享的账号、Provider 与分组反向索引。
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct RuntimeAccountDirectory {
accounts: BTreeMap<ProviderAccountId, RuntimeAccount>,
providers_with_accounts: BTreeSet<ProviderKind>,
providers_by_group: BTreeMap<AccountGroupId, BTreeSet<ProviderKind>>,
}
impl RuntimeAccountDirectory {
#[must_use]
pub fn new(accounts: BTreeMap<ProviderAccountId, RuntimeAccount>) -> Self {
let mut providers_with_accounts = BTreeSet::new();
let mut providers_by_group = BTreeMap::<AccountGroupId, BTreeSet<ProviderKind>>::new();
for account in accounts.values() {
providers_with_accounts.insert(account.provider_kind.clone());
for group_id in account.group_ids.iter() {
providers_by_group
.entry(group_id.clone())
.or_default()
.insert(account.provider_kind.clone());
}
}
Self {
accounts,
providers_with_accounts,
providers_by_group,
}
}
#[must_use]
pub fn account(&self, account_id: &ProviderAccountId) -> Option<&RuntimeAccount> {
self.accounts.get(account_id)
}
#[must_use]
pub fn providers_with_accounts(&self) -> &BTreeSet<ProviderKind> {
&self.providers_with_accounts
}
#[must_use]
pub fn providers_for_groups<'a>(
&self,
group_ids: impl IntoIterator<Item = &'a AccountGroupId>,
) -> BTreeSet<ProviderKind> {
group_ids
.into_iter()
.filter_map(|group_id| self.providers_by_group.get(group_id))
.flatten()
.cloned()
.collect()
}
}
/// 历史请求保存的账号范围种类。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AccountRoutingScopeKind {
All,
Groups,
}
impl AccountRoutingScopeKind {
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::All => "all",
Self::Groups => "groups",
}
}
}
/// 请求开始时冻结的分组 ID 与名称。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RoutingGroupSnapshot {
id: AccountGroupId,
name: String,
}
impl RoutingGroupSnapshot {
#[must_use]
pub fn new(id: AccountGroupId, name: String) -> Self {
Self { id, name }
}
#[must_use]
pub const fn id(&self) -> &AccountGroupId {
&self.id
}
#[must_use]
pub fn name(&self) -> &str {
&self.name
}
}
/// 请求历史所需的完整、稳定账号范围快照。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountRoutingSnapshot {
kind: AccountRoutingScopeKind,
groups: Arc<[RoutingGroupSnapshot]>,
}
impl AccountRoutingSnapshot {
#[must_use]
pub fn all() -> Self {
Self {
kind: AccountRoutingScopeKind::All,
groups: Arc::from([]),
}
}
#[must_use]
pub fn groups(groups: Vec<RoutingGroupSnapshot>) -> Self {
Self {
kind: AccountRoutingScopeKind::Groups,
groups: Arc::from(groups),
}
}
#[must_use]
pub const fn kind(&self) -> AccountRoutingScopeKind {
self.kind
}
#[must_use]
pub fn groups_snapshot(&self) -> &[RoutingGroupSnapshot] {
&self.groups
}
}
/// Key 持久 binding 编译出的账号权限。
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ClientRoutingScope {
AllAccounts,
Restricted {
bound_groups: Arc<[RoutingGroupSnapshot]>,
enabled_group_ids: Arc<BTreeSet<AccountGroupId>>,
provider_kinds: Arc<BTreeSet<ProviderKind>>,
},
}
impl ClientRoutingScope {
#[must_use]
pub fn all_accounts() -> Self {
Self::AllAccounts
}
pub fn restricted(
bound_groups: Vec<RoutingGroupSnapshot>,
enabled_group_ids: BTreeSet<AccountGroupId>,
provider_kinds: BTreeSet<ProviderKind>,
) -> Result<Self, RoutingError> {
if bound_groups.is_empty() {
return Err(RoutingError::InvalidAccountScope);
}
Ok(Self::Restricted {
bound_groups: Arc::from(bound_groups),
enabled_group_ids: Arc::new(enabled_group_ids),
provider_kinds: Arc::new(provider_kinds),
})
}
}
/// 一次认证随 RuntimeSnapshot 冻结的账号目录与 Key 权限。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FrozenAccountScope {
directory: Arc<RuntimeAccountDirectory>,
client_scope: ClientRoutingScope,
}
impl FrozenAccountScope {
#[must_use]
pub const fn new(
directory: Arc<RuntimeAccountDirectory>,
client_scope: ClientRoutingScope,
) -> Self {
Self {
directory,
client_scope,
}
}
#[must_use]
pub fn allows(&self, account_id: &ProviderAccountId) -> bool {
let Some(account) = self.directory.account(account_id) else {
return false;
};
match &self.client_scope {
ClientRoutingScope::AllAccounts => true,
ClientRoutingScope::Restricted {
enabled_group_ids, ..
} => account
.group_ids()
.iter()
.any(|group_id| enabled_group_ids.contains(group_id)),
}
}
#[must_use]
pub fn provider_kinds(&self) -> &BTreeSet<ProviderKind> {
match &self.client_scope {
ClientRoutingScope::AllAccounts => self.directory.providers_with_accounts(),
ClientRoutingScope::Restricted { provider_kinds, .. } => provider_kinds,
}
}
#[must_use]
pub fn routing_snapshot(&self) -> AccountRoutingSnapshot {
match &self.client_scope {
ClientRoutingScope::AllAccounts => AccountRoutingSnapshot::all(),
ClientRoutingScope::Restricted { bound_groups, .. } => {
AccountRoutingSnapshot::groups(bound_groups.to_vec())
}
}
}
#[must_use]
pub const fn directory(&self) -> &Arc<RuntimeAccountDirectory> {
&self.directory
}
}
/// 编译进二进制的 Provider adapter slug。
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct ProviderKind(String);
@@ -494,8 +775,8 @@ impl ProviderModel {
/// 本次请求选择 Provider 时使用的动态过滤事实。
#[derive(Debug, Clone, Default)]
pub struct RoutingContext {
/// 已认证 Client API Key 绑定的平台;模型名称不参与平台猜测。
pub provider_kind: Option<ProviderKind>,
/// 管理端 connection test 显式限制的 Provider;普通请求留空。
pub required_provider: Option<ProviderKind>,
pub blocked_providers: BTreeSet<ProviderKind>,
}
@@ -505,6 +786,7 @@ pub struct ProviderCandidate {
provider: ProviderKind,
upstream_model: UpstreamModelId,
emulated_features: BTreeSet<Feature>,
account_scope: Arc<FrozenAccountScope>,
}
impl ProviderCandidate {
@@ -522,6 +804,11 @@ impl ProviderCandidate {
pub const fn emulated_features(&self) -> &BTreeSet<Feature> {
&self.emulated_features
}
#[must_use]
pub const fn account_scope(&self) -> &Arc<FrozenAccountScope> {
&self.account_scope
}
}
/// 一次请求冻结的 Provider 尝试顺序。
@@ -532,6 +819,7 @@ pub struct RoutingPlan {
public_model: PublicModelId,
operation: OperationKind,
max_attempts: NonZeroU32,
account_scope: Arc<FrozenAccountScope>,
candidates: Arc<[ProviderCandidate]>,
}
@@ -561,6 +849,11 @@ impl RoutingPlan {
self.max_attempts
}
#[must_use]
pub const fn account_scope(&self) -> &Arc<FrozenAccountScope> {
&self.account_scope
}
#[must_use]
pub fn candidates(&self) -> &[ProviderCandidate] {
&self.candidates
@@ -11,7 +11,7 @@ use futures::{FutureExt as _, Stream, StreamExt as _, pin_mut, select_biased};
use futures_timer::Delay;
use crate::engine::CancellationToken;
use crate::engine::credential::{AccountSelectionPolicy, RotationStrategy};
use crate::engine::credential::{AccountSelectionPolicy, ProviderAccountId, RotationStrategy};
use crate::engine::provider::{ProviderCatalogGeneration, ProviderRegistry};
use crate::error::RoutingError;
use crate::health::{HealthProbe, HealthState};
@@ -24,8 +24,9 @@ use crate::task::{
};
use super::{
ConfigRevision, ModelCapabilities, ProviderCandidate, ProviderKind, ProviderModel,
PublicModelId, RoutingContext, RoutingPlan, UpstreamModelId,
AccountGroupId, ClientRoutingScope, ConfigRevision, FrozenAccountScope, ModelCapabilities,
ProviderCandidate, ProviderKind, ProviderModel, PublicModelId, RoutingContext,
RoutingGroupSnapshot, RoutingPlan, RuntimeAccount, RuntimeAccountDirectory, UpstreamModelId,
};
const RECONCILIATION_INTERVAL: Duration = Duration::from_secs(5);
@@ -66,7 +67,7 @@ impl SnapshotSettingsFacts {
pub struct SnapshotClientPolicyFacts {
key_id: ClientApiKeyId,
plaintext_key: PlaintextClientApiKey,
provider_kind: String,
group_ids: Vec<AccountGroupId>,
limits: RateLimits,
}
@@ -75,18 +76,67 @@ impl SnapshotClientPolicyFacts {
pub fn new(
key_id: ClientApiKeyId,
plaintext_key: PlaintextClientApiKey,
provider_kind: impl Into<String>,
group_ids: Vec<AccountGroupId>,
limits: RateLimits,
) -> Self {
Self {
key_id,
plaintext_key,
provider_kind: provider_kind.into(),
group_ids,
limits,
}
}
}
/// Store 读取到的账号分组事实。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SnapshotAccountGroupFacts {
id: AccountGroupId,
name: String,
enabled: bool,
}
impl SnapshotAccountGroupFacts {
#[must_use]
pub fn new(id: AccountGroupId, name: String, enabled: bool) -> Self {
Self { id, name, enabled }
}
}
/// Store 读取到的账号及其固有 Provider 事实。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SnapshotProviderAccountFacts {
account_id: ProviderAccountId,
provider_kind: String,
}
impl SnapshotProviderAccountFacts {
#[must_use]
pub fn new(account_id: ProviderAccountId, provider_kind: impl Into<String>) -> Self {
Self {
account_id,
provider_kind: provider_kind.into(),
}
}
}
/// Store 读取到的一条分组成员关系。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SnapshotAccountGroupMemberFacts {
group_id: AccountGroupId,
account_id: ProviderAccountId,
}
impl SnapshotAccountGroupMemberFacts {
#[must_use]
pub const fn new(group_id: AccountGroupId, account_id: ProviderAccountId) -> Self {
Self {
group_id,
account_id,
}
}
}
/// 一次一致性读取产生的全部 RuntimeSnapshot 持久事实。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SnapshotFacts {
@@ -94,6 +144,9 @@ pub struct SnapshotFacts {
observed_current_revision: ConfigRevision,
settings: SnapshotSettingsFacts,
client_policies: Vec<SnapshotClientPolicyFacts>,
account_groups: Vec<SnapshotAccountGroupFacts>,
provider_accounts: Vec<SnapshotProviderAccountFacts>,
group_memberships: Vec<SnapshotAccountGroupMemberFacts>,
}
impl SnapshotFacts {
@@ -103,12 +156,18 @@ impl SnapshotFacts {
observed_current_revision: ConfigRevision,
settings: SnapshotSettingsFacts,
client_policies: Vec<SnapshotClientPolicyFacts>,
account_groups: Vec<SnapshotAccountGroupFacts>,
provider_accounts: Vec<SnapshotProviderAccountFacts>,
group_memberships: Vec<SnapshotAccountGroupMemberFacts>,
) -> Self {
Self {
config_revision,
observed_current_revision,
settings,
client_policies,
account_groups,
provider_accounts,
group_memberships,
}
}
@@ -223,13 +282,16 @@ async fn compile_runtime_snapshot(
providers: &ProviderRegistry,
) -> Result<RuntimeSnapshot, RuntimeSnapshotCompileError> {
let provider_kinds = providers.provider_kinds().cloned().collect::<Vec<_>>();
let registered_providers = provider_kinds.iter().cloned().collect::<BTreeSet<_>>();
// 实时目录只提供公开模型与能力提示;查询失败时保留 Provider 透传语义。
// 目录查询失败表示未知;查询成功后,即使为空,也必须与“已知缺少模型”区分。
let mut provider_models = Vec::new();
let mut known_provider_catalogs = BTreeSet::new();
for provider in &provider_kinds {
let Ok(models) = providers.query_model_capabilities(provider).await else {
continue;
};
known_provider_catalogs.insert(provider.clone());
provider_models.extend(models.into_iter().map(|model| {
let compiled = ProviderModel::new(
provider.clone(),
@@ -243,6 +305,56 @@ async fn compile_runtime_snapshot(
}));
}
let mut groups = BTreeMap::new();
for group in facts.account_groups {
if group.name.trim() != group.name
|| group.name.is_empty()
|| group.name.chars().count() > 100
|| group.name.chars().any(char::is_control)
|| groups.insert(group.id.clone(), group).is_some()
{
return Err(RuntimeSnapshotCompileError::InvalidData);
}
}
let mut account_groups = BTreeMap::<ProviderAccountId, BTreeSet<AccountGroupId>>::new();
let mut accounts = BTreeMap::new();
for account in facts.provider_accounts {
let provider_kind = ProviderKind::new(account.provider_kind)
.map_err(|_| RuntimeSnapshotCompileError::InvalidData)?;
if !registered_providers.contains(&provider_kind)
|| accounts
.insert(
account.account_id.clone(),
RuntimeAccount::new(provider_kind, BTreeSet::new()),
)
.is_some()
{
return Err(RuntimeSnapshotCompileError::InvalidData);
}
account_groups.insert(account.account_id, BTreeSet::new());
}
let mut memberships = BTreeSet::new();
for membership in facts.group_memberships {
if !groups.contains_key(&membership.group_id)
|| !accounts.contains_key(&membership.account_id)
|| !memberships.insert((membership.group_id.clone(), membership.account_id.clone()))
{
return Err(RuntimeSnapshotCompileError::InvalidData);
}
account_groups
.get_mut(&membership.account_id)
.ok_or(RuntimeSnapshotCompileError::InvalidData)?
.insert(membership.group_id);
}
for (account_id, group_ids) in account_groups {
let account = accounts
.get_mut(&account_id)
.ok_or(RuntimeSnapshotCompileError::InvalidData)?;
*account = RuntimeAccount::new(account.provider_kind().clone(), group_ids);
}
let account_directory = Arc::new(RuntimeAccountDirectory::new(accounts));
let model_mappings = facts.settings.model_mappings;
let rotation_strategy = RotationStrategy::parse(facts.settings.rotation_strategy.as_str())
.ok_or(RuntimeSnapshotCompileError::InvalidData)?;
@@ -252,21 +364,48 @@ async fn compile_runtime_snapshot(
.ok_or(RuntimeSnapshotCompileError::InvalidData)?,
Duration::from_millis(facts.settings.request_interval_ms),
);
let client_policies = facts
.client_policies
.into_iter()
.map(|policy| {
let provider = ProviderKind::new(policy.provider_kind)
.map_err(|_| RuntimeSnapshotCompileError::InvalidData)?;
Ok(ClientPolicy::new(
policy.key_id,
policy.plaintext_key,
provider,
true,
policy.limits,
))
})
.collect::<Result<Vec<_>, RuntimeSnapshotCompileError>>()?;
let mut client_policies = Vec::with_capacity(facts.client_policies.len());
for policy in facts.client_policies {
let account_scope = if policy.group_ids.is_empty() {
FrozenAccountScope::new(
Arc::clone(&account_directory),
ClientRoutingScope::all_accounts(),
)
} else {
let mut seen = BTreeSet::new();
let mut bound_groups = Vec::with_capacity(policy.group_ids.len());
let mut enabled_group_ids = BTreeSet::new();
for group_id in policy.group_ids {
if !seen.insert(group_id.clone()) {
return Err(RuntimeSnapshotCompileError::InvalidData);
}
let group = groups
.get(&group_id)
.ok_or(RuntimeSnapshotCompileError::InvalidData)?;
bound_groups.push(RoutingGroupSnapshot::new(
group.id.clone(),
group.name.clone(),
));
if group.enabled {
enabled_group_ids.insert(group_id);
}
}
bound_groups.sort_by(|left, right| left.id().cmp(right.id()));
let provider_kinds = account_directory.providers_for_groups(&enabled_group_ids);
FrozenAccountScope::new(
Arc::clone(&account_directory),
ClientRoutingScope::restricted(bound_groups, enabled_group_ids, provider_kinds)
.map_err(|_| RuntimeSnapshotCompileError::InvalidData)?,
)
};
client_policies.push(ClientPolicy::new(
policy.key_id,
policy.plaintext_key,
Arc::new(account_scope),
true,
policy.limits,
));
}
RuntimeSnapshot::new(
facts.config_revision,
@@ -276,7 +415,12 @@ async fn compile_runtime_snapshot(
client_policies,
)
.map_err(|_| RuntimeSnapshotCompileError::InvalidData)
.map(|snapshot| snapshot.with_model_mappings(model_mappings))
.map(|snapshot| {
snapshot
.with_model_mappings(model_mappings)
.with_account_directory(account_directory)
.with_known_provider_catalogs(known_provider_catalogs)
})
}
/// 数据面使用的不可变配置快照。
@@ -290,6 +434,8 @@ pub struct RuntimeSnapshot {
Arc<BTreeMap<ProviderKind, BTreeMap<UpstreamModelId, super::ModelPresentation>>>,
model_mappings: Arc<BTreeMap<String, String>>,
provider_catalog_generations: Arc<BTreeMap<ProviderKind, ProviderCatalogGeneration>>,
known_provider_catalogs: Arc<BTreeSet<ProviderKind>>,
account_directory: Arc<RuntimeAccountDirectory>,
client_policies: Arc<BTreeMap<ClientApiKeyId, ClientPolicy>>,
}
@@ -312,6 +458,7 @@ impl RuntimeSnapshot {
}
}
let mut known_provider_catalogs = BTreeSet::new();
let mut model_map =
BTreeMap::<ProviderKind, BTreeMap<UpstreamModelId, ModelCapabilities>>::new();
let mut presentation_map =
@@ -323,6 +470,7 @@ impl RuntimeSnapshot {
capabilities,
presentation,
} = model;
known_provider_catalogs.insert(provider.clone());
if !provider_set.contains(&provider) {
return Err(RoutingError::NotFound {
entity: "provider",
@@ -367,6 +515,8 @@ impl RuntimeSnapshot {
provider_model_presentations: Arc::new(presentation_map),
model_mappings: Arc::new(BTreeMap::new()),
provider_catalog_generations: Arc::new(BTreeMap::new()),
known_provider_catalogs: Arc::new(known_provider_catalogs),
account_directory: Arc::new(RuntimeAccountDirectory::default()),
client_policies: Arc::new(client_policy_map),
})
}
@@ -377,6 +527,26 @@ impl RuntimeSnapshot {
self
}
#[must_use]
pub fn with_account_directory(mut self, directory: Arc<RuntimeAccountDirectory>) -> Self {
self.account_directory = directory;
self
}
#[must_use]
fn with_known_provider_catalogs(mut self, providers: BTreeSet<ProviderKind>) -> Self {
self.known_provider_catalogs = Arc::new(providers);
self
}
#[must_use]
pub fn all_account_scope(&self) -> Arc<FrozenAccountScope> {
Arc::new(FrozenAccountScope::new(
Arc::clone(&self.account_directory),
ClientRoutingScope::all_accounts(),
))
}
#[must_use]
fn with_provider_catalog_generations(
mut self,
@@ -459,8 +629,39 @@ impl RuntimeSnapshot {
.collect()
}
/// 已取得目录时按目录与全局别名判断模型可用性;目录暂不可用时沿用旧版
/// fail-open 语义,把具体 wire 可接受性留给上游。
/// 合并冻结账号范围内实际存在 Provider 的公开模型。
#[must_use]
pub fn public_models_for_scope(&self, scope: &FrozenAccountScope) -> Vec<PublicModelId> {
scope
.provider_kinds()
.iter()
.flat_map(|provider| self.public_models_for_provider(provider))
.collect::<BTreeSet<_>>()
.into_iter()
.collect()
}
/// 合并冻结账号范围内实际存在 Provider 的公开模型画像。
#[must_use]
pub fn public_model_profiles_for_scope(
&self,
scope: &FrozenAccountScope,
) -> Vec<super::PublicModelProfile> {
let mut profiles = BTreeMap::new();
for provider in scope.provider_kinds() {
for profile in self.public_model_profiles_for_provider(provider) {
profiles
.entry(profile.model().clone())
.or_insert_with(|| profile.presentation().clone());
}
}
profiles
.into_iter()
.map(|(model, presentation)| super::PublicModelProfile::new(model, presentation))
.collect()
}
/// 已取得目录时以映射后的上游模型为准;目录不可用时由 Provider 在发送前判定。
#[must_use]
pub fn contains_public_model_for_provider(
&self,
@@ -470,17 +671,25 @@ impl RuntimeSnapshot {
if !self.providers.contains(provider) {
return false;
}
if self.model_mappings.contains_key(public_model.as_str()) {
return true;
}
let upstream_model = self.mapped_model(public_model.as_str());
match self.provider_models.get(provider) {
Some(models) if !models.is_empty() => models
.keys()
.any(|model| model.as_str() == public_model.as_str()),
Some(_) | None => true,
Some(models) => models.keys().any(|model| model.as_str() == upstream_model),
None => !self.known_provider_catalogs.contains(provider),
}
}
#[must_use]
pub fn contains_public_model_for_scope(
&self,
public_model: &PublicModelId,
scope: &FrozenAccountScope,
) -> bool {
scope
.provider_kinds()
.iter()
.any(|provider| self.contains_public_model_for_provider(public_model, provider))
}
#[must_use]
pub fn mapped_model(&self, requested: &str) -> String {
let original = requested;
@@ -506,14 +715,26 @@ impl RuntimeSnapshot {
&self,
public_model: &PublicModelId,
operation: &Operation,
account_scope: Arc<FrozenAccountScope>,
context: &RoutingContext,
) -> Result<RoutingPlan, RoutingError> {
let requirements = operation.capability_requirements();
let mut candidates = Vec::new();
for provider in self.providers.iter() {
if context.required_provider.is_none() && account_scope.provider_kinds().is_empty() {
return Err(RoutingError::EmptyAccountScope);
}
let providers = context.required_provider.as_ref().map_or_else(
|| account_scope.provider_kinds().clone(),
|provider| BTreeSet::from([provider.clone()]),
);
for provider in &providers {
if !self.providers.contains(provider) {
continue;
}
if context
.provider_kind
.required_provider
.as_ref()
.is_some_and(|expected| expected != provider)
|| context.blocked_providers.contains(provider)
@@ -539,12 +760,14 @@ impl RuntimeSnapshot {
};
emulated
}
None if self.known_provider_catalogs.contains(provider) => continue,
None => BTreeSet::new(),
};
candidates.push(ProviderCandidate {
provider: provider.clone(),
upstream_model,
emulated_features,
account_scope: Arc::clone(&account_scope),
});
}
@@ -561,6 +784,7 @@ impl RuntimeSnapshot {
operation: operation.kind(),
max_attempts: NonZeroU32::new(super::MAX_REQUEST_ATTEMPTS)
.expect("constant request attempt limit is non-zero"),
account_scope,
candidates: Arc::from(candidates),
})
}
@@ -2,12 +2,14 @@ use gateway_core::engine::continuation::{
ContinuationBinding, NativeContinuationPin, PreviousResponseId,
};
use gateway_core::engine::credential::ProviderAccountId;
use gateway_core::policy::ClientApiKeyId;
use gateway_core::routing::ProviderKind;
fn pin() -> NativeContinuationPin {
NativeContinuationPin::new(
PreviousResponseId::new("response-private"),
PreviousResponseId::new("upstream-private"),
ClientApiKeyId::new("key_client").expect("valid client key"),
ProviderKind::new("openai").expect("valid provider"),
ProviderAccountId::new("acct_codex").expect("valid account"),
)
@@ -43,3 +45,8 @@ fn native_pin_should_reject_different_account() {
&ProviderAccountId::new("acct_other").expect("valid account"),
));
}
#[test]
fn native_pin_should_reject_different_client_api_key() {
assert!(!pin().matches_client(&ClientApiKeyId::new("key_other").expect("valid client key")));
}
@@ -43,8 +43,9 @@ use gateway_core::operation::{
};
use gateway_core::policy::ClientApiKeyId;
use gateway_core::routing::{
ConfigRevision, ModelCapabilities, ProviderKind, ProviderModel, PublicModelId, RoutingContext,
RoutingPlan, RuntimeSnapshot, UpstreamModelId,
AccountRoutingSnapshot, ClientRoutingScope, ConfigRevision, FrozenAccountScope,
ModelCapabilities, ProviderKind, ProviderModel, PublicModelId, RoutingContext, RoutingPlan,
RuntimeAccount, RuntimeAccountDirectory, RuntimeSnapshot, UpstreamModelId,
};
use serde_json::{Map, Value, json};
@@ -531,7 +532,32 @@ fn plan_with_policy(
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(
let directory = Arc::new(RuntimeAccountDirectory::new(
[
"acct_failed",
"acct_first",
"acct_image",
"acct_observation_mismatch",
"acct_observed",
"acct_one",
"acct_only",
"acct_other",
"acct_required",
"acct_second",
"acct_tool",
"acct_two",
"acct_wrong",
]
.into_iter()
.map(|id| {
(
ProviderAccountId::new(id).expect("account"),
RuntimeAccount::new(provider.clone(), BTreeSet::new()),
)
})
.collect(),
));
let snapshot = RuntimeSnapshot::new(
ConfigRevision::new(1).expect("config revision"),
account_selection_policy,
vec![provider.clone()],
@@ -543,15 +569,21 @@ fn plan_with_policy(
Vec::new(),
)
.expect("snapshot")
.plan(
&public_model,
operation,
&RoutingContext {
provider_kind: Some(provider),
..RoutingContext::default()
},
)
.expect("routing plan")
.with_account_directory(Arc::clone(&directory));
snapshot
.plan(
&public_model,
operation,
Arc::new(FrozenAccountScope::new(
directory,
ClientRoutingScope::all_accounts(),
)),
&RoutingContext {
required_provider: Some(provider),
..RoutingContext::default()
},
)
.expect("routing plan")
}
fn model_request(operation: &Operation, deadline: SystemTime) -> NewModelRequest {
@@ -561,6 +593,7 @@ fn model_request(operation: &Operation, deadline: SystemTime) -> NewModelRequest
client_api_key_id: Some(client_key.clone()),
client_api_key_ref: client_key,
config_revision: ConfigRevision::new(1).expect("config revision"),
routing: AccountRoutingSnapshot::all(),
protocol: "openai".to_owned(),
operation: operation.kind(),
endpoint: "responses".to_owned(),
@@ -1354,6 +1387,7 @@ fn authenticated_native_continuation_reaches_every_attempt_context() {
let continuation = NativeContinuationPin::new(
PreviousResponseId::new("previous-secret-id"),
PreviousResponseId::new("provider-native-id"),
ClientApiKeyId::new("key_client_1").expect("client key"),
ProviderKind::new("openai").expect("provider"),
ProviderAccountId::new("acct_one").expect("account"),
);
@@ -1411,6 +1445,7 @@ fn streaming_native_continuation_delivers_started_event_before_later_output() {
let continuation = NativeContinuationPin::new(
PreviousResponseId::new("previous-response"),
PreviousResponseId::new("upstream-response"),
ClientApiKeyId::new("key_client_1").expect("client key"),
ProviderKind::new("openai").expect("provider"),
ProviderAccountId::new("acct_one").expect("account"),
);
@@ -1487,6 +1522,7 @@ fn native_continuation_replays_owner_before_safely_switching_account() {
let continuation = NativeContinuationPin::new(
PreviousResponseId::new("previous-secret-id"),
PreviousResponseId::new("provider-native-id"),
ClientApiKeyId::new("key_client_1").expect("client key"),
ProviderKind::new("openai").expect("provider"),
ProviderAccountId::new("acct_one").expect("account"),
);
@@ -1552,6 +1588,7 @@ fn unavailable_native_continuation_replays_with_the_configured_account_policy()
let continuation = NativeContinuationPin::new(
PreviousResponseId::new("previous-response"),
PreviousResponseId::new("upstream-response"),
ClientApiKeyId::new("key_client_1").expect("client key"),
ProviderKind::new("openai").expect("provider"),
original.clone(),
);
@@ -2498,6 +2535,7 @@ fn native_continuation_explicit_429_is_not_retried() {
let continuation = NativeContinuationPin::new(
PreviousResponseId::new("previous-secret-id"),
PreviousResponseId::new("provider-native-id"),
ClientApiKeyId::new("key_client_1").expect("client key"),
ProviderKind::new("openai").expect("provider"),
ProviderAccountId::new("acct_first").expect("account"),
);
@@ -1,18 +1,21 @@
use std::collections::BTreeSet;
use std::collections::{BTreeMap, BTreeSet};
use std::num::NonZeroU32;
use std::sync::Arc;
use std::time::{Duration, SystemTime};
use serde_json::{Map, Value};
use gateway_core::engine::credential::{
AccountAttemptFeedback, AccountAvailability, AccountAvailabilityPolicy, AccountCandidate,
AccountAttemptFeedback, AccountCandidate, AccountConcurrencyLimit, AccountEligibilityPolicy,
AccountFeedbackStats, AccountQuotaSignals, AccountRuntimeSignals, AccountSchedulingBlocker,
AccountSelectionContext, AccountSelectionPolicy, AccountSelector, CredentialCasUpdate,
CredentialRevision, OpaqueProviderData, PlaintextCredential, PreferredAccountSelection,
ProviderAccount, ProviderAccountId, ProviderAccountIdentity, ProviderAccountUpdate,
RotationStrategy,
AccountSelectionContext, AccountSelectionPolicy, AccountSelector, AccountStatus, AccountWeight,
CredentialCasUpdate, CredentialRevision, CredentialState, OpaqueProviderData,
PlaintextCredential, PreferredAccountSelection, ProviderAccount, ProviderAccountId,
ProviderAccountIdentity, ProviderAccountUpdate, QuotaEvidence, QuotaState, RotationStrategy,
};
use gateway_core::routing::{
ClientRoutingScope, FrozenAccountScope, ProviderKind, RuntimeAccount, RuntimeAccountDirectory,
};
use gateway_core::routing::ProviderKind;
fn account(id: &str) -> ProviderAccount {
ProviderAccount::new(
@@ -24,7 +27,13 @@ fn account(id: &str) -> ProviderAccount {
CredentialRevision::new(1).expect("valid revision"),
Some(SystemTime::now() + Duration::from_secs(3600)),
)
.with_runtime_state(true, AccountAvailability::Ready)
.with_account_facts(
true,
CredentialState::Ready,
QuotaState::unknown(),
None,
None,
)
}
fn candidate(id: &str, in_flight: u32, remaining: Option<u64>) -> AccountCandidate {
@@ -35,13 +44,41 @@ 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,
rate_limited_until: None,
failure_rate_basis_points: None,
first_output_latency_ms: None,
},
}
}
fn weighted_candidate(id: &str, weight: u16, in_flight: u32) -> AccountCandidate {
let mut candidate = candidate(id, in_flight, None);
candidate.account = candidate.account.with_scheduling(
None,
AccountWeight::new(weight).expect("valid account weight"),
);
candidate
}
#[test]
fn account_scheduling_should_default_and_resolve_concurrency_override() {
let default = NonZeroU32::new(12).expect("default concurrency");
let account = account("acct_scheduling_default");
assert_eq!(account.concurrency_limit(), None);
assert_eq!(account.weight(), AccountWeight::DEFAULT);
assert_eq!(account.effective_concurrency(default), default);
let limit = AccountConcurrencyLimit::new(7).expect("concurrency override");
let weight = AccountWeight::new(80).expect("weight");
let account = account.with_scheduling(Some(limit), weight);
assert_eq!(account.concurrency_limit(), Some(limit));
assert_eq!(account.weight(), weight);
assert_eq!(account.effective_concurrency(default).get(), 7);
assert!(AccountConcurrencyLimit::new(0).is_none());
assert!(AccountWeight::new(0).is_none());
assert!(AccountWeight::new(101).is_none());
}
fn context(strategy: RotationStrategy) -> AccountSelectionContext {
AccountSelectionContext {
policy: AccountSelectionPolicy::new(
@@ -53,7 +90,8 @@ fn context(strategy: RotationStrategy) -> AccountSelectionContext {
excluded_accounts: BTreeSet::new(),
preferred_account: None,
round_robin_cursor: 0,
availability: AccountAvailabilityPolicy::Enforce,
eligibility: AccountEligibilityPolicy::Enforce,
account_scope: None,
}
}
@@ -101,43 +139,99 @@ fn provider_account_identity_debug_should_redact_upstream_ids() {
}
#[test]
fn availability_should_round_trip_all_database_values() {
fn account_fact_values_should_round_trip() {
assert_eq!(
AccountAvailability::parse(AccountAvailability::QuotaExhausted.as_str()),
Some(AccountAvailability::QuotaExhausted)
CredentialState::parse(CredentialState::Banned.as_str()),
Some(CredentialState::Banned)
);
assert_eq!(
QuotaEvidence::parse("provider_denied"),
Some(QuotaEvidence::ProviderDenied)
);
}
#[test]
fn diagnostic_selection_bypasses_all_local_account_availability() {
fn exhausted_quota_refresh_is_due_at_reset_or_provider_fallback() {
let now = SystemTime::now();
let observed_at = now - Duration::from_secs(60);
let before_reset = QuotaState::exhausted(
QuotaEvidence::ProviderDenied,
observed_at,
Some(now + Duration::from_secs(1)),
);
let after_reset = QuotaState::exhausted(
QuotaEvidence::ProviderDenied,
observed_at,
Some(now - Duration::from_secs(1)),
);
let without_reset = QuotaState::exhausted(QuotaEvidence::ProviderDenied, observed_at, None);
assert!(!before_reset.exhaustion_refresh_due(now, Duration::from_secs(30)));
assert!(after_reset.exhaustion_refresh_due(now, Duration::from_secs(30)));
assert!(without_reset.exhaustion_refresh_due(now, Duration::from_secs(30)));
assert!(!without_reset.exhaustion_refresh_due(now, Duration::from_secs(120)));
}
#[test]
fn inconclusive_quota_observation_cannot_erase_confirmed_access() {
let now = SystemTime::now();
let exhausted = QuotaState::exhausted(QuotaEvidence::ProviderDenied, now, None);
assert_eq!(
exhausted
.merge_observation(QuotaState::observed_unknown(now + Duration::from_secs(1)))
.access(),
gateway_core::engine::credential::QuotaAccessState::Exhausted
);
assert_eq!(
exhausted
.merge_observation(QuotaState::allowed(now + Duration::from_secs(1)))
.access(),
gateway_core::engine::credential::QuotaAccessState::Allowed
);
}
#[test]
fn diagnostic_selection_bypasses_all_local_account_eligibility() {
let observed_at = SystemTime::now();
let exhausted = AccountCandidate {
account: account("acct_exhausted")
.with_runtime_state(true, AccountAvailability::QuotaExhausted),
account: account("acct_exhausted").with_account_facts(
true,
CredentialState::Ready,
QuotaState::exhausted(QuotaEvidence::ProviderDenied, observed_at, None),
None,
None,
),
signals: AccountRuntimeSignals {
in_flight: 0,
last_started_at: None,
quota_reset_at: None,
quota_remaining_rank: None,
quota_limit_reached: false,
rate_limited_until: None,
failure_rate_basis_points: None,
first_output_latency_ms: None,
},
};
let disabled = AccountCandidate {
account: account("acct_disabled")
.with_runtime_state(false, AccountAvailability::QuotaExhausted),
account: account("acct_disabled").with_account_facts(
false,
CredentialState::Ready,
QuotaState::exhausted(QuotaEvidence::ProviderDenied, observed_at, None),
None,
None,
),
signals: AccountRuntimeSignals {
in_flight: 0,
last_started_at: None,
quota_reset_at: None,
quota_remaining_rank: None,
quota_limit_reached: false,
rate_limited_until: None,
failure_rate_basis_points: None,
first_output_latency_ms: None,
},
};
let mut context = context(RotationStrategy::Sticky);
context.availability = AccountAvailabilityPolicy::BypassForDiagnostic;
context.eligibility = AccountEligibilityPolicy::BypassForDiagnostic;
assert_eq!(
AccountSelector
@@ -154,10 +248,63 @@ 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);
fn selector_should_reject_a_preferred_account_outside_the_frozen_client_scope() {
let allowed = candidate("acct_allowed", 0, None);
let outside = candidate("acct_outside", 0, None);
let directory = Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([(
allowed.account.id().clone(),
RuntimeAccount::new(
ProviderKind::new("openai").expect("provider"),
BTreeSet::new(),
),
)])));
let mut context = context(RotationStrategy::Sticky);
context.preferred_account = Some(outside.account.id().clone());
context.account_scope = Some(Arc::new(FrozenAccountScope::new(
directory,
ClientRoutingScope::all_accounts(),
)));
assert!(!account.is_schedulable(SystemTime::now()));
let candidates = [outside, allowed];
let selection = AccountSelector
.select(&candidates, &context)
.expect("allowed account should remain eligible");
assert_eq!(selection.candidate().account.id().as_str(), "acct_allowed");
assert_eq!(
selection.preferred(),
PreferredAccountSelection::Blocked(AccountSchedulingBlocker::OutsideClientScope)
);
}
#[test]
fn disabled_account_should_not_be_schedulable() {
let account = account("acct_disabled").with_account_facts(
false,
CredentialState::Ready,
QuotaState::unknown(),
None,
None,
);
assert_eq!(
account.status_projection(SystemTime::now(), None).status,
AccountStatus::Disabled
);
}
#[test]
fn rate_limited_projection_carries_only_the_active_cooldown_deadline() {
let now = SystemTime::now();
let active_until = now + Duration::from_secs(60);
let account = account("acct_rate_limited");
let active = account.status_projection(now, Some(active_until));
assert_eq!(active.status, AccountStatus::RateLimited);
assert_eq!(active.rate_limited_until, Some(active_until));
let elapsed = account.status_projection(now, Some(now - Duration::from_secs(1)));
assert_eq!(elapsed.status, AccountStatus::Normal);
assert_eq!(elapsed.rate_limited_until, None);
}
#[test]
@@ -171,14 +318,20 @@ fn account_without_upstream_identity_should_stay_unknown_and_not_schedulable() {
CredentialRevision::new(1).expect("valid revision"),
Some(SystemTime::now() + Duration::from_secs(3600)),
)
.with_runtime_state(true, AccountAvailability::Ready);
.with_account_facts(
true,
CredentialState::Ready,
QuotaState::unknown(),
None,
None,
);
assert_eq!(
(
account.availability(),
account.is_schedulable(SystemTime::now())
account.credential_state(),
account.status_projection(SystemTime::now(), None).status
),
(AccountAvailability::Unknown, false)
(CredentialState::Unknown, AccountStatus::Error)
);
}
@@ -195,6 +348,70 @@ fn smart_selector_should_prefer_lower_inflight_count() {
assert_eq!(selected.candidate().account.id().as_str(), "acct_idle");
}
#[test]
fn selector_should_prioritize_the_highest_weight_for_every_rotation_strategy() {
let candidates = vec![
weighted_candidate("acct_low_weight", 1, 0),
weighted_candidate("acct_high_weight", 100, 2),
];
for strategy in [
RotationStrategy::Smart,
RotationStrategy::QuotaResetPriority,
RotationStrategy::RoundRobin,
RotationStrategy::Sticky,
] {
let selected = AccountSelector
.select(&candidates, &context(strategy))
.expect("highest weight account should be eligible");
assert_eq!(
selected.candidate().account.id().as_str(),
"acct_high_weight",
"strategy {strategy:?} must select inside the highest weight tier"
);
}
}
#[test]
fn selector_should_not_let_soft_affinity_bypass_a_higher_weight_account() {
let preferred = weighted_candidate("acct_preferred", 1, 0);
let preferred_id = preferred.account.id().clone();
let higher = weighted_candidate("acct_higher", 10, 0);
let mut selection = context(RotationStrategy::Smart);
selection.preferred_account = Some(preferred_id);
let candidates = [preferred, higher];
let selected = AccountSelector
.select(&candidates, &selection)
.expect("higher weight account should be eligible");
assert_eq!(selected.candidate().account.id().as_str(), "acct_higher");
assert_eq!(
selected.preferred(),
PreferredAccountSelection::Blocked(AccountSchedulingBlocker::LowerWeight)
);
}
#[test]
fn selector_should_fall_back_when_a_higher_weight_account_reaches_its_override() {
let mut saturated = weighted_candidate("acct_high_weight", 100, 1);
saturated.account = saturated.account.with_scheduling(
Some(AccountConcurrencyLimit::new(1).expect("concurrency override")),
AccountWeight::new(100).expect("weight"),
);
let fallback = weighted_candidate("acct_low_weight", 1, 0);
let candidates = [saturated, fallback];
let selected = AccountSelector
.select(&candidates, &context(RotationStrategy::Smart))
.expect("lower weight account should be used after saturation");
assert_eq!(
selected.candidate().account.id().as_str(),
"acct_low_weight"
);
}
#[test]
fn smart_selector_should_balance_signals_instead_of_using_lexicographic_load() {
let mut healthy = candidate("acct_healthy", 1, Some(100));
@@ -308,15 +525,11 @@ 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,
rate_limited_until: None,
failure_rate_basis_points: None,
first_output_latency_ms: None,
}
.with_provider_quota(Some(AccountQuotaSignals::new(
Some(reset_at),
Some(75),
false,
)));
.with_provider_quota(Some(AccountQuotaSignals::new(Some(reset_at), Some(75))));
assert_eq!(
(
@@ -633,21 +846,31 @@ fn credential_cas_should_reject_refresh_schedule_without_refresh_token() {
}
#[test]
fn enforced_availability_excludes_terminal_states_despite_quota_signals() {
// availability 终态(Banned/Expired/Invalid)始终排除,
// 不受 quota 信号或 skip_exhausted 放宽影响。
fn enforced_eligibility_excludes_credential_errors_despite_quota_signals() {
let mut banned = candidate("acct_banned", 0, Some(100));
banned.account = banned
.account
.with_runtime_state(true, AccountAvailability::Banned);
banned.account = banned.account.with_account_facts(
true,
CredentialState::Banned,
QuotaState::unknown(),
None,
None,
);
let mut expired = candidate("acct_expired", 0, Some(100));
expired.account = expired
.account
.with_runtime_state(true, AccountAvailability::Expired);
expired.account = expired.account.with_account_facts(
true,
CredentialState::Expired,
QuotaState::unknown(),
None,
None,
);
let mut invalid = candidate("acct_invalid", 0, Some(100));
invalid.account = invalid
.account
.with_runtime_state(true, AccountAvailability::Invalid);
invalid.account = invalid.account.with_account_facts(
true,
CredentialState::Invalid,
QuotaState::unknown(),
None,
None,
);
let healthy = candidate("acct_healthy", 0, Some(100));
let candidates = vec![banned, expired, invalid, healthy];
@@ -659,14 +882,15 @@ fn enforced_availability_excludes_terminal_states_despite_quota_signals() {
}
#[test]
fn enforced_availability_allows_quota_exhausted_with_skip_exhausted_projection() {
// QuotaExhausted 是可恢复额度耗尽,skip_exhausted=false 的
// 放宽投影(provider 层)允许其参与调度;core 的 Enforce 策略本身
// 只排除终态。这里验证 Enforce 下 QuotaExhausted 仍被排除(严格模式)。
fn enforced_eligibility_excludes_authoritative_quota_exhaustion() {
let mut exhausted = candidate("acct_exhausted", 0, Some(100));
exhausted.account = exhausted
.account
.with_runtime_state(true, AccountAvailability::QuotaExhausted);
exhausted.account = exhausted.account.with_account_facts(
true,
CredentialState::Ready,
QuotaState::exhausted(QuotaEvidence::ProviderDenied, SystemTime::now(), None),
None,
None,
);
let healthy = candidate("acct_healthy", 0, Some(100));
let candidates = vec![exhausted, healthy];
@@ -676,3 +900,31 @@ fn enforced_availability_allows_quota_exhausted_with_skip_exhausted_projection()
assert_eq!(selected.candidate().account.id().as_str(), "acct_healthy");
}
#[test]
fn elapsed_quota_reset_does_not_fabricate_recovery() {
let now = SystemTime::now();
let mut exhausted = candidate("acct_exhausted", 0, Some(100));
exhausted.account = exhausted.account.with_account_facts(
true,
CredentialState::Ready,
QuotaState::exhausted(
QuotaEvidence::ProviderDenied,
now - Duration::from_secs(2),
Some(now - Duration::from_secs(1)),
),
None,
None,
);
let healthy = candidate("acct_healthy", 0, Some(100));
assert_eq!(
exhausted.account.status_projection(now, None).status,
AccountStatus::QuotaExhausted
);
let candidates = [exhausted, healthy];
let selected = AccountSelector
.select(&candidates, &context(RotationStrategy::Smart))
.expect("healthy candidate available");
assert_eq!(selected.candidate().account.id().as_str(), "acct_healthy");
}
@@ -1,5 +1,5 @@
use std::{
collections::BTreeSet,
collections::{BTreeMap, BTreeSet},
sync::{
Arc, Mutex,
atomic::{AtomicBool, Ordering},
@@ -42,7 +42,8 @@ use gateway_core::operation::{GenerateRequest, Operation, OperationKind, Protoco
use gateway_core::policy::{ClientApiKeyId, ClientPolicy, PlaintextClientApiKey, RateLimits};
use gateway_core::routing::snapshot::RuntimeSnapshotHandle;
use gateway_core::routing::{
ConfigRevision, ModelCapabilities, ProviderKind, ProviderModel, PublicModelId, RuntimeSnapshot,
ClientRoutingScope, ConfigRevision, FrozenAccountScope, ModelCapabilities, ProviderKind,
ProviderModel, PublicModelId, RuntimeAccount, RuntimeAccountDirectory, RuntimeSnapshot,
UpstreamModelId,
};
use serde_json::json;
@@ -283,7 +284,7 @@ fn slow_circuit_store_should_time_out_and_fail_open_during_request_start() {
}
#[test]
fn stale_nonempty_catalog_should_not_block_a_new_model_during_request_start() {
fn known_catalog_should_reject_a_model_that_the_provider_did_not_publish() {
let service = DefaultExecutionService::new(
RuntimeSnapshotHandle::new(start_snapshot()),
Arc::new(TrackingExecutionStore::default()),
@@ -298,7 +299,7 @@ fn stale_nonempty_catalog_should_not_block_a_new_model_during_request_start() {
.expect("authenticated client");
let model = "gpt-future-not-in-catalog";
let started = block_on(service.start(StartExecution {
let result = block_on(service.start(StartExecution {
client,
public_model: PublicModelId::from_client_wire(model).expect("client model"),
operation: start_operation_for_model(model),
@@ -311,10 +312,80 @@ fn stale_nonempty_catalog_should_not_block_a_new_model_during_request_start() {
user_agent: None,
previous_response_id: None,
},
}))
.expect("stale model catalog must not reject the request");
}));
let Err(error) = result else {
panic!("a known provider catalog must be authoritative for model availability");
};
assert!(!started.session.is_finalized());
assert_eq!(error.kind(), GatewayErrorKind::NoAvailableProvider);
}
#[test]
fn continuation_owned_by_another_client_api_key_should_fail_closed() {
let service = DefaultExecutionService::new(
RuntimeSnapshotHandle::new(start_snapshot()),
Arc::new(TrackingExecutionStore::default()),
ProviderRegistry::default(),
Arc::new(UnusedAdmissions),
Arc::new(UnusedCircuits),
Arc::new(RejectedContinuation::OwnershipMismatch),
Arc::new(RecordingClientApiKeyUsage::default()),
);
let client = service
.authenticate("sk_start_test")
.expect("authenticated client");
let result = block_on(service.start(StartExecution {
client,
public_model: PublicModelId::new("gpt-start").expect("public model"),
operation: start_operation(),
metadata: execution_metadata_with_continuation(),
}));
let Err(error) = result else {
panic!("cross-client continuation must be rejected");
};
assert_eq!(error.kind(), GatewayErrorKind::PolicyDenied);
}
#[test]
fn invalid_continuation_record_should_not_be_forwarded_as_an_external_handle() {
let service = DefaultExecutionService::new(
RuntimeSnapshotHandle::new(start_snapshot()),
Arc::new(TrackingExecutionStore::default()),
ProviderRegistry::default(),
Arc::new(UnusedAdmissions),
Arc::new(UnusedCircuits),
Arc::new(RejectedContinuation::InvalidData),
Arc::new(RecordingClientApiKeyUsage::default()),
);
let client = service
.authenticate("sk_start_test")
.expect("authenticated client");
let result = block_on(service.start(StartExecution {
client,
public_model: PublicModelId::new("gpt-start").expect("public model"),
operation: start_operation(),
metadata: execution_metadata_with_continuation(),
}));
let Err(error) = result else {
panic!("invalid continuation state must be rejected");
};
assert_eq!(error.kind(), GatewayErrorKind::Internal);
}
fn execution_metadata_with_continuation() -> ExecutionRequestMetadata {
ExecutionRequestMetadata {
protocol: "openai".to_owned(),
endpoint: "/v1/responses".to_owned(),
transport: ClientTransport::HttpJson,
stream: false,
client_ip: None,
user_agent: None,
previous_response_id: Some(PreviousResponseId::new("response-private")),
}
}
#[derive(Default)]
@@ -526,6 +597,7 @@ struct UnusedContinuation;
impl NativeContinuationPort for UnusedContinuation {
fn resolve<'a>(
&'a self,
_: &'a ClientApiKeyId,
_: &'a PreviousResponseId,
) -> BoxFuture<'a, Result<Option<NativeContinuationPin>, NativeContinuationStoreError>> {
Box::pin(async { Ok(None) })
@@ -539,6 +611,43 @@ impl NativeContinuationPort for UnusedContinuation {
}
}
enum RejectedContinuation {
OwnershipMismatch,
InvalidData,
}
impl NativeContinuationPort for RejectedContinuation {
fn resolve<'a>(
&'a self,
_: &'a ClientApiKeyId,
_: &'a PreviousResponseId,
) -> BoxFuture<'a, Result<Option<NativeContinuationPin>, NativeContinuationStoreError>> {
Box::pin(async move {
Err(match self {
Self::OwnershipMismatch => NativeContinuationStoreError::ownership_mismatch(),
Self::InvalidData => NativeContinuationStoreError::invalid_data("invalid record"),
})
})
}
fn record<'a>(
&'a self,
_: NativeContinuationPin,
) -> BoxFuture<'a, Result<(), NativeContinuationStoreError>> {
Box::pin(async { Ok(()) })
}
}
fn account_scope(provider: &ProviderKind, account_id: &str) -> Arc<FrozenAccountScope> {
Arc::new(FrozenAccountScope::new(
Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([(
ProviderAccountId::new(account_id).expect("account ID"),
RuntimeAccount::new(provider.clone(), BTreeSet::new()),
)]))),
ClientRoutingScope::all_accounts(),
))
}
fn probe_snapshot() -> RuntimeSnapshot {
let provider = ProviderKind::new("openai").expect("provider kind");
let capabilities =
@@ -575,7 +684,7 @@ fn client_snapshot() -> RuntimeSnapshot {
vec![ClientPolicy::new(
ClientApiKeyId::new("key_usage_test").expect("client API key ID"),
PlaintextClientApiKey::new("sk_usage_test").expect("plaintext client API key"),
provider,
account_scope(&provider, "acct_usage"),
true,
RateLimits::unlimited(),
)],
@@ -603,7 +712,7 @@ fn start_snapshot() -> RuntimeSnapshot {
vec![ClientPolicy::new(
ClientApiKeyId::new("key_start_test").expect("client API key ID"),
PlaintextClientApiKey::new("sk_start_test").expect("plaintext client API key"),
provider,
account_scope(&provider, "acct_start"),
true,
RateLimits::unlimited(),
)],
@@ -1,16 +1,25 @@
use std::sync::Arc;
use gateway_core::policy::{ClientApiKeyId, ClientPolicy, PlaintextClientApiKey, RateLimits};
use gateway_core::routing::ProviderKind;
use gateway_core::routing::{ClientRoutingScope, FrozenAccountScope, RuntimeAccountDirectory};
fn plaintext(value: &str) -> PlaintextClientApiKey {
PlaintextClientApiKey::new(value).expect("valid plaintext client key")
}
fn account_scope() -> Arc<FrozenAccountScope> {
Arc::new(FrozenAccountScope::new(
Arc::new(RuntimeAccountDirectory::default()),
ClientRoutingScope::all_accounts(),
))
}
#[test]
fn disabled_client_key_should_be_denied() {
let policy = ClientPolicy::new(
ClientApiKeyId::new("key_disabled").expect("valid key ID"),
plaintext("sk_disabled_secret"),
ProviderKind::new("openai").expect("valid provider"),
account_scope(),
false,
RateLimits::unlimited(),
);
@@ -23,7 +32,7 @@ fn enabled_client_key_should_be_authorized() {
let policy = ClientPolicy::new(
ClientApiKeyId::new("key_enabled").expect("valid key ID"),
plaintext("sk_enabled_secret"),
ProviderKind::new("openai").expect("valid provider"),
account_scope(),
true,
RateLimits::unlimited(),
);
@@ -92,7 +92,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,
rate_limited_until: None,
failure_rate_basis_points: Some(125),
first_output_latency_ms: Some(250),
},
+168 -16
View File
@@ -1,15 +1,20 @@
use std::collections::{BTreeMap, BTreeSet};
use std::num::NonZeroU32;
use std::sync::Arc;
use std::time::Duration;
use gateway_core::engine::credential::{AccountSelectionPolicy, RotationStrategy};
use gateway_core::engine::credential::{
AccountSelectionPolicy, ProviderAccountId, RotationStrategy,
};
use gateway_core::operation::{
CapabilityRequirements, Feature, GenerateRequest, Operation, OperationKind, ProtocolPayload,
};
use gateway_core::policy::{ClientApiKeyId, ClientPolicy, PlaintextClientApiKey, RateLimits};
use gateway_core::routing::{
ConfigRevision, ModelCapabilities, ProviderKind, ProviderModel, PublicModelId, RoutingContext,
RuntimeSnapshot, UpstreamModelId,
AccountGroupId, AccountRoutingScopeKind, ClientRoutingScope, ConfigRevision,
FrozenAccountScope, ModelCapabilities, ProviderKind, ProviderModel, PublicModelId,
RoutingContext, RoutingGroupSnapshot, RuntimeAccount, RuntimeAccountDirectory, RuntimeSnapshot,
UpstreamModelId,
};
mod snapshot;
@@ -45,11 +50,34 @@ fn operation() -> Operation {
))
}
fn account_directory() -> Arc<RuntimeAccountDirectory> {
Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([
(
ProviderAccountId::new("acct_openai").expect("account"),
RuntimeAccount::new(
ProviderKind::new("openai").expect("provider"),
BTreeSet::new(),
),
),
(
ProviderAccountId::new("acct_xai").expect("account"),
RuntimeAccount::new(ProviderKind::new("xai").expect("provider"), BTreeSet::new()),
),
])))
}
fn account_scope() -> Arc<FrozenAccountScope> {
Arc::new(FrozenAccountScope::new(
account_directory(),
ClientRoutingScope::all_accounts(),
))
}
fn client_policy(id: &str, plaintext: &str, enabled: bool) -> ClientPolicy {
ClientPolicy::new(
ClientApiKeyId::new(id).expect("client key ID"),
PlaintextClientApiKey::new(plaintext).expect("plaintext client key"),
ProviderKind::new("openai").expect("valid provider"),
account_scope(),
enabled,
RateLimits::unlimited(),
)
@@ -70,6 +98,7 @@ fn snapshot() -> RuntimeSnapshot {
Vec::new(),
)
.expect("snapshot")
.with_account_directory(account_directory())
.with_model_mappings(BTreeMap::from([
("gpt-5.4".to_owned(), "gpt-5.5".to_owned()),
("grok-latest".to_owned(), "grok-4.5".to_owned()),
@@ -81,6 +110,92 @@ fn config_revision_should_reject_zero() {
assert!(ConfigRevision::new(0).is_err());
}
#[test]
fn account_group_id_should_match_the_database_contract_exactly() {
let valid = "grp_0123456789abcdef0123456789abcdef";
assert_eq!(
AccountGroupId::new(valid).expect("valid group ID").as_str(),
valid
);
for invalid in [
"0123456789abcdef0123456789abcdef",
"grp_0123456789abcdef0123456789abcde",
"grp_0123456789abcdef0123456789abcdef0",
"grp_0123456789abcdef0123456789abcdeF",
"grp_0123456789abcdef0123456789abcdeg",
] {
assert!(
AccountGroupId::new(invalid).is_err(),
"invalid group ID was accepted: {invalid}"
);
}
}
#[test]
fn restricted_scope_should_use_enabled_group_union_and_preserve_all_bindings_for_history() {
let enabled =
AccountGroupId::new("grp_11111111111111111111111111111111").expect("enabled group");
let disabled =
AccountGroupId::new("grp_22222222222222222222222222222222").expect("disabled group");
let provider = ProviderKind::new("openai").expect("provider");
let allowed_account = ProviderAccountId::new("acct_allowed").expect("account");
let disabled_only_account = ProviderAccountId::new("acct_disabled").expect("account");
let directory = Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([
(
allowed_account.clone(),
RuntimeAccount::new(provider.clone(), BTreeSet::from([enabled.clone()])),
),
(
disabled_only_account.clone(),
RuntimeAccount::new(provider.clone(), BTreeSet::from([disabled.clone()])),
),
])));
let scope = FrozenAccountScope::new(
directory,
ClientRoutingScope::restricted(
vec![
RoutingGroupSnapshot::new(enabled.clone(), "Enabled".to_owned()),
RoutingGroupSnapshot::new(disabled, "Disabled".to_owned()),
],
BTreeSet::from([enabled]),
BTreeSet::from([provider]),
)
.expect("restricted scope"),
);
assert!(scope.allows(&allowed_account));
assert!(!scope.allows(&disabled_only_account));
let history = scope.routing_snapshot();
assert_eq!(history.kind(), AccountRoutingScopeKind::Groups);
assert_eq!(history.groups_snapshot().len(), 2);
assert_eq!(history.groups_snapshot()[1].name(), "Disabled");
}
#[test]
fn restricted_scope_with_no_enabled_group_should_fail_closed() {
let group = AccountGroupId::new("grp_33333333333333333333333333333333").expect("group");
let account = ProviderAccountId::new("acct_group_member").expect("account");
let directory = Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([(
account.clone(),
RuntimeAccount::new(
ProviderKind::new("openai").expect("provider"),
BTreeSet::from([group.clone()]),
),
)])));
let scope = FrozenAccountScope::new(
directory,
ClientRoutingScope::restricted(
vec![RoutingGroupSnapshot::new(group, "Disabled".to_owned())],
BTreeSet::new(),
BTreeSet::new(),
)
.expect("restricted scope"),
);
assert!(scope.provider_kinds().is_empty());
assert!(!scope.allows(&account));
}
#[test]
fn snapshot_should_publish_only_enabled_plaintext_client_policies() {
let snapshot = RuntimeSnapshot::new(
@@ -98,7 +213,12 @@ fn snapshot_should_publish_only_enabled_plaintext_client_policies() {
let policies = snapshot.client_policies().collect::<Vec<_>>();
assert_eq!(policies.len(), 1);
assert_eq!(policies[0].key_id().as_str(), "key_enabled");
assert_eq!(policies[0].provider_kind().as_str(), "openai");
assert!(
policies[0]
.account_scope()
.provider_kinds()
.contains(&ProviderKind::new("openai").expect("provider"))
);
}
#[test]
@@ -137,8 +257,9 @@ fn selected_provider_should_use_global_model_mapping() {
.plan(
&PublicModelId::new("gpt-5.4").expect("model"),
&operation(),
snapshot.all_account_scope(),
&RoutingContext {
provider_kind: Some(ProviderKind::new("openai").expect("provider")),
required_provider: Some(ProviderKind::new("openai").expect("provider")),
..RoutingContext::default()
},
)
@@ -170,26 +291,37 @@ fn cyclic_model_mapping_should_fall_back_to_the_original_name() {
}
#[test]
fn global_mapping_should_apply_to_every_provider() {
fn known_provider_catalog_missing_mapped_model_should_be_filtered() {
let snapshot = snapshot();
let plan = snapshot
let error = snapshot
.plan(
&PublicModelId::new("gpt-5.4").expect("model"),
&operation(),
snapshot.all_account_scope(),
&RoutingContext {
provider_kind: Some(ProviderKind::new("xai").expect("provider")),
required_provider: Some(ProviderKind::new("xai").expect("provider")),
..RoutingContext::default()
},
)
.expect("transparent plan");
.expect_err("known xAI catalog does not contain the mapped OpenAI model");
assert_eq!(plan.candidates()[0].provider().as_str(), "xai");
assert_eq!(plan.candidates()[0].upstream_model().as_str(), "gpt-5.5");
assert!(matches!(
error,
gateway_core::error::RoutingError::NoCapableProvider { .. }
));
}
#[test]
fn unmapped_model_should_pass_through_unchanged() {
let snapshot = snapshot();
let snapshot = RuntimeSnapshot::new(
ConfigRevision::new(1).expect("revision"),
scheduling(),
vec![ProviderKind::new("openai").expect("provider")],
Vec::new(),
Vec::new(),
)
.expect("snapshot")
.with_account_directory(account_directory());
for requested in [
"future-openai-model".to_owned(),
format!("future-{}", "x".repeat(512)),
@@ -200,8 +332,9 @@ fn unmapped_model_should_pass_through_unchanged() {
.plan(
&PublicModelId::from_client_wire(requested.clone()).expect("client model"),
&operation(),
snapshot.all_account_scope(),
&RoutingContext {
provider_kind: Some(ProviderKind::new("openai").expect("provider")),
required_provider: Some(ProviderKind::new("openai").expect("provider")),
..RoutingContext::default()
},
)
@@ -229,8 +362,9 @@ fn blocked_provider_should_be_filtered() {
.plan(
&PublicModelId::new("gpt-5.5").expect("model"),
&operation(),
snapshot.all_account_scope(),
&RoutingContext {
provider_kind: Some(ProviderKind::new("openai").expect("provider")),
required_provider: Some(ProviderKind::new("openai").expect("provider")),
blocked_providers: BTreeSet::from([ProviderKind::new("openai").expect("provider")]),
},
)
@@ -262,8 +396,9 @@ fn known_unsupported_operation_should_not_be_bypassed() {
.plan(
&PublicModelId::new("gpt-known-unsupported").expect("model"),
&operation(),
snapshot.all_account_scope(),
&RoutingContext {
provider_kind: Some(ProviderKind::new("openai").expect("provider")),
required_provider: Some(ProviderKind::new("openai").expect("provider")),
..RoutingContext::default()
},
)
@@ -316,3 +451,20 @@ fn unknown_model_should_be_hidden_when_provider_catalog_is_available() {
&ProviderKind::new("openai").expect("provider"),
));
}
#[test]
fn alias_should_only_be_available_from_a_provider_with_its_mapped_model() {
let snapshot = snapshot();
let alias = PublicModelId::new("grok-latest").expect("model");
assert!(!snapshot.contains_public_model_for_provider(
&alias,
&ProviderKind::new("openai").expect("provider"),
));
assert!(
snapshot.contains_public_model_for_provider(
&alias,
&ProviderKind::new("xai").expect("provider"),
)
);
}
@@ -19,7 +19,8 @@ use gateway_core::operation::OperationKind;
use gateway_core::policy::{ClientApiKeyId, PlaintextClientApiKey, RateLimits};
use gateway_core::routing::snapshot::{
RuntimeSnapshotCompileError, RuntimeSnapshotCompiler, RuntimeSnapshotHandle,
RuntimeSnapshotPublisher, SnapshotClientPolicyFacts, SnapshotControl, SnapshotFacts,
RuntimeSnapshotPublisher, SnapshotAccountGroupFacts, SnapshotAccountGroupMemberFacts,
SnapshotClientPolicyFacts, SnapshotControl, SnapshotFacts, SnapshotProviderAccountFacts,
SnapshotRevisionStream, SnapshotSettingsFacts, SnapshotStoreError, SnapshotStorePort,
SnapshotSubscriptionError, SnapshotSubscriptionPort, runtime_revision_needs_refresh,
};
@@ -348,9 +349,12 @@ fn facts(config_revision: u64, observed_current_revision: u64) -> SnapshotFacts
vec![SnapshotClientPolicyFacts::new(
ClientApiKeyId::new("key_one").expect("key ID"),
PlaintextClientApiKey::new("sk_test").expect("plaintext key"),
"alpha",
Vec::new(),
RateLimits::unlimited(),
)],
Vec::<SnapshotAccountGroupFacts>::new(),
Vec::<SnapshotProviderAccountFacts>::new(),
Vec::<SnapshotAccountGroupMemberFacts>::new(),
)
}
+74 -28
View File
@@ -13,8 +13,9 @@ use tracing_subscriber::{Layer as _, layer::SubscriberExt as _, util::Subscriber
use crate::config::LoggingConfig;
const LOG_FILE_PREFIX: &str = "codex-proxy-rs";
/// OpenAI OAuth 恢复记录含原始 AT/RT,文件日志开启时不能被常规级别过滤掉。
const OPENAI_OAUTH_RECOVERY_LOG_TARGET: &str = "openai_oauth_recovery";
const OAUTH_RECOVERY_LOG_FILE_PREFIX: &str = "codex-proxy-rs-oath";
/// OAuth 恢复记录含原始 AT/RT,文件日志开启时必须落到独立文件。
const OAUTH_RECOVERY_LOG_TARGET: &str = "oauth_recovery";
/// non-blocking 日志 writer 的进程级守卫。
pub struct LogGuard {
@@ -36,7 +37,8 @@ pub enum LogError {
/// 按自然日、单文件大小、保留天数和文件总数初始化结构化日志。
pub fn initialize_logging(config: &LoggingConfig) -> Result<LogGuard, LogError> {
let directive = env::var("RUST_LOG").unwrap_or_else(|_| config.level.clone());
let file_filter = file_filter(&directive)?;
let file_filter = application_file_filter(&directive)?;
let recovery_file_filter = oauth_recovery_file_filter();
let stdout_filter =
EnvFilter::try_new(stdout_filter_directive(&directive, config.file.enabled))
.map_err(|_| LogError::InvalidFilter)?;
@@ -49,25 +51,37 @@ pub fn initialize_logging(config: &LoggingConfig) -> Result<LogGuard, LogError>
guards.push(guard);
writer
});
let file_writer = if config.file.enabled {
let (file_writer, recovery_file_writer) = if config.file.enabled {
let maximum_bytes = config
.file
.max_file_size_mb
.checked_mul(1024 * 1024)
.ok_or(LogError::SizeOverflow)?;
let writer = RotatingLogWriter::open(
let file_writer = RotatingLogWriter::open(
config.file.directory.clone(),
LOG_FILE_PREFIX,
maximum_bytes,
config.file.retention_days,
config.file.max_files,
)?;
let (writer, guard) = NonBlockingBuilder::default()
let recovery_file_writer = RotatingLogWriter::open(
config.file.directory.clone(),
OAUTH_RECOVERY_LOG_FILE_PREFIX,
maximum_bytes,
config.file.retention_days,
config.file.max_files,
)?;
let (file_writer, file_guard) = NonBlockingBuilder::default()
.thread_name("gateway-log-file")
.finish(writer);
guards.push(guard);
Some(writer)
.finish(file_writer);
let (recovery_file_writer, recovery_file_guard) = NonBlockingBuilder::default()
.thread_name("gateway-log-oauth-recovery-file")
.finish(recovery_file_writer);
guards.push(file_guard);
guards.push(recovery_file_guard);
(Some(file_writer), Some(recovery_file_writer))
} else {
None
(None, None)
};
let stdout_layer = stdout_writer.map(|writer| {
@@ -91,23 +105,43 @@ pub fn initialize_logging(config: &LoggingConfig) -> Result<LogGuard, LogError>
.with_span_list(true)
.with_filter(file_filter)
});
let recovery_file_layer = recovery_file_writer.map(|writer| {
tracing_subscriber::fmt::layer()
.json()
.with_writer(writer)
.with_target(true)
.with_file(true)
.with_line_number(true)
.with_thread_ids(true)
.with_thread_names(true)
.with_current_span(true)
.with_span_list(true)
.with_filter(recovery_file_filter)
});
tracing_subscriber::registry()
.with(stdout_layer)
.with(file_layer)
.with(recovery_file_layer)
.try_init()
.map_err(|_| LogError::AlreadyInitialized)?;
Ok(LogGuard { _writers: guards })
}
fn file_filter(directive: &str) -> Result<EnvFilter, LogError> {
let recovery_directive = format!("{OPENAI_OAUTH_RECOVERY_LOG_TARGET}=info")
fn application_file_filter(directive: &str) -> Result<EnvFilter, LogError> {
let recovery_directive = format!("{OAUTH_RECOVERY_LOG_TARGET}=off")
.parse()
.expect("static OpenAI OAuth recovery log directive is valid");
.expect("static OAuth recovery log directive is valid");
EnvFilter::try_new(directive)
.map(|filter| filter.add_directive(recovery_directive))
.map_err(|_| LogError::InvalidFilter)
}
fn oauth_recovery_file_filter() -> EnvFilter {
format!("off,{OAUTH_RECOVERY_LOG_TARGET}=info")
.parse()
.expect("static OAuth recovery log filter is valid")
}
fn stdout_filter_directive(directive: &str, persistent_log_enabled: bool) -> String {
if !persistent_log_enabled {
directive.to_owned()
@@ -118,6 +152,7 @@ fn stdout_filter_directive(directive: &str, persistent_log_enabled: bool) -> Str
struct RotatingLogWriter {
directory: PathBuf,
file_prefix: &'static str,
maximum_bytes: u64,
retention_days: usize,
maximum_files: usize,
@@ -130,17 +165,20 @@ struct RotatingLogWriter {
impl RotatingLogWriter {
fn open(
directory: PathBuf,
file_prefix: &'static str,
maximum_bytes: u64,
retention_days: usize,
maximum_files: usize,
) -> io::Result<Self> {
fs::create_dir_all(&directory)?;
let date = Utc::now().date_naive();
cleanup_log_files(&directory, date, retention_days, maximum_files)?;
let (segment, bytes_written) = writable_log_segment(&directory, date, maximum_bytes)?;
let file = open_log_segment(&directory, date, segment)?;
cleanup_log_files(&directory, file_prefix, date, retention_days, maximum_files)?;
let (segment, bytes_written) =
writable_log_segment(&directory, file_prefix, date, maximum_bytes)?;
let file = open_log_segment(&directory, file_prefix, date, segment)?;
Ok(Self {
directory,
file_prefix,
maximum_bytes,
retention_days,
maximum_files,
@@ -168,10 +206,11 @@ impl RotatingLogWriter {
} else {
self.segment = self.segment.saturating_add(1);
}
self.file = open_log_segment(&self.directory, self.date, self.segment)?;
self.file = open_log_segment(&self.directory, self.file_prefix, self.date, self.segment)?;
self.bytes_written = self.file.metadata()?.len();
cleanup_log_files(
&self.directory,
self.file_prefix,
self.date,
self.retention_days,
self.maximum_files,
@@ -203,10 +242,11 @@ struct ManagedLogFile {
fn writable_log_segment(
directory: &Path,
file_prefix: &str,
date: chrono::NaiveDate,
maximum_bytes: u64,
) -> io::Result<(usize, u64)> {
let latest = managed_log_files(directory)?
let latest = managed_log_files(directory, file_prefix)?
.into_iter()
.filter(|entry| entry.date == date)
.max_by_key(|entry| entry.segment);
@@ -221,26 +261,31 @@ fn writable_log_segment(
}
}
fn open_log_segment(directory: &Path, date: chrono::NaiveDate, segment: usize) -> io::Result<File> {
fn open_log_segment(
directory: &Path,
file_prefix: &str,
date: chrono::NaiveDate,
segment: usize,
) -> io::Result<File> {
OpenOptions::new()
.create(true)
.append(true)
.open(directory.join(log_file_name(date, segment)))
.open(directory.join(log_file_name(file_prefix, date, segment)))
}
fn log_file_name(date: chrono::NaiveDate, segment: usize) -> String {
fn log_file_name(file_prefix: &str, date: chrono::NaiveDate, segment: usize) -> String {
if segment == 0 {
format!("{LOG_FILE_PREFIX}.{date}.log")
format!("{file_prefix}.{date}.log")
} else {
format!("{LOG_FILE_PREFIX}.{date}.{segment}.log")
format!("{file_prefix}.{date}.{segment}.log")
}
}
fn managed_log_files(directory: &Path) -> io::Result<Vec<ManagedLogFile>> {
fn managed_log_files(directory: &Path, file_prefix: &str) -> io::Result<Vec<ManagedLogFile>> {
Ok(fs::read_dir(directory)?
.filter_map(Result::ok)
.filter_map(|entry| {
parse_log_file_name(entry.file_name().to_string_lossy().as_ref()).map(
parse_log_file_name(entry.file_name().to_string_lossy().as_ref(), file_prefix).map(
|(date, segment)| ManagedLogFile {
date,
segment,
@@ -251,9 +296,9 @@ fn managed_log_files(directory: &Path) -> io::Result<Vec<ManagedLogFile>> {
.collect())
}
fn parse_log_file_name(name: &str) -> Option<(chrono::NaiveDate, usize)> {
fn parse_log_file_name(name: &str, file_prefix: &str) -> Option<(chrono::NaiveDate, usize)> {
let body = name
.strip_prefix(&format!("{LOG_FILE_PREFIX}."))?
.strip_prefix(&format!("{file_prefix}."))?
.strip_suffix(".log")?;
let (date, segment) = body.split_once('.').map_or((body, 0), |(date, segment)| {
(date, segment.parse().ok().unwrap_or(usize::MAX))
@@ -268,6 +313,7 @@ fn parse_log_file_name(name: &str) -> Option<(chrono::NaiveDate, usize)> {
fn cleanup_log_files(
directory: &Path,
file_prefix: &str,
today: chrono::NaiveDate,
retention_days: usize,
maximum_files: usize,
@@ -277,7 +323,7 @@ fn cleanup_log_files(
.checked_sub_signed(chrono::Duration::days(retention_offset))
.unwrap_or(chrono::NaiveDate::MIN);
let mut retained = Vec::new();
for entry in managed_log_files(directory)? {
for entry in managed_log_files(directory, file_prefix)? {
if entry.date < cutoff {
fs::remove_file(entry.path)?;
} else {
+42 -9
View File
@@ -8,8 +8,13 @@ use gateway_host::system_update::SystemUpdateConfig;
const LOG_DIRECTORY_ENV: &str = "CPR_LOGGING_TEST_DIRECTORY";
const CHILD_PROCESS_ENV: &str = "CPR_LOGGING_TEST_CHILD";
const OAUTH_RECOVERY_LOG_TARGET: &str = "openai_oauth_recovery";
const APPLICATION_LOG_FILE_PREFIX: &str = "codex-proxy-rs.";
const OAUTH_RECOVERY_LOG_FILE_PREFIX: &str = "codex-proxy-rs-oath.";
const APPLICATION_LOG_TARGET: &str = "logging_test_application";
const APPLICATION_LOG_MARKER: &str = "application-file-filter-test";
const OAUTH_RECOVERY_LOG_TARGET: &str = "oauth_recovery";
const OAUTH_RECOVERY_LOG_MARKER: &str = "oauth-recovery-file-filter-test";
const OAUTH_RECOVERY_PROVIDER: &str = "openai";
#[test]
fn logging_requires_at_least_one_sink() {
@@ -42,7 +47,7 @@ fn logging_requires_at_least_one_sink() {
}
#[test]
fn oauth_recovery_file_logging_overrides_global_log_level() {
fn oauth_recovery_file_logging_is_separate_and_overrides_global_log_level() {
if env::var_os(CHILD_PROCESS_ENV).is_some() {
write_oauth_recovery_log(PathBuf::from(
env::var_os(LOG_DIRECTORY_ENV).expect("child log directory"),
@@ -54,12 +59,15 @@ fn oauth_recovery_file_logging_overrides_global_log_level() {
let output = Command::new(env::current_exe().expect("current test executable"))
.args([
"--exact",
"logging::oauth_recovery_file_logging_overrides_global_log_level",
"logging::oauth_recovery_file_logging_is_separate_and_overrides_global_log_level",
"--nocapture",
])
.env(CHILD_PROCESS_ENV, "1")
.env(LOG_DIRECTORY_ENV, directory.path())
.env("RUST_LOG", "off,openai_oauth_recovery=off")
.env(
"RUST_LOG",
"off,logging_test_application=info,oauth_recovery=off",
)
.output()
.expect("run logging child process");
@@ -70,9 +78,18 @@ fn oauth_recovery_file_logging_overrides_global_log_level() {
String::from_utf8_lossy(&output.stderr),
);
let log = read_log_directory(directory.path());
assert!(log.contains(OAUTH_RECOVERY_LOG_TARGET));
assert!(log.contains(OAUTH_RECOVERY_LOG_MARKER));
let application_log = read_log_file_set(directory.path(), APPLICATION_LOG_FILE_PREFIX);
assert!(application_log.contains(&json_target_field(APPLICATION_LOG_TARGET)));
assert!(application_log.contains(APPLICATION_LOG_MARKER));
assert!(!application_log.contains(&json_target_field(OAUTH_RECOVERY_LOG_TARGET)));
assert!(!application_log.contains(OAUTH_RECOVERY_LOG_MARKER));
let recovery_log = read_log_file_set(directory.path(), OAUTH_RECOVERY_LOG_FILE_PREFIX);
assert!(recovery_log.contains(&json_target_field(OAUTH_RECOVERY_LOG_TARGET)));
assert!(recovery_log.contains(OAUTH_RECOVERY_LOG_MARKER));
assert!(!recovery_log.contains(&json_target_field(APPLICATION_LOG_TARGET)));
assert!(!recovery_log.contains(APPLICATION_LOG_MARKER));
assert!(recovery_log.contains(&format!(r#""provider":"{OAUTH_RECOVERY_PROVIDER}""#)));
}
fn write_oauth_recovery_log(directory: PathBuf) {
@@ -104,23 +121,39 @@ fn write_oauth_recovery_log(directory: PathBuf) {
.block_on(gateway_host::initialize(config))
.expect("initialize logging");
tracing::info!(
target: APPLICATION_LOG_TARGET,
marker = APPLICATION_LOG_MARKER,
"application file test record"
);
tracing::info!(
target: OAUTH_RECOVERY_LOG_TARGET,
provider = OAUTH_RECOVERY_PROVIDER,
marker = OAUTH_RECOVERY_LOG_MARKER,
"OpenAI OAuth recovery test record"
"OAuth recovery test record"
);
drop(bundle);
}
fn read_log_directory(directory: &Path) -> String {
fn read_log_file_set(directory: &Path, file_prefix: &str) -> String {
let mut files = fs::read_dir(directory)
.expect("read log directory")
.map(|entry| entry.expect("read log entry").path())
.filter(|path| {
path.file_name()
.and_then(|name| name.to_str())
.is_some_and(|name| name.starts_with(file_prefix))
})
.collect::<Vec<_>>();
files.sort();
assert!(!files.is_empty(), "expected log file set {file_prefix}");
files
.into_iter()
.map(|path| fs::read_to_string(path).expect("read log file"))
.collect()
}
fn json_target_field(target: &str) -> String {
format!(r#""target":"{target}""#)
}
@@ -315,6 +315,26 @@ pub fn is_rate_limit_header_name(name: &str) -> bool {
|| rate_limit_id_from_header_name(&normalized).is_some()
}
/// 判断响应头是否为仅供代理本地观测的 Codex 额度信号。
///
/// 这类头会让下游官方 Codex 客户端显示账户额度提示,因此不能穿过账号隔离边界。
/// 通用限流头(例如 `retry-after`)和非额度的 Codex 头仍允许下发。
pub fn is_codex_quota_header_name(name: &str) -> bool {
let normalized = name.trim().to_ascii_lowercase();
matches!(
normalized.as_str(),
"x-codex-credits-has-credits"
| "x-codex-credits-unlimited"
| "x-codex-credits-balance"
| "x-codex-active-limit"
| "x-codex-plan-type"
| "x-codex-promo-message"
| "x-codex-rate-limit-reached-type"
| "x-codex-primary-over-secondary-limit-percent"
) || rate_limit_id_from_header_name(&normalized)
.is_some_and(|limit_id| limit_id == "codex" || limit_id.starts_with("codex_"))
}
/// 从内部 `codex.rate_limits` 事件中解析限流信息。
pub fn parse_rate_limits_event(value: &Value) -> Option<ParsedRateLimits> {
if value.get("type").and_then(Value::as_str) != Some("codex.rate_limits") {
@@ -1,7 +1,7 @@
use gateway_protocol::openai::events::{
RateLimitWindow, TokenUsage, extract_sse_usage, extract_usage, is_rate_limit_header_name,
parse_rate_limit_headers, parse_rate_limits_event, rate_limits_to_header_pairs,
retry_after_seconds_from_body,
RateLimitWindow, TokenUsage, extract_sse_usage, extract_usage, is_codex_quota_header_name,
is_rate_limit_header_name, parse_rate_limit_headers, parse_rate_limits_event,
rate_limits_to_header_pairs, retry_after_seconds_from_body,
};
use serde_json::json;
@@ -537,3 +537,21 @@ fn rate_limit_header_filter_should_accept_only_supported_domain_headers() {
[true, true, true, false, false]
);
}
#[test]
fn codex_quota_header_filter_should_only_match_client_quota_signals() {
assert_eq!(
[
"X-CoDeX-Primary-Used-Percent",
"x-codex-code-review-primary-reset-at",
"x-codex-credits-balance",
"x-codex-primary-over-secondary-limit-percent",
"retry-after",
"x-ratelimit-remaining",
"x-codex-safety-buffering-enabled",
"x-codex-turn-state",
]
.map(is_codex_quota_header_name),
[true, true, true, true, false, false, false, false]
);
}
@@ -90,6 +90,11 @@ pub async fn initialize(mut config: StoreConfig) -> StoreResult<StoreBundle> {
pool.clone(),
Some(Arc::clone(&cooldowns) as Arc<dyn ProviderCooldownPort>),
)),
Arc::new(postgres::PgAccountGroupRepository::with_runtime_state(
pool.clone(),
credential_leases.clone(),
Arc::clone(&cooldowns) as Arc<dyn ProviderCooldownPort>,
)),
Arc::new(AdminAuthStoreAdapter {
security: postgres::PgAdminSecurityAuditRepository::new(pool.clone()),
settings: postgres::PgRuntimeSettingsRepository::new(pool.clone()),
@@ -0,0 +1,730 @@
//! PostgreSQL owner for provider-neutral account groups and memberships.
use std::{
collections::{BTreeMap, BTreeSet},
str::FromStr as _,
sync::Arc,
};
use async_trait::async_trait;
use futures::future::BoxFuture;
use gateway_admin::{
model::{
MutationContext,
account_groups::{
AccountGroupAccountSummary, AccountGroupCapacity, AccountGroupColor,
AccountGroupListQuery, AccountGroupMutation, AccountGroupPage, AccountGroupRecord,
AccountGroupUsage, DeleteAccountGroup, NewAccountGroup, SetAccountGroupEnabled,
UpdateAccountGroup,
},
observability::DecimalAmount,
},
ports::store::{AccountGroupStore, AdminStoreError, AdminStoreResult},
};
use gateway_core::{
engine::credential::AccountStatus, provider_ports::ProviderCooldownPort,
routing::AccountGroupId,
};
use sqlx::{PgPool, Postgres, QueryBuilder, Row as _, Transaction};
use crate::redis::CredentialLeaseRepository as _;
use crate::{
ConflictKind, StoreError, StoreResult, admin_revision, admin_store_error, mutation_audit,
postgres_unavailable,
};
use super::{
PgProviderAccountRepository, PgRuntimeSettingsRepository, ProviderAccountRepository,
RuntimeSettingsRepository, account_status_projection, append_admin_audit_event_in_transaction,
bump_config_revision_in_transaction, load_rate_limited_until,
};
const ENTITY: &str = "account group";
/// Account group store with transactional revision and audit ownership.
#[derive(Clone)]
pub struct PgAccountGroupRepository {
pool: PgPool,
runtime_signals: Option<crate::redis::RedisCredentialLeaseRepository>,
cooldowns: Option<Arc<dyn ProviderCooldownPort>>,
}
impl PgAccountGroupRepository {
#[must_use]
pub const fn new(pool: PgPool) -> Self {
Self {
pool,
runtime_signals: None,
cooldowns: None,
}
}
#[must_use]
pub fn with_runtime_state(
pool: PgPool,
runtime_signals: crate::redis::RedisCredentialLeaseRepository,
cooldowns: Arc<dyn ProviderCooldownPort>,
) -> Self {
Self {
pool,
runtime_signals: Some(runtime_signals),
cooldowns: Some(cooldowns),
}
}
async fn current_revision(&self) -> AdminStoreResult<gateway_admin::model::Revision> {
RuntimeSettingsRepository::load_runtime_settings(&PgRuntimeSettingsRepository::new(
self.pool.clone(),
))
.await
.map_err(|error| admin_store_error(ENTITY, error))
.and_then(|settings| admin_revision(settings.config_revision))
}
async fn required_record(&self, id: &AccountGroupId) -> AdminStoreResult<AccountGroupRecord> {
let record = load_record(&self.pool, id.as_str())
.await
.map_err(|error| admin_store_error(ENTITY, error))?
.ok_or_else(|| not_found(id.as_str()))?;
let mut records = vec![record];
self.enrich_records(&mut records).await?;
records
.pop()
.ok_or_else(|| invalid_admin("account group record disappeared"))
}
async fn enrich_records(&self, records: &mut [AccountGroupRecord]) -> AdminStoreResult<()> {
enrich_group_records(
&self.pool,
self.runtime_signals.as_ref(),
self.cooldowns.as_deref(),
records,
)
.await
.map_err(|error| admin_store_error(ENTITY, error))
}
async fn mutate<F>(
&self,
audit: super::AdminAuditEvent,
mutation: F,
) -> AdminStoreResult<gateway_admin::model::Revision>
where
F: for<'a> FnOnce(&'a mut Transaction<'_, Postgres>) -> BoxFuture<'a, StoreResult<()>>,
{
let mut transaction =
self.pool.begin().await.map_err(|_| {
admin_store_error(ENTITY, unavailable("begin account group mutation"))
})?;
let result = async {
let revision = bump_config_revision_in_transaction(&mut transaction).await?;
mutation(&mut transaction).await?;
append_admin_audit_event_in_transaction(&mut transaction, audit, revision).await?;
Ok(revision)
}
.await;
match result {
Ok(revision) => {
transaction.commit().await.map_err(|_| {
admin_store_error(ENTITY, unavailable("commit account group mutation"))
})?;
admin_revision(revision)
}
Err(error) => {
transaction.rollback().await.map_err(|_| {
admin_store_error(ENTITY, unavailable("rollback account group mutation"))
})?;
Err(admin_store_error(ENTITY, error))
}
}
}
}
#[async_trait]
impl AccountGroupStore for PgAccountGroupRepository {
async fn list_account_groups(
&self,
query: AccountGroupListQuery,
) -> AdminStoreResult<AccountGroupPage> {
validate_page_query(&query)?;
let total = count_groups(&self.pool, &query)
.await
.map_err(|error| admin_store_error(ENTITY, error))?;
let offset = u64::from(query.page.saturating_sub(1)) * u64::from(query.page_size.get());
let mut statement = group_select();
push_group_filter(&mut statement, &query);
statement.push(" order by g.created_at desc, g.id desc limit ");
statement.push_bind(i64::from(query.page_size.get()));
statement.push(" offset ");
statement.push_bind(i64::try_from(offset).map_err(|_| invalid_admin("page is too large"))?);
let rows = statement
.build()
.fetch_all(&self.pool)
.await
.map_err(|_| admin_store_error(ENTITY, unavailable("list account groups")))?;
let mut items = rows
.iter()
.map(group_record)
.collect::<StoreResult<Vec<_>>>()
.map_err(|error| admin_store_error(ENTITY, error))?;
self.enrich_records(&mut items).await?;
Ok(AccountGroupPage {
config_revision: self.current_revision().await?,
items,
total,
page: query.page,
page_size: query.page_size.get(),
})
}
async fn create_account_group(
&self,
command: NewAccountGroup,
context: &MutationContext,
) -> AdminStoreResult<AccountGroupMutation> {
validate_group_fields(&command.name, command.description.as_deref())?;
let id = command.id.clone();
let audit = mutation_audit(
context,
"create",
"account_group",
id.as_str(),
vec!["name".to_owned(), "description".to_owned()],
);
let revision = self
.mutate(audit, |transaction| {
Box::pin(async move {
sqlx::query(
"insert into account_groups
(id, name, description, color, enabled, created_at, updated_at)
values ($1, $2, $3, $4, true, now(), now())",
)
.bind(command.id.as_str())
.bind(command.name)
.bind(command.description)
.bind(command.color.as_str())
.execute(&mut **transaction)
.await
.map_err(|error| map_group_write_error(error, command.id.as_str()))?;
Ok(())
})
})
.await?;
Ok(AccountGroupMutation {
config_revision: revision,
id: id.clone(),
record: Some(self.required_record(&id).await?),
})
}
async fn update_account_group(
&self,
command: UpdateAccountGroup,
context: &MutationContext,
) -> AdminStoreResult<AccountGroupMutation> {
validate_group_fields(&command.name, command.description.as_deref())?;
let id = command.id.clone();
let audit = mutation_audit(
context,
"update",
"account_group",
id.as_str(),
vec!["name".to_owned(), "description".to_owned()],
);
let revision = self
.mutate(audit, |transaction| {
Box::pin(async move {
let result = sqlx::query(
"update account_groups
set name = $2, description = $3, color = $4, updated_at = now()
where id = $1",
)
.bind(command.id.as_str())
.bind(command.name)
.bind(command.description)
.bind(command.color.as_str())
.execute(&mut **transaction)
.await
.map_err(|error| map_group_write_error(error, command.id.as_str()))?;
require_one(result.rows_affected(), command.id.as_str())
})
})
.await?;
Ok(AccountGroupMutation {
config_revision: revision,
id: id.clone(),
record: Some(self.required_record(&id).await?),
})
}
async fn set_account_group_enabled(
&self,
command: SetAccountGroupEnabled,
context: &MutationContext,
) -> AdminStoreResult<AccountGroupMutation> {
let id = command.id.clone();
let audit = mutation_audit(
context,
if command.enabled { "enable" } else { "disable" },
"account_group",
id.as_str(),
vec!["enabled".to_owned()],
);
let revision = self
.mutate(audit, |transaction| {
Box::pin(async move {
let result = sqlx::query(
"update account_groups set enabled = $2, updated_at = now() where id = $1",
)
.bind(command.id.as_str())
.bind(command.enabled)
.execute(&mut **transaction)
.await
.map_err(|_| unavailable("set account group state"))?;
require_one(result.rows_affected(), command.id.as_str())
})
})
.await?;
Ok(AccountGroupMutation {
config_revision: revision,
id: id.clone(),
record: Some(self.required_record(&id).await?),
})
}
async fn delete_account_group(
&self,
command: DeleteAccountGroup,
context: &MutationContext,
) -> AdminStoreResult<AccountGroupMutation> {
let id = command.id.clone();
let audit = mutation_audit(context, "delete", "account_group", id.as_str(), Vec::new());
let revision = self
.mutate(audit, |transaction| {
Box::pin(async move {
let result = sqlx::query(
"delete from account_groups g
where g.id = $1
and not exists (
select 1 from client_api_key_groups k where k.account_group_id = g.id
)",
)
.bind(command.id.as_str())
.execute(&mut **transaction)
.await
.map_err(|_| unavailable("delete account group"))?;
if result.rows_affected() == 1 {
return Ok(());
}
let exists = sqlx::query_scalar::<_, bool>(
"select exists(select 1 from account_groups where id = $1)",
)
.bind(command.id.as_str())
.fetch_one(&mut **transaction)
.await
.map_err(|_| unavailable("check account group delete conflict"))?;
if exists {
Err(conflict(command.id.as_str()))
} else {
Err(not_found_store(command.id.as_str()))
}
})
})
.await?;
Ok(AccountGroupMutation {
config_revision: revision,
id,
record: None,
})
}
}
fn group_select() -> QueryBuilder<Postgres> {
QueryBuilder::new(
"select g.id, g.name, g.description, g.color, g.enabled, g.created_at, g.updated_at,
coalesce(members.member_count, 0)::bigint as member_count,
coalesce(keys.client_key_count, 0)::bigint as client_key_count,
coalesce(members.provider_counts, '{}'::jsonb) as provider_counts
from account_groups g
left join lateral (
select coalesce(sum(counts.provider_count), 0)::bigint as member_count,
jsonb_object_agg(counts.provider_kind, counts.provider_count)
as provider_counts
from (
select a.provider_kind, count(*)::bigint as provider_count
from account_group_accounts gm
join provider_accounts a on a.id = gm.provider_account_id
where gm.account_group_id = g.id
group by a.provider_kind
) counts
) members on true
left join lateral (
select count(*)::bigint as client_key_count
from client_api_key_groups kg
where kg.account_group_id = g.id
) keys on true
where true",
)
}
fn push_group_filter(statement: &mut QueryBuilder<Postgres>, query: &AccountGroupListQuery) {
if let Some(search) = &query.search {
statement.push(" and lower(g.name) like ");
statement.push_bind(format!("%{}%", search.to_lowercase()));
}
if let Some(enabled) = query.enabled {
statement.push(" and g.enabled = ");
statement.push_bind(enabled);
}
}
async fn count_groups(pool: &PgPool, query: &AccountGroupListQuery) -> StoreResult<u64> {
let mut statement =
QueryBuilder::<Postgres>::new("select count(*)::bigint from account_groups g where true");
if let Some(search) = &query.search {
statement.push(" and lower(g.name) like ");
statement.push_bind(format!("%{}%", search.to_lowercase()));
}
if let Some(enabled) = query.enabled {
statement.push(" and g.enabled = ");
statement.push_bind(enabled);
}
let count = statement
.build_query_scalar::<i64>()
.fetch_one(pool)
.await
.map_err(|_| unavailable("count account groups"))?;
u64::try_from(count).map_err(|_| invalid("negative account group count"))
}
async fn load_record(pool: &PgPool, id: &str) -> StoreResult<Option<AccountGroupRecord>> {
let mut statement = group_select();
statement.push(" and g.id = ");
statement.push_bind(id.to_owned());
statement
.build()
.fetch_optional(pool)
.await
.map_err(|_| unavailable("load account group"))?
.as_ref()
.map(group_record)
.transpose()
}
fn group_record(row: &sqlx::postgres::PgRow) -> StoreResult<AccountGroupRecord> {
let provider_counts: serde_json::Value = row
.try_get("provider_counts")
.map_err(|_| invalid("invalid provider counts"))?;
let provider_counts = serde_json::from_value::<BTreeMap<String, u64>>(provider_counts)
.map_err(|_| invalid("invalid provider counts"))?;
Ok(AccountGroupRecord {
id: AccountGroupId::new(
row.try_get::<String, _>("id")
.map_err(|_| invalid("invalid id"))?,
)
.map_err(|_| invalid("invalid id"))?,
name: row.try_get("name").map_err(|_| invalid("invalid name"))?,
description: row
.try_get("description")
.map_err(|_| invalid("invalid description"))?,
color: AccountGroupColor::parse(
row.try_get::<String, _>("color")
.map_err(|_| invalid("invalid color"))?
.as_str(),
)
.ok_or_else(|| invalid("invalid color"))?,
enabled: row
.try_get("enabled")
.map_err(|_| invalid("invalid enabled"))?,
member_count: count_value(row, "member_count")?,
provider_counts,
client_key_count: count_value(row, "client_key_count")?,
account_summary: AccountGroupAccountSummary {
available: 0,
limited: 0,
total: 0,
},
capacity: AccountGroupCapacity {
used_slots: None,
total_slots: 0,
},
usage: AccountGroupUsage {
today_usd: DecimalAmount::from_str("0").map_err(|_| invalid("invalid zero cost"))?,
total_usd: DecimalAmount::from_str("0").map_err(|_| invalid("invalid zero cost"))?,
},
created_at: row
.try_get("created_at")
.map_err(|_| invalid("invalid created_at"))?,
updated_at: row
.try_get("updated_at")
.map_err(|_| invalid("invalid updated_at"))?,
})
}
async fn enrich_group_records(
pool: &PgPool,
runtime_signals: Option<&crate::redis::RedisCredentialLeaseRepository>,
cooldowns: Option<&dyn ProviderCooldownPort>,
records: &mut [AccountGroupRecord],
) -> StoreResult<()> {
if records.is_empty() {
return Ok(());
}
let settings = RuntimeSettingsRepository::load_runtime_settings(
&PgRuntimeSettingsRepository::new(pool.clone()),
)
.await?;
let accounts = PgProviderAccountRepository::new(pool.clone())
.list_provider_accounts(None, true)
.await?;
let now = chrono::Utc::now();
let cooldowns = load_rate_limited_until(cooldowns, &accounts, now.into()).await;
let mut status_by_account = BTreeMap::new();
let mut capacity_by_account = BTreeMap::new();
for account in &accounts {
let status =
account_status_projection(account, now.into(), cooldowns.get(&account.id).copied())
.status;
status_by_account.insert(account.id.as_str(), status);
capacity_by_account.insert(
account.id.as_str(),
u64::from(
account
.concurrency_limit
.map_or(settings.max_concurrent_per_account, |limit| limit.get()),
),
);
}
let group_ids = records
.iter()
.map(|record| record.id.as_str().to_owned())
.collect::<Vec<_>>();
let memberships = sqlx::query(
"select account_group_id, provider_account_id
from account_group_accounts
where account_group_id = any($1::text[])",
)
.bind(&group_ids)
.fetch_all(pool)
.await
.map_err(|_| unavailable("load account group availability"))?;
let mut accounts_by_group = BTreeMap::<String, Vec<String>>::new();
for row in memberships {
accounts_by_group
.entry(
row.try_get("account_group_id")
.map_err(|_| invalid("invalid group ID"))?,
)
.or_default()
.push(
row.try_get("provider_account_id")
.map_err(|_| invalid("invalid account ID"))?,
);
}
let normal_ids = accounts_by_group
.values()
.flatten()
.filter(|account_id| {
status_by_account.get(account_id.as_str()) == Some(&AccountStatus::Normal)
})
.cloned()
.collect::<BTreeSet<_>>()
.into_iter()
.collect::<Vec<_>>();
let used_by_account = if normal_ids.is_empty() {
Some(BTreeMap::new())
} else if let Some(runtime_signals) = runtime_signals {
runtime_signals
.credential_runtime_signals(&normal_ids)
.await
.ok()
.map(|signals| {
signals
.into_iter()
.map(|signal| (signal.resource_id, u64::from(signal.in_flight)))
.collect::<BTreeMap<_, _>>()
})
} else {
None
};
let costs = group_costs(pool, &group_ids).await?;
for record in records {
let account_ids = accounts_by_group
.get(record.id.as_str())
.map_or(&[][..], Vec::as_slice);
let available = account_ids
.iter()
.filter(|id| status_by_account.get(id.as_str()) == Some(&AccountStatus::Normal))
.count() as u64;
let total = u64::try_from(account_ids.len())
.map_err(|_| invalid("group account count overflow"))?;
record.account_summary = AccountGroupAccountSummary {
available,
limited: total.saturating_sub(available),
total,
};
record.capacity = AccountGroupCapacity {
used_slots: if available == 0 {
Some(0)
} else {
used_by_account.as_ref().map(|used| {
account_ids.iter().fold(0_u64, |sum, id| {
sum.saturating_add(used.get(id).copied().unwrap_or(0))
})
})
},
total_slots: account_ids.iter().fold(0_u64, |total, account_id| {
if status_by_account.get(account_id.as_str()) == Some(&AccountStatus::Normal) {
total.saturating_add(
capacity_by_account
.get(account_id.as_str())
.copied()
.unwrap_or(0),
)
} else {
total
}
}),
};
if let Some(usage) = costs.get(record.id.as_str()) {
record.usage = usage.clone();
}
}
Ok(())
}
async fn group_costs(
pool: &PgPool,
group_ids: &[String],
) -> StoreResult<BTreeMap<String, AccountGroupUsage>> {
let rows = sqlx::query(
"select gm.account_group_id,
coalesce(sum(mr.cost_amount) filter (
where mr.started_at >= date_trunc('day', now() at time zone 'Asia/Shanghai')
at time zone 'Asia/Shanghai'
), 0)::text as today_usd,
coalesce(sum(mr.cost_amount), 0)::text as total_usd
from account_group_accounts gm
left join model_requests mr
on mr.provider_account_ref = gm.provider_account_id
and mr.outcome = 'succeeded'
and mr.downstream_committed_at is not null
and mr.client_status_code between 200 and 399
and mr.cost_currency = 'USD'
and mr.cost_amount is not null
where gm.account_group_id = any($1::text[])
group by gm.account_group_id",
)
.bind(group_ids)
.fetch_all(pool)
.await
.map_err(|_| unavailable("load account group costs"))?;
rows.into_iter()
.map(|row| {
let group_id = row
.try_get("account_group_id")
.map_err(|_| invalid("invalid group ID"))?;
let today = row
.try_get::<String, _>("today_usd")
.map_err(|_| invalid("invalid today cost"))?;
let total = row
.try_get::<String, _>("total_usd")
.map_err(|_| invalid("invalid total cost"))?;
Ok((
group_id,
AccountGroupUsage {
today_usd: DecimalAmount::from_str(&today)
.map_err(|_| invalid("invalid today cost"))?,
total_usd: DecimalAmount::from_str(&total)
.map_err(|_| invalid("invalid total cost"))?,
},
))
})
.collect()
}
fn count_value(row: &sqlx::postgres::PgRow, field: &str) -> StoreResult<u64> {
u64::try_from(
row.try_get::<i64, _>(field)
.map_err(|_| invalid("invalid count"))?,
)
.map_err(|_| invalid("negative count"))
}
fn validate_page_query(query: &AccountGroupListQuery) -> AdminStoreResult<()> {
if query.page == 0 {
return Err(invalid_admin("page must be positive"));
}
if query.search.as_deref().is_some_and(|search| {
search.trim().is_empty() || search.len() > 256 || search.chars().any(char::is_control)
}) {
return Err(invalid_admin("invalid search"));
}
Ok(())
}
fn validate_group_fields(name: &str, description: Option<&str>) -> AdminStoreResult<()> {
if name.trim() != name
|| name.is_empty()
|| name.chars().count() > 100
|| name.chars().any(char::is_control)
|| description
.is_some_and(|value| value.len() > 4096 || value.chars().any(char::is_control))
{
return Err(invalid_admin("invalid account group fields"));
}
Ok(())
}
fn require_one(rows: u64, id: &str) -> StoreResult<()> {
if rows == 1 {
Ok(())
} else {
Err(not_found_store(id))
}
}
fn map_group_write_error(error: sqlx::Error, id: &str) -> StoreError {
if error
.as_database_error()
.is_some_and(|database| database.is_unique_violation())
{
conflict(id)
} else {
unavailable("write account group")
}
}
fn not_found(id: &str) -> AdminStoreError {
admin_store_error(ENTITY, not_found_store(id))
}
fn not_found_store(id: &str) -> StoreError {
StoreError::NotFound {
entity: ENTITY,
id: id.to_owned(),
}
}
fn conflict(id: &str) -> StoreError {
StoreError::Conflict {
entity: ENTITY,
id: id.to_owned(),
kind: ConflictKind::InvalidTransition,
}
}
fn invalid_admin(message: &str) -> AdminStoreError {
admin_store_error(ENTITY, invalid(message))
}
fn invalid(message: &str) -> StoreError {
StoreError::InvalidData {
entity: ENTITY,
message: message.to_owned(),
}
}
fn unavailable(message: &'static str) -> StoreError {
postgres_unavailable(message)
}
@@ -1,7 +1,7 @@
//! 明文 `client_api_keys` 的 PostgreSQL owner。
use std::{
collections::BTreeMap,
collections::{BTreeMap, BTreeSet},
fmt,
sync::{
Arc, Mutex,
@@ -15,6 +15,7 @@ use chrono::{DateTime, Utc};
use gateway_admin::{
model::{
MutationContext,
account_groups::AccountGroupRef as AdminAccountGroupRef,
client_keys::{
ClientKeyCursor as AdminClientKeyCursor,
ClientKeyCursorValue as AdminClientKeyCursorValue,
@@ -27,11 +28,12 @@ use gateway_admin::{
},
ports::store::{AdminStoreResult, ClientKeyStore},
};
use gateway_core::routing::ProviderKind;
use gateway_core::{
engine::execution::ClientApiKeyUsageSink,
policy::{ClientApiKeyId, PlaintextClientApiKey, RateLimits},
routing::{AccountGroupId, ProviderKind},
};
use serde::Deserialize;
use sqlx::{PgPool, Postgres, QueryBuilder, Transaction};
use crate::{
@@ -49,7 +51,7 @@ const CLIENT_API_KEY_LAST_USED_FLUSH_DELAY: Duration = Duration::from_secs(1);
pub struct ClientApiKeySnapshot {
pub id: ClientApiKeyId,
pub plaintext_key: PlaintextClientApiKey,
pub provider_kind: String,
pub group_ids: Vec<AccountGroupId>,
pub limits: RateLimits,
}
@@ -57,7 +59,7 @@ impl ClientApiKeySnapshot {
pub(crate) fn from_persisted(
id: String,
key: String,
provider_kind: String,
group_ids: Vec<String>,
max_concurrency: i64,
requests_per_minute: i64,
) -> StoreResult<Self> {
@@ -65,7 +67,12 @@ impl ClientApiKeySnapshot {
id: ClientApiKeyId::new(id).map_err(|_| invalid("persisted key ID is invalid"))?,
plaintext_key: PlaintextClientApiKey::new(key)
.map_err(|_| invalid("persisted plaintext key is invalid"))?,
provider_kind,
group_ids: group_ids
.into_iter()
.map(|id| {
AccountGroupId::new(id).map_err(|_| invalid("persisted group ID is invalid"))
})
.collect::<StoreResult<Vec<_>>>()?,
limits: RateLimits {
max_concurrency: to_u64(max_concurrency)?,
requests_per_minute: to_u64(requests_per_minute)?,
@@ -78,7 +85,6 @@ impl ClientApiKeySnapshot {
pub struct ClientApiKeySecret {
pub id: String,
pub key: String,
pub provider_kind: String,
pub enabled: bool,
pub max_concurrency: u64,
pub requests_per_minute: u64,
@@ -100,7 +106,8 @@ pub struct ClientApiKeyRecord {
pub id: String,
pub name: String,
pub label: Option<String>,
pub provider_kind: String,
pub groups: Vec<ClientApiKeyGroupRecord>,
pub provider_kinds: Vec<String>,
pub prefix: String,
pub enabled: bool,
pub max_concurrency: u64,
@@ -110,6 +117,14 @@ pub struct ClientApiKeyRecord {
pub updated_at: DateTime<Utc>,
}
#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
pub struct ClientApiKeyGroupRecord {
pub id: String,
pub name: String,
pub color: String,
pub enabled: bool,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub enum ClientApiKeySortField {
Name,
@@ -250,7 +265,7 @@ pub struct NewClientApiKey {
pub id: String,
pub name: String,
pub label: Option<String>,
pub provider_kind: String,
pub group_ids: Vec<String>,
pub key: String,
pub max_concurrency: u64,
pub requests_per_minute: u64,
@@ -271,28 +286,17 @@ impl NewClientApiKey {
pub fn validate(&self) -> StoreResult<()> {
require_nonempty(ENTITY, "id", &self.id)?;
require_nonempty(ENTITY, "name", &self.name)?;
require_nonempty(ENTITY, "provider_kind", &self.provider_kind)?;
validate_group_ids(&self.group_ids)?;
validate_key(&self.key)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UpdateClientApiKey {
pub id: String,
pub name: String,
pub label: Option<String>,
pub provider_kind: String,
pub enabled: bool,
pub max_concurrency: u64,
pub requests_per_minute: u64,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UpdateClientApiKeyDetails {
pub id: String,
pub name: String,
pub label: Option<String>,
pub provider_kind: String,
pub group_ids: Vec<String>,
pub max_concurrency: u64,
pub requests_per_minute: u64,
}
@@ -301,7 +305,7 @@ impl UpdateClientApiKeyDetails {
pub fn validate(&self) -> StoreResult<()> {
require_nonempty(ENTITY, "id", &self.id)?;
require_nonempty(ENTITY, "name", &self.name)?;
require_nonempty(ENTITY, "provider_kind", &self.provider_kind)?;
validate_group_ids(&self.group_ids)?;
to_i64(self.max_concurrency)?;
to_i64(self.requests_per_minute)?;
Ok(())
@@ -316,7 +320,6 @@ pub trait ClientApiKeyRepository: Send + Sync {
) -> StoreResult<ClientApiKeyPage>;
async fn get_client_api_key(&self, id: &str) -> StoreResult<Option<ClientApiKeyRecord>>;
async fn reveal_client_api_key(&self, id: &str) -> StoreResult<Option<ClientApiKeySecret>>;
async fn update_client_api_key(&self, key: UpdateClientApiKey) -> StoreResult<bool>;
async fn touch_client_api_keys(
&self,
touched_at: &BTreeMap<String, DateTime<Utc>>,
@@ -344,9 +347,38 @@ impl ClientApiKeyRepository for PgClientApiKeyRepository {
query.validate()?;
let total = count_client_api_keys(&self.pool, query.search.as_deref()).await?;
let mut statement = QueryBuilder::<Postgres>::new(
"select id, name, label, provider_kind, left(key, 10) as prefix, enabled, max_concurrency,
requests_per_minute, last_used_at, created_at, updated_at
from client_api_keys where true",
"select k.id, k.name, k.label, left(k.key, 10) as prefix, k.enabled,
k.max_concurrency, k.requests_per_minute, k.last_used_at, k.created_at,
k.updated_at,
coalesce(groups.groups, '[]'::jsonb) as groups,
case
when groups.binding_count = 0 then coalesce(
(select array_agg(distinct a.provider_kind order by a.provider_kind)
from provider_accounts a),
'{}'
)
else coalesce(groups.provider_kinds, '{}')
end as provider_kinds
from client_api_keys k
left join lateral (
select
(select count(*)::bigint
from client_api_key_groups kg
where kg.client_api_key_id = k.id) as binding_count,
(select jsonb_agg(jsonb_build_object(
'id', g.id, 'name', g.name, 'color', g.color, 'enabled', g.enabled
) order by g.id)
from client_api_key_groups kg
join account_groups g on g.id = kg.account_group_id
where kg.client_api_key_id = k.id) as groups,
(select array_agg(distinct a.provider_kind order by a.provider_kind)
from client_api_key_groups kg
join account_groups g on g.id = kg.account_group_id and g.enabled
join account_group_accounts gm on gm.account_group_id = g.id
join provider_accounts a on a.id = gm.provider_account_id
where kg.client_api_key_id = k.id) as provider_kinds
) groups on true
where true",
);
push_client_key_search(&mut statement, query.search.as_deref());
if let Some(cursor) = &query.cursor {
@@ -356,12 +388,12 @@ impl ClientApiKeyRepository for PgClientApiKeyRepository {
statement.push(" limit ");
statement.push_bind(i64::from(query.page_size) + 1);
let rows = statement
.build_query_as::<ClientRecordRow>()
.build()
.fetch_all(&self.pool)
.await
.map_err(|_| postgres_unavailable("list client API keys"))?;
let mut items = rows
.into_iter()
.iter()
.map(client_record_from_row)
.collect::<StoreResult<Vec<_>>>()?;
let has_more = items.len() > usize::from(query.page_size);
@@ -384,8 +416,8 @@ impl ClientApiKeyRepository for PgClientApiKeyRepository {
async fn reveal_client_api_key(&self, id: &str) -> StoreResult<Option<ClientApiKeySecret>> {
require_nonempty(ENTITY, "id", id)?;
sqlx::query_as::<_, (String, String, String, bool, i64, i64)>(
"select id, key, provider_kind, enabled, max_concurrency, requests_per_minute
sqlx::query_as::<_, (String, String, bool, i64, i64)>(
"select id, key, enabled, max_concurrency, requests_per_minute
from client_api_keys where id = $1",
)
.bind(id)
@@ -398,41 +430,48 @@ impl ClientApiKeyRepository for PgClientApiKeyRepository {
async fn get_client_api_key(&self, id: &str) -> StoreResult<Option<ClientApiKeyRecord>> {
require_nonempty(ENTITY, "id", id)?;
sqlx::query_as::<_, ClientRecordRow>(
"select id, name, label, provider_kind, left(key, 10) as prefix, enabled, max_concurrency,
requests_per_minute, last_used_at, created_at, updated_at
from client_api_keys where id = $1",
sqlx::query(
"select k.id, k.name, k.label, left(k.key, 10) as prefix, k.enabled,
k.max_concurrency, k.requests_per_minute, k.last_used_at, k.created_at,
k.updated_at, coalesce(groups.groups, '[]'::jsonb) as groups,
case
when groups.binding_count = 0 then coalesce(
(select array_agg(distinct a.provider_kind order by a.provider_kind)
from provider_accounts a),
'{}'
)
else coalesce(groups.provider_kinds, '{}')
end as provider_kinds
from client_api_keys k
left join lateral (
select
(select count(*)::bigint
from client_api_key_groups kg
where kg.client_api_key_id = k.id) as binding_count,
(select jsonb_agg(jsonb_build_object(
'id', g.id, 'name', g.name, 'color', g.color, 'enabled', g.enabled
) order by g.id)
from client_api_key_groups kg
join account_groups g on g.id = kg.account_group_id
where kg.client_api_key_id = k.id) as groups,
(select array_agg(distinct a.provider_kind order by a.provider_kind)
from client_api_key_groups kg
join account_groups g on g.id = kg.account_group_id and g.enabled
join account_group_accounts gm on gm.account_group_id = g.id
join provider_accounts a on a.id = gm.provider_account_id
where kg.client_api_key_id = k.id) as provider_kinds
) groups on true
where k.id = $1",
)
.bind(id)
.fetch_optional(&self.pool)
.await
.map_err(|_| postgres_unavailable("get client API key"))?
.as_ref()
.map(client_record_from_row)
.transpose()
}
async fn update_client_api_key(&self, key: UpdateClientApiKey) -> StoreResult<bool> {
require_nonempty(ENTITY, "id", &key.id)?;
require_nonempty(ENTITY, "name", &key.name)?;
let result = sqlx::query(
"update client_api_keys
set name = $2, label = $3, provider_kind = $4, enabled = $5, max_concurrency = $6,
requests_per_minute = $7, updated_at = now()
where id = $1",
)
.bind(key.id)
.bind(key.name)
.bind(key.label)
.bind(key.provider_kind)
.bind(key.enabled)
.bind(to_i64(key.max_concurrency)?)
.bind(to_i64(key.requests_per_minute)?)
.execute(&self.pool)
.await
.map_err(|_| postgres_unavailable("update client API key"))?;
Ok(result.rows_affected() == 1)
}
async fn touch_client_api_keys(
&self,
touched_at: &BTreeMap<String, DateTime<Utc>>,
@@ -637,7 +676,11 @@ impl ClientKeyStore for PgAdminClientKeyStore {
id: id.as_str().to_owned(),
name: command.name,
label: command.label,
provider_kind: command.provider_kind.as_str().to_owned(),
group_ids: command
.group_ids
.iter()
.map(|id| id.as_str().to_owned())
.collect(),
key: command.plaintext,
max_concurrency: command.limits.max_concurrency,
requests_per_minute: command.limits.requests_per_minute,
@@ -650,7 +693,7 @@ impl ClientKeyStore for PgAdminClientKeyStore {
[
"name",
"label",
"provider_kind",
"group_ids",
"key",
"enabled",
"max_concurrency",
@@ -679,7 +722,11 @@ impl ClientKeyStore for PgAdminClientKeyStore {
id: id.as_str().to_owned(),
name: command.name,
label: command.label,
provider_kind: command.provider_kind.as_str().to_owned(),
group_ids: command
.group_ids
.iter()
.map(|id| id.as_str().to_owned())
.collect(),
max_concurrency: command.limits.max_concurrency,
requests_per_minute: command.limits.requests_per_minute,
},
@@ -691,7 +738,7 @@ impl ClientKeyStore for PgAdminClientKeyStore {
[
"name",
"label",
"provider_kind",
"group_ids",
"max_concurrency",
"requests_per_minute",
]
@@ -828,8 +875,30 @@ fn admin_client_key_record(record: ClientApiKeyRecord) -> AdminStoreResult<Admin
.map_err(|_| admin_store_error(ENTITY, invalid("invalid client key id")))?,
name: record.name,
label: record.label,
provider_kind: ProviderKind::new(record.provider_kind)
.map_err(|_| admin_store_error(ENTITY, invalid("invalid provider kind")))?,
groups: record
.groups
.into_iter()
.map(|group| {
Ok(AdminAccountGroupRef {
id: AccountGroupId::new(group.id)
.map_err(|_| admin_store_error(ENTITY, invalid("invalid group id")))?,
name: group.name,
color: gateway_admin::model::account_groups::AccountGroupColor::parse(
&group.color,
)
.ok_or_else(|| admin_store_error(ENTITY, invalid("invalid group color")))?,
enabled: group.enabled,
})
})
.collect::<AdminStoreResult<Vec<_>>>()?,
provider_kinds: record
.provider_kinds
.into_iter()
.map(|provider| {
ProviderKind::new(provider)
.map_err(|_| admin_store_error(ENTITY, invalid("invalid provider kind")))
})
.collect::<AdminStoreResult<Vec<_>>>()?,
prefix: record.prefix,
enabled: record.enabled,
limits: RateLimits {
@@ -849,20 +918,20 @@ pub(crate) async fn insert_client_api_key_in_transaction(
key.validate()?;
sqlx::query(
"insert into client_api_keys (
id, name, label, provider_kind, key, enabled, max_concurrency, requests_per_minute,
id, name, label, key, enabled, max_concurrency, requests_per_minute,
last_used_at, created_at, updated_at
) values ($1, $2, $3, $4, $5, true, $6, $7, null, now(), now())",
) values ($1, $2, $3, $4, true, $5, $6, null, now(), now())",
)
.bind(&key.id)
.bind(&key.name)
.bind(&key.label)
.bind(&key.provider_kind)
.bind(&key.key)
.bind(to_i64(key.max_concurrency)?)
.bind(to_i64(key.requests_per_minute)?)
.execute(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("insert client API key in transaction"))?;
replace_client_api_key_groups_in_transaction(transaction, &key.id, &key.group_ids).await?;
Ok(())
}
@@ -873,20 +942,20 @@ pub(crate) async fn update_client_api_key_in_transaction(
key.validate()?;
let result = sqlx::query(
"update client_api_keys
set name = $2, label = $3, provider_kind = $4, max_concurrency = $5,
requests_per_minute = $6, updated_at = now()
set name = $2, label = $3, max_concurrency = $4,
requests_per_minute = $5, updated_at = now()
where id = $1",
)
.bind(&key.id)
.bind(&key.name)
.bind(&key.label)
.bind(&key.provider_kind)
.bind(to_i64(key.max_concurrency)?)
.bind(to_i64(key.requests_per_minute)?)
.execute(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("update client API key in transaction"))?;
require_changed(result.rows_affected(), &key.id)
require_changed(result.rows_affected(), &key.id)?;
replace_client_api_key_groups_in_transaction(transaction, &key.id, &key.group_ids).await
}
pub(crate) async fn set_client_api_key_enabled_in_transaction(
@@ -929,6 +998,59 @@ fn require_changed(rows_affected: u64, id: &str) -> StoreResult<()> {
}
}
async fn replace_client_api_key_groups_in_transaction(
transaction: &mut Transaction<'_, Postgres>,
key_id: &str,
group_ids: &[String],
) -> StoreResult<()> {
validate_group_ids(group_ids)?;
if !group_ids.is_empty() {
let count = sqlx::query_scalar::<_, i64>(
"select count(*)::bigint from account_groups where id = any($1::text[])",
)
.bind(group_ids)
.fetch_one(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("validate client API key groups"))?;
if usize::try_from(count).ok() != Some(group_ids.len()) {
return Err(StoreError::NotFound {
entity: "account group",
id: "client API key group selection".to_owned(),
});
}
}
sqlx::query("delete from client_api_key_groups where client_api_key_id = $1")
.bind(key_id)
.execute(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("delete client API key groups"))?;
if !group_ids.is_empty() {
sqlx::query(
"insert into client_api_key_groups
(client_api_key_id, account_group_id, created_at)
select $1, group_id, now() from unnest($2::text[]) group_id",
)
.bind(key_id)
.bind(group_ids)
.execute(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("insert client API key groups"))?;
}
Ok(())
}
fn validate_group_ids(group_ids: &[String]) -> StoreResult<()> {
if group_ids.len() > 1000
|| group_ids.iter().collect::<BTreeSet<_>>().len() != group_ids.len()
|| group_ids
.iter()
.any(|id| AccountGroupId::new(id.clone()).is_err())
{
return Err(invalid("group IDs are invalid or duplicated"));
}
Ok(())
}
fn validate_key(key: &str) -> StoreResult<()> {
let valid = key.len() == KEY_LENGTH
&& key.starts_with("sk_")
@@ -945,45 +1067,53 @@ fn validate_key(key: &str) -> StoreResult<()> {
}
fn client_secret_from_row(
row: (String, String, String, bool, i64, i64),
row: (String, String, bool, i64, i64),
) -> StoreResult<ClientApiKeySecret> {
Ok(ClientApiKeySecret {
id: row.0,
key: row.1,
provider_kind: row.2,
enabled: row.3,
max_concurrency: to_u64(row.4)?,
requests_per_minute: to_u64(row.5)?,
enabled: row.2,
max_concurrency: to_u64(row.3)?,
requests_per_minute: to_u64(row.4)?,
})
}
type ClientRecordRow = (
String,
String,
Option<String>,
String,
String,
bool,
i64,
i64,
Option<DateTime<Utc>>,
DateTime<Utc>,
DateTime<Utc>,
);
fn client_record_from_row(row: ClientRecordRow) -> StoreResult<ClientApiKeyRecord> {
fn client_record_from_row(row: &sqlx::postgres::PgRow) -> StoreResult<ClientApiKeyRecord> {
use sqlx::Row as _;
let groups: serde_json::Value = row
.try_get("groups")
.map_err(|_| invalid("invalid groups"))?;
Ok(ClientApiKeyRecord {
id: row.0,
name: row.1,
label: row.2,
provider_kind: row.3,
prefix: row.4,
enabled: row.5,
max_concurrency: to_u64(row.6)?,
requests_per_minute: to_u64(row.7)?,
last_used_at: row.8,
created_at: row.9,
updated_at: row.10,
id: row.try_get("id").map_err(|_| invalid("invalid id"))?,
name: row.try_get("name").map_err(|_| invalid("invalid name"))?,
label: row.try_get("label").map_err(|_| invalid("invalid label"))?,
groups: serde_json::from_value(groups).map_err(|_| invalid("invalid groups"))?,
provider_kinds: row
.try_get("provider_kinds")
.map_err(|_| invalid("invalid provider kinds"))?,
prefix: row
.try_get("prefix")
.map_err(|_| invalid("invalid prefix"))?,
enabled: row
.try_get("enabled")
.map_err(|_| invalid("invalid enabled"))?,
max_concurrency: to_u64(
row.try_get("max_concurrency")
.map_err(|_| invalid("invalid max concurrency"))?,
)?,
requests_per_minute: to_u64(
row.try_get("requests_per_minute")
.map_err(|_| invalid("invalid requests per minute"))?,
)?,
last_used_at: row
.try_get("last_used_at")
.map_err(|_| invalid("invalid last used at"))?,
created_at: row
.try_get("created_at")
.map_err(|_| invalid("invalid created at"))?,
updated_at: row
.try_get("updated_at")
.map_err(|_| invalid("invalid updated at"))?,
})
}
@@ -14,6 +14,7 @@ use gateway_core::engine::{
UpstreamSendState as CoreUpstreamSendState,
};
use gateway_core::error::{StoreError as CoreStoreError, StoreErrorKind as CoreStoreErrorKind};
use gateway_core::routing::{AccountRoutingScopeKind, AccountRoutingSnapshot};
use crate::{
ConflictKind, DecimalAmount, StoreError, StoreResult, postgres_unavailable, require_nonempty,
@@ -83,6 +84,9 @@ pub struct NewModelRequest {
pub client_api_key_id: Option<String>,
pub client_api_key_ref: String,
pub config_revision: u64,
pub routing_scope: String,
pub routing_group_refs: Vec<String>,
pub routing_group_names_snapshot: Value,
pub protocol: String,
pub operation: String,
pub endpoint: String,
@@ -109,6 +113,11 @@ impl NewModelRequest {
require_nonempty(ENTITY, "endpoint", &self.endpoint)?;
require_nonempty(ENTITY, "client_transport", &self.client_transport)?;
require_nonempty(ENTITY, "requested_model_id", &self.requested_model_id)?;
validate_routing_snapshot(
&self.routing_scope,
&self.routing_group_refs,
&self.routing_group_names_snapshot,
)?;
if self.config_revision == 0 || self.started_at > self.deadline_at {
return Err(invalid(
"revision and deadline violate the frozen constraints",
@@ -362,13 +371,15 @@ impl ModelRequestRepository for PgExecutionStore {
sqlx::query(
"insert into model_requests (
id, client_api_key_id, client_api_key_ref, config_revision, protocol,
routing_scope, routing_group_refs, routing_group_names_snapshot,
operation, endpoint, client_transport, requested_model_id,
client_ip, user_agent, reasoning_effort,
reasoning_preset, request_kind, subagent_kind, compact,
image_generation_requested, started_at, deadline_at
) values (
$1, $2, $3, $4, $5, $6, $7, $8, $9, $10::inet, $11, $12,
$13, $14, $15, $16, $17, $18, $19
$1, $2, $3, $4, $5, $6, $7, $8,
$9, $10, $11, $12, $13::inet, $14, $15,
$16, $17, $18, $19, $20, $21, $22
)",
)
.bind(request.id)
@@ -376,6 +387,9 @@ impl ModelRequestRepository for PgExecutionStore {
.bind(request.client_api_key_ref)
.bind(to_i64(request.config_revision, "config_revision")?)
.bind(request.protocol)
.bind(request.routing_scope)
.bind(request.routing_group_refs)
.bind(sqlx::types::Json(request.routing_group_names_snapshot))
.bind(request.operation)
.bind(request.endpoint)
.bind(request.client_transport)
@@ -409,6 +423,7 @@ impl ModelRequestRepository for PgExecutionStore {
sqlx::query(
"insert into model_requests (
id, client_api_key_id, client_api_key_ref, config_revision, protocol,
routing_scope, routing_group_refs, routing_group_names_snapshot,
operation, endpoint, client_transport, requested_model_id,
client_ip, user_agent, reasoning_effort,
reasoning_preset, request_kind, subagent_kind, compact,
@@ -419,13 +434,14 @@ impl ModelRequestRepository for PgExecutionStore {
upstream_model_id, upstream_transport, http_version,
attempt_count, upstream_send_state
) values (
$1, $2, $3, $4, $5, $6, $7, $8, $9, $10::inet, $11, $12,
$13, $14, $15, $16, $17, $18, $19,
$20, $21, $22,
(select name from provider_accounts where id = $21),
(select email from provider_accounts where id = $21),
(select authentication_kind from provider_accounts where id = $21),
$23, $24, $25, 1, 'not_sent'
$1, $2, $3, $4, $5, $6, $7, $8,
$9, $10, $11, $12, $13::inet, $14, $15,
$16, $17, $18, $19, $20, $21, $22,
$23, $24, $25,
(select name from provider_accounts where id = $24),
(select email from provider_accounts where id = $24),
(select authentication_kind from provider_accounts where id = $24),
$26, $27, $28, 1, 'not_sent'
)",
)
.bind(request.id)
@@ -433,6 +449,9 @@ impl ModelRequestRepository for PgExecutionStore {
.bind(request.client_api_key_ref)
.bind(to_i64(request.config_revision, "config_revision")?)
.bind(request.protocol)
.bind(request.routing_scope)
.bind(request.routing_group_refs)
.bind(sqlx::types::Json(request.routing_group_names_snapshot))
.bind(request.operation)
.bind(request.endpoint)
.bind(request.client_transport)
@@ -1045,6 +1064,8 @@ fn core_store_error(error: StoreError) -> CoreStoreError {
}
fn new_model_request_row(request: CoreNewModelRequest) -> NewModelRequest {
let (routing_scope, routing_group_refs, routing_group_names_snapshot) =
routing_snapshot_row(&request.routing);
NewModelRequest {
id: request.id.as_str().to_owned(),
client_api_key_id: request
@@ -1053,6 +1074,9 @@ fn new_model_request_row(request: CoreNewModelRequest) -> NewModelRequest {
.map(|id| id.as_str().to_owned()),
client_api_key_ref: request.client_api_key_ref.as_str().to_owned(),
config_revision: request.config_revision.get(),
routing_scope,
routing_group_refs,
routing_group_names_snapshot,
protocol: request.protocol,
operation: request.operation.as_str().to_owned(),
endpoint: request.endpoint,
@@ -1071,6 +1095,48 @@ fn new_model_request_row(request: CoreNewModelRequest) -> NewModelRequest {
}
}
fn routing_snapshot_row(snapshot: &AccountRoutingSnapshot) -> (String, Vec<String>, Value) {
let groups = snapshot.groups_snapshot();
(
snapshot.kind().as_str().to_owned(),
groups
.iter()
.map(|group| group.id().as_str().to_owned())
.collect(),
Value::Array(
groups
.iter()
.map(|group| Value::String(group.name().to_owned()))
.collect(),
),
)
}
fn validate_routing_snapshot(scope: &str, refs: &[String], names: &Value) -> StoreResult<()> {
let Some(names) = names.as_array() else {
return Err(invalid("routing group names snapshot must be an array"));
};
let valid = match scope {
value if value == AccountRoutingScopeKind::All.as_str() => {
refs.is_empty() && names.is_empty()
}
value if value == AccountRoutingScopeKind::Groups.as_str() => {
!refs.is_empty()
&& refs.len() == names.len()
&& refs.iter().all(|id| !id.is_empty())
&& names
.iter()
.all(|name| name.as_str().is_some_and(|name| !name.is_empty()))
}
_ => false,
};
if valid {
Ok(())
} else {
Err(invalid("invalid routing scope snapshot"))
}
}
fn attempt_start_row(attempt: CoreAttemptRecord) -> ModelRequestAttemptStart {
ModelRequestAttemptStart {
model_request_id: attempt.request_id.as_str().to_owned(),
@@ -7,6 +7,7 @@ use crate::{
Revision, StoreBackend, StoreError, StorePoolConfig, StoreResult, postgres_unavailable,
};
mod account_groups;
mod admin_security_audit;
mod admission_recovery;
mod backup;
@@ -20,6 +21,7 @@ mod retention;
mod runtime_settings;
mod snapshot;
pub use account_groups::*;
pub use admin_security_audit::*;
pub use admission_recovery::*;
pub use backup::*;
@@ -251,7 +253,7 @@ impl PgControlPlaneRepository {
async fn apply_targeted_mutation(
&self,
mutation: ControlPlaneMutation,
audit: AdminAuditEvent,
mut audit: AdminAuditEvent,
) -> StoreResult<Revision> {
let mut transaction = self
.pool
@@ -262,9 +264,34 @@ impl PgControlPlaneRepository {
let revision = bump_config_revision_in_transaction(&mut transaction).await?;
match mutation {
ControlPlaneMutation::CreateClientApiKey(key) => {
audit.changed_fields.push(if key.group_ids.is_empty() {
"routing_scope:all".to_owned()
} else {
"routing_scope:groups".to_owned()
});
insert_client_api_key_in_transaction(&mut transaction, &key).await?;
}
ControlPlaneMutation::UpdateClientApiKey(key) => {
let previously_restricted = sqlx::query_scalar::<_, bool>(
"select exists(
select 1 from client_api_key_groups where client_api_key_id = $1
)",
)
.bind(&key.id)
.fetch_one(&mut *transaction)
.await
.map_err(|_| {
postgres_unavailable("load client API key routing scope for audit")
})?;
if previously_restricted && key.group_ids.is_empty() {
audit
.changed_fields
.push("routing_scope:groups->all".to_owned());
} else if !previously_restricted && !key.group_ids.is_empty() {
audit
.changed_fields
.push("routing_scope:all->groups".to_owned());
}
update_client_api_key_in_transaction(&mut transaction, &key).await?;
}
ControlPlaneMutation::SetClientApiKeyEnabled { id, enabled } => {
@@ -1,12 +1,15 @@
//! Pg 观测 adapter:实现 `ObservabilityRepository` 与 `AdminObservabilityStore`。
use std::sync::Arc;
use std::time::SystemTime;
use gateway_core::provider_ports::ProviderCooldownPort;
use std::sync::Arc;
use super::*;
use crate::postgres::{
PgProviderAccountRepository, ProviderAccountRepository, ProviderAccountSummary,
account_status_projection, load_rate_limited_until,
};
use crate::redis::{CredentialLeaseRepository as _, RedisCredentialLeaseRepository};
#[derive(Clone)]
@@ -21,73 +24,48 @@ impl PgObservabilityRepository {
Self { pool, cooldowns }
}
/// 账号池指标:SQL 聚合后,仅用 Redis 429 冷却重分类仍可调度的 ready 账号。
/// 冷却中从 `active` 或 `quota_exhausted` 移到 `rate_limited`;冷却不可用时保留 SQL 分类。
async fn provider_account_metrics_with_cooldowns(
/// 从账号事实和当前冷却一次性派生 Dashboard 五态;不在 SQL 中复制状态机。
async fn account_status_snapshot(
&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);
) -> StoreResult<(ProviderAccountMetrics, Vec<ProviderAccountSummary>)> {
let accounts = PgProviderAccountRepository::new(self.pool.clone())
.list_provider_accounts(None, true)
.await?;
let now = observed_at.into();
let rate_limited_until =
load_rate_limited_until(self.cooldowns.as_deref(), &accounts, now).await;
let mut metrics = ProviderAccountMetrics {
total: u64::try_from(accounts.len()).unwrap_or(u64::MAX),
..ProviderAccountMetrics::default()
};
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;
let mut normal_accounts = Vec::new();
for account in &accounts {
let projection = account_status_projection(
account,
now,
rate_limited_until.get(&account.id).copied(),
);
match projection.status {
gateway_core::engine::credential::AccountStatus::Normal => {
metrics.normal = metrics.normal.saturating_add(1);
normal_accounts.push(account.clone());
}
quota_exhausted -= 1;
} else {
if active == 0 {
continue;
gateway_core::engine::credential::AccountStatus::QuotaExhausted => {
metrics.quota_exhausted = metrics.quota_exhausted.saturating_add(1);
}
gateway_core::engine::credential::AccountStatus::RateLimited => {
metrics.rate_limited = metrics.rate_limited.saturating_add(1);
}
gateway_core::engine::credential::AccountStatus::Disabled => {
metrics.disabled = metrics.disabled.saturating_add(1);
}
gateway_core::engine::credential::AccountStatus::Error => {
metrics.error = metrics.error.saturating_add(1);
}
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,
})
Ok((metrics, normal_accounts))
}
}
@@ -120,6 +98,7 @@ impl ObservabilityRepository for PgObservabilityRepository {
async fn dashboard_summary(
&self,
range: ObservabilityRange,
observed_at: DateTime<Utc>,
) -> StoreResult<DashboardObservation> {
let filter = UsageRecordFilter::default();
let account_usage_range = ObservabilityRange::new(
@@ -145,9 +124,7 @@ impl ObservabilityRepository for PgObservabilityRepository {
attempt_metrics(&self.pool, range, &filter),
dashboard_totals(&self.pool),
)?;
let provider_accounts = self
.provider_account_metrics_with_cooldowns(range.end)
.await?;
let (provider_accounts, _) = self.account_status_snapshot(observed_at).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),
@@ -239,10 +216,11 @@ impl AdminObservabilityStore for PgAdminObservabilityStore {
async fn dashboard_summary(
&self,
range: admin_observability::TimeRange,
observed_at: DateTime<Utc>,
) -> AdminStoreResult<admin_observability::DashboardObservation> {
let observation = self
.repository
.dashboard_summary(store_range(range)?)
.dashboard_summary(store_range(range)?, observed_at)
.await
.map_err(observability_error)?;
admin_dashboard_observation(observation)
@@ -252,28 +230,52 @@ impl AdminObservabilityStore for PgAdminObservabilityStore {
&self,
observed_at: DateTime<Utc>,
) -> AdminStoreResult<Option<admin_observability::DashboardRuntimeSlots>> {
let active_account_ids = active_provider_account_ids(&self.repository.pool, observed_at)
let (_, normal_accounts) = self
.repository
.account_status_snapshot(observed_at)
.await
.map_err(observability_error)?;
let active_accounts = u64::try_from(active_account_ids.len())
.map_err(|_| observability_error(invalid("active account count overflows u64")))?;
if active_account_ids.is_empty() {
let inherited_accounts = u64::try_from(
normal_accounts
.iter()
.filter(|account| account.concurrency_limit.is_none())
.count(),
)
.map_err(|_| observability_error(invalid("inherited account count overflows u64")))?;
let overridden_slots = normal_accounts.iter().fold(0_u64, |total, account| {
total.saturating_add(
account
.concurrency_limit
.map_or(0, |limit| u64::from(limit.get())),
)
});
if normal_accounts.is_empty() {
return Ok(Some(admin_observability::DashboardRuntimeSlots {
active_accounts,
inherited_accounts,
overridden_slots,
used_slots: Some(0),
}));
}
let normal_account_ids = normal_accounts
.iter()
.map(|account| account.id.clone())
.collect::<Vec<_>>();
let Some(runtime_signals) = &self.runtime_signals else {
return Ok(None);
return Ok(Some(admin_observability::DashboardRuntimeSlots {
inherited_accounts,
overridden_slots,
used_slots: None,
}));
};
let signals = match runtime_signals
.credential_runtime_signals(&active_account_ids)
.credential_runtime_signals(&normal_account_ids)
.await
{
Ok(signals) => signals,
Err(_) => {
return Ok(Some(admin_observability::DashboardRuntimeSlots {
active_accounts,
inherited_accounts,
overridden_slots,
used_slots: None,
}));
}
@@ -282,7 +284,8 @@ impl AdminObservabilityStore for PgAdminObservabilityStore {
total.saturating_add(u64::from(signal.in_flight))
});
Ok(Some(admin_observability::DashboardRuntimeSlots {
active_accounts,
inherited_accounts,
overridden_slots,
used_slots: Some(used_slots),
}))
}
@@ -274,16 +274,11 @@ pub(crate) const fn admin_account_pool_metrics(
) -> admin_observability::AccountPoolMetrics {
admin_observability::AccountPoolMetrics {
total: metrics.total,
enabled: metrics.enabled,
unavailable: metrics.unavailable,
active: metrics.active,
rate_limited: metrics.rate_limited,
expired: metrics.expired,
invalid: metrics.invalid,
normal: metrics.normal,
quota_exhausted: metrics.quota_exhausted,
refreshing: None,
rate_limited: metrics.rate_limited,
disabled: metrics.disabled,
banned: metrics.banned,
error: metrics.error,
}
}
@@ -297,8 +292,6 @@ pub(crate) fn admin_dashboard_account_usage(
name: usage.name,
email: usage.email,
plan_type: usage.plan_type,
enabled: usage.enabled,
availability: usage.availability,
request_count: usage.request_count,
success_count: usage.success_count,
input_tokens: usage.input_tokens,
@@ -443,6 +436,9 @@ pub(crate) fn admin_usage_record(
id: record.id,
client_api_key_ref: record.client_api_key_ref,
config_revision: record.config_revision,
routing_scope: record.routing_scope,
routing_group_refs: record.routing_group_refs,
routing_group_names_snapshot: record.routing_group_names_snapshot,
protocol: record.protocol,
operation: record.operation,
endpoint: record.endpoint,
@@ -662,6 +658,13 @@ pub(crate) fn usage_record_from_row(row: &sqlx::postgres::PgRow) -> StoreResult<
id: get(row, "id")?,
client_api_key_ref: get(row, "client_api_key_ref")?,
config_revision: unsigned(row, "config_revision")?,
routing_scope: get(row, "routing_scope")?,
routing_group_refs: get(row, "routing_group_refs")?,
routing_group_names_snapshot: get::<sqlx::types::Json<Vec<String>>>(
row,
"routing_group_names_snapshot",
)?
.0,
protocol: get(row, "protocol")?,
operation: get(row, "operation")?,
endpoint: get(row, "endpoint")?,
@@ -363,15 +363,11 @@ pub struct CalculatedUsageBillingFact {
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ProviderAccountMetrics {
pub total: u64,
pub enabled: u64,
pub unavailable: u64,
pub active: u64,
pub rate_limited: u64,
pub expired: u64,
pub invalid: u64,
pub normal: u64,
pub quota_exhausted: u64,
pub rate_limited: u64,
pub disabled: u64,
pub banned: u64,
pub error: u64,
}
#[derive(Debug, Clone, PartialEq, Eq)]
@@ -388,8 +384,6 @@ pub struct ProviderAccountUsageObservation {
pub name: String,
pub email: Option<String>,
pub plan_type: Option<String>,
pub enabled: bool,
pub availability: String,
pub request_count: u64,
pub success_count: u64,
pub input_tokens: Option<u64>,
@@ -503,6 +497,9 @@ pub struct UsageRecord {
pub id: String,
pub client_api_key_ref: String,
pub config_revision: u64,
pub routing_scope: String,
pub routing_group_refs: Vec<String>,
pub routing_group_names_snapshot: Vec<String>,
pub protocol: String,
pub operation: String,
pub endpoint: String,
@@ -689,6 +686,7 @@ pub trait ObservabilityRepository: Send + Sync {
async fn dashboard_summary(
&self,
range: ObservabilityRange,
observed_at: DateTime<Utc>,
) -> StoreResult<DashboardObservation>;
async fn dashboard_trend(
&self,
@@ -15,7 +15,7 @@ pub(crate) async fn provider_account_usage(
let mut statement = QueryBuilder::<Postgres>::new(
"select pa.id, pa.provider_kind, pa.authentication_kind, pa.name, pa.email,
pa.plan_type, pa.enabled, pa.availability,
pa.plan_type,
count(mr.id)::bigint as request_count,
count(mr.id) filter (where mr.outcome = 'succeeded')::bigint as success_count,
sum(mr.input_tokens)::bigint as input_tokens,
@@ -52,7 +52,7 @@ pub(crate) async fn provider_account_usage(
}
statement.push(
" group by pa.id, pa.provider_kind, pa.authentication_kind, pa.name, pa.email,
pa.plan_type, pa.enabled, pa.availability
pa.plan_type
order by max(mr.started_at) desc nulls last, pa.name, pa.id limit ",
);
statement.push_bind(i64::from(query.limit));
@@ -313,8 +313,6 @@ pub(crate) fn provider_account_from_row(
name: get(row, "name")?,
email: get(row, "email")?,
plan_type: get(row, "plan_type")?,
enabled: get(row, "enabled")?,
availability: get(row, "availability")?,
request_count: unsigned(row, "request_count")?,
success_count: unsigned(row, "success_count")?,
input_tokens: optional_unsigned(row, "input_tokens")?,
@@ -369,118 +369,6 @@ pub(crate) async fn request_costs(
.collect()
}
pub(crate) async fn provider_account_metrics(
pool: &PgPool,
observed_at: DateTime<Utc>,
) -> StoreResult<ProviderAccountMetrics> {
let row = sqlx::query(
"with normalized as (
select enabled, availability, quota_limit_reached,
case
when not enabled then 'disabled'
when availability = 'banned' then 'banned'
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 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
from normalized",
)
.bind(observed_at)
.fetch_one(pool)
.await
.map_err(|_| postgres_unavailable("load provider account metrics"))?;
Ok(ProviderAccountMetrics {
total: unsigned(&row, "total")?,
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")?,
})
}
pub(crate) async fn active_provider_account_ids(
pool: &PgPool,
observed_at: DateTime<Utc>,
) -> StoreResult<Vec<String>> {
sqlx::query(
"with normalized as (
select id,
case
when not enabled then 'disabled'
when availability = 'banned' then 'banned'
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'
else 'active'
end as status
from provider_accounts
)
select id
from normalized
where status = 'active'
order by id",
)
.bind(observed_at)
.fetch_all(pool)
.await
.map_err(|_| postgres_unavailable("load active provider account IDs"))?
.iter()
.map(|row| get(row, "id"))
.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,
@@ -103,7 +103,9 @@ pub(crate) fn push_completed_usage_fact_filter(query: &mut QueryBuilder<Postgres
}
pub(crate) const USAGE_RECORD_SELECT: &str =
"select mr.id, mr.client_api_key_ref, mr.config_revision, mr.protocol, mr.operation,
"select mr.id, mr.client_api_key_ref, mr.config_revision,
mr.routing_scope, mr.routing_group_refs, mr.routing_group_names_snapshot,
mr.protocol, mr.operation,
mr.endpoint, mr.client_transport, mr.requested_model_id,
mr.provider_kind, mr.provider_account_ref,
mr.provider_account_name_snapshot as provider_account_name,
@@ -2,7 +2,6 @@
use std::collections::BTreeMap;
use std::sync::Arc;
use std::time::SystemTime;
use gateway_core::provider_ports::ProviderCooldownPort;
@@ -32,34 +31,6 @@ impl PgAdminAccountStore {
}
}
/// 批量读取 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,
@@ -260,6 +231,111 @@ impl PgAdminAccountStore {
credential_revision: Some(admin_revision(rotation.credential_revision)?),
})
}
async fn account_groups_by_account(
&self,
account_ids: &[String],
) -> AdminStoreResult<BTreeMap<String, Vec<AccountGroupRef>>> {
if account_ids.is_empty() {
return Ok(BTreeMap::new());
}
let rows = sqlx::query_as::<_, (String, String, String, String, bool)>(
"select m.provider_account_id, g.id, g.name, g.color, g.enabled
from account_group_accounts m
join account_groups g on g.id = m.account_group_id
where m.provider_account_id = any($1::text[])
order by m.provider_account_id, g.id",
)
.bind(account_ids)
.fetch_all(&self.pool)
.await
.map_err(|_| {
admin_store_error(
ENTITY,
postgres_unavailable("load account group references"),
)
})?;
let mut groups = BTreeMap::<String, Vec<AccountGroupRef>>::new();
for (account_id, group_id, name, color, enabled) in rows {
let id = AccountGroupId::new(group_id).map_err(|_| {
AdminStoreError::new(
AdminStoreErrorKind::Invalid,
ENTITY,
"persisted account group ID is invalid",
)
})?;
groups.entry(account_id).or_default().push(AccountGroupRef {
id,
name,
color: gateway_admin::model::account_groups::AccountGroupColor::parse(&color)
.ok_or_else(|| {
AdminStoreError::new(
AdminStoreErrorKind::Invalid,
ENTITY,
"persisted account group color is invalid",
)
})?,
enabled,
});
}
Ok(groups)
}
async fn account_ids_matching_group(
&self,
filter: Option<&AccountGroupFilter>,
) -> AdminStoreResult<Option<BTreeSet<String>>> {
let Some(filter) = filter else {
return Ok(None);
};
let rows = match filter {
AccountGroupFilter::Group(group_id) => {
let exists = sqlx::query_scalar::<_, bool>(
"select exists(select 1 from account_groups where id = $1)",
)
.bind(group_id.as_str())
.fetch_one(&self.pool)
.await
.map_err(|_| {
admin_store_error(
ENTITY,
postgres_unavailable("validate account group filter"),
)
})?;
if !exists {
return Err(admin_store_error(
ENTITY,
StoreError::NotFound {
entity: "account group",
id: group_id.as_str().to_owned(),
},
));
}
sqlx::query_scalar::<_, String>(
"select provider_account_id from account_group_accounts
where account_group_id = $1",
)
.bind(group_id.as_str())
.fetch_all(&self.pool)
.await
}
AccountGroupFilter::Ungrouped => {
sqlx::query_scalar::<_, String>(
"select a.id from provider_accounts a
where not exists (
select 1 from account_group_accounts m
where m.provider_account_id = a.id
)",
)
.fetch_all(&self.pool)
.await
}
}
.map_err(|_| {
admin_store_error(ENTITY, postgres_unavailable("load account group filter"))
})?;
Ok(Some(rows.into_iter().collect()))
}
}
#[async_trait]
@@ -278,24 +354,26 @@ impl AccountStore for PgAdminAccountStore {
)
.map_err(|error| admin_store_error(ENTITY, error))?;
let now = Utc::now();
// 批量读取 Redis 429 冷却,供列表状态派生与筛选使用。
let rate_limited_until = self
.load_rate_limited_until(&accounts)
.await
.unwrap_or_default();
let rate_limited_until =
load_rate_limited_until(self.cooldowns.as_deref(), &accounts, now.into()).await;
let summary = admin_account_summary(&accounts, now, &rate_limited_until);
let group_matches = self
.account_ids_matching_group(query.group_filter.as_ref())
.await?;
let mut items = accounts
.into_iter()
.filter_map(|account| {
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,
status,
usage: None,
},
)
let projection = account_status_projection(&account, now.into(), until);
(account_matches_admin_query(&account, projection.status, &query)
&& group_matches
.as_ref()
.is_none_or(|matches| matches.contains(&account.id)))
.then_some(AdminAccountListItem {
account,
projection,
usage: None,
})
})
.collect::<Vec<_>>();
@@ -330,7 +408,23 @@ impl AccountStore for PgAdminAccountStore {
.into_iter()
.skip(offset)
.take(page_size)
.map(|item| admin_account_record(item.account))
.collect::<Vec<_>>();
let item_ids = items
.iter()
.map(|item| item.account.id.clone())
.collect::<Vec<_>>();
let mut groups_by_account = self.account_groups_by_account(&item_ids).await?;
let items = items
.into_iter()
.map(|item| {
let account_id = item.account.id.clone();
let mut account = admin_account_record(item.account)?;
account.groups = groups_by_account.remove(&account_id).unwrap_or_default();
Ok(AccountPageItem {
account,
projection: item.projection,
})
})
.collect::<AdminStoreResult<Vec<_>>>()?;
Ok(AccountPage {
config_revision: admin_revision(control_plane.settings.config_revision)?,
@@ -340,13 +434,35 @@ impl AccountStore for PgAdminAccountStore {
})
}
async fn load_account(&self, account_id: &str) -> AdminStoreResult<Option<AccountRecord>> {
self.accounts
async fn load_account(&self, account_id: &str) -> AdminStoreResult<Option<AccountPageItem>> {
let record = self
.accounts
.load_provider_account(account_id)
.await
.map_err(|error| admin_store_error(ENTITY, error))?
.map(|record| admin_account_record(record.summary))
.transpose()
.map_err(|error| admin_store_error(ENTITY, error))?;
let Some(record) = record else {
return Ok(None);
};
let now = Utc::now();
let rate_limited_until = load_rate_limited_until(
self.cooldowns.as_deref(),
std::slice::from_ref(&record.summary),
now.into(),
)
.await
.get(account_id)
.copied();
let projection = account_status_projection(&record.summary, now.into(), rate_limited_until);
let account_id = record.summary.id.clone();
let mut groups = self
.account_groups_by_account(std::slice::from_ref(&account_id))
.await?;
let mut account = admin_account_record(record.summary)?;
account.groups = groups.remove(&account_id).unwrap_or_default();
Ok(Some(AccountPageItem {
account,
projection,
}))
}
async fn load_account_usage(
@@ -386,9 +502,9 @@ impl AccountStore for PgAdminAccountStore {
.filter(|account| account.provider_kind == provider_kind.as_str())
.filter(|account| {
query
.availability
.credential_state
.as_ref()
.is_none_or(|expected| expected.matches(account.availability))
.is_none_or(|expected| expected.matches(account.credential_state))
})
.filter(|account| {
query
@@ -583,29 +699,134 @@ impl AccountStore for PgAdminAccountStore {
.await
}
async fn set_account_enabled(
async fn update_account(
&self,
command: SetAccountEnabled,
command: UpdateAccount,
context: &MutationContext,
) -> AdminStoreResult<AdminRevision> {
let scope = self.required_scope(&command.account_id).await?;
let action = if command.enabled { "enable" } else { "disable" };
self.accounts
.set_provider_account_enabled_admin(SetProviderAccountEnabled {
scope,
account_id: command.account_id.clone(),
) -> AdminStoreResult<AccountUpdateResult> {
let account_id = CoreProviderAccountId::new(command.account_id.clone()).map_err(|_| {
AdminStoreError::new(
AdminStoreErrorKind::Invalid,
ENTITY,
"invalid provider account ID",
)
})?;
let config_revision = self
.accounts
.batch_update_provider_accounts_admin(BatchUpdateProviderAccountsAdmin {
account_ids: vec![command.account_id.clone()],
enabled: command.enabled,
concurrency_limit: command.concurrency_limit,
weight: command.weight,
group_ids: command.group_ids,
audit: mutation_audit(
context,
action,
"update",
"provider_account",
&command.account_id,
vec!["enabled".to_owned()],
vec![
"enabled".to_owned(),
"concurrency_limit".to_owned(),
"weight".to_owned(),
"groups".to_owned(),
],
),
})
.await
.map_err(|error| admin_store_error(ENTITY, error))
.and_then(admin_revision)
.and_then(admin_revision)?;
Ok(AccountUpdateResult {
config_revision,
account_id,
})
}
async fn recover_account(
&self,
account_id: &CoreProviderAccountId,
context: &MutationContext,
) -> AdminStoreResult<AccountUpdateResult> {
if let Some(cooldowns) = self.cooldowns.as_deref() {
cooldowns.clear_all(account_id).await.map_err(|_| {
AdminStoreError::new(
AdminStoreErrorKind::Unavailable,
ENTITY,
"provider account cooldown cleanup failed",
)
})?;
}
let config_revision = self
.accounts
.recover_provider_account_admin(RecoverProviderAccount {
account_id: account_id.as_str().to_owned(),
audit: mutation_audit(
context,
"recover",
"provider_account",
account_id.as_str(),
vec!["status".to_owned(), "quota".to_owned()],
),
})
.await
.map_err(|error| admin_store_error(ENTITY, error))
.and_then(admin_revision)?;
Ok(AccountUpdateResult {
config_revision,
account_id: account_id.clone(),
})
}
async fn batch_update_accounts(
&self,
command: BatchUpdateAccounts,
context: &MutationContext,
) -> AdminStoreResult<AccountsUpdateResult> {
let account_ids = command
.account_ids
.iter()
.map(|id| {
CoreProviderAccountId::new(id.clone()).map_err(|_| {
AdminStoreError::new(
AdminStoreErrorKind::Invalid,
ENTITY,
"invalid provider account ID",
)
})
})
.collect::<AdminStoreResult<Vec<_>>>()?;
let audit_target = if command.account_ids.len() == 1 {
command.account_ids[0].clone()
} else {
"provider_accounts".to_owned()
};
let config_revision = self
.accounts
.batch_update_provider_accounts_admin(BatchUpdateProviderAccountsAdmin {
account_ids: command.account_ids,
enabled: command.enabled,
concurrency_limit: command.concurrency_limit,
weight: command.weight,
group_ids: command.group_ids,
audit: mutation_audit(
context,
"batch_update",
"provider_account",
&audit_target,
vec![
"enabled".to_owned(),
"concurrency_limit".to_owned(),
"weight".to_owned(),
"groups".to_owned(),
],
),
})
.await
.map_err(|error| admin_store_error(ENTITY, error))
.and_then(admin_revision)?;
Ok(AccountsUpdateResult {
config_revision,
account_ids,
})
}
async fn delete_accounts(
@@ -31,8 +31,10 @@ impl ProviderAccountStore for PgProviderAccountRepository {
.map(DateTime::<Utc>::from),
next_refresh_at: account.account.next_refresh_at().map(DateTime::<Utc>::from),
enabled: account.account.enabled(),
availability: account.account.availability(),
availability_observed_at: Utc::now(),
concurrency_limit: account.account.concurrency_limit(),
weight: account.account.weight(),
credential_state: account.account.credential_state(),
credential_observed_at: Utc::now(),
})
.await
.map_err(core_store_error)
@@ -164,9 +166,9 @@ impl ProviderAccountStore for PgProviderAccountRepository {
.collect::<Vec<_>>();
let rows = sqlx::query(
"select id, credential_revision, provider_quota_json, quota_observed_at, \
quota_limit_reached \
quota_access_state, quota_evidence, quota_access_observed_at, quota_reset_at \
from provider_accounts \
where id = any($1) and provider_quota_json is not null",
where id = any($1) and quota_observed_at is not null",
)
.bind(account_ids)
.fetch_all(&self.pool)
@@ -184,24 +186,49 @@ impl ProviderAccountStore for PgProviderAccountRepository {
.and_then(|value| CoreCredentialRevision::new(value).ok())
.ok_or_else(|| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?;
let quota = row
.try_get::<serde_json::Value, _>("provider_quota_json")
.ok()
.try_get::<Option<serde_json::Value>, _>("provider_quota_json")
.map_err(|_| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?
.and_then(|value| value.as_object().cloned())
.map(OpaqueProviderData::new)
.ok_or_else(|| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?;
.unwrap_or_else(|| OpaqueProviderData::new(serde_json::Map::new()));
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")
let access = row
.try_get::<String, _>("quota_access_state")
.ok()
.and_then(|value| QuotaAccessState::parse(&value))
.ok_or_else(|| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?;
let evidence = row
.try_get::<Option<String>, _>("quota_evidence")
.map_err(|_| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?;
let evidence = match evidence {
Some(value) => Some(
QuotaEvidence::parse(&value)
.ok_or_else(|| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?,
),
None => None,
};
let reset_at = row
.try_get::<Option<DateTime<Utc>>, _>("quota_reset_at")
.map_err(|_| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?;
let access_observed_at = row
.try_get::<Option<DateTime<Utc>>, _>("quota_access_observed_at")
.map_err(|_| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?;
let state = QuotaState::from_persisted(
access,
evidence,
access_observed_at.map(Into::into),
reset_at.map(Into::into),
)
.ok_or_else(|| 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),
quota,
observed_at: observed_at.into(),
state,
})
})
.collect()
@@ -211,24 +238,38 @@ impl ProviderAccountStore for PgProviderAccountRepository {
&self,
observation: QuotaObservation,
) -> Result<QuotaWriteOutcome, CoreStoreError> {
let quota = observation
.quota
.map(|value| {
JsonObject::try_from_value(
"provider_quota_json",
serde_json::Value::Object(value.into_inner()),
QUOTA_MAX_BYTES,
)
})
.transpose()
.map_err(core_store_error)?;
let quota = JsonObject::try_from_value(
"provider_quota_json",
serde_json::Value::Object(observation.quota.into_inner()),
QUOTA_MAX_BYTES,
)
.map_err(core_store_error)?;
let updated = self
.compare_and_swap_provider_quota(
observation.account_id.as_str(),
Revision::new(observation.expected_revision.get()).map_err(core_store_error)?,
quota,
observation.observed_at.map(DateTime::<Utc>::from),
observation.limit_reached,
DateTime::<Utc>::from(observation.observed_at),
observation.state,
)
.await
.map_err(core_store_error)?;
Ok(if updated {
QuotaWriteOutcome::Updated
} else {
QuotaWriteOutcome::Conflict
})
}
async fn apply_quota_access(
&self,
change: QuotaAccessChange,
) -> Result<QuotaWriteOutcome, CoreStoreError> {
let updated = self
.apply_provider_quota_access(
change.account_id.as_str(),
Revision::new(change.expected_revision.get()).map_err(core_store_error)?,
change.state,
)
.await
.map_err(core_store_error)?;
@@ -245,8 +286,9 @@ impl ProviderAccountStore for PgProviderAccountRepository {
account_id: change.account_id.as_str().to_owned(),
expected_revision: Revision::new(change.expected_revision.get())
.map_err(core_store_error)?,
availability: change.availability,
availability_observed_at: DateTime::<Utc>::from(change.observed_at),
credential_state: change.credential_state,
credential_observed_at: DateTime::<Utc>::from(change.observed_at),
error_reason: change.error_reason,
message: change.message,
})
.await
@@ -1,12 +1,12 @@
//! Admin 账号列表、用量与成本的领域映射。
use std::time::SystemTime;
use super::*;
use gateway_admin::model::accounts::AccountStatusSignals;
pub(crate) struct AdminAccountListItem {
pub(crate) account: ProviderAccountSummary,
pub(crate) status: AdminAccountStatus,
pub(crate) projection: gateway_core::engine::credential::AccountStatusProjection,
pub(crate) usage: Option<ProviderAccountUsageObservation>,
}
@@ -57,24 +57,6 @@ pub(crate) fn account_matches_admin_query(
provider_matches && search_matches && status_matches
}
pub(crate) fn admin_account_status(
account: &ProviderAccountSummary,
now: DateTime<Utc>,
rate_limited_until: Option<SystemTime>,
) -> AdminAccountStatus {
// 派生规则唯一归属 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),
quota_limit_reached: account.quota_limit_reached,
rate_limited_until,
}
.derive()
}
pub(crate) fn admin_account_summary(
accounts: &[ProviderAccountSummary],
now: DateTime<Utc>,
@@ -90,7 +72,7 @@ pub(crate) fn admin_account_summary(
};
for account in accounts {
let until = rate_limited_until.get(&account.id).copied();
match admin_account_status(account, now, until) {
match account_status_projection(account, now.into(), until).status {
AdminAccountStatus::Normal => summary.normal = summary.normal.saturating_add(1),
AdminAccountStatus::QuotaExhausted => {
summary.quota_exhausted = summary.quota_exhausted.saturating_add(1);
@@ -118,7 +100,11 @@ 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 => left.status.as_str().cmp(right.status.as_str()),
AdminAccountSortField::Status => left
.projection
.status
.sort_rank()
.cmp(&right.projection.status.sort_rank()),
AdminAccountSortField::PlanType => left.account.plan_type.cmp(&right.account.plan_type),
AdminAccountSortField::Usage => left
.usage
@@ -160,6 +146,7 @@ pub(crate) fn admin_account_record(
"persisted Provider kind is invalid",
)
})?,
groups: Vec::new(),
name: summary.name,
email: summary.email,
upstream_user_id: summary.upstream_user_id,
@@ -171,10 +158,13 @@ pub(crate) fn admin_account_record(
access_token_expires_at: summary.access_token_expires_at,
next_refresh_at: summary.next_refresh_at,
enabled: summary.enabled,
availability: summary.availability,
availability_observed_at: summary.availability_observed_at,
concurrency_limit: summary.concurrency_limit,
weight: summary.weight,
credential_state: summary.credential_state,
credential_observed_at: summary.credential_observed_at,
quota: summary.quota,
last_error_reason: summary.last_error_reason,
last_error_message: summary.last_error_message,
quota_observed_at: summary.quota_observed_at,
created_at: summary.created_at,
updated_at: summary.updated_at,
})
@@ -197,8 +187,10 @@ pub(crate) fn prepared_account(
access_token_expires_at: credential.access_token_expires_at,
next_refresh_at: credential.next_refresh_at,
enabled: credential.enabled,
availability: credential.availability,
availability_observed_at: credential.availability_observed_at,
concurrency_limit: None,
weight: AccountWeight::DEFAULT,
credential_state: credential.credential_state,
credential_observed_at: credential.credential_observed_at,
})
}
@@ -4,7 +4,6 @@ use std::{
collections::{BTreeMap, BTreeSet},
fmt,
str::FromStr,
time::SystemTime,
};
use async_trait::async_trait;
@@ -12,12 +11,15 @@ use chrono::{DateTime, TimeDelta, Utc};
use gateway_admin::{
model::{
MutationContext, Revision as AdminRevision,
account_groups::AccountGroupRef,
accounts::{
AccountCost, AccountListQuery as AdminAccountListQuery, AccountModelUsage, AccountPage,
AccountRecord, AccountRequestBucket, AccountSort as AdminAccountSort,
AccountSortField as AdminAccountSortField, AccountStatus as AdminAccountStatus,
AccountSummary, AccountUsage, AccountUsageWindowQuery, AccountUsageWindowResult,
DeleteAccounts, SetAccountEnabled, SortDirection as AdminSortDirection,
AccountCost, AccountGroupFilter, AccountListQuery as AdminAccountListQuery,
AccountModelUsage, AccountPage, AccountPageItem, AccountRecord, AccountRequestBucket,
AccountSort as AdminAccountSort, AccountSortField as AdminAccountSortField,
AccountStatus as AdminAccountStatus, AccountSummary, AccountUpdateResult, AccountUsage,
AccountUsageWindowQuery, AccountUsageWindowResult, AccountsUpdateResult,
BatchUpdateAccounts, DeleteAccounts, SortDirection as AdminSortDirection,
UpdateAccount,
},
observability::{
CostCoverage as AdminCostCoverage, DecimalAmount as AdminDecimalAmount, TimeRange,
@@ -35,15 +37,16 @@ use gateway_admin::{
use sqlx::{PgPool, Postgres, Row, Transaction};
use gateway_core::engine::credential::{
AccountAvailability, AccountStateChange, CredentialCasOutcome, CredentialCasUpdate,
CredentialRevision as CoreCredentialRevision, LoadedCredential,
NewProviderAccount as CoreNewProviderAccount, OpaqueProviderData, PlaintextCredential,
ProviderAccount as CoreProviderAccount, ProviderAccountId as CoreProviderAccountId,
ProviderAccountIdentity, ProviderAccountStore,
ProviderAccountUpdate as CoreProviderAccountUpdate, QuotaObservation, QuotaWriteOutcome,
AccountConcurrencyLimit, AccountErrorReason, AccountStateChange, AccountWeight,
CredentialCasOutcome, CredentialCasUpdate, CredentialRevision as CoreCredentialRevision,
CredentialState, LoadedCredential, NewProviderAccount as CoreNewProviderAccount,
OpaqueProviderData, PlaintextCredential, ProviderAccount as CoreProviderAccount,
ProviderAccountId as CoreProviderAccountId, ProviderAccountIdentity, ProviderAccountStore,
ProviderAccountUpdate as CoreProviderAccountUpdate, QuotaAccessChange, QuotaAccessState,
QuotaEvidence, QuotaObservation, QuotaState, QuotaWriteOutcome,
};
use gateway_core::error::{StoreError as CoreStoreError, StoreErrorKind as CoreStoreErrorKind};
use gateway_core::routing::ProviderKind;
use gateway_core::routing::{AccountGroupId, ProviderKind};
use crate::{
ConflictKind, JsonObject, Revision, StoreError, StoreResult, admin_revision, admin_store_error,
@@ -63,8 +66,10 @@ mod core_adapter;
mod mapping;
mod repository;
mod rows;
mod runtime;
pub use admin_adapter::*;
pub(crate) use mapping::*;
pub use repository::*;
pub use rows::*;
pub(crate) use runtime::*;
@@ -25,9 +25,15 @@ pub trait ProviderAccountRepository: Send + Sync {
&self,
account_id: &str,
expected_revision: Revision,
quota: Option<JsonObject>,
observed_at: Option<DateTime<Utc>>,
limit_reached: Option<bool>,
quota: JsonObject,
observed_at: DateTime<Utc>,
state: QuotaState,
) -> StoreResult<bool>;
async fn apply_provider_quota_access(
&self,
account_id: &str,
expected_revision: Revision,
state: QuotaState,
) -> StoreResult<bool>;
async fn delete_provider_account(&self, id: &str) -> StoreResult<bool>;
}
@@ -50,9 +56,14 @@ pub trait ProviderAccountAdminRepository: Send + Sync {
command: RotateProviderAccount,
) -> StoreResult<ProviderAccountAdminRotation>;
async fn set_provider_account_enabled_admin(
async fn batch_update_provider_accounts_admin(
&self,
command: SetProviderAccountEnabled,
command: BatchUpdateProviderAccountsAdmin,
) -> StoreResult<Revision>;
async fn recover_provider_account_admin(
&self,
command: RecoverProviderAccount,
) -> StoreResult<Revision>;
async fn delete_provider_accounts_admin(
@@ -93,9 +104,10 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
let rows = sqlx::query(
"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_observed_at,
quota_observed_at, quota_limit_reached, last_error_message, created_at, updated_at
access_token_expires_at, next_refresh_at, enabled, concurrency_limit, weight, credential_state,
credential_observed_at, quota_access_state, quota_evidence,
quota_access_observed_at, quota_reset_at,
quota_observed_at, last_error_reason, 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",
@@ -110,21 +122,21 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
async fn insert_provider_account(&self, account: NewProviderAccount) -> StoreResult<()> {
account.validate()?;
let availability = if account.upstream_user_id.is_some() {
account.availability
let credential_state = if account.upstream_user_id.is_some() {
account.credential_state
} else {
AccountAvailability::Unknown
CredentialState::Unknown
};
sqlx::query(
"insert into provider_accounts (
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, provider_quota_json,
availability_observed_at, quota_observed_at, created_at, updated_at
concurrency_limit, weight, credential_state, provider_quota_json,
credential_observed_at, quota_access_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, now(), greatest(now(), $15)
$14, $15, $16, null, $17, null, null, now(), greatest(now(), $17)
)",
)
.bind(account.id)
@@ -140,8 +152,10 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
.bind(account.access_token_expires_at)
.bind(account.next_refresh_at)
.bind(account.enabled)
.bind(availability.as_str())
.bind(account.availability_observed_at)
.bind(account.concurrency_limit.map(|limit| i64::from(limit.get())))
.bind(i16::try_from(account.weight.get()).map_err(|_| invalid("invalid weight"))?)
.bind(credential_state.as_str())
.bind(account.credential_observed_at)
.execute(&self.pool)
.await
.map_err(|_| postgres_unavailable("insert provider account"))?;
@@ -214,27 +228,33 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
update.validate()?;
let result = sqlx::query(
"update provider_accounts
set availability = case
set credential_state = case
when enabled and upstream_user_id is not null then $3
else availability
else credential_state
end,
availability_observed_at = case
credential_observed_at = case
when enabled and upstream_user_id is not null then $4
else availability_observed_at
else credential_observed_at
end,
last_error_reason = case
when enabled and upstream_user_id is not null then
case when $3 = 'ready' then null else $5 end
else last_error_reason
end,
last_error_message = case
when enabled and upstream_user_id is not null then
case when $3 = 'ready' then null else coalesce($5, last_error_message) end
case when $3 = 'ready' then null else $6 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 <= $4)",
and (credential_observed_at is null or credential_observed_at <= $4)",
)
.bind(update.account_id)
.bind(to_i64(update.expected_revision.get())?)
.bind(update.availability.as_str())
.bind(update.availability_observed_at)
.bind(update.credential_state.as_str())
.bind(update.credential_observed_at)
.bind(update.error_reason.map(AccountErrorReason::as_str))
.bind(update.message.as_deref())
.execute(&self.pool)
.await
@@ -259,36 +279,81 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
&self,
account_id: &str,
expected_revision: Revision,
quota: Option<JsonObject>,
observed_at: Option<DateTime<Utc>>,
limit_reached: Option<bool>,
quota: JsonObject,
observed_at: DateTime<Utc>,
state: QuotaState,
) -> StoreResult<bool> {
require_nonempty(ENTITY, "account_id", account_id)?;
if quota.is_some() != observed_at.is_some() {
return Err(invalid("quota JSON and observed_at must agree"));
}
if let Some(quota) = quota.as_ref() {
validate_object_size("provider_quota_json", quota, QUOTA_MAX_BYTES)?;
}
validate_object_size("provider_quota_json", &quota, QUOTA_MAX_BYTES)?;
let access_observed_at = state.observed_at().map(DateTime::<Utc>::from);
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()))
quota_access_state = case
when $7::timestamptz is not null
and (quota_access_observed_at is null or quota_access_observed_at <= $7)
then $5 else quota_access_state end,
quota_evidence = case
when $7::timestamptz is not null
and (quota_access_observed_at is null or quota_access_observed_at <= $7)
then $6 else quota_evidence end,
quota_access_observed_at = case
when $7::timestamptz is not null
and (quota_access_observed_at is null or quota_access_observed_at <= $7)
then $7 else quota_access_observed_at end,
quota_reset_at = case
when $7::timestamptz is not null
and (quota_access_observed_at is null or quota_access_observed_at <= $7)
then $8 else quota_reset_at end,
updated_at = greatest(now(), $4)
where id = $1 and credential_revision = $2
and ($4 is null or quota_observed_at is null or quota_observed_at <= $4)",
and (quota_observed_at is null or quota_observed_at <= $4)",
)
.bind(account_id)
.bind(to_i64(expected_revision.get())?)
.bind(quota.map(|value| value.as_value()))
.bind(quota.as_value())
.bind(observed_at)
.bind(limit_reached)
.bind(state.access().as_str())
.bind(state.evidence().map(QuotaEvidence::as_str))
.bind(access_observed_at)
.bind(state.reset_at().map(DateTime::<Utc>::from))
.execute(&self.pool)
.await
.map_err(|_| postgres_unavailable("compare and swap provider quota"))?;
Ok(result.rows_affected() == 1)
}
async fn apply_provider_quota_access(
&self,
account_id: &str,
expected_revision: Revision,
state: QuotaState,
) -> StoreResult<bool> {
require_nonempty(ENTITY, "account_id", account_id)?;
let observed_at = state
.observed_at()
.map(DateTime::<Utc>::from)
.ok_or_else(|| invalid("quota access change requires observed_at"))?;
let result = sqlx::query(
"update provider_accounts
set quota_access_observed_at = $3, quota_access_state = $4,
quota_evidence = $5, quota_reset_at = $6,
updated_at = greatest(now(), $3)
where id = $1 and credential_revision = $2
and (quota_access_observed_at is null or quota_access_observed_at <= $3)",
)
.bind(account_id)
.bind(to_i64(expected_revision.get())?)
.bind(observed_at)
.bind(state.access().as_str())
.bind(state.evidence().map(QuotaEvidence::as_str))
.bind(state.reset_at().map(DateTime::<Utc>::from))
.execute(&self.pool)
.await
.map_err(|_| postgres_unavailable("apply provider quota access"))?;
Ok(result.rows_affected() == 1)
}
async fn delete_provider_account(&self, id: &str) -> StoreResult<bool> {
require_nonempty(ENTITY, "id", id)?;
let result = sqlx::query("delete from provider_accounts where id = $1 and not enabled")
@@ -413,12 +478,12 @@ impl ProviderAccountAdminRepository for PgProviderAccountRepository {
finish_admin_transaction(transaction, result, "provider account admin rotation").await
}
async fn set_provider_account_enabled_admin(
async fn batch_update_provider_accounts_admin(
&self,
command: SetProviderAccountEnabled,
command: BatchUpdateProviderAccountsAdmin,
) -> StoreResult<Revision> {
command.scope.validate()?;
require_nonempty(ENTITY, "account_id", &command.account_id)?;
validate_batch_update_account_ids(&command.account_ids)?;
validate_batch_update_group_ids(&command.group_ids)?;
let mut transaction = self
.pool
.begin()
@@ -426,11 +491,18 @@ impl ProviderAccountAdminRepository for PgProviderAccountRepository {
.map_err(|_| postgres_unavailable("begin provider account admin state change"))?;
let result = async {
let revision = bump_config_revision_in_transaction(&mut transaction).await?;
set_provider_account_enabled_in_transaction(
update_provider_accounts_scheduling_in_transaction(
&mut transaction,
&command.scope,
&command.account_id,
&command.account_ids,
command.enabled,
command.concurrency_limit,
command.weight,
)
.await?;
replace_account_group_assignments_in_transaction(
&mut transaction,
&command.account_ids,
&command.group_ids,
)
.await?;
append_admin_audit_event_in_transaction(&mut transaction, command.audit, revision)
@@ -441,6 +513,57 @@ impl ProviderAccountAdminRepository for PgProviderAccountRepository {
finish_admin_transaction(transaction, result, "provider account admin state change").await
}
async fn recover_provider_account_admin(
&self,
command: RecoverProviderAccount,
) -> StoreResult<Revision> {
validate_admin_account_ids(std::slice::from_ref(&command.account_id))?;
let mut transaction = self
.pool
.begin()
.await
.map_err(|_| postgres_unavailable("begin provider account admin recovery"))?;
let result = async {
let revision = bump_config_revision_in_transaction(&mut transaction).await?;
let recovered = sqlx::query_scalar::<_, String>(
"update provider_accounts
set enabled = true,
credential_state = 'ready',
credential_observed_at = now(),
access_token_expires_at = case
when access_token_expires_at <= now() then null
else access_token_expires_at
end,
provider_quota_json = null,
quota_observed_at = null,
quota_access_state = 'allowed',
quota_evidence = null,
quota_access_observed_at = now(),
quota_reset_at = null,
last_error_reason = null,
last_error_message = null,
updated_at = now()
where id = $1
returning id",
)
.bind(&command.account_id)
.fetch_optional(&mut *transaction)
.await
.map_err(|_| postgres_unavailable("recover provider account state"))?
.ok_or_else(|| StoreError::NotFound {
entity: ENTITY,
id: command.account_id.clone(),
})?;
append_admin_audit_event_in_transaction(&mut transaction, command.audit, revision)
.await?;
Ok((revision, recovered))
}
.await;
finish_admin_transaction(transaction, result, "provider account admin recovery")
.await
.map(|(revision, _)| revision)
}
async fn delete_provider_accounts_admin(
&self,
command: DeleteProviderAccounts,
@@ -469,26 +592,73 @@ impl ProviderAccountAdminRepository for PgProviderAccountRepository {
}
}
async fn replace_account_group_assignments_in_transaction(
transaction: &mut Transaction<'_, Postgres>,
account_ids: &[String],
group_ids: &[AccountGroupId],
) -> StoreResult<()> {
let group_ids = group_ids
.iter()
.map(|group_id| group_id.as_str().to_owned())
.collect::<Vec<_>>();
if !group_ids.is_empty() {
let known_group_count = sqlx::query_scalar::<_, i64>(
"select count(*)::bigint from account_groups where id = any($1::text[])",
)
.bind(&group_ids)
.fetch_one(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("validate account group assignment"))?;
if usize::try_from(known_group_count).ok() != Some(group_ids.len()) {
return Err(StoreError::NotFound {
entity: "account group",
id: "one or more group IDs".to_owned(),
});
}
}
sqlx::query("delete from account_group_accounts where provider_account_id = any($1::text[])")
.bind(account_ids)
.execute(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("clear account group assignments"))?;
if group_ids.is_empty() {
return Ok(());
}
sqlx::query(
"insert into account_group_accounts
(account_group_id, provider_account_id, created_at)
select group_id, account_id, now()
from unnest($1::text[]) group_id
cross join unnest($2::text[]) account_id",
)
.bind(group_ids)
.bind(account_ids)
.execute(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("assign accounts to groups"))?;
Ok(())
}
pub(crate) async fn upsert_provider_account_in_transaction(
transaction: &mut Transaction<'_, Postgres>,
account: &NewProviderAccount,
) -> StoreResult<String> {
account.validate()?;
let availability = if account.upstream_user_id.is_some() {
account.availability
let credential_state = if account.upstream_user_id.is_some() {
account.credential_state
} else {
AccountAvailability::Unknown
CredentialState::Unknown
};
let imported_id = sqlx::query_scalar::<_, String>(
"insert into provider_accounts (
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, provider_quota_json,
availability_observed_at, quota_observed_at, created_at, updated_at
concurrency_limit, weight, credential_state, provider_quota_json,
credential_observed_at, quota_access_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, now(), greatest(now(), $15)
$14, $15, $16, null, $17, null, null, now(), greatest(now(), $17)
)
on conflict (
provider_kind,
@@ -505,11 +675,17 @@ pub(crate) async fn upsert_provider_account_in_transaction(
access_token_expires_at = excluded.access_token_expires_at,
next_refresh_at = excluded.next_refresh_at,
enabled = excluded.enabled,
availability = excluded.availability,
credential_state = excluded.credential_state,
provider_quota_json = null,
availability_observed_at = excluded.availability_observed_at,
quota_access_state = 'unknown',
quota_evidence = null,
quota_access_observed_at = null,
quota_reset_at = null,
credential_observed_at = excluded.credential_observed_at,
quota_observed_at = null,
updated_at = greatest(now(), excluded.availability_observed_at)
last_error_reason = null,
last_error_message = null,
updated_at = greatest(now(), excluded.credential_observed_at)
returning id",
)
.bind(&account.id)
@@ -525,8 +701,10 @@ pub(crate) async fn upsert_provider_account_in_transaction(
.bind(account.access_token_expires_at)
.bind(account.next_refresh_at)
.bind(account.enabled)
.bind(availability.as_str())
.bind(account.availability_observed_at)
.bind(account.concurrency_limit.map(|limit| i64::from(limit.get())))
.bind(i16::try_from(account.weight.get()).map_err(|_| invalid("invalid weight"))?)
.bind(credential_state.as_str())
.bind(account.credential_observed_at)
.fetch_optional(&mut **transaction)
.await
.map_err(|error| {
@@ -574,17 +752,17 @@ pub(crate) async fn rotate_provider_account_in_transaction(
next_refresh_at = $10,
upstream_user_id = case when $11::boolean then $12::text else upstream_user_id end,
upstream_account_id = case when $11::boolean then $13::text else upstream_account_id end,
availability = case
when not enabled then availability
when availability = 'quota_exhausted' then availability
credential_state = case
when not enabled then credential_state
when coalesce($12::text, upstream_user_id) is not null then 'ready'
else 'unknown'
end,
availability_observed_at = case
when not enabled then availability_observed_at
when availability <> 'quota_exhausted' then now()
else availability_observed_at
credential_observed_at = case
when not enabled then credential_observed_at
else now()
end,
last_error_reason = case when enabled then null else last_error_reason end,
last_error_message = case when enabled then null else last_error_message end,
updated_at = now()
where id = $1 and provider_kind = $2
and credential_revision = $3
@@ -627,23 +805,37 @@ pub(crate) async fn rotate_provider_account_in_transaction(
Revision::new(to_u64(next)?)
}
pub(crate) async fn set_provider_account_enabled_in_transaction(
pub(crate) async fn update_provider_accounts_scheduling_in_transaction(
transaction: &mut Transaction<'_, Postgres>,
scope: &ProviderAccountAdminScope,
account_id: &str,
account_ids: &[String],
enabled: bool,
concurrency_limit: Option<AccountConcurrencyLimit>,
weight: AccountWeight,
) -> StoreResult<()> {
let result = sqlx::query(
"update provider_accounts set enabled = $3, updated_at = now()
where id = $1 and provider_kind = $2",
let updated = sqlx::query_scalar::<_, String>(
"update provider_accounts
set enabled = $2, concurrency_limit = $3, weight = $4, updated_at = now()
where id = any($1::text[])
returning id",
)
.bind(account_id)
.bind(&scope.provider_kind)
.bind(account_ids)
.bind(enabled)
.execute(&mut **transaction)
.bind(concurrency_limit.map(|limit| i64::from(limit.get())))
.bind(i16::try_from(weight.get()).map_err(|_| invalid("invalid weight"))?)
.fetch_all(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("set provider account state in admin transaction"))?;
require_admin_account_changed(result.rows_affected(), account_id)
.map_err(|_| postgres_unavailable("set provider accounts state in admin transaction"))?
.into_iter()
.collect::<BTreeSet<_>>();
let expected = account_ids.iter().cloned().collect::<BTreeSet<_>>();
if updated == expected {
Ok(())
} else {
Err(StoreError::NotFound {
entity: ENTITY,
id: "one or more provider account IDs".to_owned(),
})
}
}
pub(crate) async fn delete_provider_accounts_in_transaction(
@@ -685,20 +877,6 @@ pub(crate) fn validate_credential_update(update: &ProviderCredentialUpdate) -> S
Ok(())
}
pub(crate) fn require_admin_account_changed(
rows_affected: u64,
account_id: &str,
) -> StoreResult<()> {
if rows_affected == 1 {
Ok(())
} else {
Err(StoreError::NotFound {
entity: ENTITY,
id: account_id.to_owned(),
})
}
}
pub(crate) fn validate_admin_account_ids(account_ids: &[String]) -> StoreResult<()> {
if account_ids.is_empty() || account_ids.len() > MAX_ADMIN_IMPORT_BATCH {
return Err(invalid(
@@ -715,6 +893,38 @@ pub(crate) fn validate_admin_account_ids(account_ids: &[String]) -> StoreResult<
Ok(())
}
fn validate_batch_update_account_ids(account_ids: &[String]) -> StoreResult<()> {
const MAX_BATCH_UPDATE_ACCOUNTS: usize = 1000;
if account_ids.is_empty() || account_ids.len() > MAX_BATCH_UPDATE_ACCOUNTS {
return Err(invalid(
"account batch update must contain between 1 and 1000 IDs",
));
}
let mut unique = BTreeSet::new();
for account_id in account_ids {
require_nonempty(ENTITY, "account_id", account_id)?;
if !unique.insert(account_id.as_str()) {
return Err(invalid("account batch update contains duplicate IDs"));
}
}
Ok(())
}
fn validate_batch_update_group_ids(group_ids: &[AccountGroupId]) -> StoreResult<()> {
const MAX_BATCH_UPDATE_GROUPS: usize = 1000;
if group_ids.len() > MAX_BATCH_UPDATE_GROUPS {
return Err(invalid("account batch update contains too many group IDs"));
}
let mut unique = BTreeSet::new();
if group_ids
.iter()
.any(|group_id| !unique.insert(group_id.as_str()))
{
return Err(invalid("account batch update contains duplicate group IDs"));
}
Ok(())
}
pub(crate) async fn finish_admin_transaction<T>(
transaction: Transaction<'_, Postgres>,
result: StoreResult<T>,
@@ -1,6 +1,7 @@
//! `provider_accounts` row 模型、解码与校验。
use super::*;
use gateway_core::routing::AccountGroupId;
pub(crate) const ENTITY: &str = "provider account";
pub(crate) const CREDENTIALS_MAX_BYTES: usize = 256 * 1024;
@@ -65,9 +66,29 @@ impl ProviderAccountAdminScope {
}
}
/// 解析数据库 availability 列;字符串映射本体唯一归属 gateway-core。
pub(crate) fn parse_availability(value: &str) -> StoreResult<AccountAvailability> {
AccountAvailability::parse(value).ok_or_else(|| invalid("unknown availability value"))
pub(crate) fn parse_credential_state(value: &str) -> StoreResult<CredentialState> {
CredentialState::parse(value).ok_or_else(|| invalid("unknown credential_state value"))
}
pub(crate) fn parse_quota_access_state(value: &str) -> StoreResult<QuotaAccessState> {
QuotaAccessState::parse(value).ok_or_else(|| invalid("unknown quota_access_state value"))
}
pub(crate) fn parse_quota_evidence(value: Option<String>) -> StoreResult<Option<QuotaEvidence>> {
value
.map(|value| {
QuotaEvidence::parse(&value).ok_or_else(|| invalid("unknown quota_evidence value"))
})
.transpose()
}
pub(crate) fn parse_error_reason(value: Option<String>) -> StoreResult<Option<AccountErrorReason>> {
value
.map(|value| {
AccountErrorReason::parse(&value)
.ok_or_else(|| invalid("unknown last_error_reason value"))
})
.transpose()
}
#[derive(Debug, Clone, PartialEq, Eq)]
@@ -85,12 +106,12 @@ pub struct ProviderAccountSummary {
pub access_token_expires_at: Option<DateTime<Utc>>,
pub next_refresh_at: Option<DateTime<Utc>>,
pub enabled: bool,
pub availability: AccountAvailability,
pub availability_observed_at: DateTime<Utc>,
pub quota_observed_at: Option<DateTime<Utc>>,
/// 快照级限流事实(Provider 物化)。
pub quota_limit_reached: bool,
/// 最近一次非 ready 状态变更的原始原因(上游错误原文 / 原因码);恢复 ready 后清空。
pub concurrency_limit: Option<AccountConcurrencyLimit>,
pub weight: AccountWeight,
pub credential_state: CredentialState,
pub credential_observed_at: DateTime<Utc>,
pub quota: QuotaState,
pub last_error_reason: Option<AccountErrorReason>,
pub last_error_message: Option<String>,
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
@@ -132,8 +153,10 @@ pub struct NewProviderAccount {
pub access_token_expires_at: Option<DateTime<Utc>>,
pub next_refresh_at: Option<DateTime<Utc>>,
pub enabled: bool,
pub availability: AccountAvailability,
pub availability_observed_at: DateTime<Utc>,
pub concurrency_limit: Option<AccountConcurrencyLimit>,
pub weight: AccountWeight,
pub credential_state: CredentialState,
pub credential_observed_at: DateTime<Utc>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
@@ -152,7 +175,7 @@ impl fmt::Debug for NewProviderAccount {
.field("provider_kind", &self.provider_kind)
.field("name", &self.name)
.field("enabled", &self.enabled)
.field("availability", &self.availability)
.field("credential_state", &self.credential_state)
.field("provider_credentials_json", &"[REDACTED]")
.finish_non_exhaustive()
}
@@ -256,10 +279,18 @@ impl fmt::Debug for RotateProviderAccount {
}
#[derive(Debug, Clone)]
pub struct SetProviderAccountEnabled {
pub scope: ProviderAccountAdminScope,
pub account_id: String,
pub struct BatchUpdateProviderAccountsAdmin {
pub account_ids: Vec<String>,
pub enabled: bool,
pub concurrency_limit: Option<AccountConcurrencyLimit>,
pub weight: AccountWeight,
pub group_ids: Vec<AccountGroupId>,
pub audit: AdminAuditEvent,
}
#[derive(Debug, Clone)]
pub struct RecoverProviderAccount {
pub account_id: String,
pub audit: AdminAuditEvent,
}
@@ -297,24 +328,13 @@ impl fmt::Debug for ProviderCredentialUpdate {
}
}
#[derive(Debug, Clone, PartialEq)]
pub struct ProviderAccountObservation {
pub account_id: String,
pub availability: AccountAvailability,
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)]
pub struct ProviderAccountStateUpdate {
pub account_id: String,
pub expected_revision: Revision,
pub availability: AccountAvailability,
pub availability_observed_at: DateTime<Utc>,
/// 上游原因(错误原文 / 原因码);`None` 不更新错误信息列,恢复 `Ready` 时清空。
pub credential_state: CredentialState,
pub credential_observed_at: DateTime<Utc>,
pub error_reason: Option<AccountErrorReason>,
pub message: Option<String>,
}
@@ -325,31 +345,20 @@ impl ProviderAccountStateUpdate {
}
}
impl ProviderAccountObservation {
pub fn validate(&self) -> StoreResult<()> {
require_nonempty(ENTITY, "account_id", &self.account_id)?;
if self.provider_quota_json.is_some() != self.quota_observed_at.is_some() {
return Err(invalid("quota JSON and quota_observed_at must agree"));
}
if let Some(quota) = &self.provider_quota_json {
validate_object_size("provider_quota_json", quota, QUOTA_MAX_BYTES)?;
}
Ok(())
}
}
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,
provider_quota_json, quota_limit_reached, last_error_message,
availability_observed_at, quota_observed_at, created_at, updated_at
has_refresh_token, access_token_expires_at, next_refresh_at, enabled, concurrency_limit, weight, credential_state,
provider_quota_json, quota_access_state, quota_evidence, quota_access_observed_at, quota_reset_at,
last_error_reason, last_error_message,
credential_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,
provider_quota_json, quota_limit_reached, last_error_message,
availability_observed_at, quota_observed_at, created_at, updated_at
has_refresh_token, access_token_expires_at, next_refresh_at, enabled, concurrency_limit, weight, credential_state,
provider_quota_json, quota_access_state, quota_evidence, quota_access_observed_at, quota_reset_at,
last_error_reason, last_error_message,
credential_observed_at, quota_observed_at, created_at, updated_at
from provider_accounts
where id = any($1::text[]) and provider_kind = $2
order by id";
@@ -398,7 +407,14 @@ pub(crate) fn core_account_from_summary(
summary.upstream_account_id,
summary.plan_type,
)
.with_runtime_state(summary.enabled, summary.availability)
.with_account_facts(
summary.enabled,
summary.credential_state,
summary.quota,
summary.last_error_reason,
summary.last_error_message,
)
.with_scheduling(summary.concurrency_limit, summary.weight)
.with_refresh_schedule(
summary.has_refresh_token,
summary.next_refresh_at.map(Into::into),
@@ -430,9 +446,32 @@ pub(crate) fn account_summary_from_row(
let revision = row
.try_get::<i64, _>("credential_revision")
.map_err(|_| invalid("invalid credential revision"))?;
let availability = row
.try_get::<String, _>("availability")
.map_err(|_| invalid("invalid availability"))?;
let credential_state = row
.try_get::<String, _>("credential_state")
.map_err(|_| invalid("invalid credential_state"))?;
let quota_access_state = parse_quota_access_state(&get::<String>(&row, "quota_access_state")?)?;
let quota_evidence = parse_quota_evidence(get(&row, "quota_evidence")?)?;
let quota_access_observed_at = get::<Option<DateTime<Utc>>>(&row, "quota_access_observed_at")?;
let quota_reset_at = get::<Option<DateTime<Utc>>>(&row, "quota_reset_at")?;
let quota = QuotaState::from_persisted(
quota_access_state,
quota_evidence,
quota_access_observed_at.map(Into::into),
quota_reset_at.map(Into::into),
)
.ok_or_else(|| invalid("invalid persisted quota fact"))?;
let concurrency_limit = get::<Option<i64>>(&row, "concurrency_limit")?
.map(|value| {
u32::try_from(value)
.ok()
.and_then(AccountConcurrencyLimit::new)
.ok_or_else(|| invalid("invalid concurrency_limit"))
})
.transpose()?;
let weight = u16::try_from(get::<i16>(&row, "weight")?)
.ok()
.and_then(AccountWeight::new)
.ok_or_else(|| invalid("invalid weight"))?;
Ok(ProviderAccountSummary {
id: get(&row, "id")?,
provider_kind: get(&row, "provider_kind")?,
@@ -447,10 +486,12 @@ pub(crate) fn account_summary_from_row(
access_token_expires_at: get(&row, "access_token_expires_at")?,
next_refresh_at: get(&row, "next_refresh_at")?,
enabled: get(&row, "enabled")?,
availability: parse_availability(&availability)?,
availability_observed_at: get(&row, "availability_observed_at")?,
quota_observed_at: get(&row, "quota_observed_at")?,
quota_limit_reached: get(&row, "quota_limit_reached")?,
concurrency_limit,
weight,
credential_state: parse_credential_state(&credential_state)?,
credential_observed_at: get(&row, "credential_observed_at")?,
quota,
last_error_reason: parse_error_reason(get(&row, "last_error_reason")?)?,
last_error_message: get(&row, "last_error_message")?,
created_at: get(&row, "created_at")?,
updated_at: get(&row, "updated_at")?,
@@ -0,0 +1,53 @@
//! 持久账号事实与可丢失冷却事实的统一状态投影。
use std::{collections::BTreeMap, time::SystemTime};
use futures::future::join_all;
use gateway_core::{
engine::credential::{
AccountStatusFacts, AccountStatusProjection, ProviderAccountId, resolve_account_status,
},
provider_ports::ProviderCooldownPort,
};
use super::ProviderAccountSummary;
#[must_use]
pub(crate) fn account_status_projection(
account: &ProviderAccountSummary,
now: SystemTime,
rate_limited_until: Option<SystemTime>,
) -> AccountStatusProjection {
resolve_account_status(
&AccountStatusFacts {
enabled: account.enabled,
credential_state: account.credential_state,
access_token_expires_at: account.access_token_expires_at.map(Into::into),
quota: account.quota,
rate_limited_until,
last_error_reason: account.last_error_reason,
last_error_message: account.last_error_message.clone(),
},
now,
)
}
/// 尽力读取账号当前有效的 429 冷却。
///
/// 冷却属于可丢失的账号级运行时事实:端口不可用、单条读取失败或脏账号 ID
/// 都视为没有冷却,不污染持久账号状态。凭据轮换不会解除上游对同一账号的限流。
pub(crate) async fn load_rate_limited_until(
cooldowns: Option<&dyn ProviderCooldownPort>,
accounts: &[ProviderAccountSummary],
now: SystemTime,
) -> BTreeMap<String, SystemTime> {
let Some(cooldowns) = cooldowns else {
return BTreeMap::new();
};
let reads = accounts.iter().map(|account| async move {
let account_id = ProviderAccountId::new(account.id.clone()).ok()?;
let cooldown = cooldowns.read(&account_id).await.ok().flatten()?;
(cooldown.until() > now).then(|| (account.id.clone(), cooldown.until()))
});
join_all(reads).await.into_iter().flatten().collect()
}
@@ -4,10 +4,12 @@ use std::collections::BTreeMap;
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use gateway_core::engine::credential::ProviderAccountId;
use gateway_core::routing::{
ConfigRevision,
AccountGroupId, ConfigRevision,
snapshot::{
SnapshotClientPolicyFacts, SnapshotFacts, SnapshotSettingsFacts, SnapshotStoreError,
SnapshotAccountGroupFacts, SnapshotAccountGroupMemberFacts, SnapshotClientPolicyFacts,
SnapshotFacts, SnapshotProviderAccountFacts, SnapshotSettingsFacts, SnapshotStoreError,
SnapshotStorePort,
},
};
@@ -31,7 +33,7 @@ pub struct SnapshotRuntimeSettings {
pub struct CredentialRevisionVectorEntry {
pub provider_account_id: String,
pub credential_revision: Revision,
pub availability_observed_at: DateTime<Utc>,
pub credential_observed_at: DateTime<Utc>,
pub quota_observed_at: Option<DateTime<Utc>>,
}
@@ -41,6 +43,28 @@ pub struct RuntimeSnapshotData {
pub observed_current_revision: Revision,
pub settings: SnapshotRuntimeSettings,
pub client_api_keys: Vec<ClientApiKeySnapshot>,
pub account_groups: Vec<SnapshotAccountGroupData>,
pub provider_accounts: Vec<SnapshotProviderAccountData>,
pub group_memberships: Vec<SnapshotGroupMembershipData>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SnapshotAccountGroupData {
pub id: AccountGroupId,
pub name: String,
pub enabled: bool,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SnapshotProviderAccountData {
pub id: String,
pub provider_kind: String,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SnapshotGroupMembershipData {
pub group_id: AccountGroupId,
pub account_id: String,
}
#[async_trait]
@@ -77,6 +101,9 @@ impl RuntimeSnapshotRepository for PgRuntimeSnapshotRepository {
let (config_revision, settings) = load_settings(&mut transaction).await?;
let client_api_keys = load_client_keys(&mut transaction).await?;
let account_groups = load_account_groups(&mut transaction).await?;
let provider_accounts = load_provider_accounts(&mut transaction).await?;
let group_memberships = load_group_memberships(&mut transaction).await?;
transaction
.commit()
.await
@@ -89,6 +116,9 @@ impl RuntimeSnapshotRepository for PgRuntimeSnapshotRepository {
observed_current_revision,
settings,
client_api_keys,
account_groups,
provider_accounts,
group_memberships,
})
}
@@ -108,7 +138,7 @@ impl RuntimeSnapshotRepository for PgRuntimeSnapshotRepository {
async fn credential_revision_vector(&self) -> StoreResult<Vec<CredentialRevisionVectorEntry>> {
let rows = sqlx::query_as::<_, (String, i64, DateTime<Utc>, Option<DateTime<Utc>>)>(
"select id, credential_revision, availability_observed_at, quota_observed_at
"select id, credential_revision, credential_observed_at, quota_observed_at
from provider_accounts order by id",
)
.fetch_all(&self.pool)
@@ -119,7 +149,7 @@ impl RuntimeSnapshotRepository for PgRuntimeSnapshotRepository {
Ok(CredentialRevisionVectorEntry {
provider_account_id: row.0,
credential_revision: revision_from_i64(row.1)?,
availability_observed_at: row.2,
credential_observed_at: row.2,
quota_observed_at: row.3,
})
})
@@ -151,16 +181,44 @@ impl SnapshotStorePort for PgRuntimeSnapshotRepository {
SnapshotClientPolicyFacts::new(
key.id,
key.plaintext_key,
key.provider_kind,
key.group_ids,
key.limits,
)
})
.collect();
let account_groups = data
.account_groups
.into_iter()
.map(|group| SnapshotAccountGroupFacts::new(group.id, group.name, group.enabled))
.collect();
let provider_accounts = data
.provider_accounts
.into_iter()
.map(|account| {
ProviderAccountId::new(account.id)
.map(|id| SnapshotProviderAccountFacts::new(id, account.provider_kind))
.map_err(|_| SnapshotStoreError::unavailable())
})
.collect::<Result<Vec<_>, _>>()?;
let group_memberships = data
.group_memberships
.into_iter()
.map(|membership| {
ProviderAccountId::new(membership.account_id)
.map(|account_id| {
SnapshotAccountGroupMemberFacts::new(membership.group_id, account_id)
})
.map_err(|_| SnapshotStoreError::unavailable())
})
.collect::<Result<Vec<_>, _>>()?;
Ok(SnapshotFacts::new(
config_revision,
observed_current_revision,
settings,
client_policies,
account_groups,
provider_accounts,
group_memberships,
))
})
}
@@ -224,9 +282,16 @@ async fn load_settings(
async fn load_client_keys(
transaction: &mut Transaction<'_, Postgres>,
) -> StoreResult<Vec<ClientApiKeySnapshot>> {
let rows = sqlx::query_as::<_, (String, String, String, i64, i64)>(
"select id, key, provider_kind, max_concurrency, requests_per_minute
from client_api_keys where enabled order by id",
let rows = sqlx::query_as::<_, (String, String, Vec<String>, i64, i64)>(
"select k.id, k.key,
coalesce(array_agg(kg.account_group_id order by kg.account_group_id)
filter (where kg.account_group_id is not null), '{}') as group_ids,
k.max_concurrency, k.requests_per_minute
from client_api_keys k
left join client_api_key_groups kg on kg.client_api_key_id = k.id
where k.enabled
group by k.id
order by k.id",
)
.fetch_all(&mut **transaction)
.await
@@ -236,6 +301,63 @@ async fn load_client_keys(
.collect()
}
async fn load_account_groups(
transaction: &mut Transaction<'_, Postgres>,
) -> StoreResult<Vec<SnapshotAccountGroupData>> {
let rows = sqlx::query_as::<_, (String, String, bool)>(
"select id, name, enabled from account_groups order by id",
)
.fetch_all(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("load snapshot account groups"))?;
rows.into_iter()
.map(|(id, name, enabled)| {
Ok(SnapshotAccountGroupData {
id: AccountGroupId::new(id).map_err(|_| invalid("invalid account group id"))?,
name,
enabled,
})
})
.collect()
}
async fn load_provider_accounts(
transaction: &mut Transaction<'_, Postgres>,
) -> StoreResult<Vec<SnapshotProviderAccountData>> {
sqlx::query_as::<_, (String, String)>(
"select id, provider_kind from provider_accounts order by id",
)
.fetch_all(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("load snapshot provider accounts"))
.map(|rows| {
rows.into_iter()
.map(|(id, provider_kind)| SnapshotProviderAccountData { id, provider_kind })
.collect()
})
}
async fn load_group_memberships(
transaction: &mut Transaction<'_, Postgres>,
) -> StoreResult<Vec<SnapshotGroupMembershipData>> {
let rows = sqlx::query_as::<_, (String, String)>(
"select account_group_id, provider_account_id
from account_group_accounts order by account_group_id, provider_account_id",
)
.fetch_all(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("load snapshot group memberships"))?;
rows.into_iter()
.map(|(group_id, account_id)| {
Ok(SnapshotGroupMembershipData {
group_id: AccountGroupId::new(group_id)
.map_err(|_| invalid("invalid membership group id"))?,
account_id,
})
})
.collect()
}
fn revision_from_i64(value: i64) -> StoreResult<Revision> {
Revision::new(to_u64(value)?)
}
@@ -9,6 +9,7 @@ use async_trait::async_trait;
use chrono::{DateTime, Utc};
use futures::future::join_all;
use gateway_core::engine::credential::{AccountRuntimeSignals, ProviderAccountId};
use gateway_core::policy::ClientApiKeyId;
use gateway_core::provider_ports::{
ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest,
ProviderRefreshCapacityRequest, ProviderSchedulingLeaseRequest, ProviderSchedulingState,
@@ -329,10 +330,13 @@ impl RedisCredentialLeaseRepository {
}))
}
/// 原子推进跨进程共享的 Provider 调度游标。
pub async fn advance_scheduling_cursor(&self, provider_kind: &str) -> StoreResult<u64> {
require_nonempty("provider scheduling cursor", "provider_kind", provider_kind)?;
let key = self.scheduling_cursor_key(provider_kind)?;
/// 原子推进跨进程共享、按 Client Key 与 Provider 隔离的调度游标。
pub async fn advance_scheduling_cursor(
&self,
client_api_key_id: &ClientApiKeyId,
provider_kind: &ProviderKind,
) -> StoreResult<u64> {
let key = self.scheduling_cursor_key(client_api_key_id, provider_kind)?;
let mut connection = self.connection.clone();
let cursor = Script::new(ADVANCE_SCHEDULING_CURSOR_SCRIPT)
.key(key)
@@ -369,10 +373,17 @@ impl RedisCredentialLeaseRepository {
))
}
fn scheduling_cursor_key(&self, provider_kind: &str) -> StoreResult<String> {
let fingerprint = resource_fingerprint("provider scheduling cursor", provider_kind)?;
fn scheduling_cursor_key(
&self,
client_api_key_id: &ClientApiKeyId,
provider_kind: &ProviderKind,
) -> StoreResult<String> {
let fingerprint = resource_fingerprint(
"provider scheduling cursor",
&format!("{}\0{}", client_api_key_id.as_str(), provider_kind.as_str()),
)?;
Ok(format!(
"{}:scheduler:provider:{{{fingerprint}}}:cursor",
"{}:scheduler:client-provider:{{{fingerprint}}}:cursor",
self.namespace
))
}
@@ -475,10 +486,11 @@ impl RedisProviderLeaseCoordinator {
async fn next_scheduling_cursor(
&self,
client_api_key_id: &ClientApiKeyId,
provider_kind: &ProviderKind,
) -> Result<u64, ProviderStoreError> {
self.repository
.advance_scheduling_cursor(provider_kind.as_str())
.advance_scheduling_cursor(client_api_key_id, provider_kind)
.await
.map_err(|_| provider_unavailable("advance scheduling cursor"))
}
@@ -512,7 +524,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,
rate_limited_until: None,
failure_rate_basis_points: None,
first_output_latency_ms: None,
},
@@ -589,12 +601,15 @@ impl RedisProviderLeaseCoordinator {
impl ProviderLeasePort for RedisProviderLeaseCoordinator {
fn load_state<'a>(
&'a self,
client_api_key_id: &'a ClientApiKeyId,
provider_kind: &'a ProviderKind,
accounts: &'a [ProviderAccountId],
) -> futures::future::BoxFuture<'a, Result<ProviderSchedulingState, ProviderStoreError>> {
Box::pin(async move {
let signals = self.load_signals(accounts).await?;
let round_robin_cursor = self.next_scheduling_cursor(provider_kind).await?;
let round_robin_cursor = self
.next_scheduling_cursor(client_api_key_id, provider_kind)
.await?;
Ok(ProviderSchedulingState::new(signals, round_robin_cursor))
})
}
@@ -1,4 +1,4 @@
//! Provider account availability 的可重建 Redis cache fence。
//! Provider credential state 的可重建 Redis cache fence。
use std::time::Duration;
@@ -7,7 +7,7 @@ use chrono::{DateTime, Utc};
use redis::{Script, aio::ConnectionManager};
use gateway_core::engine::credential::{
AccountAvailability, CredentialRevision, OpaqueProviderData, ProviderAccountId,
CredentialRevision, CredentialState, OpaqueProviderData, ProviderAccountId,
};
use gateway_core::provider_ports::{
ProviderCatalogCacheKey, ProviderCatalogCachePort, ProviderCredentialState,
@@ -24,7 +24,7 @@ if current >= tonumber(ARGV[1]) then return 0 end
redis.call('HSET', KEYS[1],
'revision', ARGV[1],
'enabled', ARGV[2],
'availability', ARGV[3],
'credential_state', ARGV[3],
'observed_at_ms', ARGV[4])
redis.call('PEXPIRE', KEYS[1], ARGV[5])
return 1
@@ -44,7 +44,7 @@ pub struct CredentialStateCache {
pub provider_account_id: String,
pub revision: Revision,
pub enabled: bool,
pub availability: String,
pub credential_state: String,
pub observed_at: DateTime<Utc>,
}
@@ -55,8 +55,8 @@ impl CredentialStateCache {
"provider_account_id",
&self.provider_account_id,
)?;
if AccountAvailability::parse(&self.availability).is_none() {
return Err(invalid("availability is not registered"));
if CredentialState::parse(&self.credential_state).is_none() {
return Err(invalid("credential_state is not registered"));
}
Ok(())
}
@@ -172,7 +172,7 @@ impl CredentialStateRepository for RedisCredentialStateRepository {
.key(self.key(&state.provider_account_id)?)
.arg(state.revision.get())
.arg(u8::from(state.enabled))
.arg(&state.availability)
.arg(&state.credential_state)
.arg(state.observed_at.timestamp_millis())
.arg(ttl_ms)
.invoke_async::<i64>(&mut connection)
@@ -195,12 +195,12 @@ impl CredentialStateRepository for RedisCredentialStateRepository {
.arg(self.key(provider_account_id)?)
.arg("revision")
.arg("enabled")
.arg("availability")
.arg("credential_state")
.arg("observed_at_ms")
.query_async(&mut connection)
.await
.map_err(|_| redis_unavailable("read credential state"))?;
let (Some(revision), Some(enabled), Some(availability), Some(observed_at_ms)) = values
let (Some(revision), Some(enabled), Some(credential_state), Some(observed_at_ms)) = values
else {
return Ok(None);
};
@@ -208,7 +208,7 @@ impl CredentialStateRepository for RedisCredentialStateRepository {
provider_account_id: provider_account_id.to_owned(),
revision: Revision::new(revision)?,
enabled: enabled == 1,
availability,
credential_state,
observed_at: DateTime::from_timestamp_millis(observed_at_ms)
.ok_or_else(|| invalid("observed_at is invalid"))?,
};
@@ -299,7 +299,7 @@ impl ProviderCredentialStatePort for RedisCredentialStateRepository {
revision: Revision::new(state.credential_revision().get())
.map_err(|_| provider_invalid("encode credential state"))?,
enabled: state.enabled(),
availability: state.availability().as_str().to_owned(),
credential_state: state.credential_state().as_str().to_owned(),
observed_at,
};
CredentialStateRepository::cache_credential_state(
@@ -329,13 +329,13 @@ impl ProviderCredentialStatePort for RedisCredentialStateRepository {
.map_err(|_| provider_invalid("decode credential state"))?;
let revision = CredentialRevision::new(state.revision.get())
.map_err(|_| provider_invalid("decode credential state"))?;
let availability = AccountAvailability::parse(&state.availability)
let credential_state = CredentialState::parse(&state.credential_state)
.ok_or_else(|| provider_invalid("decode credential state"))?;
Ok(ProviderCredentialState::new(
account_id,
revision,
state.enabled,
availability,
credential_state,
state.observed_at.into(),
))
})
@@ -1,8 +1,9 @@
//! 旧版 Responses continuation 的 Redis 会话亲和记录。
//! Responses continuation 的 Redis 会话亲和记录。
//!
//! 该存储只保存可丢失的 response → Provider/account pin;PostgreSQL 审计记录
//! 不是 continuation 的可用性前提。键只含 response ID 的不可逆指纹,Provider
//! 私有状态作为不透明 JSON 保存,Store 不解释其内容。
//! 不是 continuation 的可用性前提。键只含 response ID 的不可逆指纹;Client Key
//! 所有权写入 payload 并在读取时严格校验,使跨 Key 复用能够 fail closed。
//! Provider 私有状态作为不透明 JSON 保存,Store 不解释其内容。
use std::time::{Duration, SystemTime, UNIX_EPOCH};
@@ -13,6 +14,7 @@ use gateway_core::engine::continuation::{
};
use gateway_core::engine::credential::ProviderAccountId;
use gateway_core::operation::ProviderSessionState;
use gateway_core::policy::ClientApiKeyId;
use gateway_core::routing::ProviderKind;
use redis::aio::ConnectionManager;
use serde::{Deserialize, Serialize};
@@ -72,8 +74,9 @@ pub struct RedisNativeContinuationRepository {
}
#[derive(Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct ContinuationWire {
client_api_key_id: String,
upstream_response_id: String,
provider: String,
account: String,
@@ -86,13 +89,13 @@ impl RedisNativeContinuationRepository {
pub fn new(connection: ConnectionManager, key_namespace: &str) -> StoreResult<Self> {
Ok(Self {
connection,
namespace: format!("{}:response-affinity:v1", namespace(key_namespace)?),
namespace: format!("{}:response-affinity:v2", namespace(key_namespace)?),
})
}
fn entry_fingerprint(&self, response_id: &str) -> Result<String, NativeContinuationStoreError> {
resource_fingerprint("native continuation", response_id)
.map_err(|error| NativeContinuationStoreError::new(error.to_string()))
.map_err(|error| NativeContinuationStoreError::invalid_data(error.to_string()))
}
fn entry_key(&self, fingerprint: &str) -> String {
@@ -111,6 +114,7 @@ impl RedisNativeContinuationRepository {
impl NativeContinuationPort for RedisNativeContinuationRepository {
fn resolve<'a>(
&'a self,
client_api_key_id: &'a ClientApiKeyId,
previous_response_id: &'a PreviousResponseId,
) -> BoxFuture<'a, Result<Option<NativeContinuationPin>, NativeContinuationStoreError>> {
Box::pin(async move {
@@ -124,10 +128,10 @@ impl NativeContinuationPort for RedisNativeContinuationRepository {
.query_async::<Option<String>>(&mut connection)
.await
.map_err(|error| {
NativeContinuationStoreError::new(format!("redis GET failed: {error}"))
NativeContinuationStoreError::unavailable(format!("redis GET failed: {error}"))
})?;
payload
.map(|payload| decode_pin(previous_response_id, &payload))
.map(|payload| decode_pin(client_api_key_id, previous_response_id, &payload))
.transpose()
})
}
@@ -143,20 +147,24 @@ impl NativeContinuationPort for RedisNativeContinuationRepository {
let fingerprint = self.entry_fingerprint(pin.previous_response_id().as_str())?;
let payload = encode_pin(&pin)?;
if payload.len() > MAX_RECORD_BYTES {
return Err(NativeContinuationStoreError::new(format!(
return Err(NativeContinuationStoreError::invalid_data(format!(
"pin payload is too large ({} bytes, limit {MAX_RECORD_BYTES})",
payload.len()
)));
}
let now_millis = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map_err(|_| NativeContinuationStoreError::new("system clock predates Unix epoch"))?
.map_err(|_| {
NativeContinuationStoreError::unavailable("system clock predates Unix epoch")
})?
.as_millis();
let now_millis = u64::try_from(now_millis).map_err(|_| {
NativeContinuationStoreError::new("system time is outside supported range")
NativeContinuationStoreError::unavailable("system time is outside supported range")
})?;
let ttl_millis = u64::try_from(CONTINUATION_TTL.as_millis()).map_err(|_| {
NativeContinuationStoreError::new("continuation TTL is outside supported range")
NativeContinuationStoreError::invalid_data(
"continuation TTL is outside supported range",
)
})?;
let cutoff = now_millis.saturating_sub(ttl_millis);
let ttl_seconds = CONTINUATION_TTL.as_secs();
@@ -175,7 +183,7 @@ impl NativeContinuationPort for RedisNativeContinuationRepository {
.invoke_async::<i64>(&mut connection)
.await
.map_err(|error| {
NativeContinuationStoreError::new(format!(
NativeContinuationStoreError::unavailable(format!(
"redis affinity script failed: {error}"
))
})?;
@@ -189,11 +197,12 @@ fn encode_pin(pin: &NativeContinuationPin) -> Result<String, NativeContinuationS
.session_state()
.is_some_and(|state| state.provider() != pin.provider().as_str())
{
return Err(NativeContinuationStoreError::new(
return Err(NativeContinuationStoreError::invalid_data(
"session state provider does not match pin provider",
));
}
serde_json::to_string(&ContinuationWire {
client_api_key_id: pin.client_api_key_id().as_str().to_owned(),
upstream_response_id: pin.upstream_response_id().as_str().to_owned(),
provider: pin.provider().as_str().to_owned(),
account: pin.account().as_str().to_owned(),
@@ -201,34 +210,44 @@ fn encode_pin(pin: &NativeContinuationPin) -> Result<String, NativeContinuationS
session_state: pin.session_state().cloned(),
})
.map_err(|error| {
NativeContinuationStoreError::new(format!("pin serialization failed: {error}"))
NativeContinuationStoreError::invalid_data(format!("pin serialization failed: {error}"))
})
}
fn decode_pin(
expected_client_api_key_id: &ClientApiKeyId,
previous_response_id: &PreviousResponseId,
payload: &str,
) -> Result<NativeContinuationPin, NativeContinuationStoreError> {
let wire: ContinuationWire = serde_json::from_str(payload)
.map_err(|error| NativeContinuationStoreError::new(format!("{error}")))?;
.map_err(|error| NativeContinuationStoreError::invalid_data(format!("{error}")))?;
let provider = ProviderKind::new(wire.provider).map_err(|_| {
NativeContinuationStoreError::new("pin payload contains an invalid provider")
NativeContinuationStoreError::invalid_data("pin payload contains an invalid provider")
})?;
let account = ProviderAccountId::new(wire.account).map_err(|_| {
NativeContinuationStoreError::new("pin payload contains an invalid account id")
NativeContinuationStoreError::invalid_data("pin payload contains an invalid account id")
})?;
let client_api_key_id = ClientApiKeyId::new(wire.client_api_key_id).map_err(|_| {
NativeContinuationStoreError::invalid_data(
"pin payload contains an invalid client API key id",
)
})?;
if client_api_key_id != *expected_client_api_key_id {
return Err(NativeContinuationStoreError::ownership_mismatch());
}
let upstream_response_id = PreviousResponseId::new(wire.upstream_response_id);
let scope = parse_scope(&wire.scope)?;
let mut pin = NativeContinuationPin::new(
previous_response_id.clone(),
upstream_response_id,
client_api_key_id,
provider.clone(),
account,
)
.with_scope(scope);
if let Some(state) = wire.session_state {
if state.provider() != provider.as_str() {
return Err(NativeContinuationStoreError::new(
return Err(NativeContinuationStoreError::invalid_data(
"pin payload session state provider does not match pin provider",
));
}
@@ -248,7 +267,7 @@ fn parse_scope(scope: &str) -> Result<NativeContinuationScope, NativeContinuatio
match scope {
"persisted" => Ok(NativeContinuationScope::Persisted),
"connection_local" => Ok(NativeContinuationScope::ConnectionLocal),
_ => Err(NativeContinuationStoreError::new(
_ => Err(NativeContinuationStoreError::invalid_data(
"pin payload contains an unsupported continuation scope",
)),
}
@@ -0,0 +1,299 @@
use std::collections::BTreeMap;
use gateway_admin::{
model::{
MutationActor, MutationContext, PageSize,
account_groups::{
AccountGroupColor, AccountGroupListQuery, DeleteAccountGroup, NewAccountGroup,
},
client_keys::NewClientKey,
client_keys::UpdateClientKey,
},
ports::store::{AccountGroupStore, AdminStoreErrorKind, ClientKeyStore},
};
use gateway_core::{
policy::{ClientApiKeyId, RateLimits},
routing::AccountGroupId,
};
use gateway_store::postgres::{PgAccountGroupRepository, PgAdminClientKeyStore};
use super::TestDatabase;
const MIXED_GROUP: &str = "grp_00000000000000000000000000000001";
const EMPTY_GROUP: &str = "grp_00000000000000000000000000000002";
#[tokio::test]
async fn groups_aggregate_cross_provider_members_and_key_bindings_without_multiplication() {
let Some(database) = TestDatabase::create("account_group_aggregate").await else {
return;
};
seed_account(
&database.pool,
"acct_group_openai",
"openai",
"OpenAI Account",
)
.await;
seed_account(&database.pool, "acct_group_xai", "xai", "xAI Account").await;
sqlx::query(
"update provider_accounts set concurrency_limit = 4 where id = 'acct_group_openai'",
)
.execute(&database.pool)
.await
.expect("set account concurrency override");
let groups = PgAccountGroupRepository::new(database.pool.clone());
let keys = PgAdminClientKeyStore::new(database.pool.clone());
let mixed_group = group_id(MIXED_GROUP);
let empty_group = group_id(EMPTY_GROUP);
groups
.create_account_group(
NewAccountGroup {
id: mixed_group.clone(),
name: "Mixed Production".to_owned(),
description: Some("cross-provider".to_owned()),
color: group_color("#2563EBFF"),
},
&context("create-mixed"),
)
.await
.expect("create mixed account group");
groups
.create_account_group(
NewAccountGroup {
id: empty_group.clone(),
name: "Empty Pool".to_owned(),
description: None,
color: group_color("#06B6D4CC"),
},
&context("create-empty"),
)
.await
.expect("create empty account group");
assign_accounts(
&database.pool,
MIXED_GROUP,
&["acct_group_openai", "acct_group_xai"],
)
.await;
for (id, group_ids) in [
("key_group_one", vec![mixed_group.clone()]),
("key_group_two", vec![mixed_group.clone()]),
("key_empty_pool", vec![empty_group.clone()]),
("key_all_accounts", Vec::new()),
] {
keys.create_client_key(new_key(id, group_ids), &context(id))
.await
.expect("create scoped client key");
}
let page = groups
.list_account_groups(AccountGroupListQuery {
page: 1,
page_size: PageSize::new(20).expect("page size"),
search: None,
enabled: None,
})
.await
.expect("list account groups");
assert_eq!(page.total, 2);
let by_id = page
.items
.into_iter()
.map(|group| (group.id.to_string(), group))
.collect::<BTreeMap<_, _>>();
let mixed = by_id.get(MIXED_GROUP).expect("mixed group");
assert_eq!(mixed.member_count, 2);
assert_eq!(
mixed.provider_counts,
BTreeMap::from([("openai".to_owned(), 1), ("xai".to_owned(), 1)])
);
assert_eq!(mixed.client_key_count, 2);
assert_eq!(mixed.account_summary.available, 2);
assert_eq!(mixed.account_summary.limited, 0);
assert_eq!(mixed.account_summary.total, 2);
assert_eq!(mixed.capacity.used_slots, None);
assert_eq!(mixed.capacity.total_slots, 7);
assert_eq!(mixed.usage.today_usd.as_str(), "0");
assert_eq!(mixed.usage.total_usd.as_str(), "0");
let empty = by_id.get(EMPTY_GROUP).expect("empty group");
assert_eq!(empty.member_count, 0);
assert!(empty.provider_counts.is_empty());
assert_eq!(empty.client_key_count, 1);
let all_key = keys
.reveal_client_key(&client_key_id("key_all_accounts"))
.await
.expect("reveal all-accounts key")
.expect("all-accounts key exists");
assert!(all_key.record.groups.is_empty());
assert_eq!(
all_key
.record
.provider_kinds
.iter()
.map(|kind| kind.as_str())
.collect::<Vec<_>>(),
["openai", "xai"]
);
let empty_pool_key = keys
.reveal_client_key(&client_key_id("key_empty_pool"))
.await
.expect("reveal empty-pool key")
.expect("empty-pool key exists");
assert_eq!(empty_pool_key.record.groups.len(), 1);
assert!(empty_pool_key.record.provider_kinds.is_empty());
let (scope_revision, widened) = keys
.update_client_key(
UpdateClientKey {
id: client_key_id("key_group_one"),
name: "key_group_one".to_owned(),
label: None,
group_ids: Vec::new(),
limits: RateLimits::unlimited(),
},
&context("widen-group-key"),
)
.await
.expect("widen restricted key to all accounts");
assert!(widened.groups.is_empty());
let scope_audit: Vec<String> = sqlx::query_scalar(
"select changed_fields from admin_audit_events
where admin_request_id = 'widen-group-key'",
)
.fetch_one(&database.pool)
.await
.expect("load scope widening audit");
assert!(scope_audit.contains(&"routing_scope:groups->all".to_owned()));
assert_eq!(current_revision(&database.pool).await, scope_revision.get());
let (restricted_revision, restricted) = keys
.update_client_key(
UpdateClientKey {
id: client_key_id("key_group_one"),
name: "key_group_one".to_owned(),
label: None,
group_ids: vec![empty_group],
limits: RateLimits::unlimited(),
},
&context("restrict-all-key"),
)
.await
.expect("restrict all-accounts key to groups");
assert_eq!(restricted.groups.len(), 1);
let restricted_audit: Vec<String> = sqlx::query_scalar(
"select changed_fields from admin_audit_events
where admin_request_id = 'restrict-all-key'",
)
.fetch_one(&database.pool)
.await
.expect("load scope restriction audit");
assert!(restricted_audit.contains(&"routing_scope:all->groups".to_owned()));
assert_eq!(
current_revision(&database.pool).await,
restricted_revision.get()
);
let revision_before_delete = current_revision(&database.pool).await;
let audit_before_delete = audit_count(&database.pool).await;
let error = groups
.delete_account_group(
DeleteAccountGroup { id: mixed_group },
&context("delete-referenced"),
)
.await
.expect_err("referenced group must not be deleted");
assert_eq!(error.kind(), AdminStoreErrorKind::Conflict);
assert_eq!(
current_revision(&database.pool).await,
revision_before_delete
);
assert_eq!(audit_count(&database.pool).await, audit_before_delete);
database.close().await;
}
fn new_key(id: &str, group_ids: Vec<AccountGroupId>) -> NewClientKey {
let marker = char::from(id.as_bytes().last().copied().unwrap_or(b'k'));
NewClientKey {
id: client_key_id(id),
name: id.to_owned(),
label: None,
group_ids,
limits: RateLimits::unlimited(),
plaintext: format!("sk_{}", marker.to_string().repeat(43)),
}
}
fn group_id(value: &str) -> AccountGroupId {
AccountGroupId::new(value).expect("valid account group ID")
}
fn group_color(value: &str) -> AccountGroupColor {
AccountGroupColor::parse(value).expect("valid account group color")
}
fn client_key_id(value: &str) -> ClientApiKeyId {
ClientApiKeyId::new(value).expect("valid client key ID")
}
fn context(request_id: &str) -> MutationContext {
MutationContext {
actor: MutationActor::System,
request_id: request_id.to_owned(),
}
}
async fn seed_account(pool: &sqlx::PgPool, id: &str, provider: &str, name: &str) {
sqlx::query(
"insert into provider_accounts (
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,
credential_state, credential_observed_at, created_at, updated_at
) values (
$1, $2, $3, null, $1 || '-user', null, null, 'oauth', '{}'::jsonb, 1,
false, null, null, true, 'ready', now(), now(), now()
)",
)
.bind(id)
.bind(provider)
.bind(name)
.execute(pool)
.await
.expect("seed provider account");
}
async fn assign_accounts(pool: &sqlx::PgPool, group_id: &str, account_ids: &[&str]) {
for account_id in account_ids {
sqlx::query(
"insert into account_group_accounts (
account_group_id, provider_account_id, created_at
)
values ($1, $2, now())",
)
.bind(group_id)
.bind(account_id)
.execute(pool)
.await
.expect("seed account group membership");
}
}
async fn current_revision(pool: &sqlx::PgPool) -> u64 {
let value =
sqlx::query_scalar::<_, i64>("select config_revision from runtime_settings where id = 1")
.fetch_one(pool)
.await
.expect("load config revision");
u64::try_from(value).expect("positive config revision")
}
async fn audit_count(pool: &sqlx::PgPool) -> u64 {
let value = sqlx::query_scalar::<_, i64>("select count(*) from admin_audit_events")
.fetch_one(pool)
.await
.expect("count audit rows");
u64::try_from(value).expect("non-negative audit count")
}
@@ -93,10 +93,12 @@ async fn seed_request(
"insert into model_requests (
id, client_api_key_ref, config_revision, protocol, operation, endpoint,
client_transport, requested_model_id, outcome,
started_at, deadline_at, completed_at
started_at, deadline_at, completed_at,
routing_scope, routing_group_refs, routing_group_names_snapshot
) values (
$1, 'key-recovery', 1, 'openai', 'responses', '/v1/responses',
'http_sse', 'coding', $2, $3, $4, $5
'http_sse', 'coding', $2, $3, $4, $5,
'all', '{}'::text[], '[]'::jsonb
)",
)
.bind(id)
@@ -23,7 +23,7 @@ fn client_key_requires_the_frozen_plaintext_format() {
id: "key-1".to_owned(),
name: "default".to_owned(),
label: None,
provider_kind: "openai".to_owned(),
group_ids: Vec::new(),
key: format!("sk_{}", "a".repeat(43)),
max_concurrency: 0,
requests_per_minute: 0,
@@ -44,9 +44,9 @@ async fn client_key_list_uses_safe_keyset_search_and_filtered_total() {
);
sqlx::query(
"insert into client_api_keys (
id, name, label, provider_kind, key, enabled, max_concurrency, requests_per_minute,
id, name, label, key, enabled, max_concurrency, requests_per_minute,
created_at, updated_at
) values ($1, $2, $2, 'openai', $3, true, 0, 0,
) values ($1, $2, $2, $3, true, 0, 0,
now() - ($4::bigint * interval '1 minute'), now())",
)
.bind(id)
@@ -186,9 +186,9 @@ async fn client_key_database_sort_is_stable_and_keeps_null_last_used_at_last() {
{
sqlx::query(
"insert into client_api_keys (
id, name, provider_kind, key, enabled, max_concurrency, requests_per_minute,
id, name, key, enabled, max_concurrency, requests_per_minute,
last_used_at, created_at, updated_at
) values ($1, $2, 'openai', $3, $4, 0, 0, $5::timestamptz, $6::timestamptz,
) values ($1, $2, $3, $4, 0, 0, $5::timestamptz, $6::timestamptz,
$6::timestamptz)",
)
.bind(id)
@@ -272,9 +272,9 @@ async fn client_key_usage_touch_should_preserve_the_latest_observed_timestamp()
};
sqlx::query(
"insert into client_api_keys (
id, name, provider_kind, key, enabled, max_concurrency, requests_per_minute,
id, name, key, enabled, max_concurrency, requests_per_minute,
created_at, updated_at
) values ('key_usage_touch', 'usage', 'openai', $1, true, 0, 0, now(), now())",
) values ('key_usage_touch', 'usage', $1, true, 0, 0, now(), now())",
)
.bind(format!("sk_{}", "t".repeat(43)))
.execute(&database.pool)
@@ -315,9 +315,9 @@ async fn dedicated_reveal_returns_plaintext_without_debug_exposure() {
let plaintext = format!("sk_{}", "r".repeat(43));
sqlx::query(
"insert into client_api_keys (
id, name, provider_kind, key, enabled, max_concurrency, requests_per_minute,
id, name, key, enabled, max_concurrency, requests_per_minute,
created_at, updated_at
) values ('key_reveal', 'reveal', 'openai', $1, true, 1, 2, now(), now())",
) values ('key_reveal', 'reveal', $1, true, 1, 2, now(), now())",
)
.bind(&plaintext)
.execute(&database.pool)
@@ -340,7 +340,7 @@ fn client_key_debug_redacts_plaintext() {
id: "key-1".to_owned(),
name: "default".to_owned(),
label: None,
provider_kind: "openai".to_owned(),
group_ids: Vec::new(),
key: secret.clone(),
max_concurrency: 0,
requests_per_minute: 0,
@@ -29,6 +29,9 @@ fn model_request_rejects_mismatched_client_key_live_id() {
endpoint: "/v1/responses".to_owned(),
client_transport: "http_sse".to_owned(),
requested_model_id: "coding".to_owned(),
routing_scope: "all".to_owned(),
routing_group_refs: Vec::new(),
routing_group_names_snapshot: serde_json::json!([]),
client_ip: None,
user_agent: None,
reasoning_effort: None,
@@ -55,8 +58,8 @@ async fn merged_first_attempt_insert_should_match_sequential_semantics() {
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, enabled, availability,
availability_observed_at, created_at, updated_at
has_refresh_token, access_token_expires_at, enabled, credential_state,
credential_observed_at, created_at, updated_at
) values (
'acct_merged', 'openai', 'merged', null, 'user-merged', null, null, 'oauth',
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'ready', $1, $1, $1
@@ -76,6 +79,9 @@ async fn merged_first_attempt_insert_should_match_sequential_semantics() {
endpoint: "/v1/responses".to_owned(),
client_transport: "http_sse".to_owned(),
requested_model_id: "coding".to_owned(),
routing_scope: "all".to_owned(),
routing_group_refs: Vec::new(),
routing_group_names_snapshot: serde_json::json!([]),
client_ip: None,
user_agent: None,
reasoning_effort: None,
@@ -153,6 +159,66 @@ async fn merged_first_attempt_insert_should_match_sequential_semantics() {
database.close().await;
}
#[tokio::test]
async fn model_request_persists_group_routing_snapshot_without_live_group_foreign_keys() {
let Some(database) = TestDatabase::create("execution_group_routing_history").await else {
return;
};
let repository = PgExecutionStore::new(database.pool.clone());
let started_at = Utc::now();
repository
.insert_model_request(NewModelRequest {
id: "req_group_history".to_owned(),
client_api_key_id: None,
client_api_key_ref: "key_group_history".to_owned(),
config_revision: 7,
routing_scope: "groups".to_owned(),
routing_group_refs: vec![
"grp_00000000000000000000000000000001".to_owned(),
"grp_00000000000000000000000000000002".to_owned(),
],
routing_group_names_snapshot: serde_json::json!(["Production Pool", "Overflow Pool"]),
protocol: "openai".to_owned(),
operation: "responses".to_owned(),
endpoint: "/v1/responses".to_owned(),
client_transport: "http_sse".to_owned(),
requested_model_id: "coding".to_owned(),
client_ip: None,
user_agent: None,
reasoning_effort: None,
reasoning_preset: None,
request_kind: None,
subagent_kind: None,
compact: false,
image_generation_requested: false,
started_at,
deadline_at: started_at + Duration::seconds(30),
})
.await
.expect("insert grouped request history");
let stored: (String, Vec<String>, serde_json::Value) = sqlx::query_as(
"select routing_scope, routing_group_refs, routing_group_names_snapshot
from model_requests where id = 'req_group_history'",
)
.fetch_one(&database.pool)
.await
.expect("load grouped request history");
assert_eq!(stored.0, "groups");
assert_eq!(
stored.1,
[
"grp_00000000000000000000000000000001",
"grp_00000000000000000000000000000002",
]
);
assert_eq!(
stored.2,
serde_json::json!(["Production Pool", "Overflow Pool"])
);
database.close().await;
}
#[tokio::test]
async fn downstream_commit_should_atomically_record_http_status_once() {
let Some(database) = TestDatabase::create("execution_commit_status").await else {
@@ -483,9 +549,11 @@ async fn seed_running_request(pool: &sqlx::PgPool, id: &str) -> Result<(), sqlx:
"insert into model_requests (
id, client_api_key_ref, config_revision, protocol, operation, endpoint,
client_transport, requested_model_id, provider_kind, provider_account_ref, cost_source,
started_at, deadline_at
started_at, deadline_at,
routing_scope, routing_group_refs, routing_group_names_snapshot
) values ($1, 'key_status', 1, 'openai_responses', 'generate', '/v1/responses',
'http_json', 'status-model', 'openai', 'acct_status', 'unavailable', now(), now() + interval '1 minute')",
'http_json', 'status-model', 'openai', 'acct_status', 'unavailable', now(), now() + interval '1 minute',
'all', '{}'::text[], '[]'::jsonb)",
)
.bind(id)
.execute(pool)
@@ -7,6 +7,7 @@ use sqlx::{
};
use uuid::Uuid;
mod account_groups;
mod admin_security_audit;
mod admission_recovery;
mod backup;
@@ -146,6 +147,29 @@ async fn connect_and_migrate_should_apply_all_migrations_once_and_reopen_cleanly
.fetch_one(&second)
.await
.expect("check removed raw response ID index");
let legacy_key_provider_column_exists = sqlx::query_scalar::<_, bool>(
"select exists (
select 1 from information_schema.columns
where table_schema = 'public'
and table_name = 'client_api_keys'
and column_name = 'provider_kind'
)",
)
.fetch_one(&second)
.await
.expect("check removed client key provider column");
let routing_history_columns = sqlx::query_scalar::<_, String>(
"select column_name from information_schema.columns
where table_schema = 'public'
and table_name = 'model_requests'
and column_name in (
'routing_scope', 'routing_group_refs', 'routing_group_names_snapshot'
)
order by column_name",
)
.fetch_all(&second)
.await
.expect("load routing history columns");
second.close().await;
sqlx::raw_sql(sqlx::AssertSqlSafe(format!(
@@ -156,7 +180,7 @@ 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, 10);
assert_eq!(first_table_count, 13);
assert_eq!(
migration_count,
i64::try_from(TEST_MIGRATOR.iter().count())
@@ -164,6 +188,15 @@ async fn connect_and_migrate_should_apply_all_migrations_once_and_reopen_cleanly
);
assert_eq!(response_id_types, ["bytea", "bytea"]);
assert!(!raw_response_id_index_exists);
assert!(!legacy_key_provider_column_exists);
assert_eq!(
routing_history_columns,
[
"routing_group_names_snapshot",
"routing_group_refs",
"routing_scope",
]
);
}
#[test]
@@ -134,26 +134,27 @@ async fn dashboard_account_metrics_should_partition_account_statuses() {
"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,
availability_observed_at, created_at, updated_at
access_token_expires_at, enabled, credential_state,
quota_access_state, quota_evidence, quota_access_observed_at,
credential_observed_at, created_at, updated_at
) values
('acct_available', 'openai', 'available', 'user-available', 'oauth',
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'ready',
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'ready', 'allowed', null, $1,
$1, $1, $1),
('acct_expired', 'openai', 'expired', 'user-expired', 'oauth',
'{}'::jsonb, 1, false, $1 - interval '1 day', true, 'ready',
'{}'::jsonb, 1, false, $1 - interval '1 day', true, 'ready', 'allowed', null, $1,
$1, $1, $1),
('acct_rate_cooldown', 'openai', 'rate-cooldown', 'user-rate-cooldown', 'oauth',
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'ready',
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'ready', 'allowed', null, $1,
$1, $1, $1),
('acct_usage_limit', 'xai', 'usage-limit', 'user-usage-limit', 'oauth',
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'quota_exhausted',
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'ready', 'exhausted', 'usage_limit_reached', $1,
$1, $1, $1),
('acct_banned', 'xai', 'banned', 'user-banned', 'oauth',
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'banned',
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'banned', 'allowed', null, $1,
$1, $1, $1),
('acct_disabled', 'xai', 'disabled', 'user-disabled', 'oauth',
'{}'::jsonb, 1, false, $1 + interval '1 day', false, 'ready',
'{}'::jsonb, 1, false, $1 + interval '1 day', false, 'ready', 'allowed', null, $1,
$1, $1, $1)",
)
.bind(now)
@@ -165,7 +166,7 @@ async fn dashboard_account_metrics_should_partition_account_statuses() {
.expect("dashboard range");
let metrics = repository
.dashboard_summary(range)
.dashboard_summary(range, now)
.await
.expect("dashboard summary")
.provider_accounts;
@@ -173,22 +174,22 @@ async fn dashboard_account_metrics_should_partition_account_statuses() {
assert_eq!(
(
metrics.total,
metrics.enabled,
metrics.active,
metrics.unavailable,
metrics.expired,
metrics.normal,
metrics.quota_exhausted,
metrics.rate_limited,
metrics.disabled,
metrics.banned,
metrics.error,
),
(6, 5, 2, 4, 1, 1, 1, 1),
(6, 2, 1, 0, 1, 2),
);
// 首页按调度可用性聚合:配额耗尽也属于不可用。
assert_eq!(
metrics.unavailable,
metrics.quota_exhausted + metrics.expired + metrics.disabled + metrics.banned
metrics.total,
metrics.normal
+ metrics.quota_exhausted
+ metrics.rate_limited
+ metrics.disabled
+ metrics.error
);
assert_eq!(metrics.total, metrics.active + metrics.unavailable);
database.close().await;
}
@@ -202,35 +203,36 @@ async fn dashboard_account_metrics_with_cooldowns_should_only_reclassify_eligibl
"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
access_token_expires_at, enabled, credential_state,
quota_access_state, quota_evidence, quota_access_observed_at,
credential_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,
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'ready', 'allowed', null, $1,
$1, $1, $1),
('acct_metrics_active_cooling', 'openai', 'active-cooling', 'user-active-cooling', 'oauth',
'{}'::jsonb, 2, false, $1 + interval '1 day', true, 'ready', false,
'{}'::jsonb, 2, false, $1 + interval '1 day', true, 'ready', 'allowed', null, $1,
$1, $1, $1),
('acct_metrics_ready_quota', 'openai', 'ready-quota', 'user-ready-quota', 'oauth',
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'ready', true,
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'ready', 'exhausted', 'provider_denied', $1,
$1, $1, $1),
('acct_metrics_cooling', 'openai', 'cooling', 'user-cooling', 'oauth',
'{}'::jsonb, 3, false, $1 + interval '1 day', true, 'ready', true,
'{}'::jsonb, 3, false, $1 + interval '1 day', true, 'ready', 'exhausted', 'provider_denied', $1,
$1, $1, $1),
('acct_metrics_stale_cooldown', 'openai', 'stale-cooldown', 'user-stale-cooldown', 'oauth',
'{}'::jsonb, 4, false, $1 + interval '1 day', true, 'ready', true,
'{}'::jsonb, 4, false, $1 + interval '1 day', true, 'ready', 'allowed', null, $1,
$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,
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'ready', 'exhausted', 'provider_denied', $1,
$1, $1, $1),
('acct_metrics_expired', 'openai', 'expired', 'user-expired', 'oauth',
'{}'::jsonb, 5, false, $1 - interval '1 second', true, 'ready', true,
'{}'::jsonb, 5, false, $1 - interval '1 second', true, 'ready', 'exhausted', 'provider_denied', $1,
$1, $1, $1),
('acct_metrics_unknown', 'openai', 'unknown', 'user-unknown', 'oauth',
'{}'::jsonb, 6, false, $1 + interval '1 day', true, 'unknown', false,
'{}'::jsonb, 6, false, $1 + interval '1 day', true, 'unknown', 'unknown', null, null,
$1, $1, $1),
('acct_metrics_disabled', 'openai', 'disabled', 'user-disabled', 'oauth',
'{}'::jsonb, 1, false, $1 + interval '1 day', false, 'ready', false,
'{}'::jsonb, 1, false, $1 + interval '1 day', false, 'ready', 'allowed', null, $1,
$1, $1, $1)",
)
.bind(now)
@@ -250,7 +252,7 @@ async fn dashboard_account_metrics_with_cooldowns_should_only_reclassify_eligibl
.expect("dashboard range");
let metrics = repository
.dashboard_summary(range)
.dashboard_summary(range, now)
.await
.expect("dashboard summary")
.provider_accounts;
@@ -258,36 +260,21 @@ async fn dashboard_account_metrics_with_cooldowns_should_only_reclassify_eligibl
assert_eq!(
(
metrics.total,
metrics.enabled,
metrics.active,
metrics.normal,
metrics.rate_limited,
metrics.quota_exhausted,
metrics.expired,
metrics.invalid,
metrics.disabled,
metrics.banned,
metrics.unavailable,
metrics.error,
),
(9, 8, 1, 2, 3, 1, 1, 1, 0, 8),
(9, 1, 2, 3, 1, 2),
);
assert_eq!(
metrics.total,
metrics.active
metrics.normal
+ 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
+ metrics.error
);
database.close().await;
}
@@ -429,7 +416,7 @@ async fn admin_observability_adapter_preserves_utc_queries_metrics_costs_and_det
let store = PgAdminObservabilityStore::new(database.pool.clone(), None, None);
let dashboard = store
.dashboard_summary(range)
.dashboard_summary(range, now)
.await
.expect("admin dashboard summary");
assert_eq!(dashboard.range, range);
@@ -563,6 +550,12 @@ async fn admin_observability_adapter_preserves_utc_queries_metrics_costs_and_det
.await
.expect("fully filtered usage page");
assert_eq!(filtered.total, 1);
assert_eq!(filtered.items[0].routing_scope, "groups");
assert_eq!(filtered.items[0].routing_group_refs, ["grp_history"]);
assert_eq!(
filtered.items[0].routing_group_names_snapshot,
["Historical group"]
);
assert_eq!(filtered.items[0].service_tier.as_deref(), Some("priority"));
assert_eq!(
filtered.items[0].upstream_request_id.as_deref(),
@@ -597,6 +590,12 @@ async fn admin_observability_adapter_preserves_utc_queries_metrics_costs_and_det
admin_observability::RequestOutcome::Succeeded
);
assert_eq!(detail.request.service_tier.as_deref(), Some("priority"));
assert_eq!(detail.request.routing_scope, "groups");
assert_eq!(detail.request.routing_group_refs, ["grp_history"]);
assert_eq!(
detail.request.routing_group_names_snapshot,
["Historical group"]
);
assert_eq!(detail.attempts.len(), 1);
assert_eq!(
detail.attempts[0].outcome,
@@ -727,7 +726,7 @@ async fn dashboard_summary_totals_include_history_outside_selected_range() {
let repository = PgObservabilityRepository::new(database.pool.clone(), None);
let dashboard = repository
.dashboard_summary(range)
.dashboard_summary(range, now)
.await
.expect("dashboard summary");
assert_eq!(dashboard.requests.request_count, 2);
@@ -766,7 +765,7 @@ async fn observability_queries_preserve_request_account_cost_and_diagnostic_fact
let repository = PgObservabilityRepository::new(database.pool.clone(), None);
let dashboard = repository
.dashboard_summary(range)
.dashboard_summary(range, now)
.await
.expect("dashboard summary");
assert_eq!(dashboard.requests.request_count, 3);
@@ -1013,8 +1012,8 @@ async fn seed_observability_facts(
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, enabled, availability,
availability_observed_at, created_at, updated_at
has_refresh_token, access_token_expires_at, enabled, credential_state,
credential_observed_at, created_at, updated_at
) values (
'acct_observe', 'openai', 'primary', 'account@example.invalid',
'user-observe', null, 'pro', 'oauth', '{}'::jsonb, 1, false, $1 + interval '1 day',
@@ -1034,7 +1033,8 @@ async fn seed_observability_facts(
provider_account_authentication_kind_snapshot,
attempt_count, upstream_send_state, outcome, client_status_code,
input_tokens, output_tokens, total_tokens, cost_source, latency_ms,
started_at, deadline_at, completed_at
started_at, deadline_at, completed_at,
routing_scope, routing_group_refs, routing_group_names_snapshot
) values (
'req_observe_uncommitted', 'key_observe', 1, 'openai', 'responses', '/v1/responses',
'http_sse', 'public-model', 'openai', 'acct_observe',
@@ -1042,7 +1042,8 @@ async fn seed_observability_facts(
'primary', 'account@example.invalid', 'oauth',
1, 'sent', 'succeeded', 200,
900, 900, 1800, 'unavailable', 650,
$1 - interval '15 minutes', $1 + interval '10 minutes', $1 - interval '14 minutes'
$1 - interval '15 minutes', $1 + interval '10 minutes', $1 - interval '14 minutes',
'all', '{}'::text[], '[]'::jsonb
)",
)
.bind(now)
@@ -1064,7 +1065,8 @@ async fn seed_observability_facts(
image_input_tokens, image_output_tokens, total_tokens,
image_generation_requested, image_generation_succeeded,
cost_source, cost_amount, cost_currency, first_token_ms, latency_ms,
started_at, deadline_at, completed_at
started_at, deadline_at, completed_at,
routing_scope, routing_group_refs, routing_group_names_snapshot
) values (
'req_observe_success', 'key_observe', 1, 'openai', 'responses', '/v1/responses',
'http_sse', 'public-model', 'openai', 'acct_observe',
@@ -1075,7 +1077,8 @@ async fn seed_observability_facts(
'resp_observe_success', 'upstream_req_success', 'upstream_resp_success',
100, 20, 40, 3, 5, 31, 9, 120, true, true,
'provider_reported', 1.25, 'USD', 120, 900,
$1 - interval '20 minutes', $1 + interval '10 minutes', $1 - interval '19 minutes'
$1 - interval '20 minutes', $1 + interval '10 minutes', $1 - interval '19 minutes',
'groups', array['grp_history'], jsonb_build_array('Historical group')
)",
)
.bind(now)
@@ -1093,7 +1096,8 @@ async fn seed_observability_facts(
error_kind, provider_error_code, error_message, retry_after_ms,
input_tokens, cached_tokens, image_generation_requested,
image_generation_succeeded, cost_source, latency_ms,
started_at, deadline_at, completed_at
started_at, deadline_at, completed_at,
routing_scope, routing_group_refs, routing_group_names_snapshot
) values (
'req_observe_failed', 'key_observe', 1, 'openai', 'responses', '/v1/responses',
'http_sse', 'public-model', 'openai', 'acct_observe',
@@ -1102,7 +1106,8 @@ async fn seed_observability_facts(
'primary', 'account@example.invalid', 'oauth',
'sent', 'failed', 502, 429, 'rate_limited', 'rate_limit',
'upstream limited', 1000, 0, 0, true, false, 'unavailable', 700,
$1 - interval '10 minutes', $1 + interval '20 minutes', $1 - interval '9 minutes'
$1 - interval '10 minutes', $1 + interval '20 minutes', $1 - interval '9 minutes',
'all', '{}'::text[], '[]'::jsonb
)",
)
.bind(now)
@@ -1144,7 +1149,8 @@ async fn seed_calculated_billing_facts(
provider_account_authentication_kind_snapshot,
upstream_send_state, downstream_committed_at, outcome, client_status_code,
input_tokens, output_tokens, cached_tokens, cache_write_tokens, total_tokens, service_tier,
cost_source, cost_amount, cost_currency, started_at, deadline_at, completed_at
cost_source, cost_amount, cost_currency, started_at, deadline_at, completed_at,
routing_scope, routing_group_refs, routing_group_names_snapshot
) values (
'req_observe_calculated', 'key_observe', 1, 'openai', 'responses', '/v1/responses',
'http_sse', 'public-model', 'openai', 'acct_observe', 'acct_observe', 'gpt-5.5',
@@ -1152,7 +1158,8 @@ async fn seed_calculated_billing_facts(
'primary', 'account@example.invalid', 'oauth',
'sent', $1 - interval '29 minutes', 'succeeded', 200,
800, 200, 0, 0, 1000, 'priority', 'calculated', 1.25, 'USD',
$1 - interval '30 minutes', $1 + interval '10 minutes', $1 - interval '29 minutes'
$1 - interval '30 minutes', $1 + interval '10 minutes', $1 - interval '29 minutes',
'all', '{}'::text[], '[]'::jsonb
)",
)
.bind(now)
@@ -1167,7 +1174,7 @@ async fn seed_calculated_billing_facts(
provider_account_authentication_kind_snapshot,
upstream_send_state, outcome, client_status_code, input_tokens, output_tokens,
total_tokens, cost_source, cost_amount, cost_currency, started_at, deadline_at,
completed_at
completed_at, routing_scope, routing_group_refs, routing_group_names_snapshot
) values (
'req_observe_calculated_uncommitted', 'key_observe', 1, 'openai', 'responses',
'/v1/responses', 'http_sse', 'public-model', 'openai', 'acct_observe',
@@ -1175,7 +1182,8 @@ async fn seed_calculated_billing_facts(
'primary', 'account@example.invalid', 'oauth',
'sent', 'succeeded', 200, 800, 200,
1000, 'calculated', 1.25, 'USD',
$1 - interval '40 minutes', $1 + interval '10 minutes', $1 - interval '39 minutes'
$1 - interval '40 minutes', $1 + interval '10 minutes', $1 - interval '39 minutes',
'all', '{}'::text[], '[]'::jsonb
)",
)
.bind(now)
File diff suppressed because it is too large Load Diff
@@ -108,7 +108,7 @@ async fn insert_refreshable_account(
"insert into provider_accounts (
id, provider_kind, name, upstream_user_id, authentication_kind,
provider_credentials_json, has_refresh_token, access_token_expires_at,
next_refresh_at, availability, availability_observed_at, created_at, updated_at
next_refresh_at, credential_state, credential_observed_at, created_at, updated_at
) values ($1, 'openai', $1, $1, 'oauth', '{}'::jsonb, true, $2, $3,
'ready', now(), now(), now())",
)
@@ -10,14 +10,14 @@ fn snapshot_client_policy_contains_only_common_limits() {
let policy = ClientApiKeySnapshot {
id: ClientApiKeyId::new("key-1").expect("client key ID"),
plaintext_key: PlaintextClientApiKey::new("sk_snapshot_secret").expect("plaintext key"),
provider_kind: "openai".to_owned(),
group_ids: Vec::new(),
limits: RateLimits {
max_concurrency: 3,
requests_per_minute: 60,
},
};
assert_eq!(policy.limits.max_concurrency, 3);
assert_eq!(policy.provider_kind, "openai");
assert!(policy.group_ids.is_empty());
assert!(!format!("{policy:?}").contains("sk_snapshot_secret"));
}
@@ -29,9 +29,8 @@ async fn runtime_snapshot_loads_enabled_plaintext_key_without_debug_exposure() {
let plaintext = format!("sk_{}", "s".repeat(43));
sqlx::query(
"insert into client_api_keys (
id, name, provider_kind, key, enabled, max_concurrency, requests_per_minute,
created_at, updated_at
) values ('key_snapshot', 'snapshot', 'xai', $1, true, 2, 60, now(), now())",
id, name, key, enabled, max_concurrency, requests_per_minute, created_at, updated_at
) values ('key_snapshot', 'snapshot', $1, true, 2, 60, now(), now())",
)
.bind(&plaintext)
.execute(&database.pool)
@@ -42,7 +41,7 @@ async fn runtime_snapshot_loads_enabled_plaintext_key_without_debug_exposure() {
.await
.expect("load runtime snapshot");
assert_eq!(snapshot.client_api_keys.len(), 1);
assert_eq!(snapshot.client_api_keys[0].provider_kind, "xai");
assert!(snapshot.client_api_keys[0].group_ids.is_empty());
assert_eq!(
snapshot.client_api_keys[0].plaintext_key.expose_for_auth(),
plaintext
@@ -665,6 +665,9 @@ fn new_request(id: &str, started_at: DateTime<Utc>) -> NewModelRequest {
endpoint: "/v1/responses".to_owned(),
client_transport: "http_sse".to_owned(),
requested_model_id: "coding".to_owned(),
routing_scope: "all".to_owned(),
routing_group_refs: Vec::new(),
routing_group_names_snapshot: serde_json::json!([]),
client_ip: None,
user_agent: None,
reasoning_effort: None,
@@ -722,8 +725,8 @@ async fn seed_account(
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, enabled, availability,
availability_observed_at, created_at, updated_at
has_refresh_token, access_token_expires_at, enabled, credential_state,
credential_observed_at, created_at, updated_at
) values (
$1, 'openai', $2, $3, 'user-' || $1, null, null, $4,
'{}'::jsonb, 1, false, $5 + interval '1 day', true, 'ready',
@@ -756,8 +759,8 @@ async fn finalize_request(pool: &PgPool, id: &str, started_at: DateTime<Utc>) {
async fn seed_api_key(pool: &PgPool, id: &str, name: &str, now: DateTime<Utc>) {
sqlx::query(
"insert into client_api_keys (id, name, provider_kind, key, enabled, created_at, updated_at)
values ($1, $2, 'openai', 'sk_' || $3, true, $4, $4)",
"insert into client_api_keys (id, name, key, enabled, created_at, updated_at)
values ($1, $2, 'sk_' || $3, true, $4, $4)",
)
.bind(id)
.bind(name)
@@ -211,30 +211,41 @@ async fn scheduling_lease_ttl_and_cache_loss_rebuild_empty_runtime_signal() {
}
#[tokio::test]
async fn scheduling_cursor_is_shared_across_gateway_processes() {
async fn scheduling_cursor_is_shared_across_processes_and_isolated_by_client_and_provider() {
let Some((repository, mut connection, namespace)) = repository().await else {
return;
};
let peer = RedisCredentialLeaseRepository::new(connection.clone(), &namespace)
.expect("peer repository");
let primary = gateway_core::policy::ClientApiKeyId::new("key_primary").expect("client key");
let secondary = gateway_core::policy::ClientApiKeyId::new("key_secondary").expect("client key");
let openai = gateway_core::routing::ProviderKind::new("openai").expect("provider");
let xai = gateway_core::routing::ProviderKind::new("xai").expect("provider");
assert_eq!(
repository
.advance_scheduling_cursor("route_openai_real")
.advance_scheduling_cursor(&primary, &openai)
.await
.expect("first cursor"),
0
);
assert_eq!(
peer.advance_scheduling_cursor("route_openai_real")
peer.advance_scheduling_cursor(&primary, &openai)
.await
.expect("shared cursor"),
1
);
assert_eq!(
peer.advance_scheduling_cursor("route_openai_backup")
peer.advance_scheduling_cursor(&secondary, &openai)
.await
.expect("independent cursor"),
.expect("client-isolated cursor"),
0
);
assert_eq!(
peer.advance_scheduling_cursor(&primary, &xai)
.await
.expect("provider-isolated cursor"),
0
);
@@ -243,12 +254,13 @@ async fn scheduling_cursor_is_shared_across_gateway_processes() {
.query_async::<Vec<String>>(&mut connection)
.await
.expect("list cursor keys");
assert_eq!(keys.len(), 2);
assert!(
keys.iter().all(|key| {
!key.contains("route_openai_real") && !key.contains("route_openai_backup")
})
);
assert_eq!(keys.len(), 3);
assert!(keys.iter().all(|key| {
!key.contains("key_primary")
&& !key.contains("key_secondary")
&& !key.contains("openai")
&& !key.contains("xai")
}));
redis::cmd("DEL")
.arg(keys)
.query_async::<i64>(&mut connection)

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