mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-07 08:05:32 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
e7c2f025cc | ||
|
|
4ee294b208 | ||
|
|
93f1a82d3f | ||
|
|
c46f98cd2b | ||
|
|
5c69b1b082 | ||
|
|
6a78fb9232 | ||
|
|
02a03df79a | ||
|
|
2aca8fb79b | ||
|
|
891b227b69 | ||
|
|
04096dddb7 | ||
|
|
3a9b1c479f | ||
|
|
944d3f299e | ||
|
|
56451601fb | ||
|
|
4d53ef6068 | ||
|
|
a62bb196f7 | ||
|
|
7ad0ca48c5 | ||
|
|
e25639c11b | ||
|
|
8abac080a2 | ||
|
|
37cd059531 | ||
|
|
c46910df22 | ||
|
|
8cfd7f53be | ||
|
|
75e22d7a5b | ||
|
|
6b2d335bc5 | ||
|
|
270b1540cb | ||
|
|
660214d9a5 | ||
|
|
2231fee768 | ||
|
|
a2e60c1c6e | ||
|
|
fc388bb799 | ||
|
|
62a35c3960 | ||
|
|
82653544ac | ||
|
|
f10f825404 | ||
|
|
a268f5ca66 | ||
|
|
832e2e8bdd | ||
|
|
be703a0c5f | ||
|
|
c2c86928c5 | ||
|
|
bda21ccc61 | ||
|
|
72c56680b9 | ||
|
|
a6808499f0 | ||
|
|
fdb954b3d1 | ||
|
|
a069001f01 | ||
|
|
d3a87d11ad |
@@ -79,7 +79,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@db488ddef3bf6cb639b32c2e9a7c0a7ea8271d28 # v4.37.8
|
||||
uses: github/codeql-action/upload-sarif@cdf488f595d80d6e07e03d4674febd5ab45fa938 # v4.37.9
|
||||
with:
|
||||
sarif_file: trivy-results.sarif
|
||||
category: runtime-container
|
||||
|
||||
@@ -4,6 +4,8 @@ on:
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
paths-ignore:
|
||||
- 'release/version.yaml'
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
|
||||
@@ -163,6 +163,15 @@ jobs:
|
||||
ref: ${{ needs.release-info.outputs.git_sha }}
|
||||
upload_frontend: true
|
||||
|
||||
container:
|
||||
needs: [release-info]
|
||||
permissions:
|
||||
contents: read
|
||||
security-events: write
|
||||
uses: ./.github/workflows/_container.yml
|
||||
with:
|
||||
ref: ${{ needs.release-info.outputs.git_sha }}
|
||||
|
||||
build-docker-binary:
|
||||
name: build-docker-binary (${{ matrix.os }}-${{ matrix.arch }})
|
||||
needs: release-info
|
||||
@@ -497,7 +506,7 @@ jobs:
|
||||
exit-code: '1'
|
||||
|
||||
publish:
|
||||
needs: [release-info, package-assets, build-docker-image]
|
||||
needs: [release-info, package-assets, build-docker-image, container]
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 20
|
||||
permissions:
|
||||
|
||||
Generated
+45
-30
@@ -49,13 +49,13 @@ checksum = "03918c3dbd7701a85c6b9887732e2921175f26c350b4563841d0958c21d57e6d"
|
||||
|
||||
[[package]]
|
||||
name = "argon2"
|
||||
version = "0.5.3"
|
||||
version = "0.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3c3610892ee6e0cbce8ae2700349fcf8f98adb0dbfbee85aec3c9179d29cc072"
|
||||
checksum = "134c52ddac6d63c576bef8168db10c83c49c26444ecbc68060fef078925a901c"
|
||||
dependencies = [
|
||||
"base64ct",
|
||||
"blake2",
|
||||
"cpufeatures 0.2.17",
|
||||
"cpufeatures 0.3.0",
|
||||
"password-hash",
|
||||
]
|
||||
|
||||
@@ -116,7 +116,7 @@ checksum = "82f6aeea286b8eb4dd3431a1be1b59d290ace00f5bfd8e2a159bc2a05e2c1667"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 3.0.3",
|
||||
"syn 3.0.4",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -206,9 +206,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "aws-sdk-s3"
|
||||
version = "1.142.0"
|
||||
version = "1.144.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f9e15a5c55e05f4b0b7e483160b3c85cccdf77cff02c95504f3e71d460855cd2"
|
||||
checksum = "30dc8bf6baaf7d46336a0ca2c69f223d9b90d7a801fb3e28f7ea17b00dc6b1de"
|
||||
dependencies = [
|
||||
"arc-swap",
|
||||
"aws-credential-types",
|
||||
@@ -331,9 +331,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "aws-smithy-http-client"
|
||||
version = "1.3.0"
|
||||
version = "1.4.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3c1c8a04cb31ba74d0115af5a890bb8c0d48fba64b52812fa13929a6ef0cc83c"
|
||||
checksum = "ebfd138fac0337cee7516c352757ea73b9f2266e57d0bcb5bc70e9547e45aef1"
|
||||
dependencies = [
|
||||
"aws-smithy-async",
|
||||
"aws-smithy-runtime-api",
|
||||
@@ -375,9 +375,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "aws-smithy-runtime"
|
||||
version = "1.13.1"
|
||||
version = "1.14.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "483b858ff67522011c4786310c5cd8fd88d0be7ea3d5f1a48328446300c4269e"
|
||||
checksum = "b82e438d30e02a825d363bd639a9efaed68a8089d86101054b0081e7e0d3e606"
|
||||
dependencies = [
|
||||
"aws-smithy-async",
|
||||
"aws-smithy-http",
|
||||
@@ -599,11 +599,11 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "blake2"
|
||||
version = "0.10.6"
|
||||
version = "0.11.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "46502ad458c9a52b69d4d4d32775c788b7a1b85e8bc9d482d92250fc0e3f8efe"
|
||||
checksum = "5b5d4d889834ee8ecfc0f8426ad30faf7cdcb10f741a8e6d7224d95325479f6f"
|
||||
dependencies = [
|
||||
"digest 0.10.7",
|
||||
"digest 0.11.3",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -740,7 +740,7 @@ dependencies = [
|
||||
"provider-xai",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"syn 3.0.3",
|
||||
"syn 3.0.4",
|
||||
"tempfile",
|
||||
"thiserror",
|
||||
"tokio",
|
||||
@@ -1013,7 +1013,6 @@ checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292"
|
||||
dependencies = [
|
||||
"block-buffer 0.10.4",
|
||||
"crypto-common 0.1.7",
|
||||
"subtle",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1318,7 +1317,7 @@ checksum = "9fb9654ba8355388abeb8dcb4fc62f511300867002afc858860463bdd9fe0c44"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 3.0.3",
|
||||
"syn 3.0.4",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1363,6 +1362,7 @@ dependencies = [
|
||||
"argon2",
|
||||
"async-trait",
|
||||
"base64 0.23.1",
|
||||
"bytes",
|
||||
"chrono",
|
||||
"chrono-tz",
|
||||
"cron",
|
||||
@@ -1586,6 +1586,11 @@ name = "hashbrown"
|
||||
version = "0.17.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a"
|
||||
dependencies = [
|
||||
"allocator-api2",
|
||||
"equivalent",
|
||||
"foldhash",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hashlink"
|
||||
@@ -2046,11 +2051,11 @@ checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad"
|
||||
|
||||
[[package]]
|
||||
name = "lru"
|
||||
version = "0.16.4"
|
||||
version = "0.18.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7f66e8d5d03f609abc3a39e6f08e4164ebf1447a732906d39eb9b99b7919ef39"
|
||||
checksum = "0d317b4b9eb398e6acce275758ec6125535505e7a146fb1a9b8bda2451b0ff4c"
|
||||
dependencies = [
|
||||
"hashbrown 0.16.1",
|
||||
"hashbrown 0.17.1",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -2232,13 +2237,12 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "password-hash"
|
||||
version = "0.5.0"
|
||||
version = "0.6.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "346f04948ba92c43e8469c1ee6736c7563d71012b17d40745260fe106aac2166"
|
||||
checksum = "aab41826031698d6ffcd9cff78ef56ef998e39dc7e5067cdfebe373842d4723b"
|
||||
dependencies = [
|
||||
"base64ct",
|
||||
"rand_core 0.6.4",
|
||||
"subtle",
|
||||
"getrandom 0.4.3",
|
||||
"phc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -2262,6 +2266,17 @@ version = "2.3.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220"
|
||||
|
||||
[[package]]
|
||||
name = "phc"
|
||||
version = "0.6.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "44dc769b75f93afdddd8c7fa12d685292ddeff1e66f7f0f3a234cf1818afe892"
|
||||
dependencies = [
|
||||
"base64ct",
|
||||
"ctutils",
|
||||
"getrandom 0.4.3",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "phf"
|
||||
version = "0.11.3"
|
||||
@@ -2982,7 +2997,7 @@ checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 3.0.3",
|
||||
"syn 3.0.4",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -3407,9 +3422,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "syn"
|
||||
version = "3.0.3"
|
||||
version = "3.0.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "53e9bae58849f64dfa4f5d5ae372c8341f7305f82a3868709269343628b659a3"
|
||||
checksum = "e6275cddf4610d1775e6d1fe9469b2e77d0f39fd98fb7450901b821e0c53649f"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
@@ -3477,7 +3492,7 @@ checksum = "bc04cd3e1236dd4a98afca4569f2deb3f120e5422a4023be2cb683f8486292af"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 3.0.3",
|
||||
"syn 3.0.4",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -3913,9 +3928,9 @@ checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be"
|
||||
|
||||
[[package]]
|
||||
name = "uuid"
|
||||
version = "1.24.1"
|
||||
version = "1.26.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2cefc03fd367c0c6d4305de1b312cf00248c4114f4a0418ce6a6af769e3b0bd9"
|
||||
checksum = "b5772d71c9be8a8a6ac2117d949c5b224c1b72241bb611d9a3012edcf8af7812"
|
||||
dependencies = [
|
||||
"getrandom 0.4.3",
|
||||
"js-sys",
|
||||
|
||||
@@ -144,7 +144,7 @@ fn config_loader_should_reject_removed_fingerprint_section() {
|
||||
#[test]
|
||||
fn config_loader_should_reject_invalid_codex_cli_version() {
|
||||
assert_rejected(valid_config().replace(
|
||||
"codex_version: '0.147.0-alpha.6.6'",
|
||||
"codex_version: '0.151.0-alpha.7.2'",
|
||||
"codex_version: 'latest'",
|
||||
));
|
||||
}
|
||||
@@ -165,7 +165,7 @@ fn config_loader_should_reject_zero_server_port() {
|
||||
|
||||
#[test]
|
||||
fn config_loader_should_reject_invalid_desktop_profile_fields() {
|
||||
assert_rejected(valid_config().replace("desktop_build: '6415'", "desktop_build: 'build'"));
|
||||
assert_rejected(valid_config().replace("desktop_build: '7377'", "desktop_build: 'build'"));
|
||||
}
|
||||
|
||||
fn assert_rejected(config: String) {
|
||||
|
||||
@@ -12,9 +12,10 @@ name = "main"
|
||||
path = "tests/main.rs"
|
||||
|
||||
[dependencies]
|
||||
argon2 = "0.5.3"
|
||||
argon2 = "0.6.0"
|
||||
async-trait = "0.1.89"
|
||||
base64 = "0.23.0"
|
||||
bytes = "1.11.1"
|
||||
chrono = "0.4.45"
|
||||
futures = "0.3.32"
|
||||
gateway-core = { path = "../gateway-core" }
|
||||
@@ -28,7 +29,7 @@ tracing = "0.1.44"
|
||||
uuid = { version = "1.23.3", features = ["v7"] }
|
||||
cron = "0.17"
|
||||
chrono-tz = "0.10"
|
||||
tokio = { version = "1.52", features = ["rt", "time", "macros"] }
|
||||
tokio = { version = "1.52", features = ["rt", "time", "macros", "sync"] }
|
||||
|
||||
[dev-dependencies]
|
||||
serde_json = { version = "1.0.150", features = ["arbitrary_precision", "preserve_order"] }
|
||||
|
||||
@@ -261,6 +261,7 @@ pub async fn initialize(
|
||||
|
||||
let accounts = Arc::new(DefaultAccountsService::new(
|
||||
store.accounts(),
|
||||
store.account_runtime(),
|
||||
registry.clone(),
|
||||
snapshot.clone(),
|
||||
probe.clone(),
|
||||
@@ -282,6 +283,7 @@ pub async fn initialize(
|
||||
accounts,
|
||||
account_groups: Arc::new(DefaultAccountGroupService::new(
|
||||
store.account_groups(),
|
||||
store.account_runtime(),
|
||||
snapshot.clone(),
|
||||
)),
|
||||
client_keys: Arc::new(DefaultClientKeyService::new(
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use gateway_core::routing::AccountGroupId;
|
||||
use gateway_core::{engine::credential::AccountStatusFacts, routing::AccountGroupId};
|
||||
|
||||
use super::{PageSize, Revision, observability::DecimalAmount};
|
||||
|
||||
@@ -45,7 +45,17 @@ pub struct AccountGroupCapacity {
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AccountGroupUsage {
|
||||
pub today_usd: DecimalAmount,
|
||||
pub total_usd: DecimalAmount,
|
||||
/// 当前 usage retention 窗口内、按请求发生时 routing group 快照归属的累计成本。
|
||||
pub retained_total_usd: DecimalAmount,
|
||||
}
|
||||
|
||||
/// 当前页账号组 membership 对应的持久账号事实;运行态在 Admin query service 合并。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AccountGroupMemberFact {
|
||||
pub group_id: AccountGroupId,
|
||||
pub account_id: String,
|
||||
pub status: AccountStatusFacts,
|
||||
pub total_slots: u64,
|
||||
}
|
||||
|
||||
/// Lightweight group reference embedded in account and client-key views.
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
//! 多 Provider 账号目录与连接测试的公共事实。
|
||||
|
||||
use std::pin::Pin;
|
||||
use std::{collections::BTreeMap, pin::Pin};
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use futures::Stream;
|
||||
@@ -56,6 +56,14 @@ pub struct AccountListQuery {
|
||||
pub sort: Option<AccountSort>,
|
||||
}
|
||||
|
||||
/// Admin query service 从运行态存储取得的当前账号冷却快照。
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct AccountRuntimeSnapshot {
|
||||
pub rate_limited_until: BTreeMap<String, DateTime<Utc>>,
|
||||
/// `None` 表示实时 lease 存储不可用;`Some` 中未出现的账号当前使用量为零。
|
||||
pub in_flight: Option<BTreeMap<String, u64>>,
|
||||
}
|
||||
|
||||
/// Optional account membership filter.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum AccountGroupFilter {
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
//! 用量、成本、健康与错误诊断的 UTC 语义事实。
|
||||
|
||||
use std::{num::NonZeroU32, str::FromStr};
|
||||
use std::str::FromStr;
|
||||
|
||||
use chrono::{DateTime, TimeDelta, Timelike as _, Utc};
|
||||
|
||||
@@ -38,29 +38,6 @@ impl TimeRange {
|
||||
}
|
||||
}
|
||||
|
||||
/// 从一开始的观测页码。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct PageNumber(NonZeroU32);
|
||||
|
||||
impl PageNumber {
|
||||
#[must_use]
|
||||
pub const fn new(value: NonZeroU32) -> Self {
|
||||
Self(value)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn get(self) -> u32 {
|
||||
self.0.get()
|
||||
}
|
||||
}
|
||||
|
||||
/// 观测记录的稳定键集游标。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ObservabilityCursor {
|
||||
pub observed_at: DateTime<Utc>,
|
||||
pub stable_id: String,
|
||||
}
|
||||
|
||||
/// 请求结果状态。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum RequestOutcome {
|
||||
@@ -146,8 +123,7 @@ pub struct UsageFilter {
|
||||
pub struct UsageQuery {
|
||||
pub range: TimeRange,
|
||||
pub filter: UsageFilter,
|
||||
pub cursor: Option<ObservabilityCursor>,
|
||||
pub page: PageNumber,
|
||||
pub current_page: u32,
|
||||
pub page_size: PageSize,
|
||||
}
|
||||
|
||||
@@ -159,6 +135,7 @@ pub struct OpsErrorFilter {
|
||||
pub provider_kind: Option<String>,
|
||||
pub provider_account_ref: Option<String>,
|
||||
pub operation: Option<String>,
|
||||
pub endpoint: Option<String>,
|
||||
pub model: Option<String>,
|
||||
pub transport: Option<String>,
|
||||
pub attempt_index: Option<u32>,
|
||||
@@ -174,8 +151,7 @@ pub struct OpsErrorFilter {
|
||||
pub struct OpsErrorQuery {
|
||||
pub range: TimeRange,
|
||||
pub filter: OpsErrorFilter,
|
||||
pub cursor: Option<ObservabilityCursor>,
|
||||
pub page: PageNumber,
|
||||
pub current_page: u32,
|
||||
pub page_size: PageSize,
|
||||
}
|
||||
|
||||
@@ -545,13 +521,60 @@ pub struct DashboardTotals {
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct DashboardObservation {
|
||||
pub range: TimeRange,
|
||||
pub requests: RequestMetrics,
|
||||
pub attempts: AttemptMetrics,
|
||||
pub totals: DashboardTotals,
|
||||
pub provider_accounts: AccountPoolMetrics,
|
||||
pub trend: Vec<RequestMetricPoint>,
|
||||
pub account_usage: Vec<DashboardAccountUsage>,
|
||||
pub recent_requests: Vec<UsageRecord>,
|
||||
pub recent_requests: Vec<UsageListRecord>,
|
||||
}
|
||||
|
||||
/// 使用记录表格的窄读模型。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct UsageListRecord {
|
||||
pub id: String,
|
||||
pub endpoint: String,
|
||||
pub client_transport: String,
|
||||
pub requested_model_id: Option<String>,
|
||||
pub provider_kind: Option<String>,
|
||||
pub provider_account_ref: Option<String>,
|
||||
pub provider_account_name: Option<String>,
|
||||
pub provider_account_email: Option<String>,
|
||||
pub provider_account_authentication_kind: Option<String>,
|
||||
pub upstream_model_id: Option<String>,
|
||||
pub upstream_transport: Option<String>,
|
||||
pub service_tier: Option<String>,
|
||||
pub input_tokens: Option<u64>,
|
||||
pub output_tokens: Option<u64>,
|
||||
pub cached_tokens: Option<u64>,
|
||||
pub cache_write_tokens: Option<u64>,
|
||||
pub reasoning_tokens: Option<u64>,
|
||||
pub image_input_tokens: Option<u64>,
|
||||
pub image_output_tokens: Option<u64>,
|
||||
pub total_tokens: Option<u64>,
|
||||
pub cost_source: String,
|
||||
pub cost_amount: Option<DecimalAmount>,
|
||||
pub cost_currency: Option<String>,
|
||||
pub billing: Option<UsageBilling>,
|
||||
pub transport_decision_wait_ms: Option<u64>,
|
||||
pub connect_ms: Option<u64>,
|
||||
pub headers_ms: Option<u64>,
|
||||
pub first_event_ms: Option<u64>,
|
||||
pub first_reasoning_ms: Option<u64>,
|
||||
pub first_text_ms: Option<u64>,
|
||||
pub first_token_ms: Option<u64>,
|
||||
pub provider_processing_ms: Option<u64>,
|
||||
pub latency_ms: Option<u64>,
|
||||
pub admission_decision_ms: Option<u64>,
|
||||
pub account_selection_wait_ms: Option<u64>,
|
||||
pub capacity_used_slots: Option<u64>,
|
||||
pub capacity_total_slots: Option<u64>,
|
||||
pub client_ip: Option<String>,
|
||||
pub user_agent: Option<String>,
|
||||
pub reasoning_effort: Option<String>,
|
||||
pub reasoning_preset: Option<String>,
|
||||
pub subagent_kind: Option<String>,
|
||||
pub compact: bool,
|
||||
pub started_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
/// 一次完整模型请求的公共观测记录。
|
||||
@@ -635,9 +658,10 @@ pub struct UsageRecord {
|
||||
/// 用量分页结果。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct UsagePage {
|
||||
pub items: Vec<UsageRecord>,
|
||||
pub items: Vec<UsageListRecord>,
|
||||
pub current_page: u32,
|
||||
pub page_size: u16,
|
||||
pub total: u64,
|
||||
pub next_cursor: Option<ObservabilityCursor>,
|
||||
}
|
||||
|
||||
/// 请求中的一次上游尝试或运维事件。
|
||||
@@ -739,6 +763,10 @@ pub struct OpsError {
|
||||
pub client_api_key_ref: Option<String>,
|
||||
pub component: String,
|
||||
pub operation: String,
|
||||
pub protocol: Option<String>,
|
||||
pub client_transport: Option<String>,
|
||||
pub requested_model_id: Option<String>,
|
||||
pub service_tier: Option<String>,
|
||||
pub endpoint: Option<String>,
|
||||
pub provider_kind: Option<String>,
|
||||
pub provider_account_ref: Option<String>,
|
||||
@@ -756,6 +784,13 @@ pub struct OpsError {
|
||||
pub latency_ms: Option<u64>,
|
||||
pub message: String,
|
||||
pub occurrence_count: u32,
|
||||
pub client_ip: Option<String>,
|
||||
pub user_agent: Option<String>,
|
||||
pub reasoning_effort: Option<String>,
|
||||
pub reasoning_preset: Option<String>,
|
||||
pub request_kind: Option<String>,
|
||||
pub subagent_kind: Option<String>,
|
||||
pub compact: Option<bool>,
|
||||
pub occurred_at: DateTime<Utc>,
|
||||
pub stable_sort_id: String,
|
||||
}
|
||||
@@ -764,8 +799,9 @@ pub struct OpsError {
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct OpsErrorPage {
|
||||
pub items: Vec<OpsError>,
|
||||
pub current_page: u32,
|
||||
pub page_size: u16,
|
||||
pub total: u64,
|
||||
pub next_cursor: Option<ObservabilityCursor>,
|
||||
}
|
||||
|
||||
/// 仪表盘趋势指标。
|
||||
|
||||
@@ -1,9 +1,11 @@
|
||||
//! Provider 管理能力交换的中立 Command 与 Result。
|
||||
|
||||
use std::fmt;
|
||||
use std::{fmt, pin::Pin};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use bytes::Bytes;
|
||||
use chrono::{DateTime, NaiveDate, Utc};
|
||||
use futures::Stream;
|
||||
use gateway_core::{
|
||||
engine::credential::{OpaqueProviderData, ProviderAccountId, ProviderAccountIdentity},
|
||||
routing::{ProviderKind, UpstreamModelId},
|
||||
@@ -679,6 +681,35 @@ pub struct ProviderProfileStatistics {
|
||||
pub activity_insights: ProviderProfileActivityInsights,
|
||||
}
|
||||
|
||||
/// Provider 头像正文的中立字节流;公共层不解释 MIME,也不限制总字节数。
|
||||
pub type ProviderProfileAvatarStream =
|
||||
Pin<Box<dyn Stream<Item = Result<Bytes, ProviderProfileAvatarStreamError>> + Send + 'static>>;
|
||||
|
||||
/// Provider 头像流在响应开始后的通用失败,不携带上游 URL 或正文。
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
#[error("provider profile avatar stream failed")]
|
||||
pub struct ProviderProfileAvatarStreamError;
|
||||
|
||||
/// Provider 已校验并打开的官方头像响应。
|
||||
pub struct ProviderProfileAvatar {
|
||||
pub content_type: Option<String>,
|
||||
pub content_length: Option<u64>,
|
||||
pub etag: Option<String>,
|
||||
pub body: ProviderProfileAvatarStream,
|
||||
}
|
||||
|
||||
impl fmt::Debug for ProviderProfileAvatar {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter
|
||||
.debug_struct("ProviderProfileAvatar")
|
||||
.field("content_type", &self.content_type)
|
||||
.field("content_length", &self.content_length)
|
||||
.field("etag", &self.etag)
|
||||
.field("body", &"<stream>")
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
/// Provider 返回的一张安全主动额度重置卡。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ProviderResetCredit {
|
||||
|
||||
@@ -17,8 +17,8 @@ use crate::model::provider_credentials::{
|
||||
PendingAuthorizationMutation, PrepareCredentialImport, PrepareCredentialRefresh,
|
||||
PrepareCredentialRotation, PreparedAuthorizationCommit, PreparedCredentialImport,
|
||||
PreparedCredentialRotation, ProviderExport, ProviderExportCredentialInput, ProviderModels,
|
||||
ProviderProfileStatistics, ProviderQuota, ProviderQuotaRequest, ProviderResetCreditResult,
|
||||
ProviderResetCredits,
|
||||
ProviderProfileAvatar, ProviderProfileStatistics, ProviderQuota, ProviderQuotaRequest,
|
||||
ProviderResetCreditResult, ProviderResetCredits,
|
||||
};
|
||||
|
||||
/// Provider 管理失败的稳定分类。
|
||||
@@ -157,6 +157,14 @@ pub trait ProviderAdmin: Send + Sync {
|
||||
Err(ProviderAdminError::new(ProviderAdminErrorKind::Unsupported))
|
||||
}
|
||||
|
||||
/// 打开 Provider 官方头像字节流;不支持该能力的 Provider 使用默认拒绝。
|
||||
async fn profile_avatar(
|
||||
&self,
|
||||
_account_id: &ProviderAccountId,
|
||||
) -> Result<ProviderProfileAvatar, ProviderAdminError> {
|
||||
Err(ProviderAdminError::new(ProviderAdminErrorKind::Unsupported))
|
||||
}
|
||||
|
||||
/// 查询 Provider 主动额度重置卡;不支持该能力的 Provider 使用默认拒绝。
|
||||
async fn reset_credits(
|
||||
&self,
|
||||
|
||||
@@ -11,13 +11,13 @@ use super::backup::BackupStorePorts;
|
||||
use crate::model::{
|
||||
MutationContext, Revision,
|
||||
account_groups::{
|
||||
AccountGroupListQuery, AccountGroupMutation, AccountGroupPage, DeleteAccountGroup,
|
||||
NewAccountGroup, SetAccountGroupEnabled, UpdateAccountGroup,
|
||||
AccountGroupListQuery, AccountGroupMemberFact, AccountGroupMutation, AccountGroupPage,
|
||||
DeleteAccountGroup, NewAccountGroup, SetAccountGroupEnabled, UpdateAccountGroup,
|
||||
},
|
||||
accounts::{
|
||||
AccountListQuery, AccountPage, AccountPageItem, AccountUpdateResult, AccountUsage,
|
||||
AccountUsageWindowQuery, AccountUsageWindowResult, AccountsUpdateResult,
|
||||
BatchUpdateAccounts, DeleteAccounts, UpdateAccount,
|
||||
AccountListQuery, AccountPage, AccountPageItem, AccountRuntimeSnapshot,
|
||||
AccountUpdateResult, AccountUsage, AccountUsageWindowQuery, AccountUsageWindowResult,
|
||||
AccountsUpdateResult, BatchUpdateAccounts, DeleteAccounts, UpdateAccount,
|
||||
},
|
||||
auth::{AdminAuditEvent, AdminSession},
|
||||
client_keys::{
|
||||
@@ -86,9 +86,17 @@ pub type AdminStoreResult<T> = Result<T, AdminStoreError>;
|
||||
/// 账号目录与公共账号写操作。
|
||||
#[async_trait]
|
||||
pub trait AccountStore: Send + Sync {
|
||||
async fn list_accounts(&self, query: AccountListQuery) -> AdminStoreResult<AccountPage>;
|
||||
async fn list_accounts(
|
||||
&self,
|
||||
query: AccountListQuery,
|
||||
runtime: AccountRuntimeSnapshot,
|
||||
) -> AdminStoreResult<AccountPage>;
|
||||
|
||||
async fn load_account(&self, account_id: &str) -> AdminStoreResult<Option<AccountPageItem>>;
|
||||
async fn load_account(
|
||||
&self,
|
||||
account_id: &str,
|
||||
runtime: AccountRuntimeSnapshot,
|
||||
) -> AdminStoreResult<Option<AccountPageItem>>;
|
||||
|
||||
async fn load_account_usage(
|
||||
&self,
|
||||
@@ -174,6 +182,17 @@ pub trait AccountStore: Send + Sync {
|
||||
) -> AdminStoreResult<()>;
|
||||
}
|
||||
|
||||
/// 可丢失账号运行态的管理读端口;跨存储编排由 Admin application service 拥有。
|
||||
#[async_trait]
|
||||
pub trait AccountRuntimeStore: Send + Sync {
|
||||
async fn active_rate_limits(&self) -> AdminStoreResult<AccountRuntimeSnapshot>;
|
||||
|
||||
async fn account_runtime(
|
||||
&self,
|
||||
account_ids: &[String],
|
||||
) -> AdminStoreResult<AccountRuntimeSnapshot>;
|
||||
}
|
||||
|
||||
/// 管理员密码、会话和安全审计。
|
||||
#[async_trait]
|
||||
pub trait AuthStore: Send + Sync {
|
||||
@@ -240,6 +259,11 @@ pub trait AccountGroupStore: Send + Sync {
|
||||
query: AccountGroupListQuery,
|
||||
) -> AdminStoreResult<AccountGroupPage>;
|
||||
|
||||
async fn load_account_group_members(
|
||||
&self,
|
||||
group_ids: &[gateway_core::routing::AccountGroupId],
|
||||
) -> AdminStoreResult<Vec<AccountGroupMemberFact>>;
|
||||
|
||||
async fn create_account_group(
|
||||
&self,
|
||||
command: NewAccountGroup,
|
||||
@@ -346,13 +370,35 @@ pub trait SettingsStore: Send + Sync {
|
||||
) -> AdminStoreResult<AdminApiKeyMutation>;
|
||||
}
|
||||
|
||||
/// 账号目录、运行态与分组所需的 Store 能力集合。
|
||||
#[derive(Clone)]
|
||||
pub struct AdminAccountStorePorts {
|
||||
accounts: Arc<dyn AccountStore>,
|
||||
runtime: Arc<dyn AccountRuntimeStore>,
|
||||
groups: Arc<dyn AccountGroupStore>,
|
||||
}
|
||||
|
||||
impl AdminAccountStorePorts {
|
||||
#[must_use]
|
||||
pub fn new(
|
||||
accounts: Arc<dyn AccountStore>,
|
||||
runtime: Arc<dyn AccountRuntimeStore>,
|
||||
groups: Arc<dyn AccountGroupStore>,
|
||||
) -> Self {
|
||||
Self {
|
||||
accounts,
|
||||
runtime,
|
||||
groups,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 管理用例所需能力的封闭集合。
|
||||
///
|
||||
/// 字段保持私有,每个 getter 只交出一种明确能力。该类型不提供通用拆包入口。
|
||||
#[derive(Clone)]
|
||||
pub struct AdminStorePorts {
|
||||
accounts: Arc<dyn AccountStore>,
|
||||
account_groups: Arc<dyn AccountGroupStore>,
|
||||
accounts: AdminAccountStorePorts,
|
||||
auth: Arc<dyn AuthStore>,
|
||||
client_keys: Arc<dyn ClientKeyStore>,
|
||||
observability: Arc<dyn ObservabilityStore>,
|
||||
@@ -363,8 +409,7 @@ pub struct AdminStorePorts {
|
||||
impl AdminStorePorts {
|
||||
#[must_use]
|
||||
pub fn new(
|
||||
accounts: Arc<dyn AccountStore>,
|
||||
account_groups: Arc<dyn AccountGroupStore>,
|
||||
accounts: AdminAccountStorePorts,
|
||||
auth: Arc<dyn AuthStore>,
|
||||
client_keys: Arc<dyn ClientKeyStore>,
|
||||
observability: Arc<dyn ObservabilityStore>,
|
||||
@@ -373,7 +418,6 @@ impl AdminStorePorts {
|
||||
) -> Self {
|
||||
Self {
|
||||
accounts,
|
||||
account_groups,
|
||||
auth,
|
||||
client_keys,
|
||||
observability,
|
||||
@@ -384,12 +428,17 @@ impl AdminStorePorts {
|
||||
|
||||
#[must_use]
|
||||
pub fn accounts(&self) -> Arc<dyn AccountStore> {
|
||||
self.accounts.clone()
|
||||
self.accounts.accounts.clone()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn account_runtime(&self) -> Arc<dyn AccountRuntimeStore> {
|
||||
self.accounts.runtime.clone()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn account_groups(&self) -> Arc<dyn AccountGroupStore> {
|
||||
self.account_groups.clone()
|
||||
self.accounts.groups.clone()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
|
||||
@@ -1,20 +1,29 @@
|
||||
//! Account group management use cases.
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::{
|
||||
collections::{BTreeMap, BTreeSet},
|
||||
sync::Arc,
|
||||
time::SystemTime,
|
||||
};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use gateway_core::routing::{AccountGroupId, snapshot::SnapshotControl};
|
||||
use gateway_core::{
|
||||
engine::credential::{AccountStatus, resolve_account_status},
|
||||
routing::{AccountGroupId, snapshot::SnapshotControl},
|
||||
};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::{
|
||||
model::{
|
||||
AdminError, MutationContext,
|
||||
account_groups::{
|
||||
AccountGroupListQuery, AccountGroupMutation, AccountGroupPage, CreateAccountGroup,
|
||||
DeleteAccountGroup, NewAccountGroup, SetAccountGroupEnabled, UpdateAccountGroup,
|
||||
AccountGroupAccountSummary, AccountGroupCapacity, AccountGroupListQuery,
|
||||
AccountGroupMemberFact, AccountGroupMutation, AccountGroupPage, AccountGroupRecord,
|
||||
CreateAccountGroup, DeleteAccountGroup, NewAccountGroup, SetAccountGroupEnabled,
|
||||
UpdateAccountGroup,
|
||||
},
|
||||
},
|
||||
ports::store::AccountGroupStore,
|
||||
ports::store::{AccountGroupStore, AccountRuntimeStore},
|
||||
};
|
||||
|
||||
use super::{map_store_error, publish_committed};
|
||||
@@ -47,6 +56,7 @@ pub trait AccountGroupService: Send + Sync {
|
||||
|
||||
pub(crate) struct DefaultAccountGroupService {
|
||||
store: Arc<dyn AccountGroupStore>,
|
||||
runtime: Arc<dyn AccountRuntimeStore>,
|
||||
snapshot: Arc<dyn SnapshotControl>,
|
||||
}
|
||||
|
||||
@@ -54,16 +64,54 @@ impl DefaultAccountGroupService {
|
||||
#[must_use]
|
||||
pub(crate) fn new(
|
||||
store: Arc<dyn AccountGroupStore>,
|
||||
runtime: Arc<dyn AccountRuntimeStore>,
|
||||
snapshot: Arc<dyn SnapshotControl>,
|
||||
) -> Self {
|
||||
Self { store, snapshot }
|
||||
Self {
|
||||
store,
|
||||
runtime,
|
||||
snapshot,
|
||||
}
|
||||
}
|
||||
|
||||
async fn enrich_records(&self, records: &mut [AccountGroupRecord]) -> Result<(), AdminError> {
|
||||
if records.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
let group_ids = records
|
||||
.iter()
|
||||
.map(|record| record.id.clone())
|
||||
.collect::<Vec<_>>();
|
||||
let members = self
|
||||
.store
|
||||
.load_account_group_members(&group_ids)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "account group members"))?;
|
||||
let account_ids = members
|
||||
.iter()
|
||||
.map(|member| member.account_id.clone())
|
||||
.collect::<BTreeSet<_>>()
|
||||
.into_iter()
|
||||
.collect::<Vec<_>>();
|
||||
let runtime = match self.runtime.account_runtime(&account_ids).await {
|
||||
Ok(runtime) => runtime,
|
||||
Err(error) => {
|
||||
tracing::warn!(error = %error, "account group runtime projection is unavailable");
|
||||
Default::default()
|
||||
}
|
||||
};
|
||||
project_group_runtime(records, &members, &runtime);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
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"))?;
|
||||
let mut result = result.map_err(|error| map_store_error(error, "account group"))?;
|
||||
if let Some(record) = result.record.as_mut() {
|
||||
self.enrich_records(std::slice::from_mut(record)).await?;
|
||||
}
|
||||
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
|
||||
Ok(result)
|
||||
}
|
||||
@@ -72,10 +120,13 @@ impl DefaultAccountGroupService {
|
||||
#[async_trait]
|
||||
impl AccountGroupService for DefaultAccountGroupService {
|
||||
async fn list(&self, query: AccountGroupListQuery) -> Result<AccountGroupPage, AdminError> {
|
||||
self.store
|
||||
let mut page = self
|
||||
.store
|
||||
.list_account_groups(query)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "account group"))
|
||||
.map_err(|error| map_store_error(error, "account group"))?;
|
||||
self.enrich_records(&mut page.items).await?;
|
||||
Ok(page)
|
||||
}
|
||||
|
||||
async fn create(
|
||||
@@ -128,3 +179,63 @@ impl AccountGroupService for DefaultAccountGroupService {
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
fn project_group_runtime(
|
||||
records: &mut [AccountGroupRecord],
|
||||
members: &[AccountGroupMemberFact],
|
||||
runtime: &crate::model::accounts::AccountRuntimeSnapshot,
|
||||
) {
|
||||
let now = SystemTime::now();
|
||||
let mut status_by_account = BTreeMap::new();
|
||||
let mut slots_by_account = BTreeMap::new();
|
||||
let mut accounts_by_group = BTreeMap::<&str, Vec<&str>>::new();
|
||||
for member in members {
|
||||
let mut status = member.status.clone();
|
||||
status.rate_limited_until = runtime
|
||||
.rate_limited_until
|
||||
.get(&member.account_id)
|
||||
.copied()
|
||||
.map(Into::into);
|
||||
status_by_account
|
||||
.entry(member.account_id.as_str())
|
||||
.or_insert_with(|| resolve_account_status(&status, now).status);
|
||||
slots_by_account
|
||||
.entry(member.account_id.as_str())
|
||||
.or_insert(member.total_slots);
|
||||
accounts_by_group
|
||||
.entry(member.group_id.as_str())
|
||||
.or_default()
|
||||
.push(member.account_id.as_str());
|
||||
}
|
||||
for record in records {
|
||||
let account_ids = accounts_by_group
|
||||
.get(record.id.as_str())
|
||||
.map_or(&[][..], Vec::as_slice);
|
||||
let available_ids = account_ids
|
||||
.iter()
|
||||
.copied()
|
||||
.filter(|account_id| status_by_account.get(account_id) == Some(&AccountStatus::Normal))
|
||||
.collect::<Vec<_>>();
|
||||
let total = u64::try_from(account_ids.len()).unwrap_or(u64::MAX);
|
||||
let available = u64::try_from(available_ids.len()).unwrap_or(u64::MAX);
|
||||
record.account_summary = AccountGroupAccountSummary {
|
||||
available,
|
||||
limited: total.saturating_sub(available),
|
||||
total,
|
||||
};
|
||||
record.capacity = AccountGroupCapacity {
|
||||
used_slots: if available_ids.is_empty() {
|
||||
Some(0)
|
||||
} else {
|
||||
runtime.in_flight.as_ref().map(|in_flight| {
|
||||
available_ids.iter().fold(0_u64, |sum, account_id| {
|
||||
sum.saturating_add(in_flight.get(*account_id).copied().unwrap_or(0))
|
||||
})
|
||||
})
|
||||
},
|
||||
total_slots: available_ids.iter().fold(0_u64, |sum, account_id| {
|
||||
sum.saturating_add(slots_by_account.get(*account_id).copied().unwrap_or(0))
|
||||
}),
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
@@ -25,11 +25,15 @@ use crate::{
|
||||
provider_credentials::{
|
||||
AccountDirectoryItem, AccountDirectoryPage, AccountExportBundle, AccountRefreshResult,
|
||||
ConsumeProviderResetCredit, PrepareCredentialRefresh, ProviderModels,
|
||||
ProviderProfileStatistics, ProviderQuota, ProviderQuotaRequest, ProviderQuotaWindow,
|
||||
ProviderResetCreditResult, ProviderResetCredits, QuotaLocalUsageAttribution,
|
||||
ProviderProfileAvatar, ProviderProfileStatistics, ProviderQuota, ProviderQuotaRequest,
|
||||
ProviderQuotaWindow, ProviderResetCreditResult, ProviderResetCredits,
|
||||
QuotaLocalUsageAttribution,
|
||||
},
|
||||
},
|
||||
ports::{provider::ProviderAdminRegistry, store::AccountStore},
|
||||
ports::{
|
||||
provider::ProviderAdminRegistry,
|
||||
store::{AccountRuntimeStore, AccountStore},
|
||||
},
|
||||
};
|
||||
|
||||
use super::{
|
||||
@@ -87,6 +91,13 @@ pub trait AccountsService: Send + Sync {
|
||||
Err(AdminError::invalid("当前 Provider 不支持账号统计"))
|
||||
}
|
||||
|
||||
async fn profile_avatar(
|
||||
&self,
|
||||
_account_id: &ProviderAccountId,
|
||||
) -> Result<ProviderProfileAvatar, AdminError> {
|
||||
Err(AdminError::invalid("当前 Provider 不支持账号头像"))
|
||||
}
|
||||
|
||||
async fn reset_credits(
|
||||
&self,
|
||||
_context: &MutationContext,
|
||||
@@ -118,6 +129,7 @@ pub trait AccountsService: Send + Sync {
|
||||
|
||||
pub(crate) struct DefaultAccountsService {
|
||||
accounts: Arc<dyn AccountStore>,
|
||||
account_runtime: Arc<dyn AccountRuntimeStore>,
|
||||
providers: ProviderAdminRegistry,
|
||||
snapshot: Arc<dyn SnapshotControl>,
|
||||
probe: Arc<dyn AccountProbe>,
|
||||
@@ -129,12 +141,14 @@ impl DefaultAccountsService {
|
||||
#[must_use]
|
||||
pub(crate) fn new(
|
||||
accounts: Arc<dyn AccountStore>,
|
||||
account_runtime: Arc<dyn AccountRuntimeStore>,
|
||||
providers: ProviderAdminRegistry,
|
||||
snapshot: Arc<dyn SnapshotControl>,
|
||||
probe: Arc<dyn AccountProbe>,
|
||||
) -> Self {
|
||||
Self {
|
||||
accounts,
|
||||
account_runtime,
|
||||
providers,
|
||||
snapshot,
|
||||
probe,
|
||||
@@ -158,8 +172,13 @@ impl DefaultAccountsService {
|
||||
&self,
|
||||
account_id: &ProviderAccountId,
|
||||
) -> Result<AccountPageItem, AdminError> {
|
||||
let runtime = self
|
||||
.account_runtime
|
||||
.account_runtime(&[account_id.as_str().to_owned()])
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "account runtime"))?;
|
||||
self.accounts
|
||||
.load_account(account_id.as_str())
|
||||
.load_account(account_id.as_str(), runtime)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "provider account"))?
|
||||
.ok_or_else(|| AdminError::not_found("Provider 账号不存在"))
|
||||
@@ -269,9 +288,14 @@ impl DefaultAccountsService {
|
||||
#[async_trait]
|
||||
impl AccountsService for DefaultAccountsService {
|
||||
async fn list(&self, query: AccountListQuery) -> Result<AccountDirectoryPage, AdminError> {
|
||||
let runtime = self
|
||||
.account_runtime
|
||||
.active_rate_limits()
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "account runtime"))?;
|
||||
let page = self
|
||||
.accounts
|
||||
.list_accounts(query)
|
||||
.list_accounts(query, runtime)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "account directory"))?;
|
||||
let now = Utc::now();
|
||||
@@ -550,6 +574,17 @@ impl AccountsService for DefaultAccountsService {
|
||||
.map_err(|error| map_provider_error(error, "provider profile statistics"))
|
||||
}
|
||||
|
||||
async fn profile_avatar(
|
||||
&self,
|
||||
account_id: &ProviderAccountId,
|
||||
) -> Result<ProviderProfileAvatar, AdminError> {
|
||||
let (_, provider) = self.provider_for_account(account_id).await?;
|
||||
provider
|
||||
.profile_avatar(account_id)
|
||||
.await
|
||||
.map_err(|error| map_provider_error(error, "provider profile avatar"))
|
||||
}
|
||||
|
||||
async fn reset_credits(
|
||||
&self,
|
||||
context: &MutationContext,
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use argon2::{Argon2, PasswordHash, PasswordHasher, PasswordVerifier, password_hash::SaltString};
|
||||
use argon2::{Argon2, PasswordHash, PasswordHasher, PasswordVerifier};
|
||||
use async_trait::async_trait;
|
||||
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
|
||||
use chrono::{Duration, Utc};
|
||||
@@ -195,9 +195,8 @@ impl AuthService for DefaultAuthService {
|
||||
}
|
||||
|
||||
fn hash_admin_password(password: &str) -> Result<String, AdminError> {
|
||||
let salt = SaltString::generate(&mut OsRng);
|
||||
Argon2::default()
|
||||
.hash_password(password.as_bytes(), &salt)
|
||||
.hash_password(password.as_bytes())
|
||||
.map(|hash| hash.to_string())
|
||||
.map_err(|_| AdminError::internal("管理员密码哈希失败"))
|
||||
}
|
||||
|
||||
@@ -1,9 +1,10 @@
|
||||
//! 观测查询、趋势、费用覆盖与健康阈值规则。
|
||||
|
||||
use std::{
|
||||
collections::{BTreeMap, btree_map::Entry},
|
||||
collections::{BTreeMap, HashMap, btree_map::Entry},
|
||||
str::FromStr,
|
||||
sync::Arc,
|
||||
time::{Duration as StdDuration, Instant},
|
||||
};
|
||||
|
||||
use async_trait::async_trait;
|
||||
@@ -39,6 +40,58 @@ const HEALTH_TIMELINE_SLOTS: i64 = 24 * 4;
|
||||
const HEALTH_TIMELINE_MIN_SAMPLE_SIZE: u64 = 10;
|
||||
const HEALTH_TIMELINE_UNAVAILABLE_FAILURE_THRESHOLD: u64 = 3;
|
||||
const HEALTH_TIMELINE_STABLE_RELIABILITY: f64 = 99.0;
|
||||
const DASHBOARD_QUERY_BUCKET_SECONDS: i64 = 2;
|
||||
const DASHBOARD_QUERY_CACHE_RETENTION: StdDuration = StdDuration::from_secs(2);
|
||||
|
||||
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
|
||||
struct DashboardQueryKey {
|
||||
start: DateTime<Utc>,
|
||||
end_bucket: i64,
|
||||
}
|
||||
|
||||
impl DashboardQueryKey {
|
||||
fn new(range: TimeRange) -> Self {
|
||||
Self {
|
||||
start: range.start,
|
||||
end_bucket: range
|
||||
.end
|
||||
.timestamp()
|
||||
.div_euclid(DASHBOARD_QUERY_BUCKET_SECONDS),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct DashboardQueryEntry {
|
||||
created_at: Instant,
|
||||
result: Arc<tokio::sync::OnceCell<Result<DashboardResult, AdminError>>>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct DashboardQueryCache {
|
||||
entries: tokio::sync::Mutex<HashMap<DashboardQueryKey, DashboardQueryEntry>>,
|
||||
}
|
||||
|
||||
impl DashboardQueryCache {
|
||||
async fn result_cell(
|
||||
&self,
|
||||
key: DashboardQueryKey,
|
||||
) -> Arc<tokio::sync::OnceCell<Result<DashboardResult, AdminError>>> {
|
||||
let now = Instant::now();
|
||||
let mut entries = self.entries.lock().await;
|
||||
entries.retain(|_, entry| {
|
||||
now.duration_since(entry.created_at) < DASHBOARD_QUERY_CACHE_RETENTION
|
||||
});
|
||||
Arc::clone(
|
||||
&entries
|
||||
.entry(key)
|
||||
.or_insert_with(|| DashboardQueryEntry {
|
||||
created_at: now,
|
||||
result: Arc::new(tokio::sync::OnceCell::new()),
|
||||
})
|
||||
.result,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// API 消费的观测控制面服务。
|
||||
#[async_trait]
|
||||
@@ -75,6 +128,7 @@ pub(crate) struct DefaultObservabilityService {
|
||||
store: Arc<dyn ObservabilityStore>,
|
||||
settings: Arc<dyn SettingsStore>,
|
||||
providers: ProviderAdminRegistry,
|
||||
dashboard_queries: DashboardQueryCache,
|
||||
}
|
||||
|
||||
impl DefaultObservabilityService {
|
||||
@@ -88,16 +142,13 @@ impl DefaultObservabilityService {
|
||||
store,
|
||||
settings,
|
||||
providers,
|
||||
dashboard_queries: DashboardQueryCache::default(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ObservabilityService for DefaultObservabilityService {
|
||||
async fn dashboard_summary(
|
||||
async fn load_dashboard_summary(
|
||||
&self,
|
||||
range: TimeRange,
|
||||
kind: TrendKind,
|
||||
) -> Result<DashboardResult, AdminError> {
|
||||
let observed_at = Utc::now();
|
||||
let (mut observation, settings, runtime_slots) = futures::try_join!(
|
||||
@@ -106,7 +157,7 @@ impl ObservabilityService for DefaultObservabilityService {
|
||||
self.store.dashboard_runtime_slots(observed_at),
|
||||
)
|
||||
.map_err(|error| map_store_error(error, "dashboard"))?;
|
||||
self.enrich_billing(&mut observation.recent_requests);
|
||||
self.enrich_list_billing(&mut observation.recent_requests);
|
||||
self.enrich_dashboard_quotas(&mut observation.account_usage)
|
||||
.await;
|
||||
let today_start = china_day_start(observation.range.end);
|
||||
@@ -119,11 +170,18 @@ impl ObservabilityService for DefaultObservabilityService {
|
||||
observation.totals.cached_tokens,
|
||||
observation.totals.input_tokens,
|
||||
);
|
||||
let average_first_token_latency_ms = average(
|
||||
observation.requests.first_token_latency_sum_ms,
|
||||
observation.requests.first_token_latency_count,
|
||||
);
|
||||
let trend = trend(kind, observation.trend.clone())?;
|
||||
let (first_token_latency_sum_ms, first_token_latency_count) = observation
|
||||
.trend
|
||||
.iter()
|
||||
.fold((0_u64, 0_u64), |(sum, count), point| {
|
||||
(
|
||||
sum.saturating_add(point.metrics.first_token_latency_sum_ms),
|
||||
count.saturating_add(point.metrics.first_token_latency_count),
|
||||
)
|
||||
});
|
||||
let average_first_token_latency_ms =
|
||||
average(first_token_latency_sum_ms, first_token_latency_count);
|
||||
let trend = trend(TrendKind::Usage, observation.trend.clone())?;
|
||||
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);
|
||||
@@ -161,6 +219,26 @@ impl ObservabilityService for DefaultObservabilityService {
|
||||
wire_profiles,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ObservabilityService for DefaultObservabilityService {
|
||||
async fn dashboard_summary(
|
||||
&self,
|
||||
range: TimeRange,
|
||||
kind: TrendKind,
|
||||
) -> Result<DashboardResult, AdminError> {
|
||||
let cell = self
|
||||
.dashboard_queries
|
||||
.result_cell(DashboardQueryKey::new(range))
|
||||
.await;
|
||||
let mut result = cell
|
||||
.get_or_init(|| self.load_dashboard_summary(range))
|
||||
.await
|
||||
.clone()?;
|
||||
result.trend = trend(kind, result.observation.trend.clone())?;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
async fn dashboard_trend(
|
||||
&self,
|
||||
@@ -181,7 +259,7 @@ impl ObservabilityService for DefaultObservabilityService {
|
||||
.list_usage_records(query)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "usage records"))?;
|
||||
self.enrich_billing(&mut page.items);
|
||||
self.enrich_list_billing(&mut page.items);
|
||||
Ok(page)
|
||||
}
|
||||
|
||||
@@ -194,7 +272,7 @@ impl ObservabilityService for DefaultObservabilityService {
|
||||
.usage_record_detail(request_id)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "usage record"))?;
|
||||
self.enrich_billing(std::slice::from_mut(&mut detail.request));
|
||||
self.enrich_detail_billing(std::slice::from_mut(&mut detail.request));
|
||||
Ok(detail)
|
||||
}
|
||||
|
||||
@@ -659,7 +737,39 @@ impl DefaultObservabilityService {
|
||||
|
||||
/// 逐条尽力把可校验的总额升级为完整分解;单条脏数据(非法 Provider kind、
|
||||
/// 不支持的来源或费用规则失败)只保留该条已存的总额,不影响整页返回。
|
||||
fn enrich_billing(&self, records: &mut [crate::model::observability::UsageRecord]) {
|
||||
fn enrich_list_billing(&self, records: &mut [crate::model::observability::UsageListRecord]) {
|
||||
for record in records {
|
||||
let Some(UsageBilling::Total { source, total }) = record.billing.as_ref() else {
|
||||
continue;
|
||||
};
|
||||
if !matches!(source.as_str(), "calculated" | "provider_reported") {
|
||||
continue;
|
||||
}
|
||||
let (Some(provider), Some(upstream_model_id)) = (
|
||||
record.provider_kind.as_deref(),
|
||||
record.upstream_model_id.as_ref(),
|
||||
) else {
|
||||
continue;
|
||||
};
|
||||
let Ok(provider_kind) = ProviderKind::new(provider.to_owned()) else {
|
||||
continue;
|
||||
};
|
||||
let input = ProviderBillingInput {
|
||||
upstream_model_id: upstream_model_id.clone(),
|
||||
service_tier: record.service_tier.clone(),
|
||||
input_tokens: record.input_tokens,
|
||||
output_tokens: record.output_tokens,
|
||||
cached_tokens: record.cached_tokens,
|
||||
cache_write_tokens: record.cache_write_tokens,
|
||||
total: total.clone(),
|
||||
};
|
||||
if let Ok(Some(breakdown)) = self.providers.calculated_billing(&provider_kind, &input) {
|
||||
record.billing = Some(UsageBilling::Calculated(Box::new(breakdown)));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn enrich_detail_billing(&self, records: &mut [crate::model::observability::UsageRecord]) {
|
||||
for record in records {
|
||||
let Some(UsageBilling::Total { source, total }) = record.billing.as_ref() else {
|
||||
continue;
|
||||
|
||||
@@ -0,0 +1,236 @@
|
||||
use std::{
|
||||
collections::BTreeMap,
|
||||
str::FromStr as _,
|
||||
sync::{Arc, Mutex},
|
||||
};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{Duration, Utc};
|
||||
use gateway_admin::{
|
||||
model::{
|
||||
MutationContext, PageSize, Revision,
|
||||
account_groups::{
|
||||
AccountGroupAccountSummary, AccountGroupCapacity, AccountGroupColor,
|
||||
AccountGroupListQuery, AccountGroupMemberFact, AccountGroupMutation, AccountGroupPage,
|
||||
AccountGroupRecord, AccountGroupUsage, DeleteAccountGroup, NewAccountGroup,
|
||||
SetAccountGroupEnabled, UpdateAccountGroup,
|
||||
},
|
||||
accounts::AccountRuntimeSnapshot,
|
||||
observability::DecimalAmount,
|
||||
},
|
||||
ports::store::{
|
||||
AccountGroupStore, AccountRuntimeStore, AdminStoreError, AdminStoreErrorKind,
|
||||
AdminStoreResult,
|
||||
},
|
||||
};
|
||||
use gateway_core::{
|
||||
engine::credential::{AccountStatusFacts, CredentialState, QuotaState},
|
||||
routing::AccountGroupId,
|
||||
};
|
||||
|
||||
use super::AdminHarness;
|
||||
|
||||
const GROUP_ID: &str = "grp_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa";
|
||||
|
||||
#[tokio::test]
|
||||
async fn group_query_service_enriches_only_current_page_members_with_runtime_facts() {
|
||||
let groups = Arc::new(FakeGroupStore::default());
|
||||
let runtime = Arc::new(FakeRuntimeStore::default());
|
||||
let service = AdminHarness::new()
|
||||
.account_groups(groups.clone())
|
||||
.account_runtime(runtime.clone())
|
||||
.build()
|
||||
.await;
|
||||
|
||||
let page = service
|
||||
.account_groups()
|
||||
.list(AccountGroupListQuery {
|
||||
page: 1,
|
||||
page_size: PageSize::new(20).expect("page size"),
|
||||
search: None,
|
||||
enabled: None,
|
||||
})
|
||||
.await
|
||||
.expect("list projected account groups");
|
||||
|
||||
assert_eq!(
|
||||
groups
|
||||
.requested_groups
|
||||
.lock()
|
||||
.expect("requested groups")
|
||||
.as_slice(),
|
||||
[GROUP_ID]
|
||||
);
|
||||
assert_eq!(
|
||||
runtime
|
||||
.requested_accounts
|
||||
.lock()
|
||||
.expect("requested accounts")
|
||||
.as_slice(),
|
||||
["acct_available", "acct_limited"]
|
||||
);
|
||||
let group = &page.items[0];
|
||||
assert_eq!(
|
||||
group.account_summary,
|
||||
AccountGroupAccountSummary {
|
||||
available: 1,
|
||||
limited: 1,
|
||||
total: 2,
|
||||
}
|
||||
);
|
||||
assert_eq!(
|
||||
group.capacity,
|
||||
AccountGroupCapacity {
|
||||
used_slots: Some(2),
|
||||
total_slots: 4,
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct FakeGroupStore {
|
||||
requested_groups: Mutex<Vec<String>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AccountGroupStore for FakeGroupStore {
|
||||
async fn list_account_groups(
|
||||
&self,
|
||||
query: AccountGroupListQuery,
|
||||
) -> AdminStoreResult<AccountGroupPage> {
|
||||
Ok(AccountGroupPage {
|
||||
config_revision: Revision::new(1).expect("revision"),
|
||||
items: vec![group_record()],
|
||||
total: 1,
|
||||
page: query.page,
|
||||
page_size: query.page_size.get(),
|
||||
})
|
||||
}
|
||||
|
||||
async fn load_account_group_members(
|
||||
&self,
|
||||
group_ids: &[AccountGroupId],
|
||||
) -> AdminStoreResult<Vec<AccountGroupMemberFact>> {
|
||||
*self.requested_groups.lock().expect("requested groups") = group_ids
|
||||
.iter()
|
||||
.map(|group_id| group_id.as_str().to_owned())
|
||||
.collect();
|
||||
Ok(vec![member("acct_available", 4), member("acct_limited", 3)])
|
||||
}
|
||||
|
||||
async fn create_account_group(
|
||||
&self,
|
||||
_: NewAccountGroup,
|
||||
_: &MutationContext,
|
||||
) -> AdminStoreResult<AccountGroupMutation> {
|
||||
Err(unused())
|
||||
}
|
||||
|
||||
async fn update_account_group(
|
||||
&self,
|
||||
_: UpdateAccountGroup,
|
||||
_: &MutationContext,
|
||||
) -> AdminStoreResult<AccountGroupMutation> {
|
||||
Err(unused())
|
||||
}
|
||||
|
||||
async fn set_account_group_enabled(
|
||||
&self,
|
||||
_: SetAccountGroupEnabled,
|
||||
_: &MutationContext,
|
||||
) -> AdminStoreResult<AccountGroupMutation> {
|
||||
Err(unused())
|
||||
}
|
||||
|
||||
async fn delete_account_group(
|
||||
&self,
|
||||
_: DeleteAccountGroup,
|
||||
_: &MutationContext,
|
||||
) -> AdminStoreResult<AccountGroupMutation> {
|
||||
Err(unused())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct FakeRuntimeStore {
|
||||
requested_accounts: Mutex<Vec<String>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AccountRuntimeStore for FakeRuntimeStore {
|
||||
async fn active_rate_limits(&self) -> AdminStoreResult<AccountRuntimeSnapshot> {
|
||||
Ok(AccountRuntimeSnapshot::default())
|
||||
}
|
||||
|
||||
async fn account_runtime(
|
||||
&self,
|
||||
account_ids: &[String],
|
||||
) -> AdminStoreResult<AccountRuntimeSnapshot> {
|
||||
*self.requested_accounts.lock().expect("requested accounts") = account_ids.to_vec();
|
||||
Ok(AccountRuntimeSnapshot {
|
||||
rate_limited_until: BTreeMap::from([(
|
||||
"acct_limited".to_owned(),
|
||||
Utc::now() + Duration::minutes(5),
|
||||
)]),
|
||||
in_flight: Some(BTreeMap::from([("acct_available".to_owned(), 2)])),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn group_record() -> AccountGroupRecord {
|
||||
let now = Utc::now();
|
||||
AccountGroupRecord {
|
||||
id: group_id(),
|
||||
name: "Primary".to_owned(),
|
||||
description: None,
|
||||
color: AccountGroupColor::parse("#2563EBFF").expect("color"),
|
||||
enabled: true,
|
||||
member_count: 2,
|
||||
provider_counts: BTreeMap::from([("openai".to_owned(), 2)]),
|
||||
client_key_count: 1,
|
||||
account_summary: AccountGroupAccountSummary {
|
||||
available: 0,
|
||||
limited: 0,
|
||||
total: 0,
|
||||
},
|
||||
capacity: AccountGroupCapacity {
|
||||
used_slots: None,
|
||||
total_slots: 0,
|
||||
},
|
||||
usage: AccountGroupUsage {
|
||||
today_usd: DecimalAmount::from_str("1").expect("today usage"),
|
||||
retained_total_usd: DecimalAmount::from_str("3").expect("retained usage"),
|
||||
},
|
||||
created_at: now,
|
||||
updated_at: now,
|
||||
}
|
||||
}
|
||||
|
||||
fn member(account_id: &str, total_slots: u64) -> AccountGroupMemberFact {
|
||||
AccountGroupMemberFact {
|
||||
group_id: group_id(),
|
||||
account_id: account_id.to_owned(),
|
||||
status: AccountStatusFacts {
|
||||
enabled: true,
|
||||
credential_state: CredentialState::Ready,
|
||||
access_token_expires_at: None,
|
||||
quota: QuotaState::default(),
|
||||
rate_limited_until: None,
|
||||
last_error_reason: None,
|
||||
last_error_message: None,
|
||||
},
|
||||
total_slots,
|
||||
}
|
||||
}
|
||||
|
||||
fn group_id() -> AccountGroupId {
|
||||
AccountGroupId::new(GROUP_ID).expect("group ID")
|
||||
}
|
||||
|
||||
fn unused() -> AdminStoreError {
|
||||
AdminStoreError::new(
|
||||
AdminStoreErrorKind::Unavailable,
|
||||
"account group",
|
||||
"unused test operation",
|
||||
)
|
||||
}
|
||||
@@ -26,8 +26,8 @@ use gateway_admin::{
|
||||
AdminError, MutationContext, Revision,
|
||||
accounts::{
|
||||
AccountConnectionTestEvent, AccountListQuery, AccountPage, AccountPageItem,
|
||||
AccountRecord, AccountSummary, AccountUpdateResult, AccountUsage,
|
||||
AccountUsageWindowQuery, AccountUsageWindowResult, AccountsUpdateResult,
|
||||
AccountRecord, AccountRuntimeSnapshot, AccountSummary, AccountUpdateResult,
|
||||
AccountUsage, AccountUsageWindowQuery, AccountUsageWindowResult, AccountsUpdateResult,
|
||||
BatchUpdateAccounts, DeleteAccounts, UpdateAccount,
|
||||
},
|
||||
observability::TimeRange,
|
||||
@@ -507,7 +507,11 @@ impl FakeAccountStore {
|
||||
|
||||
#[async_trait]
|
||||
impl AccountStore for FakeAccountStore {
|
||||
async fn list_accounts(&self, _: AccountListQuery) -> AdminStoreResult<AccountPage> {
|
||||
async fn list_accounts(
|
||||
&self,
|
||||
_: AccountListQuery,
|
||||
_: AccountRuntimeSnapshot,
|
||||
) -> AdminStoreResult<AccountPage> {
|
||||
self.record("store.list_accounts");
|
||||
let accounts = self.accounts.lock().expect("accounts").clone();
|
||||
let total = accounts.len() as u64;
|
||||
@@ -526,7 +530,11 @@ impl AccountStore for FakeAccountStore {
|
||||
})
|
||||
}
|
||||
|
||||
async fn load_account(&self, account_id: &str) -> AdminStoreResult<Option<AccountPageItem>> {
|
||||
async fn load_account(
|
||||
&self,
|
||||
account_id: &str,
|
||||
_: AccountRuntimeSnapshot,
|
||||
) -> AdminStoreResult<Option<AccountPageItem>> {
|
||||
self.record("store.load_account");
|
||||
// probe 后的账号状态覆盖只对同一 id 生效;其余按账号列表查询。
|
||||
let account = self
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
mod account_groups;
|
||||
mod accounts;
|
||||
mod auth;
|
||||
mod backup;
|
||||
@@ -21,12 +22,12 @@ use gateway_admin::{
|
||||
model::{
|
||||
MutationContext, Revision,
|
||||
account_groups::{
|
||||
AccountGroupListQuery, AccountGroupMutation, AccountGroupPage, DeleteAccountGroup,
|
||||
NewAccountGroup, SetAccountGroupEnabled, UpdateAccountGroup,
|
||||
AccountGroupListQuery, AccountGroupMemberFact, AccountGroupMutation, AccountGroupPage,
|
||||
DeleteAccountGroup, NewAccountGroup, SetAccountGroupEnabled, UpdateAccountGroup,
|
||||
},
|
||||
accounts::{
|
||||
AccountListQuery, AccountPage, AccountUpdateResult, AccountUsage,
|
||||
AccountUsageWindowQuery, AccountUsageWindowResult, AccountsUpdateResult,
|
||||
AccountListQuery, AccountPage, AccountRuntimeSnapshot, AccountUpdateResult,
|
||||
AccountUsage, AccountUsageWindowQuery, AccountUsageWindowResult, AccountsUpdateResult,
|
||||
BatchUpdateAccounts, DeleteAccounts, UpdateAccount,
|
||||
},
|
||||
auth::{AdminAuditEvent, AdminSession},
|
||||
@@ -56,8 +57,9 @@ use gateway_admin::{
|
||||
backup::BackupStorePorts,
|
||||
provider::{ProviderAdmin, ProviderAdminError, ProviderAdminErrorKind},
|
||||
store::{
|
||||
AccountGroupStore, AccountStore, AdminStoreError, AdminStoreErrorKind, AdminStorePorts,
|
||||
AdminStoreResult, AuthStore, ClientKeyStore, ObservabilityStore, SettingsStore,
|
||||
AccountGroupStore, AccountRuntimeStore, AccountStore, AdminAccountStorePorts,
|
||||
AdminStoreError, AdminStoreErrorKind, AdminStorePorts, AdminStoreResult, AuthStore,
|
||||
ClientKeyStore, ObservabilityStore, SettingsStore,
|
||||
},
|
||||
system::{
|
||||
SystemOperationError, SystemOperationErrorKind, SystemOperations,
|
||||
@@ -79,6 +81,8 @@ pub(super) struct AdminHarness {
|
||||
default_password: String,
|
||||
session_ttl_minutes: u64,
|
||||
accounts: Arc<dyn AccountStore>,
|
||||
account_runtime: Arc<dyn AccountRuntimeStore>,
|
||||
account_groups: Arc<dyn AccountGroupStore>,
|
||||
auth: Arc<dyn AuthStore>,
|
||||
client_keys: Arc<dyn ClientKeyStore>,
|
||||
observability: Arc<dyn ObservabilityStore>,
|
||||
@@ -96,6 +100,8 @@ impl AdminHarness {
|
||||
default_password: "strong-test-password".to_owned(),
|
||||
session_ttl_minutes: 60,
|
||||
accounts: unavailable.clone(),
|
||||
account_runtime: unavailable.clone(),
|
||||
account_groups: Arc::new(UnavailableAccountGroupStore),
|
||||
auth: Arc::new(BootstrapAuthStore::default()),
|
||||
client_keys: unavailable.clone(),
|
||||
observability: unavailable.clone(),
|
||||
@@ -125,6 +131,16 @@ impl AdminHarness {
|
||||
self
|
||||
}
|
||||
|
||||
pub(super) fn account_runtime(mut self, store: Arc<dyn AccountRuntimeStore>) -> Self {
|
||||
self.account_runtime = store;
|
||||
self
|
||||
}
|
||||
|
||||
pub(super) fn account_groups(mut self, store: Arc<dyn AccountGroupStore>) -> Self {
|
||||
self.account_groups = store;
|
||||
self
|
||||
}
|
||||
|
||||
pub(super) fn auth(mut self, store: Arc<dyn AuthStore>) -> Self {
|
||||
self.auth = store;
|
||||
self
|
||||
@@ -175,8 +191,11 @@ impl AdminHarness {
|
||||
default_password: InitialAdminPassword::new(self.default_password),
|
||||
},
|
||||
AdminStorePorts::new(
|
||||
self.accounts,
|
||||
Arc::new(UnavailableAccountGroupStore),
|
||||
AdminAccountStorePorts::new(
|
||||
self.accounts,
|
||||
self.account_runtime,
|
||||
self.account_groups,
|
||||
),
|
||||
self.auth,
|
||||
self.client_keys,
|
||||
self.observability,
|
||||
@@ -252,6 +271,13 @@ impl AccountGroupStore for UnavailableAccountGroupStore {
|
||||
Err(unavailable("account groups"))
|
||||
}
|
||||
|
||||
async fn load_account_group_members(
|
||||
&self,
|
||||
_: &[gateway_core::routing::AccountGroupId],
|
||||
) -> AdminStoreResult<Vec<AccountGroupMemberFact>> {
|
||||
Err(unavailable("account group members"))
|
||||
}
|
||||
|
||||
async fn create_account_group(
|
||||
&self,
|
||||
_: NewAccountGroup,
|
||||
@@ -287,13 +313,18 @@ impl AccountGroupStore for UnavailableAccountGroupStore {
|
||||
|
||||
#[async_trait]
|
||||
impl AccountStore for UnavailableStore {
|
||||
async fn list_accounts(&self, _: AccountListQuery) -> AdminStoreResult<AccountPage> {
|
||||
async fn list_accounts(
|
||||
&self,
|
||||
_: AccountListQuery,
|
||||
_: AccountRuntimeSnapshot,
|
||||
) -> AdminStoreResult<AccountPage> {
|
||||
Err(unavailable("accounts"))
|
||||
}
|
||||
|
||||
async fn load_account(
|
||||
&self,
|
||||
_: &str,
|
||||
_: AccountRuntimeSnapshot,
|
||||
) -> AdminStoreResult<Option<gateway_admin::model::accounts::AccountPageItem>> {
|
||||
Err(unavailable("account"))
|
||||
}
|
||||
@@ -410,6 +441,17 @@ impl AccountStore for UnavailableStore {
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AccountRuntimeStore for UnavailableStore {
|
||||
async fn active_rate_limits(&self) -> AdminStoreResult<AccountRuntimeSnapshot> {
|
||||
Ok(AccountRuntimeSnapshot::default())
|
||||
}
|
||||
|
||||
async fn account_runtime(&self, _: &[String]) -> AdminStoreResult<AccountRuntimeSnapshot> {
|
||||
Ok(AccountRuntimeSnapshot::default())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ClientKeyStore for UnavailableStore {
|
||||
async fn list_client_keys(&self, _: ClientKeyListQuery) -> AdminStoreResult<ClientKeyPage> {
|
||||
|
||||
@@ -1,7 +1,10 @@
|
||||
use std::{
|
||||
num::NonZeroU32,
|
||||
str::FromStr as _,
|
||||
sync::{Arc, Mutex},
|
||||
sync::{
|
||||
Arc, Mutex,
|
||||
atomic::{AtomicUsize, Ordering},
|
||||
},
|
||||
time::Duration as StdDuration,
|
||||
};
|
||||
|
||||
use async_trait::async_trait;
|
||||
@@ -14,10 +17,10 @@ use gateway_admin::{
|
||||
observability::{
|
||||
AccountPoolMetrics, AttemptMetrics, CostCoverage, CurrencyCost, DashboardObservation,
|
||||
DashboardRuntimeSlots, DiagnosticDimension, DiagnosticObservation, Granularity,
|
||||
HealthStatus, LatencyPercentiles, OpsErrorPage, OpsErrorQuery, PageNumber,
|
||||
PercentileMilliseconds, RequestMetricPoint, RequestMetrics, RequestOutcome, TimeRange,
|
||||
TrendKind, UsageBilling, UsageCalculatedBillingFact, UsageDetail, UsageFilter,
|
||||
UsageOverview, UsagePage, UsageQuery, UsageRecord, china_day_start,
|
||||
HealthStatus, LatencyPercentiles, OpsErrorPage, OpsErrorQuery, PercentileMilliseconds,
|
||||
RequestMetricPoint, RequestMetrics, TimeRange, TrendKind, UsageBilling,
|
||||
UsageCalculatedBillingFact, UsageDetail, UsageFilter, UsageListRecord, UsageOverview,
|
||||
UsagePage, UsageQuery, china_day_start,
|
||||
},
|
||||
settings::{
|
||||
AdminApiKey, AdminApiKeyMutation, ReplaceRuntimeSettings, RotationStrategy,
|
||||
@@ -105,6 +108,40 @@ async fn health_timeline_should_keep_exactly_china_day_quarter_hour_slots() {
|
||||
}));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dashboard_summary_should_derive_first_token_average_from_trend() {
|
||||
let now = Utc::now();
|
||||
let current_slot = quarter_hour_start(now);
|
||||
let store = Arc::new(FixtureObservabilityStore::new(observation_range(now)));
|
||||
store.replace_trend(vec![
|
||||
health_metric_point(
|
||||
current_slot - Duration::minutes(15),
|
||||
RequestMetrics {
|
||||
first_token_latency_sum_ms: 120,
|
||||
first_token_latency_count: 2,
|
||||
..RequestMetrics::default()
|
||||
},
|
||||
),
|
||||
health_metric_point(
|
||||
current_slot,
|
||||
RequestMetrics {
|
||||
first_token_latency_sum_ms: 180,
|
||||
first_token_latency_count: 1,
|
||||
..RequestMetrics::default()
|
||||
},
|
||||
),
|
||||
]);
|
||||
let services = observability_services(store).await;
|
||||
|
||||
let summary = services
|
||||
.observability()
|
||||
.dashboard_summary(observation_range(now), TrendKind::Usage)
|
||||
.await
|
||||
.expect("dashboard summary");
|
||||
|
||||
assert_eq!(summary.average_first_token_latency_ms, Some(100));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn health_timeline_should_match_legacy_status_precedence_and_thresholds() {
|
||||
let now = Utc::now();
|
||||
@@ -160,11 +197,15 @@ async fn health_timeline_should_match_legacy_status_precedence_and_thresholds()
|
||||
),
|
||||
];
|
||||
|
||||
for (metrics, expected_status, expected_reliability) in fixtures {
|
||||
for (index, (metrics, expected_status, expected_reliability)) in
|
||||
fixtures.into_iter().enumerate()
|
||||
{
|
||||
store.replace_trend(vec![health_metric_point(current_slot, metrics)]);
|
||||
let range_end =
|
||||
now + Duration::seconds(i64::try_from(index).expect("fixture index fits i64") * 2);
|
||||
let timeline = services
|
||||
.observability()
|
||||
.dashboard_summary(observation_range(now), TrendKind::Usage)
|
||||
.dashboard_summary(observation_range(range_end), TrendKind::Usage)
|
||||
.await
|
||||
.expect("dashboard summary")
|
||||
.health_timeline;
|
||||
@@ -220,6 +261,53 @@ async fn dashboard_summary_should_share_one_current_observation_time_across_runt
|
||||
assert!(summary_observed_at.is_some_and(|value| value > historical_end));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dashboard_summary_should_coalesce_concurrent_requests_in_one_short_window() {
|
||||
let end_bucket = Utc::now().timestamp().div_euclid(2) * 2;
|
||||
let first_end = DateTime::from_timestamp(end_bucket, 100_000_000).expect("first end");
|
||||
let second_end = DateTime::from_timestamp(end_bucket, 900_000_000).expect("second end");
|
||||
let first_range = TimeRange::new(china_day_start(first_end), first_end).expect("first range");
|
||||
let second_range =
|
||||
TimeRange::new(china_day_start(second_end), second_end).expect("second range");
|
||||
let store = Arc::new(FixtureObservabilityStore::new(first_range));
|
||||
store.set_dashboard_delay(StdDuration::from_millis(20));
|
||||
let services = observability_services(store.clone()).await;
|
||||
let observability = services.observability();
|
||||
|
||||
let (first, second) = tokio::join!(
|
||||
observability.dashboard_summary(first_range, TrendKind::Usage),
|
||||
observability.dashboard_summary(second_range, TrendKind::Errors),
|
||||
);
|
||||
|
||||
first.expect("first dashboard summary");
|
||||
second.expect("second dashboard summary");
|
||||
assert_eq!(store.dashboard_summary_calls(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dashboard_summary_should_reload_after_the_short_window_changes() {
|
||||
let end_bucket = Utc::now().timestamp().div_euclid(2) * 2;
|
||||
let first_end = DateTime::from_timestamp(end_bucket, 100_000_000).expect("first end");
|
||||
let second_end = first_end + Duration::seconds(2);
|
||||
let first_range = TimeRange::new(china_day_start(first_end), first_end).expect("first range");
|
||||
let second_range =
|
||||
TimeRange::new(china_day_start(second_end), second_end).expect("second range");
|
||||
let store = Arc::new(FixtureObservabilityStore::new(first_range));
|
||||
let services = observability_services(store.clone()).await;
|
||||
let observability = services.observability();
|
||||
|
||||
observability
|
||||
.dashboard_summary(first_range, TrendKind::Usage)
|
||||
.await
|
||||
.expect("first dashboard summary");
|
||||
observability
|
||||
.dashboard_summary(second_range, TrendKind::Usage)
|
||||
.await
|
||||
.expect("second dashboard summary");
|
||||
|
||||
assert_eq!(store.dashboard_summary_calls(), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn observability_services_should_calculate_usage_insights_and_diagnostic_shares() {
|
||||
let now = Utc::now();
|
||||
@@ -457,7 +545,9 @@ struct FixtureObservabilityStore {
|
||||
runtime_slots: Mutex<Option<DashboardRuntimeSlots>>,
|
||||
summary_observed_at: Mutex<Option<DateTime<Utc>>>,
|
||||
slots_observed_at: Mutex<Option<DateTime<Utc>>>,
|
||||
usage_records: Mutex<Vec<UsageRecord>>,
|
||||
usage_records: Mutex<Vec<UsageListRecord>>,
|
||||
dashboard_delay: Mutex<StdDuration>,
|
||||
dashboard_summary_calls: AtomicUsize,
|
||||
}
|
||||
|
||||
impl FixtureObservabilityStore {
|
||||
@@ -476,9 +566,19 @@ impl FixtureObservabilityStore {
|
||||
summary_observed_at: Mutex::new(None),
|
||||
slots_observed_at: Mutex::new(None),
|
||||
usage_records: Mutex::new(Vec::new()),
|
||||
dashboard_delay: Mutex::new(StdDuration::ZERO),
|
||||
dashboard_summary_calls: AtomicUsize::new(0),
|
||||
}
|
||||
}
|
||||
|
||||
fn set_dashboard_delay(&self, delay: StdDuration) {
|
||||
*self.dashboard_delay.lock().expect("dashboard delay") = delay;
|
||||
}
|
||||
|
||||
fn dashboard_summary_calls(&self) -> usize {
|
||||
self.dashboard_summary_calls.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
fn replace_trend(&self, trend: Vec<RequestMetricPoint>) {
|
||||
*self.trend.lock().expect("trend") = trend;
|
||||
}
|
||||
@@ -512,7 +612,7 @@ impl FixtureObservabilityStore {
|
||||
)
|
||||
}
|
||||
|
||||
fn replace_usage_records(&self, records: Vec<UsageRecord>) {
|
||||
fn replace_usage_records(&self, records: Vec<UsageListRecord>) {
|
||||
*self.usage_records.lock().expect("usage records") = records;
|
||||
}
|
||||
}
|
||||
@@ -524,14 +624,17 @@ impl ObservabilityStore for FixtureObservabilityStore {
|
||||
range: TimeRange,
|
||||
observed_at: DateTime<Utc>,
|
||||
) -> AdminStoreResult<DashboardObservation> {
|
||||
self.dashboard_summary_calls.fetch_add(1, Ordering::Relaxed);
|
||||
*self
|
||||
.summary_observed_at
|
||||
.lock()
|
||||
.expect("summary observed at") = Some(observed_at);
|
||||
let dashboard_delay = *self.dashboard_delay.lock().expect("dashboard delay");
|
||||
if !dashboard_delay.is_zero() {
|
||||
tokio::time::sleep(dashboard_delay).await;
|
||||
}
|
||||
Ok(DashboardObservation {
|
||||
range,
|
||||
requests: RequestMetrics::default(),
|
||||
attempts: AttemptMetrics::default(),
|
||||
totals: Default::default(),
|
||||
provider_accounts: AccountPoolMetrics::default(),
|
||||
trend: self.trend.lock().expect("trend").clone(),
|
||||
@@ -572,13 +675,14 @@ impl ObservabilityStore for FixtureObservabilityStore {
|
||||
.clone())
|
||||
}
|
||||
|
||||
async fn list_usage_records(&self, _: UsageQuery) -> AdminStoreResult<UsagePage> {
|
||||
async fn list_usage_records(&self, query: UsageQuery) -> AdminStoreResult<UsagePage> {
|
||||
let items = self.usage_records.lock().expect("usage records").clone();
|
||||
let total = u64::try_from(items.len()).unwrap_or(u64::MAX);
|
||||
Ok(UsagePage {
|
||||
items,
|
||||
current_page: query.current_page,
|
||||
page_size: query.page_size.get(),
|
||||
total,
|
||||
next_cursor: None,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -680,8 +784,7 @@ fn usage_query(now: DateTime<Utc>) -> UsageQuery {
|
||||
UsageQuery {
|
||||
range: observation_range(now),
|
||||
filter: UsageFilter::default(),
|
||||
cursor: None,
|
||||
page: PageNumber::new(NonZeroU32::new(1).expect("page number")),
|
||||
current_page: 1,
|
||||
page_size: PageSize::new(50).expect("page size"),
|
||||
}
|
||||
}
|
||||
@@ -691,16 +794,9 @@ fn total_record(
|
||||
provider_kind: Option<&str>,
|
||||
source: &str,
|
||||
now: DateTime<Utc>,
|
||||
) -> UsageRecord {
|
||||
UsageRecord {
|
||||
) -> UsageListRecord {
|
||||
UsageListRecord {
|
||||
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(),
|
||||
client_transport: "http_sse".to_owned(),
|
||||
requested_model_id: Some("gpt-5.5".to_owned()),
|
||||
@@ -711,23 +807,7 @@ fn total_record(
|
||||
provider_account_authentication_kind: None,
|
||||
upstream_model_id: Some("gpt-5.5".to_owned()),
|
||||
upstream_transport: None,
|
||||
http_version: None,
|
||||
websocket_pool: None,
|
||||
service_tier: None,
|
||||
provider_metadata_json: None,
|
||||
attempt_count: 1,
|
||||
upstream_send_state: "sent".to_owned(),
|
||||
downstream_committed_at: None,
|
||||
outcome: RequestOutcome::Succeeded,
|
||||
client_status_code: Some(200),
|
||||
upstream_status_code: Some(200),
|
||||
client_response_id: None,
|
||||
upstream_request_id: None,
|
||||
upstream_response_id: None,
|
||||
error_kind: None,
|
||||
provider_error_code: None,
|
||||
error_message: None,
|
||||
retry_after_ms: None,
|
||||
input_tokens: Some(800),
|
||||
output_tokens: Some(200),
|
||||
cached_tokens: Some(0),
|
||||
@@ -764,14 +844,9 @@ fn total_record(
|
||||
user_agent: None,
|
||||
reasoning_effort: None,
|
||||
reasoning_preset: None,
|
||||
request_kind: None,
|
||||
subagent_kind: None,
|
||||
compact: false,
|
||||
image_generation_requested: false,
|
||||
image_generation_succeeded: None,
|
||||
started_at: now,
|
||||
deadline_at: now,
|
||||
completed_at: Some(now),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -24,7 +24,7 @@ secrecy = "0.10"
|
||||
serde = { version = "1.0.228", features = ["derive"] }
|
||||
serde_json = { version = "1.0.150", features = ["arbitrary_precision", "preserve_order"] }
|
||||
thiserror = "2.0.18"
|
||||
tokio = { version = "1.52.3", default-features = false, features = ["macros", "time"] }
|
||||
tokio = { version = "1.52.3", default-features = false, features = ["macros", "rt", "sync", "time"] }
|
||||
tower-http = { version = "0.7.0", features = ["cors", "fs", "request-id", "timeout", "trace"] }
|
||||
tracing = "0.1.44"
|
||||
url = "2.5.8"
|
||||
@@ -33,7 +33,7 @@ uuid = { version = "1.23.3", features = ["v7"] }
|
||||
[dev-dependencies]
|
||||
async-trait = "0.1.89"
|
||||
axum = { version = "0.8.9", default-features = false, features = ["http1", "tokio"] }
|
||||
tokio = { version = "1.52.3", default-features = false, features = ["macros", "net", "rt-multi-thread"] }
|
||||
tokio = { version = "1.52.3", default-features = false, features = ["macros", "net", "rt-multi-thread", "test-util"] }
|
||||
tokio-tungstenite = "0.29.0"
|
||||
tower = { version = "0.5.3", features = ["util"] }
|
||||
|
||||
|
||||
@@ -127,7 +127,7 @@ struct AccountGroupCapacityView {
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct AccountGroupUsageView {
|
||||
today_usd: String,
|
||||
total_usd: String,
|
||||
retained_total_usd: String,
|
||||
}
|
||||
|
||||
impl From<AccountGroupRecord> for AccountGroupView {
|
||||
@@ -168,7 +168,7 @@ fn capacity_view(capacity: AccountGroupCapacity) -> AccountGroupCapacityView {
|
||||
fn usage_view(usage: AccountGroupUsage) -> AccountGroupUsageView {
|
||||
AccountGroupUsageView {
|
||||
today_usd: usage.today_usd.to_string(),
|
||||
total_usd: usage.total_usd.to_string(),
|
||||
retained_total_usd: usage.retained_total_usd.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -4,8 +4,9 @@ use std::{collections::BTreeSet, convert::Infallible, fmt};
|
||||
|
||||
use axum::{
|
||||
Router,
|
||||
body::Body,
|
||||
extract::State,
|
||||
http::{HeaderValue, StatusCode, header::CACHE_CONTROL},
|
||||
http::{HeaderName, HeaderValue, StatusCode, header},
|
||||
response::{
|
||||
IntoResponse, Response,
|
||||
sse::{Event, KeepAlive, Sse},
|
||||
@@ -28,10 +29,10 @@ use gateway_admin::model::{
|
||||
AuthorizationStarted, CompleteAuthorization, ConsumeProviderResetCredit,
|
||||
CredentialDeletion, CredentialDeletionResult, CredentialImportResult, CredentialMutation,
|
||||
CredentialMutationResult, ImportCredentials, ProviderDocument, ProviderModels,
|
||||
ProviderProfileActivityInsights, ProviderProfileDailyUsage, ProviderProfileInvocation,
|
||||
ProviderProfileStatistics, ProviderProfileStatisticsSummary, ProviderQuota,
|
||||
ProviderQuotaWindow, ProviderResetCredit, ProviderResetCreditResult, ProviderResetCredits,
|
||||
RotateCredential, StartAuthorization,
|
||||
ProviderProfileActivityInsights, ProviderProfileAvatar, ProviderProfileDailyUsage,
|
||||
ProviderProfileInvocation, ProviderProfileStatistics, ProviderProfileStatisticsSummary,
|
||||
ProviderQuota, ProviderQuotaWindow, ProviderResetCredit, ProviderResetCreditResult,
|
||||
ProviderResetCredits, RotateCredential, StartAuthorization,
|
||||
},
|
||||
};
|
||||
use gateway_core::{
|
||||
@@ -60,6 +61,7 @@ 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;
|
||||
const MAX_AVATAR_VERSION_BYTES: usize = 32;
|
||||
|
||||
/// 账号列表查询参数。
|
||||
#[derive(Debug, Clone, Default, Deserialize)]
|
||||
@@ -422,6 +424,33 @@ impl AccountIdQuery {
|
||||
}
|
||||
}
|
||||
|
||||
/// 头像 GET 的固定 query;`version` 只参与浏览器缓存键,不进入上游请求。
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
pub struct AccountProfileAvatarQuery {
|
||||
pub account_id: String,
|
||||
pub version: Option<String>,
|
||||
}
|
||||
|
||||
impl AccountProfileAvatarQuery {
|
||||
pub fn validate(&self) -> Result<(), WireValidationError> {
|
||||
require_account_id(&self.account_id, "accountId")?;
|
||||
if self.version.as_deref().is_some_and(|version| {
|
||||
version.is_empty()
|
||||
|| version.len() > MAX_AVATAR_VERSION_BYTES
|
||||
|| !version.bytes().all(|byte| byte.is_ascii_alphanumeric())
|
||||
}) {
|
||||
return Err(WireValidationError::new("version"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn into_id(self) -> Result<ProviderAccountId, WireValidationError> {
|
||||
self.validate()?;
|
||||
ProviderAccountId::new(self.account_id).map_err(|_| WireValidationError::new("accountId"))
|
||||
}
|
||||
}
|
||||
|
||||
/// 敏感导出的固定 query;IDs 使用逗号分隔,禁止隐式导出全部账号。
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
@@ -1291,6 +1320,10 @@ where
|
||||
"/api/admin/accounts/profile-statistics",
|
||||
get(account_profile_statistics::<S>),
|
||||
)
|
||||
.route(
|
||||
"/api/admin/accounts/profile-avatar",
|
||||
get(account_profile_avatar::<S>),
|
||||
)
|
||||
.route(
|
||||
"/api/admin/accounts/reset-credits",
|
||||
get(account_reset_credits::<S>).post(consume_account_reset_credit::<S>),
|
||||
@@ -1400,7 +1433,7 @@ where
|
||||
let mut response = AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data)).into_response();
|
||||
response
|
||||
.headers_mut()
|
||||
.insert(CACHE_CONTROL, HeaderValue::from_static("no-store"));
|
||||
.insert(header::CACHE_CONTROL, HeaderValue::from_static("no-store"));
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
@@ -1646,6 +1679,69 @@ where
|
||||
))
|
||||
}
|
||||
|
||||
async fn account_profile_avatar<S>(
|
||||
_auth: AdminAuth,
|
||||
State(state): State<S>,
|
||||
AdminQuery(query): AdminQuery<AccountProfileAvatarQuery>,
|
||||
) -> Result<Response, AdminError>
|
||||
where
|
||||
S: AdminSessionState + Send + Sync,
|
||||
{
|
||||
let account_id = query.into_id().map_err(map_wire_error)?;
|
||||
let avatar = state
|
||||
.admin_services()
|
||||
.accounts()
|
||||
.profile_avatar(&account_id)
|
||||
.await
|
||||
.map_err(map_service_error)?;
|
||||
Ok(profile_avatar_response(avatar))
|
||||
}
|
||||
|
||||
/// 将 Provider 头像流投影为受保护的同源 HTTP 响应。
|
||||
#[must_use]
|
||||
pub fn profile_avatar_response(avatar: ProviderProfileAvatar) -> Response {
|
||||
let ProviderProfileAvatar {
|
||||
content_type,
|
||||
content_length,
|
||||
etag,
|
||||
body,
|
||||
} = avatar;
|
||||
let mut response = Response::new(Body::from_stream(body));
|
||||
*response.status_mut() = StatusCode::OK;
|
||||
let headers = response.headers_mut();
|
||||
headers.insert(
|
||||
header::CONTENT_TYPE,
|
||||
content_type
|
||||
.and_then(|value| HeaderValue::from_str(&value).ok())
|
||||
.unwrap_or_else(|| HeaderValue::from_static("application/octet-stream")),
|
||||
);
|
||||
if let Some(value) =
|
||||
content_length.and_then(|length| HeaderValue::from_str(&length.to_string()).ok())
|
||||
{
|
||||
headers.insert(header::CONTENT_LENGTH, value);
|
||||
}
|
||||
if let Some(value) = etag.and_then(|value| HeaderValue::from_str(&value).ok()) {
|
||||
headers.insert(header::ETAG, value);
|
||||
}
|
||||
headers.insert(
|
||||
header::CACHE_CONTROL,
|
||||
HeaderValue::from_static("private, max-age=3600"),
|
||||
);
|
||||
headers.insert(
|
||||
HeaderName::from_static("x-content-type-options"),
|
||||
HeaderValue::from_static("nosniff"),
|
||||
);
|
||||
headers.insert(
|
||||
HeaderName::from_static("cross-origin-resource-policy"),
|
||||
HeaderValue::from_static("same-origin"),
|
||||
);
|
||||
headers.insert(
|
||||
HeaderName::from_static("content-security-policy"),
|
||||
HeaderValue::from_static("sandbox; default-src 'none'"),
|
||||
);
|
||||
response
|
||||
}
|
||||
|
||||
async fn refresh_account_quota<S>(
|
||||
_auth: AdminAuth,
|
||||
State(state): State<S>,
|
||||
|
||||
@@ -58,6 +58,7 @@ async fn admin_not_found() -> AdminError {
|
||||
async fn no_store(mut response: Response) -> Response {
|
||||
response
|
||||
.headers_mut()
|
||||
.insert(header::CACHE_CONTROL, HeaderValue::from_static("no-store"));
|
||||
.entry(header::CACHE_CONTROL)
|
||||
.or_insert(HeaderValue::from_static("no-store"));
|
||||
response
|
||||
}
|
||||
|
||||
@@ -1,9 +1,8 @@
|
||||
//! Dashboard、用量与诊断查询的 wire 映射和固定路由。
|
||||
|
||||
use std::{collections::BTreeMap, num::NonZeroU32};
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use axum::{Router, extract::State, http::StatusCode, response::IntoResponse, routing::get};
|
||||
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
|
||||
use chrono::{DateTime, Duration, Utc};
|
||||
use gateway_admin::model::{PageSize as DomainPageSize, observability as domain};
|
||||
use serde::{Deserialize, Serialize};
|
||||
@@ -11,7 +10,7 @@ use serde_json::Value;
|
||||
|
||||
use super::presenter::{format_compact_number, format_decimal_currency, format_number};
|
||||
use super::{
|
||||
AdminAuth, AdminEnvelope, AdminError, AdminQuery, AdminResponse, AdminSessionState, PageMeta,
|
||||
AdminAuth, AdminEnvelope, AdminError, AdminQuery, AdminResponse, AdminSessionState,
|
||||
WireValidationError, wire::map_admin_service_error,
|
||||
};
|
||||
|
||||
|
||||
@@ -146,6 +146,43 @@ pub(crate) fn token_details(record: &domain::UsageRecord) -> TokenDetailsView {
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn usage_list_token_details(record: &domain::UsageListRecord) -> TokenDetailsView {
|
||||
TokenDetailsView {
|
||||
input_tokens: record.input_tokens,
|
||||
output_tokens: record.output_tokens,
|
||||
cached_tokens: record.cached_tokens,
|
||||
cache_write_tokens: record.cache_write_tokens,
|
||||
reasoning_tokens: record.reasoning_tokens,
|
||||
image_input_tokens: record.image_input_tokens,
|
||||
image_output_tokens: record.image_output_tokens,
|
||||
total_tokens: record.total_tokens,
|
||||
input_tokens_display: record
|
||||
.input_tokens
|
||||
.map_or_else(|| "-".to_owned(), format_number),
|
||||
output_tokens_display: record
|
||||
.output_tokens
|
||||
.map_or_else(|| "-".to_owned(), format_number),
|
||||
cached_tokens_display: record
|
||||
.cached_tokens
|
||||
.map_or_else(|| "-".to_owned(), format_compact_number),
|
||||
cache_write_tokens_display: record
|
||||
.cache_write_tokens
|
||||
.map_or_else(|| "-".to_owned(), format_compact_number),
|
||||
reasoning_tokens_display: record
|
||||
.reasoning_tokens
|
||||
.map_or_else(|| "-".to_owned(), format_number),
|
||||
image_input_tokens_display: record
|
||||
.image_input_tokens
|
||||
.map_or_else(|| "-".to_owned(), format_number),
|
||||
image_output_tokens_display: record
|
||||
.image_output_tokens
|
||||
.map_or_else(|| "-".to_owned(), format_number),
|
||||
total_tokens_display: record
|
||||
.total_tokens
|
||||
.map_or_else(|| "-".to_owned(), format_number),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn format_money(cost: &domain::CurrencyCost) -> String {
|
||||
format_decimal_currency(cost.amount.as_str(), &cost.currency)
|
||||
}
|
||||
@@ -207,6 +244,57 @@ pub(crate) fn billing_view(billing: Option<&domain::UsageBilling>) -> Option<Bil
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn usage_list_record_view(record: domain::UsageListRecord) -> UsageListRecordView {
|
||||
let token_details = usage_list_token_details(&record);
|
||||
let billing = billing_view(record.billing.as_ref());
|
||||
let created_at_display = china_datetime(&record.started_at);
|
||||
let model = record
|
||||
.upstream_model_id
|
||||
.clone()
|
||||
.or_else(|| record.requested_model_id.clone());
|
||||
UsageListRecordView {
|
||||
id: record.id,
|
||||
provider: record.provider_kind,
|
||||
authentication_kind: record.provider_account_authentication_kind,
|
||||
account_id: record.provider_account_ref,
|
||||
account_email: record.provider_account_email,
|
||||
account_name: record.provider_account_name,
|
||||
route: record.endpoint,
|
||||
model,
|
||||
requested_model: record.requested_model_id,
|
||||
upstream_model: record.upstream_model_id,
|
||||
service_tier: record.service_tier,
|
||||
client_transport: record.client_transport,
|
||||
upstream_transport: record.upstream_transport,
|
||||
reasoning_effort: record.reasoning_effort,
|
||||
reasoning_preset: record.reasoning_preset,
|
||||
subagent_kind: record.subagent_kind,
|
||||
compact: record.compact,
|
||||
token_details,
|
||||
billing,
|
||||
latency_details: UsageLatencyDetailsView {
|
||||
admission_decision_ms: record.admission_decision_ms,
|
||||
account_selection_wait_ms: record.account_selection_wait_ms,
|
||||
capacity_used_slots: record.capacity_used_slots,
|
||||
capacity_total_slots: record.capacity_total_slots,
|
||||
transport_decision_wait_ms: record.transport_decision_wait_ms,
|
||||
ws_connect_ms: record.connect_ms,
|
||||
upstream_headers_ms: record.headers_ms,
|
||||
first_event_ms: record.first_event_ms,
|
||||
first_reasoning_ms: record.first_reasoning_ms,
|
||||
first_text_ms: record.first_text_ms,
|
||||
first_token_ms: record.first_token_ms,
|
||||
openai_processing_ms: record.provider_processing_ms,
|
||||
},
|
||||
first_token_latency_ms: record.first_token_ms,
|
||||
latency_ms: record.latency_ms,
|
||||
created_at: record.started_at,
|
||||
created_at_display,
|
||||
client_ip: record.client_ip,
|
||||
user_agent: record.user_agent,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn usage_record_view(record: domain::UsageRecord) -> UsageRecordView {
|
||||
let tokens = token_details(&record);
|
||||
let billing = billing_view(record.billing.as_ref());
|
||||
@@ -250,10 +338,6 @@ pub(crate) fn usage_record_view(record: domain::UsageRecord) -> UsageRecordView
|
||||
.upstream_model_id
|
||||
.clone()
|
||||
.or_else(|| record.requested_model_id.clone());
|
||||
let transport = record
|
||||
.upstream_transport
|
||||
.clone()
|
||||
.or_else(|| Some(record.client_transport.clone()));
|
||||
let metadata = provider_metadata_fields(record.provider_metadata_json.as_deref());
|
||||
UsageRecordView {
|
||||
id: record.id.clone(),
|
||||
@@ -274,7 +358,8 @@ pub(crate) fn usage_record_view(record: domain::UsageRecord) -> UsageRecordView
|
||||
upstream_model: record.upstream_model_id,
|
||||
service_tier: record.service_tier,
|
||||
status_code,
|
||||
transport,
|
||||
client_transport: record.client_transport,
|
||||
upstream_transport: record.upstream_transport,
|
||||
protocol: record.protocol.clone(),
|
||||
http_version: record.http_version.clone(),
|
||||
client_status_code: record.client_status_code.map(i64::from),
|
||||
@@ -349,6 +434,8 @@ pub(crate) fn provider_metadata_fields(value: Option<&str>) -> BTreeMap<String,
|
||||
"requestKind",
|
||||
"subagentKind",
|
||||
"transport",
|
||||
"clientTransport",
|
||||
"upstreamTransport",
|
||||
"httpVersion",
|
||||
"clientStatusCode",
|
||||
"upstreamStatusCode",
|
||||
@@ -428,31 +515,13 @@ pub(crate) fn usage_detail_view(detail: domain::UsageDetail) -> UsageRecordDetai
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn page_meta(page: u32, page_size: u16, total: u64) -> PageMeta {
|
||||
let page_size_u64 = u64::from(page_size);
|
||||
let total_pages = total.saturating_add(page_size_u64 - 1) / page_size_u64;
|
||||
PageMeta::new(
|
||||
page,
|
||||
u32::from(page_size),
|
||||
total,
|
||||
u32::try_from(total_pages).unwrap_or(u32::MAX),
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn usage_page_view(
|
||||
page: domain::UsagePage,
|
||||
page_number: u32,
|
||||
page_size: u16,
|
||||
) -> Result<PageData<UsageRecordView>, WireValidationError> {
|
||||
Ok(PageData {
|
||||
items: page.items.into_iter().map(usage_record_view).collect(),
|
||||
page: page_meta(page_number, page_size, page.total),
|
||||
next_cursor: page
|
||||
.next_cursor
|
||||
.as_ref()
|
||||
.map(encode_observability_cursor)
|
||||
.transpose()?,
|
||||
})
|
||||
pub(crate) fn usage_page_view(page: domain::UsagePage) -> PageData<UsageListRecordView> {
|
||||
PageData {
|
||||
items: page.items.into_iter().map(usage_list_record_view).collect(),
|
||||
current_page: page.current_page,
|
||||
page_size: page.page_size,
|
||||
total: page.total,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn ops_error_view(error: domain::OpsError) -> OpsErrorView {
|
||||
@@ -462,24 +531,45 @@ pub(crate) fn ops_error_view(error: domain::OpsError) -> OpsErrorView {
|
||||
.or(error.provider_account_name.as_ref())
|
||||
.or(error.provider_account_ref.as_ref())
|
||||
.cloned();
|
||||
let model = error
|
||||
.upstream_model_id
|
||||
.clone()
|
||||
.or_else(|| error.requested_model_id.clone());
|
||||
OpsErrorView {
|
||||
id: error.event_id,
|
||||
request_id: error.request_id,
|
||||
client_api_key_id: error.client_api_key_ref,
|
||||
kind: error.component.clone(),
|
||||
operation: error.operation,
|
||||
protocol: error.protocol,
|
||||
client_transport: error.client_transport,
|
||||
provider: error.provider_kind,
|
||||
authentication_kind: error.provider_account_authentication_kind,
|
||||
account_id: error.provider_account_ref,
|
||||
account_name: error.provider_account_name,
|
||||
account_email: error.provider_account_email,
|
||||
route: error.endpoint.unwrap_or_else(|| "—".to_owned()),
|
||||
model: error.upstream_model_id,
|
||||
model,
|
||||
requested_model: error.requested_model_id,
|
||||
upstream_model: error.upstream_model_id,
|
||||
service_tier: error.service_tier,
|
||||
client_status_code: error.client_status_code.map(i64::from),
|
||||
upstream_status_code: error.upstream_status_code.map(i64::from),
|
||||
transport: error.upstream_transport,
|
||||
attempt_index: error.attempt_index,
|
||||
failure_class: error.failure_kind,
|
||||
provider_error_code: error.provider_error_code,
|
||||
occurrence_count: error.occurrence_count,
|
||||
response_id: error.client_response_id,
|
||||
upstream_request_id: error.upstream_request_id,
|
||||
latency_ms: error.latency_ms,
|
||||
client_ip: error.client_ip,
|
||||
user_agent: error.user_agent,
|
||||
reasoning_effort: error.reasoning_effort,
|
||||
reasoning_preset: error.reasoning_preset,
|
||||
request_kind: error.request_kind,
|
||||
subagent_kind: error.subagent_kind,
|
||||
compact: error.compact,
|
||||
message: error.message,
|
||||
metadata: OpsErrorMetadataView {
|
||||
source: error.source,
|
||||
@@ -492,20 +582,13 @@ pub(crate) fn ops_error_view(error: domain::OpsError) -> OpsErrorView {
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn ops_page_view(
|
||||
page: domain::OpsErrorPage,
|
||||
page_number: u32,
|
||||
page_size: u16,
|
||||
) -> Result<PageData<OpsErrorView>, WireValidationError> {
|
||||
Ok(PageData {
|
||||
pub(crate) fn ops_page_view(page: domain::OpsErrorPage) -> PageData<OpsErrorView> {
|
||||
PageData {
|
||||
items: page.items.into_iter().map(ops_error_view).collect(),
|
||||
page: page_meta(page_number, page_size, page.total),
|
||||
next_cursor: page
|
||||
.next_cursor
|
||||
.as_ref()
|
||||
.map(encode_observability_cursor)
|
||||
.transpose()?,
|
||||
})
|
||||
current_page: page.current_page,
|
||||
page_size: page.page_size,
|
||||
total: page.total,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn trend_point_view(point: domain::TrendPoint) -> TrendPointView {
|
||||
@@ -834,7 +917,10 @@ pub(crate) fn dashboard_view(
|
||||
health_timeline: health_timeline_view(health_timeline),
|
||||
wire_profiles: wire_profiles.into_iter().map(wire_profile_view).collect(),
|
||||
account_usage: account_usage_views,
|
||||
usage_records: recent_requests.into_iter().map(usage_record_view).collect(),
|
||||
usage_records: recent_requests
|
||||
.into_iter()
|
||||
.map(usage_list_record_view)
|
||||
.collect(),
|
||||
pool_summary: DashboardPoolSummaryView {
|
||||
total: provider_accounts.total,
|
||||
normal: provider_accounts.normal,
|
||||
|
||||
@@ -2,14 +2,10 @@
|
||||
|
||||
use super::*;
|
||||
|
||||
use base64::Engine as _;
|
||||
|
||||
/// 观测列表默认页大小。
|
||||
pub const DEFAULT_PAGE_SIZE: u16 = 50;
|
||||
/// 观测列表允许的最大页大小。
|
||||
pub const MAX_PAGE_SIZE: u16 = 100;
|
||||
/// 游标编码在 HTTP 层解码前允许的最大长度。
|
||||
pub const MAX_CURSOR_BYTES: usize = 1024;
|
||||
|
||||
/// Dashboard 查询参数。
|
||||
#[derive(Clone, Default, Deserialize)]
|
||||
@@ -31,9 +27,8 @@ impl DashboardQuery {
|
||||
#[derive(Clone, Default, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
pub struct UsageQuery {
|
||||
pub page: Option<u32>,
|
||||
pub current_page: Option<u32>,
|
||||
pub page_size: Option<u16>,
|
||||
pub cursor: Option<String>,
|
||||
pub kind: Option<String>,
|
||||
pub outcome: Option<String>,
|
||||
pub client_api_key_id: Option<String>,
|
||||
@@ -53,22 +48,17 @@ pub struct UsageQuery {
|
||||
}
|
||||
|
||||
impl UsageQuery {
|
||||
/// 校验分页字段,不改变仓储层的分页 owner。
|
||||
pub fn validate_page(&self) -> Result<(u32, u16), WireValidationError> {
|
||||
let page = self.page.unwrap_or(1);
|
||||
if page == 0 {
|
||||
return Err(WireValidationError::new("page"));
|
||||
/// 校验 Element Plus 风格的页码分页字段。
|
||||
pub fn validate_pagination(&self) -> Result<(u32, u16), WireValidationError> {
|
||||
let current_page = self.current_page.unwrap_or(1);
|
||||
if current_page == 0 {
|
||||
return Err(WireValidationError::new("currentPage"));
|
||||
}
|
||||
let page_size = self.page_size.unwrap_or(DEFAULT_PAGE_SIZE);
|
||||
if page_size == 0 || page_size > MAX_PAGE_SIZE {
|
||||
return Err(WireValidationError::new("pageSize"));
|
||||
}
|
||||
Ok((page, page_size))
|
||||
}
|
||||
|
||||
/// 校验游标的 wire 边界;编码和排序语义由应用层负责。
|
||||
pub fn validate_cursor(&self) -> Result<(), WireValidationError> {
|
||||
validate_cursor(self.cursor.as_deref())
|
||||
Ok((current_page, page_size))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -110,9 +100,8 @@ impl DiagnosticsQuery {
|
||||
#[derive(Clone, Default, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
pub struct OpsQuery {
|
||||
pub page: Option<u32>,
|
||||
pub current_page: Option<u32>,
|
||||
pub page_size: Option<u16>,
|
||||
pub cursor: Option<String>,
|
||||
pub kind: Option<String>,
|
||||
pub client_api_key_id: Option<String>,
|
||||
pub provider: Option<String>,
|
||||
@@ -134,22 +123,17 @@ pub struct OpsQuery {
|
||||
}
|
||||
|
||||
impl OpsQuery {
|
||||
/// 校验分页字段。
|
||||
pub fn validate_page(&self) -> Result<(u32, u16), WireValidationError> {
|
||||
let page = self.page.unwrap_or(1);
|
||||
if page == 0 {
|
||||
return Err(WireValidationError::new("page"));
|
||||
/// 校验 Element Plus 风格的页码分页字段。
|
||||
pub fn validate_pagination(&self) -> Result<(u32, u16), WireValidationError> {
|
||||
let current_page = self.current_page.unwrap_or(1);
|
||||
if current_page == 0 {
|
||||
return Err(WireValidationError::new("currentPage"));
|
||||
}
|
||||
let page_size = self.page_size.unwrap_or(DEFAULT_PAGE_SIZE);
|
||||
if page_size == 0 || page_size > MAX_PAGE_SIZE {
|
||||
return Err(WireValidationError::new("pageSize"));
|
||||
}
|
||||
Ok((page, page_size))
|
||||
}
|
||||
|
||||
/// 校验游标的 wire 边界。
|
||||
pub fn validate_cursor(&self) -> Result<(), WireValidationError> {
|
||||
validate_cursor(self.cursor.as_deref())
|
||||
Ok((current_page, page_size))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -258,13 +242,6 @@ pub(crate) fn require_text(value: &str, field: &'static str) -> Result<(), WireV
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn validate_cursor(value: Option<&str>) -> Result<(), WireValidationError> {
|
||||
if value.is_some_and(|value| value.is_empty() || value.len() > MAX_CURSOR_BYTES) {
|
||||
return Err(WireValidationError::new("cursor"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) fn trimmed(value: Option<&str>) -> Option<&str> {
|
||||
value.map(str::trim).filter(|value| !value.is_empty())
|
||||
}
|
||||
@@ -345,37 +322,20 @@ pub(crate) fn usage_filter(query: &UsageQuery) -> Result<domain::UsageFilter, Wi
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn usage_command(
|
||||
query: &UsageQuery,
|
||||
) -> Result<(domain::UsageQuery, u32, u16), WireValidationError> {
|
||||
let (page, page_size) = query.validate_page()?;
|
||||
query.validate_cursor()?;
|
||||
let page_number = NonZeroU32::new(page)
|
||||
.map(domain::PageNumber::new)
|
||||
.ok_or_else(|| WireValidationError::new("page"))?;
|
||||
pub(crate) fn usage_command(query: &UsageQuery) -> Result<domain::UsageQuery, WireValidationError> {
|
||||
let (current_page, page_size) = query.validate_pagination()?;
|
||||
let page_size_value =
|
||||
DomainPageSize::new(page_size).map_err(|_| WireValidationError::new("pageSize"))?;
|
||||
Ok((
|
||||
domain::UsageQuery {
|
||||
range: usage_range(query.start_time.as_deref(), query.end_time.as_deref())?,
|
||||
filter: usage_filter(query)?,
|
||||
cursor: decode_observability_cursor(query.cursor.as_deref())?,
|
||||
page: page_number,
|
||||
page_size: page_size_value,
|
||||
},
|
||||
page,
|
||||
page_size,
|
||||
))
|
||||
Ok(domain::UsageQuery {
|
||||
range: usage_range(query.start_time.as_deref(), query.end_time.as_deref())?,
|
||||
filter: usage_filter(query)?,
|
||||
current_page,
|
||||
page_size: page_size_value,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn ops_command(
|
||||
query: &OpsQuery,
|
||||
) -> Result<(domain::OpsErrorQuery, u32, u16), WireValidationError> {
|
||||
let (page, page_size) = query.validate_page()?;
|
||||
query.validate_cursor()?;
|
||||
let page_number = NonZeroU32::new(page)
|
||||
.map(domain::PageNumber::new)
|
||||
.ok_or_else(|| WireValidationError::new("page"))?;
|
||||
pub(crate) fn ops_command(query: &OpsQuery) -> Result<domain::OpsErrorQuery, WireValidationError> {
|
||||
let (current_page, page_size) = query.validate_pagination()?;
|
||||
let page_size_value =
|
||||
DomainPageSize::new(page_size).map_err(|_| WireValidationError::new("pageSize"))?;
|
||||
let status_code = parse_status(
|
||||
@@ -384,66 +344,27 @@ pub(crate) fn ops_command(
|
||||
.or(query.client_status_code)
|
||||
.or(query.status_code),
|
||||
)?;
|
||||
Ok((
|
||||
domain::OpsErrorQuery {
|
||||
range: usage_range(query.start_time.as_deref(), query.end_time.as_deref())?,
|
||||
filter: domain::OpsErrorFilter {
|
||||
client_api_key_ref: non_empty(query.client_api_key_id.clone()),
|
||||
request_id: non_empty(query.request_id.clone()),
|
||||
provider_kind: non_empty(query.provider.clone()),
|
||||
provider_account_ref: non_empty(query.account_id.clone()),
|
||||
operation: non_empty(query.route.clone()).or_else(|| non_empty(query.kind.clone())),
|
||||
model: non_empty(query.model.clone()),
|
||||
transport: non_empty(query.transport.clone()),
|
||||
attempt_index: parse_attempt_index(query.attempt_index)?,
|
||||
response_id: non_empty(query.response_id.clone()),
|
||||
upstream_request_id: non_empty(query.upstream_request_id.clone()),
|
||||
failure_kind: non_empty(query.failure_class.clone()),
|
||||
status_code,
|
||||
search: non_empty(query.search.clone()),
|
||||
},
|
||||
cursor: decode_observability_cursor(query.cursor.as_deref())?,
|
||||
page: page_number,
|
||||
page_size: page_size_value,
|
||||
Ok(domain::OpsErrorQuery {
|
||||
range: usage_range(query.start_time.as_deref(), query.end_time.as_deref())?,
|
||||
filter: domain::OpsErrorFilter {
|
||||
client_api_key_ref: non_empty(query.client_api_key_id.clone()),
|
||||
request_id: non_empty(query.request_id.clone()),
|
||||
provider_kind: non_empty(query.provider.clone()),
|
||||
provider_account_ref: non_empty(query.account_id.clone()),
|
||||
operation: non_empty(query.kind.clone()),
|
||||
endpoint: non_empty(query.route.clone()),
|
||||
model: non_empty(query.model.clone()),
|
||||
transport: non_empty(query.transport.clone()),
|
||||
attempt_index: parse_attempt_index(query.attempt_index)?,
|
||||
response_id: non_empty(query.response_id.clone()),
|
||||
upstream_request_id: non_empty(query.upstream_request_id.clone()),
|
||||
failure_kind: non_empty(query.failure_class.clone()),
|
||||
status_code,
|
||||
search: non_empty(query.search.clone()),
|
||||
},
|
||||
page,
|
||||
page_size,
|
||||
))
|
||||
}
|
||||
|
||||
pub(crate) fn decode_observability_cursor(
|
||||
value: Option<&str>,
|
||||
) -> Result<Option<domain::ObservabilityCursor>, WireValidationError> {
|
||||
value
|
||||
.map(|encoded| {
|
||||
if encoded.is_empty() || encoded.len() > MAX_CURSOR_BYTES {
|
||||
return Err(WireValidationError::new("cursor"));
|
||||
}
|
||||
let bytes = URL_SAFE_NO_PAD
|
||||
.decode(encoded)
|
||||
.map_err(|_| WireValidationError::new("cursor"))?;
|
||||
let wire: CursorWire =
|
||||
serde_json::from_slice(&bytes).map_err(|_| WireValidationError::new("cursor"))?;
|
||||
if wire.stable_id.trim().is_empty() {
|
||||
return Err(WireValidationError::new("cursor"));
|
||||
}
|
||||
Ok(domain::ObservabilityCursor {
|
||||
observed_at: wire.observed_at,
|
||||
stable_id: wire.stable_id,
|
||||
})
|
||||
})
|
||||
.transpose()
|
||||
}
|
||||
|
||||
pub(crate) fn encode_observability_cursor(
|
||||
cursor: &domain::ObservabilityCursor,
|
||||
) -> Result<String, WireValidationError> {
|
||||
let bytes = serde_json::to_vec(&CursorWire {
|
||||
observed_at: cursor.observed_at,
|
||||
stable_id: cursor.stable_id.clone(),
|
||||
current_page,
|
||||
page_size: page_size_value,
|
||||
})
|
||||
.map_err(|_| WireValidationError::new("cursor"))?;
|
||||
Ok(URL_SAFE_NO_PAD.encode(bytes))
|
||||
}
|
||||
|
||||
pub(crate) fn non_empty(value: Option<String>) -> Option<String> {
|
||||
|
||||
@@ -84,14 +84,14 @@ pub(crate) async fn usage_records<S>(
|
||||
where
|
||||
S: AdminSessionState + Send + Sync,
|
||||
{
|
||||
let (command, page, page_size) = usage_command(&query).map_err(map_wire_error)?;
|
||||
let command = usage_command(&query).map_err(map_wire_error)?;
|
||||
let result = state
|
||||
.admin_services()
|
||||
.observability()
|
||||
.usage_records(command)
|
||||
.await
|
||||
.map_err(map_service_error)?;
|
||||
let data = usage_page_view(result, page, page_size).map_err(|_| AdminError::internal())?;
|
||||
let data = usage_page_view(result);
|
||||
Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data)))
|
||||
}
|
||||
|
||||
@@ -200,13 +200,13 @@ pub(crate) async fn ops_errors<S>(
|
||||
where
|
||||
S: AdminSessionState + Send + Sync,
|
||||
{
|
||||
let (command, page, page_size) = ops_command(&query).map_err(map_wire_error)?;
|
||||
let command = ops_command(&query).map_err(map_wire_error)?;
|
||||
let result = state
|
||||
.admin_services()
|
||||
.observability()
|
||||
.ops_errors(command)
|
||||
.await
|
||||
.map_err(map_service_error)?;
|
||||
let data = ops_page_view(result, page, page_size).map_err(|_| AdminError::internal())?;
|
||||
let data = ops_page_view(result);
|
||||
Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data)))
|
||||
}
|
||||
|
||||
@@ -2,21 +2,14 @@
|
||||
|
||||
use super::*;
|
||||
|
||||
/// 观测列表游标的稳定 wire 形状。
|
||||
#[derive(Clone, Serialize, serde::Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
pub struct CursorWire {
|
||||
pub observed_at: DateTime<Utc>,
|
||||
pub stable_id: String,
|
||||
}
|
||||
|
||||
/// 观测列表响应数据。
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct PageData<T> {
|
||||
pub items: Vec<T>,
|
||||
pub page: PageMeta,
|
||||
pub next_cursor: Option<String>,
|
||||
pub current_page: u32,
|
||||
pub page_size: u16,
|
||||
pub total: u64,
|
||||
}
|
||||
|
||||
/// Token 详情展示。
|
||||
@@ -77,7 +70,39 @@ pub struct BillingView {
|
||||
pub multiplier_display: String,
|
||||
}
|
||||
|
||||
/// 单条逻辑请求展示。
|
||||
/// 使用记录表格的窄展示。
|
||||
#[derive(Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct UsageListRecordView {
|
||||
pub id: String,
|
||||
pub provider: Option<String>,
|
||||
pub authentication_kind: Option<String>,
|
||||
pub account_id: Option<String>,
|
||||
pub account_email: Option<String>,
|
||||
pub account_name: Option<String>,
|
||||
pub route: String,
|
||||
pub model: Option<String>,
|
||||
pub requested_model: Option<String>,
|
||||
pub upstream_model: Option<String>,
|
||||
pub service_tier: Option<String>,
|
||||
pub client_transport: String,
|
||||
pub upstream_transport: Option<String>,
|
||||
pub reasoning_effort: Option<String>,
|
||||
pub reasoning_preset: Option<String>,
|
||||
pub subagent_kind: Option<String>,
|
||||
pub compact: bool,
|
||||
pub token_details: TokenDetailsView,
|
||||
pub billing: Option<BillingView>,
|
||||
pub latency_details: UsageLatencyDetailsView,
|
||||
pub first_token_latency_ms: Option<u64>,
|
||||
pub latency_ms: Option<u64>,
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub created_at_display: String,
|
||||
pub client_ip: Option<String>,
|
||||
pub user_agent: Option<String>,
|
||||
}
|
||||
|
||||
/// 单条逻辑请求详情展示。
|
||||
#[derive(Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct UsageRecordView {
|
||||
@@ -99,7 +124,8 @@ pub struct UsageRecordView {
|
||||
pub upstream_model: Option<String>,
|
||||
pub service_tier: Option<String>,
|
||||
pub status_code: Option<i64>,
|
||||
pub transport: Option<String>,
|
||||
pub client_transport: String,
|
||||
pub upstream_transport: Option<String>,
|
||||
pub protocol: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub http_version: Option<String>,
|
||||
@@ -525,7 +551,7 @@ pub struct DashboardDataView {
|
||||
pub health_timeline: HealthTimelineView,
|
||||
pub wire_profiles: Vec<DashboardWireProfileView>,
|
||||
pub account_usage: Vec<DashboardAccountUsageView>,
|
||||
pub usage_records: Vec<UsageRecordView>,
|
||||
pub usage_records: Vec<UsageListRecordView>,
|
||||
pub pool_summary: DashboardPoolSummaryView,
|
||||
pub capacity_info: DashboardCapacityInfoView,
|
||||
pub rotation_strategy: String,
|
||||
@@ -732,19 +758,36 @@ pub struct OpsErrorView {
|
||||
pub request_id: Option<String>,
|
||||
pub client_api_key_id: Option<String>,
|
||||
pub kind: String,
|
||||
pub operation: String,
|
||||
pub protocol: Option<String>,
|
||||
pub client_transport: Option<String>,
|
||||
pub provider: Option<String>,
|
||||
pub authentication_kind: Option<String>,
|
||||
pub account_id: Option<String>,
|
||||
pub account_name: Option<String>,
|
||||
pub account_email: Option<String>,
|
||||
pub route: String,
|
||||
pub model: Option<String>,
|
||||
pub requested_model: Option<String>,
|
||||
pub upstream_model: Option<String>,
|
||||
pub service_tier: Option<String>,
|
||||
pub client_status_code: Option<i64>,
|
||||
pub upstream_status_code: Option<i64>,
|
||||
pub transport: Option<String>,
|
||||
pub attempt_index: Option<u32>,
|
||||
pub failure_class: String,
|
||||
pub provider_error_code: Option<String>,
|
||||
pub occurrence_count: u32,
|
||||
pub response_id: Option<String>,
|
||||
pub upstream_request_id: Option<String>,
|
||||
pub latency_ms: Option<u64>,
|
||||
pub client_ip: Option<String>,
|
||||
pub user_agent: Option<String>,
|
||||
pub reasoning_effort: Option<String>,
|
||||
pub reasoning_preset: Option<String>,
|
||||
pub request_kind: Option<String>,
|
||||
pub subagent_kind: Option<String>,
|
||||
pub compact: Option<bool>,
|
||||
pub message: String,
|
||||
pub metadata: OpsErrorMetadataView,
|
||||
pub created_at: DateTime<Utc>,
|
||||
|
||||
@@ -455,10 +455,11 @@ fn passthrough_header_name(name: &str, connection_headers: &[String]) -> bool {
|
||||
| "x-openai-project"
|
||||
// installation ID 始终由当前 lease 重建。
|
||||
| "x-codex-installation-id"
|
||||
// 上游指纹必须由运行时画像统一生成,客户端 originator/User-Agent
|
||||
// 不能透传覆盖,避免不同下游客户端暴露不一致的设备指纹。
|
||||
// 上游指纹必须由运行时画像统一生成,客户端 originator/User-Agent/version
|
||||
// 不能作为不透明头透传覆盖,避免不同下游客户端暴露不一致的设备指纹。
|
||||
| "originator"
|
||||
| "user-agent"
|
||||
| "version"
|
||||
// 设备 attestation/integrity 头只能由官方客户端或网关自身生成,
|
||||
// 客户端注入的 x-oai-attestation / X-OAI-IS 不得透传上游。
|
||||
| "x-oai-attestation"
|
||||
|
||||
@@ -1,699 +0,0 @@
|
||||
//! OpenAI Responses WebSocket adapter。
|
||||
|
||||
use std::net::{IpAddr, SocketAddr};
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::{
|
||||
extract::{
|
||||
Extension, State,
|
||||
connect_info::ConnectInfo,
|
||||
ws::{CloseFrame, Message, WebSocket, WebSocketUpgrade, close_code},
|
||||
},
|
||||
http::{HeaderMap, HeaderName, StatusCode},
|
||||
response::{IntoResponse, Response},
|
||||
};
|
||||
use futures::StreamExt;
|
||||
use gateway_core::engine::execution::{AuthenticatedClient, ClientTransport, StartedExecution};
|
||||
use gateway_core::engine::{CommitRequirement, CoordinatedEvent, EngineError};
|
||||
use gateway_core::error::{GatewayError, GatewayErrorKind};
|
||||
use gateway_core::event::ProviderResponseHeader;
|
||||
use gateway_core::lifecycle::{ConnectionGuard, ConnectionLifecycle};
|
||||
use gateway_core::operation::ProviderSessionState;
|
||||
use serde_json::{Map, Value, json};
|
||||
use thiserror::Error;
|
||||
|
||||
use crate::{
|
||||
ApiState,
|
||||
openai::{
|
||||
auth::{authenticate_client, authentication_error_response},
|
||||
error::{gateway_error_contract, gateway_error_from_engine, runtime_unavailable_response},
|
||||
service::OpenAiService,
|
||||
},
|
||||
};
|
||||
|
||||
use super::{
|
||||
DecodedResponsesRequest, OpenAiResponsesEncoder, ProtocolErrorBody, RequestDecodeError,
|
||||
http::{PendingExecution, request_client_context},
|
||||
request::OpenAiRequestHeaders,
|
||||
};
|
||||
|
||||
const ACTIVE_RESPONSE_VIOLATION: &str = "Only one response.create may be active per connection";
|
||||
const TEXT_FRAMES_ONLY: &str = "Responses WebSocket accepts text frames only";
|
||||
|
||||
/// 将已认证的 `GET /v1/responses` 升级为 Responses WebSocket。
|
||||
pub(crate) async fn responses_websocket(
|
||||
State(state): State<ApiState>,
|
||||
connect_info: Option<Extension<ConnectInfo<SocketAddr>>>,
|
||||
headers: HeaderMap,
|
||||
websocket: WebSocketUpgrade,
|
||||
) -> Response {
|
||||
let service = state.openai().clone();
|
||||
let client = match authenticate_client(&service, &headers) {
|
||||
Ok(client) => client,
|
||||
Err(error) => return authentication_error_response(error),
|
||||
};
|
||||
let (client_ip, user_agent) = request_client_context(
|
||||
&headers,
|
||||
connect_info.map(|Extension(ConnectInfo(address))| address),
|
||||
);
|
||||
ResponsesWebSocketAdapter::new(service).upgrade_with_client_context(
|
||||
websocket,
|
||||
client,
|
||||
client_ip,
|
||||
user_agent,
|
||||
OpenAiRequestHeaders::from_headers(&headers),
|
||||
)
|
||||
}
|
||||
|
||||
/// 已鉴权 Responses WebSocket 升级边界。
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct ResponsesWebSocketAdapter {
|
||||
service: OpenAiService,
|
||||
}
|
||||
|
||||
impl ResponsesWebSocketAdapter {
|
||||
/// 绑定应用提供的唯一 OpenAI 客户端服务端口。
|
||||
#[must_use]
|
||||
pub const fn new(service: OpenAiService) -> Self {
|
||||
Self { service }
|
||||
}
|
||||
|
||||
fn upgrade_with_client_context(
|
||||
&self,
|
||||
websocket: WebSocketUpgrade,
|
||||
client: AuthenticatedClient,
|
||||
client_ip: Option<IpAddr>,
|
||||
user_agent: Option<String>,
|
||||
request_headers: OpenAiRequestHeaders,
|
||||
) -> Response {
|
||||
let connection_guard = match self.service.try_register_connection() {
|
||||
Ok(guard) => guard,
|
||||
Err(_) => return runtime_unavailable_response().into_response(),
|
||||
};
|
||||
let session = ResponsesWebSocketSession {
|
||||
service: self.service.clone(),
|
||||
client,
|
||||
connection_id: self.service.next_request_id().replacen("req_", "ws_", 1),
|
||||
client_ip,
|
||||
user_agent,
|
||||
request_headers,
|
||||
lifecycle: self.service.lifecycle(),
|
||||
_connection_guard: connection_guard,
|
||||
};
|
||||
websocket
|
||||
// 覆盖 axum/tungstenite 的私有 64 MiB message 与 16 MiB frame 默认值。
|
||||
// Responses JSON 的协议可接受性由上游决定,代理不另设 wire 长度上限。
|
||||
.max_message_size(usize::MAX)
|
||||
.max_frame_size(usize::MAX)
|
||||
.on_upgrade(move |socket| async move {
|
||||
serve_responses_websocket(socket, session).await;
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
struct ResponsesWebSocketSession {
|
||||
service: OpenAiService,
|
||||
client: AuthenticatedClient,
|
||||
connection_id: String,
|
||||
client_ip: Option<IpAddr>,
|
||||
user_agent: Option<String>,
|
||||
request_headers: OpenAiRequestHeaders,
|
||||
lifecycle: Arc<dyn ConnectionLifecycle>,
|
||||
_connection_guard: Box<dyn ConnectionGuard>,
|
||||
}
|
||||
|
||||
async fn serve_responses_websocket(mut socket: WebSocket, session: ResponsesWebSocketSession) {
|
||||
tracing::info!(
|
||||
websocket_connection_id = %session.connection_id,
|
||||
"Responses WebSocket connected"
|
||||
);
|
||||
let mut request_count = 0_u64;
|
||||
let mut replay = ConnectionReplaySnapshot::default();
|
||||
|
||||
let cancellation = session.lifecycle.cancellation();
|
||||
loop {
|
||||
let message = tokio::select! {
|
||||
() = cancellation.cancelled() => break,
|
||||
message = socket.next() => message,
|
||||
};
|
||||
let Some(message) = message else {
|
||||
break;
|
||||
};
|
||||
let payload = match message {
|
||||
Ok(Message::Text(payload)) => payload,
|
||||
Ok(Message::Ping(_) | Message::Pong(_)) => continue,
|
||||
Ok(Message::Close(_)) => break,
|
||||
Ok(Message::Binary(_)) => {
|
||||
close_for_protocol_violation(&mut socket, TEXT_FRAMES_ONLY).await;
|
||||
break;
|
||||
}
|
||||
Err(error) => {
|
||||
tracing::info!(
|
||||
websocket_connection_id = %session.connection_id,
|
||||
error = %error,
|
||||
"Responses WebSocket receive failed"
|
||||
);
|
||||
break;
|
||||
}
|
||||
};
|
||||
request_count = request_count.saturating_add(1);
|
||||
let correlation_id = session.service.next_request_id();
|
||||
let decoded =
|
||||
match decode_response_create_with_context(payload.as_str(), &session.request_headers) {
|
||||
Ok(decoded) => {
|
||||
decoded.with_client_context(session.client_ip, session.user_agent.clone())
|
||||
}
|
||||
Err(error) => {
|
||||
tracing::info!(
|
||||
websocket_connection_id = %session.connection_id,
|
||||
request_id = %correlation_id,
|
||||
error = %error,
|
||||
"Responses WebSocket request rejected"
|
||||
);
|
||||
if send_protocol_error(
|
||||
&mut socket,
|
||||
StatusCode::BAD_REQUEST,
|
||||
error.protocol_body(),
|
||||
&correlation_id,
|
||||
)
|
||||
.await
|
||||
== ForwardOutcome::Disconnect
|
||||
{
|
||||
break;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let decoded = replay.prepare(decoded);
|
||||
let started = match session
|
||||
.service
|
||||
.start_response(
|
||||
session.client.clone(),
|
||||
decoded,
|
||||
ClientTransport::WebSocket,
|
||||
"/v1/responses",
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(started) => started,
|
||||
Err(error) => {
|
||||
if send_gateway_error(&mut socket, &error, &correlation_id).await
|
||||
== ForwardOutcome::Disconnect
|
||||
{
|
||||
break;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
if forward_execution(&mut socket, started, &mut replay).await == ForwardOutcome::Disconnect
|
||||
{
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
tracing::info!(
|
||||
websocket_connection_id = %session.connection_id,
|
||||
request_count,
|
||||
"Responses WebSocket disconnected"
|
||||
);
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, PartialEq, Eq)]
|
||||
enum ForwardOutcome {
|
||||
Continue,
|
||||
Disconnect,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct ConnectionReplaySnapshot {
|
||||
last_response_id: Option<String>,
|
||||
provider_state: Option<ProviderSessionState>,
|
||||
}
|
||||
|
||||
impl ConnectionReplaySnapshot {
|
||||
fn prepare(&self, request: DecodedResponsesRequest) -> DecodedResponsesRequest {
|
||||
match (
|
||||
request.metadata().continuation().previous_response_id(),
|
||||
self.last_response_id.as_deref(),
|
||||
self.provider_state.as_ref(),
|
||||
) {
|
||||
(Some(previous), Some(last), Some(state)) if previous == last => {
|
||||
request.with_provider_session_state(state.clone())
|
||||
}
|
||||
_ => request,
|
||||
}
|
||||
}
|
||||
|
||||
fn commit(&mut self, response_id: String, provider_state: Option<ProviderSessionState>) {
|
||||
self.last_response_id = Some(response_id);
|
||||
self.provider_state = provider_state;
|
||||
}
|
||||
}
|
||||
|
||||
async fn forward_execution(
|
||||
socket: &mut WebSocket,
|
||||
started: StartedExecution,
|
||||
replay: &mut ConnectionReplaySnapshot,
|
||||
) -> ForwardOutcome {
|
||||
let request_id = started.request_id.to_string();
|
||||
let streaming = started.stream;
|
||||
let mut execution = PendingExecution::new(started.session);
|
||||
if !streaming {
|
||||
let error = GatewayError::new(
|
||||
GatewayErrorKind::Internal,
|
||||
"WebSocket execution was not initialized as a stream",
|
||||
);
|
||||
return send_gateway_error(socket, &error, &request_id).await;
|
||||
}
|
||||
let first = loop {
|
||||
match next_active_input(socket, &mut execution).await {
|
||||
ActiveInput::Event(Ok(Some(event))) => break event,
|
||||
ActiveInput::Event(Ok(None)) => {
|
||||
let error = GatewayError::new(
|
||||
GatewayErrorKind::Internal,
|
||||
"gateway response ended before its first event",
|
||||
);
|
||||
return send_gateway_error(socket, &error, &request_id).await;
|
||||
}
|
||||
ActiveInput::Event(Err(error)) => {
|
||||
let error = gateway_error_from_engine(&error);
|
||||
return send_gateway_error(socket, &error, &request_id).await;
|
||||
}
|
||||
ActiveInput::Control => continue,
|
||||
ActiveInput::Disconnect => return ForwardOutcome::Disconnect,
|
||||
ActiveInput::ProtocolViolation => {
|
||||
close_for_protocol_violation(socket, ACTIVE_RESPONSE_VIOLATION).await;
|
||||
return ForwardOutcome::Disconnect;
|
||||
}
|
||||
}
|
||||
};
|
||||
let requirement = first.commit_requirement();
|
||||
let mut first = first.into_provider_events();
|
||||
if requirement != CommitRequirement::CommitBeforeDelivery {
|
||||
let error = GatewayError::new(
|
||||
GatewayErrorKind::Internal,
|
||||
"gateway first event did not require commit",
|
||||
);
|
||||
return send_gateway_error(socket, &error, &request_id).await;
|
||||
}
|
||||
let mut encoder = OpenAiResponsesEncoder::new();
|
||||
let mut provider_state = None;
|
||||
let mut first_messages = Vec::new();
|
||||
for event in &mut first {
|
||||
if let Some(update) = event.take_session_update() {
|
||||
provider_state = Some(update);
|
||||
}
|
||||
first_messages.extend(encoder.push_websocket(event));
|
||||
}
|
||||
if first_messages.is_empty() {
|
||||
let error = GatewayError::new(
|
||||
GatewayErrorKind::Internal,
|
||||
"gateway commit batch encoded no output",
|
||||
);
|
||||
return send_gateway_error(socket, &error, &request_id).await;
|
||||
}
|
||||
let Some(response_session) = execution.session_mut() else {
|
||||
return ForwardOutcome::Disconnect;
|
||||
};
|
||||
let response_headers = response_session.response_headers().to_vec();
|
||||
if let Err(error) = response_session.commit_downstream(None).await {
|
||||
let error = gateway_error_from_engine(&error);
|
||||
return send_gateway_error(socket, &error, &request_id).await;
|
||||
}
|
||||
if encoder.is_completed() {
|
||||
if let Err(outcome) = confirm_completed_execution(socket, &mut execution, &request_id).await
|
||||
{
|
||||
return outcome;
|
||||
}
|
||||
commit_connection_replay(replay, &encoder, provider_state.take());
|
||||
if !send_text(
|
||||
socket,
|
||||
response_metadata_event(&request_id, &response_headers),
|
||||
)
|
||||
.await
|
||||
|| !send_messages(socket, first_messages).await
|
||||
{
|
||||
return ForwardOutcome::Disconnect;
|
||||
}
|
||||
execution.disarm();
|
||||
return ForwardOutcome::Continue;
|
||||
}
|
||||
if !send_text(
|
||||
socket,
|
||||
response_metadata_event(&request_id, &response_headers),
|
||||
)
|
||||
.await
|
||||
|| !send_messages(socket, first_messages).await
|
||||
{
|
||||
return ForwardOutcome::Disconnect;
|
||||
}
|
||||
|
||||
loop {
|
||||
match next_active_input(socket, &mut execution).await {
|
||||
ActiveInput::Event(Ok(Some(delivery))) => {
|
||||
let requirement = delivery.commit_requirement();
|
||||
let mut events = delivery.into_provider_events();
|
||||
if requirement != CommitRequirement::AlreadyCommitted {
|
||||
let error = GatewayError::new(
|
||||
GatewayErrorKind::Internal,
|
||||
"gateway requested another downstream commit",
|
||||
);
|
||||
return send_gateway_error(socket, &error, &request_id).await;
|
||||
}
|
||||
let mut messages = Vec::new();
|
||||
for event in &mut events {
|
||||
if let Some(update) = event.take_session_update() {
|
||||
provider_state = Some(update);
|
||||
}
|
||||
messages.extend(encoder.push_websocket(event));
|
||||
}
|
||||
if encoder.is_completed() {
|
||||
if let Err(outcome) =
|
||||
confirm_completed_execution(socket, &mut execution, &request_id).await
|
||||
{
|
||||
return outcome;
|
||||
}
|
||||
commit_connection_replay(replay, &encoder, provider_state.take());
|
||||
if !send_messages(socket, messages).await {
|
||||
return ForwardOutcome::Disconnect;
|
||||
}
|
||||
execution.disarm();
|
||||
return ForwardOutcome::Continue;
|
||||
}
|
||||
if !send_messages(socket, messages).await {
|
||||
return ForwardOutcome::Disconnect;
|
||||
}
|
||||
}
|
||||
ActiveInput::Event(Ok(None)) => {
|
||||
if execution
|
||||
.session_mut()
|
||||
.is_some_and(|session| session.is_finalized())
|
||||
{
|
||||
commit_connection_replay(replay, &encoder, provider_state.take());
|
||||
execution.disarm();
|
||||
return ForwardOutcome::Continue;
|
||||
}
|
||||
let error = GatewayError::new(
|
||||
GatewayErrorKind::Internal,
|
||||
"gateway response ended without finalizing the execution",
|
||||
);
|
||||
return send_gateway_error(socket, &error, &request_id).await;
|
||||
}
|
||||
ActiveInput::Event(Err(error)) => {
|
||||
if encoder.has_wire_failure() {
|
||||
execution.disarm();
|
||||
return ForwardOutcome::Continue;
|
||||
}
|
||||
let error = gateway_error_from_engine(&error);
|
||||
return send_gateway_error(socket, &error, &request_id).await;
|
||||
}
|
||||
ActiveInput::Control => {}
|
||||
ActiveInput::Disconnect => return ForwardOutcome::Disconnect,
|
||||
ActiveInput::ProtocolViolation => {
|
||||
close_for_protocol_violation(socket, ACTIVE_RESPONSE_VIOLATION).await;
|
||||
return ForwardOutcome::Disconnect;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn commit_connection_replay(
|
||||
replay: &mut ConnectionReplaySnapshot,
|
||||
encoder: &OpenAiResponsesEncoder,
|
||||
provider_state: Option<ProviderSessionState>,
|
||||
) {
|
||||
if let Some(response_id) = encoder.response_id() {
|
||||
replay.commit(response_id.to_owned(), provider_state);
|
||||
}
|
||||
}
|
||||
|
||||
async fn confirm_completed_execution(
|
||||
socket: &mut WebSocket,
|
||||
execution: &mut PendingExecution,
|
||||
request_id: &str,
|
||||
) -> Result<(), ForwardOutcome> {
|
||||
loop {
|
||||
match next_active_input(socket, execution).await {
|
||||
ActiveInput::Event(Ok(None))
|
||||
if execution
|
||||
.session_mut()
|
||||
.is_some_and(|session| session.is_finalized()) =>
|
||||
{
|
||||
return Ok(());
|
||||
}
|
||||
ActiveInput::Event(Ok(None)) => {
|
||||
let error = GatewayError::new(
|
||||
GatewayErrorKind::Internal,
|
||||
"gateway response was not finalized after its terminal event",
|
||||
);
|
||||
return Err(send_gateway_error(socket, &error, request_id).await);
|
||||
}
|
||||
ActiveInput::Event(Ok(Some(_))) => {
|
||||
let error = GatewayError::new(
|
||||
GatewayErrorKind::Internal,
|
||||
"gateway response continued after its terminal event",
|
||||
);
|
||||
return Err(send_gateway_error(socket, &error, request_id).await);
|
||||
}
|
||||
ActiveInput::Event(Err(error)) => {
|
||||
let error = gateway_error_from_engine(&error);
|
||||
return Err(send_gateway_error(socket, &error, request_id).await);
|
||||
}
|
||||
ActiveInput::Control => {}
|
||||
ActiveInput::Disconnect => return Err(ForwardOutcome::Disconnect),
|
||||
ActiveInput::ProtocolViolation => {
|
||||
close_for_protocol_violation(socket, ACTIVE_RESPONSE_VIOLATION).await;
|
||||
return Err(ForwardOutcome::Disconnect);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
enum ActiveInput {
|
||||
Event(Result<Option<CoordinatedEvent>, EngineError>),
|
||||
Control,
|
||||
Disconnect,
|
||||
ProtocolViolation,
|
||||
}
|
||||
|
||||
async fn next_active_input(
|
||||
socket: &mut WebSocket,
|
||||
execution: &mut PendingExecution,
|
||||
) -> ActiveInput {
|
||||
let Some(session) = execution.session_mut() else {
|
||||
return ActiveInput::Disconnect;
|
||||
};
|
||||
tokio::select! {
|
||||
event = session.next_event() => ActiveInput::Event(event),
|
||||
message = socket.next() => match message {
|
||||
Some(Ok(Message::Ping(_) | Message::Pong(_))) => ActiveInput::Control,
|
||||
Some(Ok(Message::Close(_))) | None | Some(Err(_)) => ActiveInput::Disconnect,
|
||||
Some(Ok(Message::Text(_) | Message::Binary(_))) => ActiveInput::ProtocolViolation,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// 使用连接级请求头和 Provider 上下文解码官方 `response.create` 文本帧。
|
||||
///
|
||||
/// 缺省 `stream` 等价于 WebSocket 固有的流式语义;显式 `false` 会被拒绝。
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// 帧不是合法 JSON object、消息类型错误、显式关闭 stream,或 Responses 请求
|
||||
/// 无法映射到 canonical operation 时返回不包含正文内容的稳定错误。
|
||||
pub fn decode_response_create_with_context(
|
||||
payload: &str,
|
||||
request_headers: &OpenAiRequestHeaders,
|
||||
) -> Result<DecodedResponsesRequest, ResponseCreateFrameError> {
|
||||
decode_response_create_inner(payload, request_headers)
|
||||
}
|
||||
|
||||
fn decode_response_create_inner(
|
||||
payload: &str,
|
||||
request_headers: &OpenAiRequestHeaders,
|
||||
) -> Result<DecodedResponsesRequest, ResponseCreateFrameError> {
|
||||
let Value::Object(mut body) = serde_json::from_str::<Value>(payload)
|
||||
.map_err(|_| ResponseCreateFrameError::InvalidJson)?
|
||||
else {
|
||||
return Err(ResponseCreateFrameError::ExpectedObject);
|
||||
};
|
||||
match body.remove("type") {
|
||||
Some(Value::String(message_type)) if message_type == "response.create" => {}
|
||||
_ => return Err(ResponseCreateFrameError::UnsupportedType),
|
||||
}
|
||||
if matches!(body.get("stream"), Some(value) if value.as_bool() != Some(true)) {
|
||||
return Err(ResponseCreateFrameError::StreamingRequired);
|
||||
}
|
||||
super::request::decode_request_object(body, false, request_headers)
|
||||
.map_err(ResponseCreateFrameError::Request)
|
||||
}
|
||||
|
||||
/// `response.create` 帧的稳定安全错误。
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Error)]
|
||||
pub enum ResponseCreateFrameError {
|
||||
/// 文本不是合法 JSON。
|
||||
#[error("response.create frame must be valid JSON")]
|
||||
InvalidJson,
|
||||
/// 顶层不是 object。
|
||||
#[error("response.create frame must be a JSON object")]
|
||||
ExpectedObject,
|
||||
/// `type` 缺失或不是 `response.create`。
|
||||
#[error("unsupported Responses WebSocket message type")]
|
||||
UnsupportedType,
|
||||
/// WebSocket 请求显式声明 `stream=false`。
|
||||
#[error("Responses WebSocket requests require stream=true")]
|
||||
StreamingRequired,
|
||||
/// 内层 Responses 请求无法映射到 canonical operation。
|
||||
#[error(transparent)]
|
||||
Request(RequestDecodeError),
|
||||
}
|
||||
|
||||
impl ResponseCreateFrameError {
|
||||
fn protocol_body(&self) -> ProtocolErrorBody {
|
||||
match self {
|
||||
Self::Request(error) => error.protocol_body(),
|
||||
Self::InvalidJson => RequestDecodeError::MalformedJson.protocol_body(),
|
||||
Self::ExpectedObject => RequestDecodeError::ExpectedObject.protocol_body(),
|
||||
Self::UnsupportedType => RequestDecodeError::InvalidValue {
|
||||
field: "type".to_owned(),
|
||||
}
|
||||
.protocol_body(),
|
||||
Self::StreamingRequired => RequestDecodeError::InvalidValue {
|
||||
field: "stream".to_owned(),
|
||||
}
|
||||
.protocol_body(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn response_metadata_event(
|
||||
request_id: &str,
|
||||
response_headers: &[ProviderResponseHeader],
|
||||
) -> String {
|
||||
let connection_options = super::response_connection_options(response_headers);
|
||||
let mut headers = Map::new();
|
||||
for header in response_headers {
|
||||
if !super::response_header_is_forwardable(header.name(), &connection_options) {
|
||||
continue;
|
||||
}
|
||||
let Ok(name) = HeaderName::from_bytes(header.name().as_bytes()) else {
|
||||
continue;
|
||||
};
|
||||
let Ok(value) = std::str::from_utf8(header.value()) else {
|
||||
continue;
|
||||
};
|
||||
// 官方 response.metadata 使用 string map,无法表达同名多值;后值保持既有覆盖语义。
|
||||
headers.insert(name.as_str().to_owned(), Value::String(value.to_owned()));
|
||||
}
|
||||
headers
|
||||
.entry("x-request-id".to_owned())
|
||||
.or_insert_with(|| Value::String(request_id.to_owned()));
|
||||
json!({
|
||||
"type": "response.metadata",
|
||||
"headers": headers,
|
||||
})
|
||||
.to_string()
|
||||
}
|
||||
|
||||
async fn send_gateway_error(
|
||||
socket: &mut WebSocket,
|
||||
error: &GatewayError,
|
||||
request_id: &str,
|
||||
) -> ForwardOutcome {
|
||||
let (status, default_type, default_code) = gateway_error_contract(error.kind());
|
||||
if send_error_event(
|
||||
socket,
|
||||
status,
|
||||
error.client_error_type().unwrap_or(default_type),
|
||||
error.client_error_code().unwrap_or(default_code),
|
||||
error.client_message(),
|
||||
None,
|
||||
request_id,
|
||||
)
|
||||
.await
|
||||
{
|
||||
ForwardOutcome::Continue
|
||||
} else {
|
||||
ForwardOutcome::Disconnect
|
||||
}
|
||||
}
|
||||
|
||||
async fn send_protocol_error(
|
||||
socket: &mut WebSocket,
|
||||
status: StatusCode,
|
||||
body: ProtocolErrorBody,
|
||||
request_id: &str,
|
||||
) -> ForwardOutcome {
|
||||
let error = body.error;
|
||||
if send_error_event(
|
||||
socket,
|
||||
status,
|
||||
error.kind,
|
||||
error.code,
|
||||
&error.message,
|
||||
error.param.as_deref(),
|
||||
request_id,
|
||||
)
|
||||
.await
|
||||
{
|
||||
ForwardOutcome::Continue
|
||||
} else {
|
||||
ForwardOutcome::Disconnect
|
||||
}
|
||||
}
|
||||
|
||||
async fn send_error_event(
|
||||
socket: &mut WebSocket,
|
||||
status: StatusCode,
|
||||
error_type: &str,
|
||||
code: &str,
|
||||
message: &str,
|
||||
param: Option<&str>,
|
||||
request_id: &str,
|
||||
) -> bool {
|
||||
let mut error = Map::new();
|
||||
error.insert("type".to_owned(), Value::String(error_type.to_owned()));
|
||||
error.insert("code".to_owned(), Value::String(code.to_owned()));
|
||||
error.insert("message".to_owned(), Value::String(message.to_owned()));
|
||||
if let Some(param) = param {
|
||||
error.insert("param".to_owned(), Value::String(param.to_owned()));
|
||||
}
|
||||
send_text(
|
||||
socket,
|
||||
json!({
|
||||
"type": "error",
|
||||
"status": status.as_u16(),
|
||||
"error": error,
|
||||
"request_id": request_id,
|
||||
})
|
||||
.to_string(),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn send_messages(socket: &mut WebSocket, messages: Vec<String>) -> bool {
|
||||
for message in messages {
|
||||
if !send_text(socket, message).await {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
async fn send_text(socket: &mut WebSocket, payload: String) -> bool {
|
||||
if let Err(error) = socket.send(Message::Text(payload.into())).await {
|
||||
tracing::info!(error = %error, "Responses WebSocket send failed");
|
||||
return false;
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
async fn close_for_protocol_violation(socket: &mut WebSocket, reason: &'static str) {
|
||||
let _ = socket
|
||||
.send(Message::Close(Some(CloseFrame {
|
||||
code: close_code::POLICY,
|
||||
reason: reason.into(),
|
||||
})))
|
||||
.await;
|
||||
}
|
||||
@@ -0,0 +1,550 @@
|
||||
//! 下游客户端 WebSocket 的单 owner pump 与有界收发边界。
|
||||
|
||||
use std::{
|
||||
fmt,
|
||||
sync::{
|
||||
Arc,
|
||||
atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering},
|
||||
},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use axum::extract::ws::{CloseFrame, Message, WebSocket, close_code};
|
||||
use futures::{Sink, SinkExt, Stream, StreamExt};
|
||||
use gateway_core::engine::CancellationToken;
|
||||
use thiserror::Error;
|
||||
use tokio::{
|
||||
sync::{mpsc, oneshot},
|
||||
task::JoinHandle,
|
||||
time::{Instant, timeout},
|
||||
};
|
||||
|
||||
const OUTBOUND_COMMAND_BUFFER: usize = 32;
|
||||
const INBOUND_EVENT_BUFFER: usize = 32;
|
||||
const DOWNSTREAM_WRITE_TIMEOUT: Duration = Duration::from_secs(5 * 60);
|
||||
const CONNECTION_MAX_AGE: Duration = Duration::from_secs(60 * 60);
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
pub(super) struct ConnectionConfig {
|
||||
pub(super) write_timeout: Duration,
|
||||
pub(super) max_age: Duration,
|
||||
}
|
||||
|
||||
impl ConnectionConfig {
|
||||
pub(super) const PRODUCTION: Self = Self {
|
||||
write_timeout: DOWNSTREAM_WRITE_TIMEOUT,
|
||||
max_age: CONNECTION_MAX_AGE,
|
||||
};
|
||||
}
|
||||
|
||||
/// 业务层可观察的客户端输入;Ping/Pong 始终由 pump 消费。
|
||||
pub(super) enum ConnectionEvent {
|
||||
Text(String),
|
||||
Binary,
|
||||
Expired,
|
||||
Exited(PumpExitReason),
|
||||
}
|
||||
|
||||
/// 下游写入阶段;名称刻意使用 write,而不是暗示客户端已消费的 delivery。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(super) enum FramePhase {
|
||||
Metadata,
|
||||
First,
|
||||
Data,
|
||||
Terminal,
|
||||
FirstAndTerminal,
|
||||
Error,
|
||||
ConnectionLimit,
|
||||
Close,
|
||||
}
|
||||
|
||||
impl FramePhase {
|
||||
const fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::Metadata => "metadata",
|
||||
Self::First => "first",
|
||||
Self::Data => "data",
|
||||
Self::Terminal => "terminal",
|
||||
Self::FirstAndTerminal => "first_and_terminal",
|
||||
Self::Error => "error",
|
||||
Self::ConnectionLimit => "connection_limit",
|
||||
Self::Close => "close",
|
||||
}
|
||||
}
|
||||
|
||||
const fn is_milestone(self) -> bool {
|
||||
matches!(
|
||||
self,
|
||||
Self::First
|
||||
| Self::Terminal
|
||||
| Self::FirstAndTerminal
|
||||
| Self::Error
|
||||
| Self::ConnectionLimit
|
||||
| Self::Close
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(super) struct WriteContext {
|
||||
request_id: Option<Arc<str>>,
|
||||
phase: FramePhase,
|
||||
}
|
||||
|
||||
impl WriteContext {
|
||||
pub(super) fn request(request_id: &Arc<str>, phase: FramePhase) -> Self {
|
||||
Self {
|
||||
request_id: Some(Arc::clone(request_id)),
|
||||
phase,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) const fn connection(phase: FramePhase) -> Self {
|
||||
Self {
|
||||
request_id: None,
|
||||
phase,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(super) enum PumpExitReason {
|
||||
ClientClose,
|
||||
PeerEof,
|
||||
ReadError,
|
||||
WriteError,
|
||||
WriteTimeout,
|
||||
LifecycleShutdown,
|
||||
ConnectionMaxAge,
|
||||
CoordinatorDropped,
|
||||
InboundOverload,
|
||||
ServerClose,
|
||||
PumpStopped,
|
||||
}
|
||||
|
||||
impl PumpExitReason {
|
||||
pub(super) const fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::ClientClose => "client_close",
|
||||
Self::PeerEof => "peer_eof",
|
||||
Self::ReadError => "read_error",
|
||||
Self::WriteError => "write_error",
|
||||
Self::WriteTimeout => "write_timeout",
|
||||
Self::LifecycleShutdown => "lifecycle_shutdown",
|
||||
Self::ConnectionMaxAge => "connection_max_age",
|
||||
Self::CoordinatorDropped => "coordinator_dropped",
|
||||
Self::InboundOverload => "inbound_overload",
|
||||
Self::ServerClose => "server_close",
|
||||
Self::PumpStopped => "pump_stopped",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub(super) enum ConnectionWriteError {
|
||||
#[error("downstream WebSocket pump is closed")]
|
||||
Closed,
|
||||
#[error("downstream WebSocket write timed out after {timeout:?}")]
|
||||
Timeout { timeout: Duration },
|
||||
#[error("downstream WebSocket transport write failed: {message}")]
|
||||
Transport { message: String },
|
||||
}
|
||||
|
||||
struct ConnectionCommand {
|
||||
message: Message,
|
||||
context: WriteContext,
|
||||
acknowledged: oneshot::Sender<Result<(), ConnectionWriteError>>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct ConnectionStats {
|
||||
command_queue_high_water: AtomicUsize,
|
||||
command_backpressure_count: AtomicU64,
|
||||
ping_received_count: AtomicU64,
|
||||
pong_written_count: AtomicU64,
|
||||
}
|
||||
|
||||
impl ConnectionStats {
|
||||
fn observe_command_queue(&self, sender: &mpsc::Sender<ConnectionCommand>) {
|
||||
let queued = OUTBOUND_COMMAND_BUFFER
|
||||
.saturating_sub(sender.capacity())
|
||||
.saturating_add(1)
|
||||
.min(OUTBOUND_COMMAND_BUFFER);
|
||||
self.command_queue_high_water
|
||||
.fetch_max(queued, Ordering::Relaxed);
|
||||
if sender.capacity() == 0 {
|
||||
self.command_backpressure_count
|
||||
.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) struct ResponsesWebSocketConnection {
|
||||
connection_id: Arc<str>,
|
||||
opened_at: Instant,
|
||||
expired: Arc<AtomicBool>,
|
||||
commands: Option<mpsc::Sender<ConnectionCommand>>,
|
||||
incoming: mpsc::Receiver<ConnectionEvent>,
|
||||
pump_task: Option<JoinHandle<()>>,
|
||||
stats: Arc<ConnectionStats>,
|
||||
config: ConnectionConfig,
|
||||
exit_reason: Option<PumpExitReason>,
|
||||
}
|
||||
|
||||
impl ResponsesWebSocketConnection {
|
||||
pub(super) fn new(
|
||||
socket: WebSocket,
|
||||
connection_id: String,
|
||||
cancellation: CancellationToken,
|
||||
) -> Self {
|
||||
spawn_connection(
|
||||
socket,
|
||||
Arc::<str>::from(connection_id),
|
||||
cancellation,
|
||||
ConnectionConfig::PRODUCTION,
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn id(&self) -> &str {
|
||||
&self.connection_id
|
||||
}
|
||||
|
||||
pub(super) fn is_expired(&self) -> bool {
|
||||
self.expired.load(Ordering::Acquire)
|
||||
}
|
||||
|
||||
pub(super) fn age(&self) -> Duration {
|
||||
self.opened_at.elapsed()
|
||||
}
|
||||
|
||||
pub(super) async fn next_event(&mut self) -> Option<ConnectionEvent> {
|
||||
let event = self.incoming.recv().await;
|
||||
if let Some(ConnectionEvent::Exited(reason)) = event.as_ref() {
|
||||
self.exit_reason.get_or_insert(*reason);
|
||||
} else if event.is_none() {
|
||||
self.exit_reason.get_or_insert(PumpExitReason::PumpStopped);
|
||||
}
|
||||
event
|
||||
}
|
||||
|
||||
pub(super) async fn send_text(
|
||||
&mut self,
|
||||
payload: String,
|
||||
context: WriteContext,
|
||||
) -> Result<(), ConnectionWriteError> {
|
||||
self.send(Message::Text(payload.into()), context).await
|
||||
}
|
||||
|
||||
pub(super) async fn close_policy(
|
||||
&mut self,
|
||||
reason: &'static str,
|
||||
request_id: Option<&Arc<str>>,
|
||||
) {
|
||||
let context = request_id.map_or_else(
|
||||
|| WriteContext::connection(FramePhase::Close),
|
||||
|request_id| WriteContext::request(request_id, FramePhase::Close),
|
||||
);
|
||||
self.close(
|
||||
CloseFrame {
|
||||
code: close_code::POLICY,
|
||||
reason: reason.into(),
|
||||
},
|
||||
context,
|
||||
PumpExitReason::ServerClose,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
pub(super) async fn close_for_connection_limit(&mut self, reason: &'static str) {
|
||||
self.close(
|
||||
CloseFrame {
|
||||
code: close_code::NORMAL,
|
||||
reason: reason.into(),
|
||||
},
|
||||
WriteContext::connection(FramePhase::Close),
|
||||
PumpExitReason::ConnectionMaxAge,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
async fn close(
|
||||
&mut self,
|
||||
frame: CloseFrame,
|
||||
context: WriteContext,
|
||||
exit_reason: PumpExitReason,
|
||||
) {
|
||||
if self
|
||||
.send(Message::Close(Some(frame)), context)
|
||||
.await
|
||||
.is_ok()
|
||||
{
|
||||
self.exit_reason.get_or_insert(exit_reason);
|
||||
}
|
||||
}
|
||||
|
||||
async fn send(
|
||||
&mut self,
|
||||
message: Message,
|
||||
context: WriteContext,
|
||||
) -> Result<(), ConnectionWriteError> {
|
||||
let Some(commands) = self.commands.as_ref().cloned() else {
|
||||
return Err(ConnectionWriteError::Closed);
|
||||
};
|
||||
self.stats.observe_command_queue(&commands);
|
||||
let (acknowledged, acknowledgement) = oneshot::channel();
|
||||
let started_at = Instant::now();
|
||||
let request_id = context.request_id.clone();
|
||||
let phase = context.phase;
|
||||
let command = ConnectionCommand {
|
||||
message,
|
||||
context,
|
||||
acknowledged,
|
||||
};
|
||||
let write = async move {
|
||||
commands
|
||||
.send(command)
|
||||
.await
|
||||
.map_err(|_| ConnectionWriteError::Closed)?;
|
||||
acknowledgement
|
||||
.await
|
||||
.map_err(|_| ConnectionWriteError::Closed)?
|
||||
};
|
||||
let result = match timeout(self.config.write_timeout, write).await {
|
||||
Ok(result) => result,
|
||||
Err(_) => Err(ConnectionWriteError::Timeout {
|
||||
timeout: self.config.write_timeout,
|
||||
}),
|
||||
};
|
||||
let duration = started_at.elapsed();
|
||||
match &result {
|
||||
Ok(()) if phase.is_milestone() => tracing::info!(
|
||||
websocket_connection_id = %self.connection_id,
|
||||
request_id = request_id.as_deref().unwrap_or(""),
|
||||
frame_phase = phase.as_str(),
|
||||
write_duration_ms = duration.as_millis(),
|
||||
"Responses WebSocket frame write succeeded"
|
||||
),
|
||||
Ok(()) => tracing::debug!(
|
||||
websocket_connection_id = %self.connection_id,
|
||||
request_id = request_id.as_deref().unwrap_or(""),
|
||||
frame_phase = phase.as_str(),
|
||||
write_duration_ms = duration.as_millis(),
|
||||
"Responses WebSocket frame write succeeded"
|
||||
),
|
||||
Err(error) => tracing::info!(
|
||||
websocket_connection_id = %self.connection_id,
|
||||
request_id = request_id.as_deref().unwrap_or(""),
|
||||
frame_phase = phase.as_str(),
|
||||
write_duration_ms = duration.as_millis(),
|
||||
error = %error,
|
||||
"Responses WebSocket frame write failed"
|
||||
),
|
||||
}
|
||||
if let Err(error) = &result {
|
||||
let reason = match error {
|
||||
ConnectionWriteError::Timeout { .. } => PumpExitReason::WriteTimeout,
|
||||
ConnectionWriteError::Transport { .. } => PumpExitReason::WriteError,
|
||||
ConnectionWriteError::Closed => PumpExitReason::PumpStopped,
|
||||
};
|
||||
self.terminate(reason);
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
fn terminate(&mut self, reason: PumpExitReason) {
|
||||
self.exit_reason.get_or_insert(reason);
|
||||
self.commands.take();
|
||||
if let Some(task) = self.pump_task.take() {
|
||||
task.abort();
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn log_summary(&self, request_count: u64) {
|
||||
tracing::info!(
|
||||
websocket_connection_id = %self.connection_id,
|
||||
request_count,
|
||||
connection_age_ms = self.age().as_millis(),
|
||||
pump_exit_reason = self
|
||||
.exit_reason
|
||||
.map_or("unknown", PumpExitReason::as_str),
|
||||
command_queue_high_water = self
|
||||
.stats
|
||||
.command_queue_high_water
|
||||
.load(Ordering::Relaxed),
|
||||
command_backpressure_count = self
|
||||
.stats
|
||||
.command_backpressure_count
|
||||
.load(Ordering::Relaxed),
|
||||
ping_received_count = self.stats.ping_received_count.load(Ordering::Relaxed),
|
||||
pong_written_count = self.stats.pong_written_count.load(Ordering::Relaxed),
|
||||
"Responses WebSocket disconnected"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for ResponsesWebSocketConnection {
|
||||
fn drop(&mut self) {
|
||||
self.commands.take();
|
||||
if let Some(task) = self.pump_task.take() {
|
||||
task.abort();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn spawn_connection<S, E>(
|
||||
socket: S,
|
||||
connection_id: Arc<str>,
|
||||
cancellation: CancellationToken,
|
||||
config: ConnectionConfig,
|
||||
) -> ResponsesWebSocketConnection
|
||||
where
|
||||
S: Stream<Item = Result<Message, E>> + Sink<Message, Error = E> + Unpin + Send + 'static,
|
||||
E: fmt::Display + Send + 'static,
|
||||
{
|
||||
let opened_at = Instant::now();
|
||||
let expired = Arc::new(AtomicBool::new(false));
|
||||
let stats = Arc::new(ConnectionStats::default());
|
||||
let (command_tx, command_rx) = mpsc::channel(OUTBOUND_COMMAND_BUFFER);
|
||||
let (incoming_tx, incoming_rx) = mpsc::channel(INBOUND_EVENT_BUFFER);
|
||||
let pump_task = tokio::spawn(run_pump(
|
||||
socket,
|
||||
command_rx,
|
||||
incoming_tx,
|
||||
Arc::clone(&connection_id),
|
||||
cancellation,
|
||||
opened_at,
|
||||
Arc::clone(&expired),
|
||||
Arc::clone(&stats),
|
||||
config,
|
||||
));
|
||||
ResponsesWebSocketConnection {
|
||||
connection_id,
|
||||
opened_at,
|
||||
expired,
|
||||
commands: Some(command_tx),
|
||||
incoming: incoming_rx,
|
||||
pump_task: Some(pump_task),
|
||||
stats,
|
||||
config,
|
||||
exit_reason: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[expect(clippy::too_many_arguments)]
|
||||
async fn run_pump<S, E>(
|
||||
mut socket: S,
|
||||
mut commands: mpsc::Receiver<ConnectionCommand>,
|
||||
incoming: mpsc::Sender<ConnectionEvent>,
|
||||
connection_id: Arc<str>,
|
||||
cancellation: CancellationToken,
|
||||
opened_at: Instant,
|
||||
expired: Arc<AtomicBool>,
|
||||
stats: Arc<ConnectionStats>,
|
||||
config: ConnectionConfig,
|
||||
) where
|
||||
S: Stream<Item = Result<Message, E>> + Sink<Message, Error = E> + Unpin,
|
||||
E: fmt::Display,
|
||||
{
|
||||
let deadline = tokio::time::sleep_until(opened_at + config.max_age);
|
||||
tokio::pin!(deadline);
|
||||
let mut deadline_elapsed = false;
|
||||
let reason = loop {
|
||||
tokio::select! {
|
||||
biased;
|
||||
() = cancellation.cancelled() => break PumpExitReason::LifecycleShutdown,
|
||||
() = &mut deadline, if !deadline_elapsed => {
|
||||
deadline_elapsed = true;
|
||||
expired.store(true, Ordering::Release);
|
||||
if let Err(reason) = emit_incoming(&incoming, ConnectionEvent::Expired) {
|
||||
break reason;
|
||||
}
|
||||
}
|
||||
command = commands.recv() => {
|
||||
let Some(command) = command else {
|
||||
break PumpExitReason::CoordinatorDropped;
|
||||
};
|
||||
let ConnectionCommand {
|
||||
message,
|
||||
context,
|
||||
acknowledged,
|
||||
} = command;
|
||||
let closing = matches!(message, Message::Close(_));
|
||||
let result = socket.send(message).await.map_err(|error| {
|
||||
let message = error.to_string();
|
||||
tracing::debug!(
|
||||
websocket_connection_id = %connection_id,
|
||||
request_id = context.request_id.as_deref().unwrap_or(""),
|
||||
frame_phase = context.phase.as_str(),
|
||||
error = %message,
|
||||
"Responses WebSocket pump transport write failed"
|
||||
);
|
||||
ConnectionWriteError::Transport { message }
|
||||
});
|
||||
let failed = result.is_err();
|
||||
let _ = acknowledged.send(result);
|
||||
if failed {
|
||||
break PumpExitReason::WriteError;
|
||||
}
|
||||
if closing {
|
||||
break PumpExitReason::ServerClose;
|
||||
}
|
||||
}
|
||||
message = socket.next() => match message {
|
||||
Some(Ok(Message::Text(payload))) => {
|
||||
if let Err(reason) = emit_incoming(&incoming, ConnectionEvent::Text(payload.to_string())) {
|
||||
break reason;
|
||||
}
|
||||
}
|
||||
Some(Ok(Message::Binary(_))) => {
|
||||
if let Err(reason) = emit_incoming(&incoming, ConnectionEvent::Binary) {
|
||||
break reason;
|
||||
}
|
||||
}
|
||||
Some(Ok(Message::Ping(payload))) => {
|
||||
stats.ping_received_count.fetch_add(1, Ordering::Relaxed);
|
||||
match timeout(config.write_timeout, socket.send(Message::Pong(payload))).await {
|
||||
Ok(Ok(())) => {
|
||||
stats.pong_written_count.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
Ok(Err(error)) => {
|
||||
tracing::info!(
|
||||
websocket_connection_id = %connection_id,
|
||||
error = %error,
|
||||
"Responses WebSocket Pong write failed"
|
||||
);
|
||||
break PumpExitReason::WriteError;
|
||||
}
|
||||
Err(_) => break PumpExitReason::WriteTimeout,
|
||||
}
|
||||
}
|
||||
Some(Ok(Message::Pong(_))) => {}
|
||||
Some(Ok(Message::Close(_))) => break PumpExitReason::ClientClose,
|
||||
Some(Err(error)) => {
|
||||
tracing::info!(
|
||||
websocket_connection_id = %connection_id,
|
||||
error = %error,
|
||||
"Responses WebSocket receive failed"
|
||||
);
|
||||
break PumpExitReason::ReadError;
|
||||
}
|
||||
None => break PumpExitReason::PeerEof,
|
||||
},
|
||||
}
|
||||
};
|
||||
let _ = incoming.try_send(ConnectionEvent::Exited(reason));
|
||||
tracing::debug!(
|
||||
websocket_connection_id = %connection_id,
|
||||
pump_exit_reason = reason.as_str(),
|
||||
"Responses WebSocket pump exited"
|
||||
);
|
||||
}
|
||||
|
||||
fn emit_incoming(
|
||||
incoming: &mpsc::Sender<ConnectionEvent>,
|
||||
event: ConnectionEvent,
|
||||
) -> Result<(), PumpExitReason> {
|
||||
incoming.try_send(event).map_err(|error| match error {
|
||||
mpsc::error::TrySendError::Full(_) => PumpExitReason::InboundOverload,
|
||||
mpsc::error::TrySendError::Closed(_) => PumpExitReason::CoordinatorDropped,
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,482 @@
|
||||
//! 已启动 Responses execution 到客户端 WebSocket wire event 的串行转发。
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::http::StatusCode;
|
||||
use gateway_core::engine::execution::StartedExecution;
|
||||
use gateway_core::engine::{CommitRequirement, CoordinatedEvent, EngineError};
|
||||
use gateway_core::error::{GatewayError, GatewayErrorKind};
|
||||
use gateway_core::operation::ProviderSessionState;
|
||||
use tokio::time::Instant;
|
||||
|
||||
use crate::openai::error::{gateway_error_contract, gateway_error_from_engine};
|
||||
|
||||
use super::{
|
||||
super::{DecodedResponsesRequest, OpenAiResponsesEncoder, PendingExecution, ProtocolErrorBody},
|
||||
connection::{ConnectionEvent, FramePhase, ResponsesWebSocketConnection, WriteContext},
|
||||
protocol::{error_event, response_metadata_event},
|
||||
};
|
||||
|
||||
const ACTIVE_RESPONSE_VIOLATION: &str = "Only one response.create may be active per connection";
|
||||
|
||||
#[derive(Clone, Copy, PartialEq, Eq)]
|
||||
pub(super) enum ForwardOutcome {
|
||||
Continue,
|
||||
Disconnect,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub(super) struct ConnectionReplaySnapshot {
|
||||
last_response_id: Option<String>,
|
||||
provider_state: Option<ProviderSessionState>,
|
||||
}
|
||||
|
||||
impl ConnectionReplaySnapshot {
|
||||
pub(super) fn prepare(&self, request: DecodedResponsesRequest) -> DecodedResponsesRequest {
|
||||
match (
|
||||
request.metadata().continuation().previous_response_id(),
|
||||
self.last_response_id.as_deref(),
|
||||
self.provider_state.as_ref(),
|
||||
) {
|
||||
(Some(previous), Some(last), Some(state)) if previous == last => {
|
||||
request.with_provider_session_state(state.clone())
|
||||
}
|
||||
_ => request,
|
||||
}
|
||||
}
|
||||
|
||||
fn commit(&mut self, response_id: String, provider_state: Option<ProviderSessionState>) {
|
||||
self.last_response_id = Some(response_id);
|
||||
self.provider_state = provider_state;
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn forward_execution(
|
||||
connection: &mut ResponsesWebSocketConnection,
|
||||
started: StartedExecution,
|
||||
replay: &mut ConnectionReplaySnapshot,
|
||||
) -> ForwardOutcome {
|
||||
let request_id = Arc::<str>::from(started.request_id.to_string());
|
||||
let streaming = started.stream;
|
||||
let mut execution = PendingExecution::new(started.session);
|
||||
if !streaming {
|
||||
let error = GatewayError::new(
|
||||
GatewayErrorKind::Internal,
|
||||
"WebSocket execution was not initialized as a stream",
|
||||
);
|
||||
return send_gateway_error(connection, &error, &request_id).await;
|
||||
}
|
||||
let first = loop {
|
||||
match next_active_input(connection, &mut execution).await {
|
||||
ActiveInput::Event(Ok(Some(event))) => break event,
|
||||
ActiveInput::Event(Ok(None)) => {
|
||||
let error = GatewayError::new(
|
||||
GatewayErrorKind::Internal,
|
||||
"gateway response ended before its first event",
|
||||
);
|
||||
return send_gateway_error(connection, &error, &request_id).await;
|
||||
}
|
||||
ActiveInput::Event(Err(error)) => {
|
||||
let error = gateway_error_from_engine(&error);
|
||||
return send_gateway_error(connection, &error, &request_id).await;
|
||||
}
|
||||
ActiveInput::Control => continue,
|
||||
ActiveInput::Disconnect => return ForwardOutcome::Disconnect,
|
||||
ActiveInput::ProtocolViolation => {
|
||||
connection
|
||||
.close_policy(ACTIVE_RESPONSE_VIOLATION, Some(&request_id))
|
||||
.await;
|
||||
return ForwardOutcome::Disconnect;
|
||||
}
|
||||
}
|
||||
};
|
||||
let requirement = first.commit_requirement();
|
||||
let mut first = first.into_provider_events();
|
||||
if requirement != CommitRequirement::CommitBeforeDelivery {
|
||||
let error = GatewayError::new(
|
||||
GatewayErrorKind::Internal,
|
||||
"gateway first event did not require commit",
|
||||
);
|
||||
return send_gateway_error(connection, &error, &request_id).await;
|
||||
}
|
||||
let mut encoder = OpenAiResponsesEncoder::new();
|
||||
let mut provider_state = None;
|
||||
let mut first_messages = Vec::new();
|
||||
for event in &mut first {
|
||||
if let Some(update) = event.take_session_update() {
|
||||
provider_state = Some(update);
|
||||
}
|
||||
first_messages.extend(encoder.push_websocket(event));
|
||||
}
|
||||
if first_messages.is_empty() {
|
||||
let error = GatewayError::new(
|
||||
GatewayErrorKind::Internal,
|
||||
"gateway commit batch encoded no output",
|
||||
);
|
||||
return send_gateway_error(connection, &error, &request_id).await;
|
||||
}
|
||||
let Some(response_session) = execution.session_mut() else {
|
||||
return ForwardOutcome::Disconnect;
|
||||
};
|
||||
let response_headers = response_session.response_headers().to_vec();
|
||||
if let Err(error) = response_session.commit_downstream(None).await {
|
||||
let error = gateway_error_from_engine(&error);
|
||||
return send_gateway_error(connection, &error, &request_id).await;
|
||||
}
|
||||
let mut first_frame_written = false;
|
||||
if encoder.is_completed() {
|
||||
if let Err(outcome) =
|
||||
confirm_completed_execution(connection, &mut execution, &request_id).await
|
||||
{
|
||||
return outcome;
|
||||
}
|
||||
let provider_terminal_at = Instant::now();
|
||||
commit_connection_replay(replay, &encoder, provider_state.take());
|
||||
let writes_succeeded = send_metadata(
|
||||
connection,
|
||||
&request_id,
|
||||
response_metadata_event(&request_id, &response_headers),
|
||||
)
|
||||
.await
|
||||
&& send_messages(
|
||||
connection,
|
||||
first_messages,
|
||||
&request_id,
|
||||
&mut first_frame_written,
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
if !writes_succeeded {
|
||||
log_terminal_write_failure(connection, &request_id, provider_terminal_at.elapsed());
|
||||
return ForwardOutcome::Disconnect;
|
||||
}
|
||||
log_terminal_write_success(connection, &request_id, provider_terminal_at.elapsed());
|
||||
execution.disarm();
|
||||
return ForwardOutcome::Continue;
|
||||
}
|
||||
if !send_metadata(
|
||||
connection,
|
||||
&request_id,
|
||||
response_metadata_event(&request_id, &response_headers),
|
||||
)
|
||||
.await
|
||||
|| !send_messages(
|
||||
connection,
|
||||
first_messages,
|
||||
&request_id,
|
||||
&mut first_frame_written,
|
||||
encoder.has_wire_failure(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
return ForwardOutcome::Disconnect;
|
||||
}
|
||||
|
||||
loop {
|
||||
match next_active_input(connection, &mut execution).await {
|
||||
ActiveInput::Event(Ok(Some(delivery))) => {
|
||||
let requirement = delivery.commit_requirement();
|
||||
let mut events = delivery.into_provider_events();
|
||||
if requirement != CommitRequirement::AlreadyCommitted {
|
||||
let error = GatewayError::new(
|
||||
GatewayErrorKind::Internal,
|
||||
"gateway requested another downstream commit",
|
||||
);
|
||||
return send_gateway_error(connection, &error, &request_id).await;
|
||||
}
|
||||
let mut messages = Vec::new();
|
||||
for event in &mut events {
|
||||
if let Some(update) = event.take_session_update() {
|
||||
provider_state = Some(update);
|
||||
}
|
||||
messages.extend(encoder.push_websocket(event));
|
||||
}
|
||||
if encoder.is_completed() {
|
||||
if let Err(outcome) =
|
||||
confirm_completed_execution(connection, &mut execution, &request_id).await
|
||||
{
|
||||
return outcome;
|
||||
}
|
||||
let provider_terminal_at = Instant::now();
|
||||
commit_connection_replay(replay, &encoder, provider_state.take());
|
||||
if !send_messages(
|
||||
connection,
|
||||
messages,
|
||||
&request_id,
|
||||
&mut first_frame_written,
|
||||
true,
|
||||
)
|
||||
.await
|
||||
{
|
||||
log_terminal_write_failure(
|
||||
connection,
|
||||
&request_id,
|
||||
provider_terminal_at.elapsed(),
|
||||
);
|
||||
return ForwardOutcome::Disconnect;
|
||||
}
|
||||
log_terminal_write_success(
|
||||
connection,
|
||||
&request_id,
|
||||
provider_terminal_at.elapsed(),
|
||||
);
|
||||
execution.disarm();
|
||||
return ForwardOutcome::Continue;
|
||||
}
|
||||
let terminal_batch = encoder.has_wire_failure();
|
||||
if !send_messages(
|
||||
connection,
|
||||
messages,
|
||||
&request_id,
|
||||
&mut first_frame_written,
|
||||
terminal_batch,
|
||||
)
|
||||
.await
|
||||
{
|
||||
return ForwardOutcome::Disconnect;
|
||||
}
|
||||
}
|
||||
ActiveInput::Event(Ok(None)) => {
|
||||
if execution
|
||||
.session_mut()
|
||||
.is_some_and(|session| session.is_finalized())
|
||||
{
|
||||
commit_connection_replay(replay, &encoder, provider_state.take());
|
||||
execution.disarm();
|
||||
return ForwardOutcome::Continue;
|
||||
}
|
||||
let error = GatewayError::new(
|
||||
GatewayErrorKind::Internal,
|
||||
"gateway response ended without finalizing the execution",
|
||||
);
|
||||
return send_gateway_error(connection, &error, &request_id).await;
|
||||
}
|
||||
ActiveInput::Event(Err(error)) => {
|
||||
if encoder.has_wire_failure() {
|
||||
execution.disarm();
|
||||
return ForwardOutcome::Continue;
|
||||
}
|
||||
let error = gateway_error_from_engine(&error);
|
||||
return send_gateway_error(connection, &error, &request_id).await;
|
||||
}
|
||||
ActiveInput::Control => {}
|
||||
ActiveInput::Disconnect => return ForwardOutcome::Disconnect,
|
||||
ActiveInput::ProtocolViolation => {
|
||||
connection
|
||||
.close_policy(ACTIVE_RESPONSE_VIOLATION, Some(&request_id))
|
||||
.await;
|
||||
return ForwardOutcome::Disconnect;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn commit_connection_replay(
|
||||
replay: &mut ConnectionReplaySnapshot,
|
||||
encoder: &OpenAiResponsesEncoder,
|
||||
provider_state: Option<ProviderSessionState>,
|
||||
) {
|
||||
if let Some(response_id) = encoder.response_id() {
|
||||
replay.commit(response_id.to_owned(), provider_state);
|
||||
}
|
||||
}
|
||||
|
||||
async fn confirm_completed_execution(
|
||||
connection: &mut ResponsesWebSocketConnection,
|
||||
execution: &mut PendingExecution,
|
||||
request_id: &Arc<str>,
|
||||
) -> Result<(), ForwardOutcome> {
|
||||
loop {
|
||||
match next_active_input(connection, execution).await {
|
||||
ActiveInput::Event(Ok(None))
|
||||
if execution
|
||||
.session_mut()
|
||||
.is_some_and(|session| session.is_finalized()) =>
|
||||
{
|
||||
return Ok(());
|
||||
}
|
||||
ActiveInput::Event(Ok(None)) => {
|
||||
let error = GatewayError::new(
|
||||
GatewayErrorKind::Internal,
|
||||
"gateway response was not finalized after its terminal event",
|
||||
);
|
||||
return Err(send_gateway_error(connection, &error, request_id).await);
|
||||
}
|
||||
ActiveInput::Event(Ok(Some(_))) => {
|
||||
let error = GatewayError::new(
|
||||
GatewayErrorKind::Internal,
|
||||
"gateway response continued after its terminal event",
|
||||
);
|
||||
return Err(send_gateway_error(connection, &error, request_id).await);
|
||||
}
|
||||
ActiveInput::Event(Err(error)) => {
|
||||
let error = gateway_error_from_engine(&error);
|
||||
return Err(send_gateway_error(connection, &error, request_id).await);
|
||||
}
|
||||
ActiveInput::Control => {}
|
||||
ActiveInput::Disconnect => return Err(ForwardOutcome::Disconnect),
|
||||
ActiveInput::ProtocolViolation => {
|
||||
connection
|
||||
.close_policy(ACTIVE_RESPONSE_VIOLATION, Some(request_id))
|
||||
.await;
|
||||
return Err(ForwardOutcome::Disconnect);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
enum ActiveInput {
|
||||
Event(Result<Option<CoordinatedEvent>, EngineError>),
|
||||
Control,
|
||||
Disconnect,
|
||||
ProtocolViolation,
|
||||
}
|
||||
|
||||
async fn next_active_input(
|
||||
connection: &mut ResponsesWebSocketConnection,
|
||||
execution: &mut PendingExecution,
|
||||
) -> ActiveInput {
|
||||
let Some(session) = execution.session_mut() else {
|
||||
return ActiveInput::Disconnect;
|
||||
};
|
||||
tokio::select! {
|
||||
event = session.next_event() => ActiveInput::Event(event),
|
||||
message = connection.next_event() => match message {
|
||||
Some(ConnectionEvent::Expired) => ActiveInput::Control,
|
||||
Some(ConnectionEvent::Text(_) | ConnectionEvent::Binary) => {
|
||||
ActiveInput::ProtocolViolation
|
||||
}
|
||||
Some(ConnectionEvent::Exited(_)) | None => ActiveInput::Disconnect,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn send_gateway_error(
|
||||
connection: &mut ResponsesWebSocketConnection,
|
||||
error: &GatewayError,
|
||||
request_id: &Arc<str>,
|
||||
) -> ForwardOutcome {
|
||||
let (status, default_type, default_code) = gateway_error_contract(error.kind());
|
||||
send_error(
|
||||
connection,
|
||||
status,
|
||||
error.client_error_type().unwrap_or(default_type),
|
||||
error.client_error_code().unwrap_or(default_code),
|
||||
error.client_message(),
|
||||
None,
|
||||
request_id,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(super) async fn send_protocol_error(
|
||||
connection: &mut ResponsesWebSocketConnection,
|
||||
status: StatusCode,
|
||||
body: ProtocolErrorBody,
|
||||
request_id: &Arc<str>,
|
||||
) -> ForwardOutcome {
|
||||
let error = body.error;
|
||||
send_error(
|
||||
connection,
|
||||
status,
|
||||
error.kind,
|
||||
error.code,
|
||||
&error.message,
|
||||
error.param.as_deref(),
|
||||
request_id,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn send_error(
|
||||
connection: &mut ResponsesWebSocketConnection,
|
||||
status: StatusCode,
|
||||
error_type: &str,
|
||||
code: &str,
|
||||
message: &str,
|
||||
param: Option<&str>,
|
||||
request_id: &Arc<str>,
|
||||
) -> ForwardOutcome {
|
||||
if connection
|
||||
.send_text(
|
||||
error_event(status, error_type, code, message, param, Some(request_id)),
|
||||
WriteContext::request(request_id, FramePhase::Error),
|
||||
)
|
||||
.await
|
||||
.is_ok()
|
||||
{
|
||||
ForwardOutcome::Continue
|
||||
} else {
|
||||
ForwardOutcome::Disconnect
|
||||
}
|
||||
}
|
||||
|
||||
async fn send_metadata(
|
||||
connection: &mut ResponsesWebSocketConnection,
|
||||
request_id: &Arc<str>,
|
||||
message: String,
|
||||
) -> bool {
|
||||
connection
|
||||
.send_text(
|
||||
message,
|
||||
WriteContext::request(request_id, FramePhase::Metadata),
|
||||
)
|
||||
.await
|
||||
.is_ok()
|
||||
}
|
||||
|
||||
async fn send_messages(
|
||||
connection: &mut ResponsesWebSocketConnection,
|
||||
messages: Vec<String>,
|
||||
request_id: &Arc<str>,
|
||||
first_frame_written: &mut bool,
|
||||
terminal_batch: bool,
|
||||
) -> bool {
|
||||
let message_count = messages.len();
|
||||
for (index, message) in messages.into_iter().enumerate() {
|
||||
let is_terminal = terminal_batch && index + 1 == message_count;
|
||||
let phase = match (*first_frame_written, is_terminal) {
|
||||
(false, true) => FramePhase::FirstAndTerminal,
|
||||
(false, false) => FramePhase::First,
|
||||
(true, true) => FramePhase::Terminal,
|
||||
(true, false) => FramePhase::Data,
|
||||
};
|
||||
if connection
|
||||
.send_text(message, WriteContext::request(request_id, phase))
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
return false;
|
||||
}
|
||||
*first_frame_written = true;
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
fn log_terminal_write_success(
|
||||
connection: &ResponsesWebSocketConnection,
|
||||
request_id: &Arc<str>,
|
||||
provider_terminal_to_write: std::time::Duration,
|
||||
) {
|
||||
tracing::info!(
|
||||
websocket_connection_id = connection.id(),
|
||||
request_id = %request_id,
|
||||
provider_terminal_to_terminal_write_ms = provider_terminal_to_write.as_millis(),
|
||||
terminal_frame_written = true,
|
||||
"Provider execution terminal was written to the WebSocket transport"
|
||||
);
|
||||
}
|
||||
|
||||
fn log_terminal_write_failure(
|
||||
connection: &ResponsesWebSocketConnection,
|
||||
request_id: &Arc<str>,
|
||||
provider_terminal_to_write: std::time::Duration,
|
||||
) {
|
||||
tracing::info!(
|
||||
websocket_connection_id = connection.id(),
|
||||
request_id = %request_id,
|
||||
provider_terminal_to_terminal_write_ms = provider_terminal_to_write.as_millis(),
|
||||
provider_succeeded_but_terminal_write_failed = true,
|
||||
"Provider execution succeeded before the terminal WebSocket frame write failed"
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,248 @@
|
||||
//! OpenAI Responses WebSocket 的升级、连接与串行 session 编排。
|
||||
|
||||
mod connection;
|
||||
mod forward;
|
||||
mod protocol;
|
||||
|
||||
use std::{
|
||||
net::{IpAddr, SocketAddr},
|
||||
sync::Arc,
|
||||
};
|
||||
|
||||
use axum::{
|
||||
extract::{
|
||||
Extension, State,
|
||||
connect_info::ConnectInfo,
|
||||
ws::{WebSocket, WebSocketUpgrade},
|
||||
},
|
||||
http::{HeaderMap, StatusCode},
|
||||
response::{IntoResponse, Response},
|
||||
};
|
||||
use gateway_core::{
|
||||
engine::execution::{AuthenticatedClient, ClientTransport},
|
||||
lifecycle::{ConnectionGuard, ConnectionLifecycle},
|
||||
};
|
||||
|
||||
use crate::{
|
||||
ApiState,
|
||||
openai::{
|
||||
auth::{authenticate_client, authentication_error_response},
|
||||
error::runtime_unavailable_response,
|
||||
service::OpenAiService,
|
||||
},
|
||||
};
|
||||
|
||||
use super::{http::request_client_context, request::OpenAiRequestHeaders};
|
||||
use connection::{ConnectionEvent, FramePhase, ResponsesWebSocketConnection, WriteContext};
|
||||
use forward::{
|
||||
ConnectionReplaySnapshot, ForwardOutcome, forward_execution, send_gateway_error,
|
||||
send_protocol_error,
|
||||
};
|
||||
use protocol::connection_limit_event;
|
||||
pub use protocol::{ResponseCreateFrameError, decode_response_create_with_context};
|
||||
|
||||
const TEXT_FRAMES_ONLY: &str = "Responses WebSocket accepts text frames only";
|
||||
const CONNECTION_LIMIT_CLOSE_REASON: &str = "Responses websocket connection limit reached";
|
||||
|
||||
/// 将已认证的 `GET /v1/responses` 升级为 Responses WebSocket。
|
||||
pub(crate) async fn responses_websocket(
|
||||
State(state): State<ApiState>,
|
||||
connect_info: Option<Extension<ConnectInfo<SocketAddr>>>,
|
||||
headers: HeaderMap,
|
||||
websocket: WebSocketUpgrade,
|
||||
) -> Response {
|
||||
let service = state.openai().clone();
|
||||
let client = match authenticate_client(&service, &headers) {
|
||||
Ok(client) => client,
|
||||
Err(error) => return authentication_error_response(error),
|
||||
};
|
||||
let (client_ip, user_agent) = request_client_context(
|
||||
&headers,
|
||||
connect_info.map(|Extension(ConnectInfo(address))| address),
|
||||
);
|
||||
ResponsesWebSocketAdapter::new(service).upgrade_with_client_context(
|
||||
websocket,
|
||||
client,
|
||||
client_ip,
|
||||
user_agent,
|
||||
OpenAiRequestHeaders::from_headers(&headers),
|
||||
)
|
||||
}
|
||||
|
||||
/// 已鉴权 Responses WebSocket 升级边界。
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct ResponsesWebSocketAdapter {
|
||||
service: OpenAiService,
|
||||
}
|
||||
|
||||
impl ResponsesWebSocketAdapter {
|
||||
/// 绑定应用提供的唯一 OpenAI 客户端服务端口。
|
||||
#[must_use]
|
||||
pub const fn new(service: OpenAiService) -> Self {
|
||||
Self { service }
|
||||
}
|
||||
|
||||
fn upgrade_with_client_context(
|
||||
&self,
|
||||
websocket: WebSocketUpgrade,
|
||||
client: AuthenticatedClient,
|
||||
client_ip: Option<IpAddr>,
|
||||
user_agent: Option<String>,
|
||||
request_headers: OpenAiRequestHeaders,
|
||||
) -> Response {
|
||||
let connection_guard = match self.service.try_register_connection() {
|
||||
Ok(guard) => guard,
|
||||
Err(_) => return runtime_unavailable_response().into_response(),
|
||||
};
|
||||
let session = ResponsesWebSocketSession {
|
||||
service: self.service.clone(),
|
||||
client,
|
||||
connection_id: self.service.next_request_id().replacen("req_", "ws_", 1),
|
||||
client_ip,
|
||||
user_agent,
|
||||
request_headers,
|
||||
lifecycle: self.service.lifecycle(),
|
||||
connection_guard,
|
||||
};
|
||||
websocket
|
||||
// 覆盖 axum/tungstenite 的私有 64 MiB message 与 16 MiB frame 默认值。
|
||||
// Responses JSON 的协议可接受性由上游决定,代理不另设 wire 长度上限。
|
||||
.max_message_size(usize::MAX)
|
||||
.max_frame_size(usize::MAX)
|
||||
.on_upgrade(move |socket| async move {
|
||||
serve_responses_websocket(socket, session).await;
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
struct ResponsesWebSocketSession {
|
||||
service: OpenAiService,
|
||||
client: AuthenticatedClient,
|
||||
connection_id: String,
|
||||
client_ip: Option<IpAddr>,
|
||||
user_agent: Option<String>,
|
||||
request_headers: OpenAiRequestHeaders,
|
||||
lifecycle: Arc<dyn ConnectionLifecycle>,
|
||||
connection_guard: Box<dyn ConnectionGuard>,
|
||||
}
|
||||
|
||||
async fn serve_responses_websocket(socket: WebSocket, session: ResponsesWebSocketSession) {
|
||||
let ResponsesWebSocketSession {
|
||||
service,
|
||||
client,
|
||||
connection_id,
|
||||
client_ip,
|
||||
user_agent,
|
||||
request_headers,
|
||||
lifecycle,
|
||||
connection_guard,
|
||||
} = session;
|
||||
tracing::info!(
|
||||
websocket_connection_id = %connection_id,
|
||||
"Responses WebSocket connected"
|
||||
);
|
||||
let _connection_guard = connection_guard;
|
||||
let cancellation = lifecycle.cancellation();
|
||||
let mut connection = ResponsesWebSocketConnection::new(socket, connection_id, cancellation);
|
||||
let mut request_count = 0_u64;
|
||||
let mut replay = ConnectionReplaySnapshot::default();
|
||||
|
||||
loop {
|
||||
let Some(event) = connection.next_event().await else {
|
||||
break;
|
||||
};
|
||||
let payload = match event {
|
||||
ConnectionEvent::Text(payload) => payload,
|
||||
ConnectionEvent::Binary => {
|
||||
connection.close_policy(TEXT_FRAMES_ONLY, None).await;
|
||||
break;
|
||||
}
|
||||
ConnectionEvent::Expired => {
|
||||
expire_connection(&mut connection).await;
|
||||
break;
|
||||
}
|
||||
ConnectionEvent::Exited(_) => break,
|
||||
};
|
||||
if connection.is_expired() {
|
||||
expire_connection(&mut connection).await;
|
||||
break;
|
||||
}
|
||||
request_count = request_count.saturating_add(1);
|
||||
let correlation_id = Arc::<str>::from(service.next_request_id());
|
||||
let decoded = match decode_response_create_with_context(&payload, &request_headers) {
|
||||
Ok(decoded) => decoded.with_client_context(client_ip, user_agent.clone()),
|
||||
Err(error) => {
|
||||
tracing::info!(
|
||||
websocket_connection_id = connection.id(),
|
||||
request_id = %correlation_id,
|
||||
error = %error,
|
||||
"Responses WebSocket request rejected"
|
||||
);
|
||||
if send_protocol_error(
|
||||
&mut connection,
|
||||
StatusCode::BAD_REQUEST,
|
||||
error.protocol_body(),
|
||||
&correlation_id,
|
||||
)
|
||||
.await
|
||||
== ForwardOutcome::Disconnect
|
||||
{
|
||||
break;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let decoded = replay.prepare(decoded);
|
||||
// deadline 与 Text 可能同时就绪;在任何上游执行开始前再次封住该竞争窗口。
|
||||
if connection.is_expired() {
|
||||
expire_connection(&mut connection).await;
|
||||
break;
|
||||
}
|
||||
let started = match service
|
||||
.start_response(
|
||||
client.clone(),
|
||||
decoded,
|
||||
ClientTransport::WebSocket,
|
||||
"/v1/responses",
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(started) => started,
|
||||
Err(error) => {
|
||||
if send_gateway_error(&mut connection, &error, &correlation_id).await
|
||||
== ForwardOutcome::Disconnect
|
||||
{
|
||||
break;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
if forward_execution(&mut connection, started, &mut replay).await
|
||||
== ForwardOutcome::Disconnect
|
||||
{
|
||||
break;
|
||||
}
|
||||
if connection.is_expired() {
|
||||
expire_connection(&mut connection).await;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
connection.log_summary(request_count);
|
||||
}
|
||||
|
||||
async fn expire_connection(connection: &mut ResponsesWebSocketConnection) {
|
||||
if connection
|
||||
.send_text(
|
||||
connection_limit_event(),
|
||||
WriteContext::connection(FramePhase::ConnectionLimit),
|
||||
)
|
||||
.await
|
||||
.is_ok()
|
||||
{
|
||||
connection
|
||||
.close_for_connection_limit(CONNECTION_LIMIT_CLOSE_REASON)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,151 @@
|
||||
//! Responses WebSocket 入站请求与下行事件的纯协议映射。
|
||||
|
||||
use axum::http::{HeaderName, StatusCode};
|
||||
use gateway_core::event::ProviderResponseHeader;
|
||||
use serde_json::{Map, Value, json};
|
||||
use thiserror::Error;
|
||||
|
||||
use super::super::{
|
||||
DecodedResponsesRequest, ProtocolErrorBody, RequestDecodeError, request::OpenAiRequestHeaders,
|
||||
};
|
||||
|
||||
/// 使用连接级请求头和 Provider 上下文解码官方 `response.create` 文本帧。
|
||||
///
|
||||
/// 缺省 `stream` 等价于 WebSocket 固有的流式语义;显式 `false` 会被拒绝。
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// 帧不是合法 JSON object、消息类型错误、显式关闭 stream,或 Responses 请求
|
||||
/// 无法映射到 canonical operation 时返回不包含正文内容的稳定错误。
|
||||
pub fn decode_response_create_with_context(
|
||||
payload: &str,
|
||||
request_headers: &OpenAiRequestHeaders,
|
||||
) -> Result<DecodedResponsesRequest, ResponseCreateFrameError> {
|
||||
decode_response_create_inner(payload, request_headers)
|
||||
}
|
||||
|
||||
fn decode_response_create_inner(
|
||||
payload: &str,
|
||||
request_headers: &OpenAiRequestHeaders,
|
||||
) -> Result<DecodedResponsesRequest, ResponseCreateFrameError> {
|
||||
let Value::Object(mut body) = serde_json::from_str::<Value>(payload)
|
||||
.map_err(|_| ResponseCreateFrameError::InvalidJson)?
|
||||
else {
|
||||
return Err(ResponseCreateFrameError::ExpectedObject);
|
||||
};
|
||||
match body.remove("type") {
|
||||
Some(Value::String(message_type)) if message_type == "response.create" => {}
|
||||
_ => return Err(ResponseCreateFrameError::UnsupportedType),
|
||||
}
|
||||
if matches!(body.get("stream"), Some(value) if value.as_bool() != Some(true)) {
|
||||
return Err(ResponseCreateFrameError::StreamingRequired);
|
||||
}
|
||||
super::super::request::decode_request_object(body, false, request_headers)
|
||||
.map_err(ResponseCreateFrameError::Request)
|
||||
}
|
||||
|
||||
/// `response.create` 帧的稳定安全错误。
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Error)]
|
||||
pub enum ResponseCreateFrameError {
|
||||
/// 文本不是合法 JSON。
|
||||
#[error("response.create frame must be valid JSON")]
|
||||
InvalidJson,
|
||||
/// 顶层不是 object。
|
||||
#[error("response.create frame must be a JSON object")]
|
||||
ExpectedObject,
|
||||
/// `type` 缺失或不是 `response.create`。
|
||||
#[error("unsupported Responses WebSocket message type")]
|
||||
UnsupportedType,
|
||||
/// WebSocket 请求显式声明 `stream=false`。
|
||||
#[error("Responses WebSocket requests require stream=true")]
|
||||
StreamingRequired,
|
||||
/// 内层 Responses 请求无法映射到 canonical operation。
|
||||
#[error(transparent)]
|
||||
Request(RequestDecodeError),
|
||||
}
|
||||
|
||||
impl ResponseCreateFrameError {
|
||||
pub(super) fn protocol_body(&self) -> ProtocolErrorBody {
|
||||
match self {
|
||||
Self::Request(error) => error.protocol_body(),
|
||||
Self::InvalidJson => RequestDecodeError::MalformedJson.protocol_body(),
|
||||
Self::ExpectedObject => RequestDecodeError::ExpectedObject.protocol_body(),
|
||||
Self::UnsupportedType => RequestDecodeError::InvalidValue {
|
||||
field: "type".to_owned(),
|
||||
}
|
||||
.protocol_body(),
|
||||
Self::StreamingRequired => RequestDecodeError::InvalidValue {
|
||||
field: "stream".to_owned(),
|
||||
}
|
||||
.protocol_body(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn response_metadata_event(
|
||||
request_id: &str,
|
||||
response_headers: &[ProviderResponseHeader],
|
||||
) -> String {
|
||||
let connection_options = super::super::response_connection_options(response_headers);
|
||||
let mut headers = Map::new();
|
||||
for header in response_headers {
|
||||
if !super::super::response_header_is_forwardable(header.name(), &connection_options) {
|
||||
continue;
|
||||
}
|
||||
let Ok(name) = HeaderName::from_bytes(header.name().as_bytes()) else {
|
||||
continue;
|
||||
};
|
||||
let Ok(value) = std::str::from_utf8(header.value()) else {
|
||||
continue;
|
||||
};
|
||||
// 官方 response.metadata 使用 string map,无法表达同名多值;后值保持既有覆盖语义。
|
||||
headers.insert(name.as_str().to_owned(), Value::String(value.to_owned()));
|
||||
}
|
||||
headers
|
||||
.entry("x-request-id".to_owned())
|
||||
.or_insert_with(|| Value::String(request_id.to_owned()));
|
||||
json!({
|
||||
"type": "response.metadata",
|
||||
"headers": headers,
|
||||
})
|
||||
.to_string()
|
||||
}
|
||||
|
||||
pub(super) fn error_event(
|
||||
status: StatusCode,
|
||||
error_type: &str,
|
||||
code: &str,
|
||||
message: &str,
|
||||
param: Option<&str>,
|
||||
request_id: Option<&str>,
|
||||
) -> String {
|
||||
let mut error = Map::new();
|
||||
error.insert("type".to_owned(), Value::String(error_type.to_owned()));
|
||||
error.insert("code".to_owned(), Value::String(code.to_owned()));
|
||||
error.insert("message".to_owned(), Value::String(message.to_owned()));
|
||||
if let Some(param) = param {
|
||||
error.insert("param".to_owned(), Value::String(param.to_owned()));
|
||||
}
|
||||
let mut event = Map::new();
|
||||
event.insert("type".to_owned(), Value::String("error".to_owned()));
|
||||
event.insert("status".to_owned(), Value::Number(status.as_u16().into()));
|
||||
event.insert("error".to_owned(), Value::Object(error));
|
||||
if let Some(request_id) = request_id {
|
||||
event.insert(
|
||||
"request_id".to_owned(),
|
||||
Value::String(request_id.to_owned()),
|
||||
);
|
||||
}
|
||||
Value::Object(event).to_string()
|
||||
}
|
||||
|
||||
pub(super) fn connection_limit_event() -> String {
|
||||
error_event(
|
||||
StatusCode::BAD_REQUEST,
|
||||
"invalid_request_error",
|
||||
"websocket_connection_limit_reached",
|
||||
"Responses websocket connection limit reached (60 minutes). Create a new websocket connection to continue.",
|
||||
None,
|
||||
None,
|
||||
)
|
||||
}
|
||||
@@ -40,7 +40,10 @@ async fn list_route_should_keep_camel_case_group_and_page_wire() {
|
||||
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_eq!(
|
||||
value["data"]["items"][0]["usage"]["retainedTotalUsd"],
|
||||
"5.5"
|
||||
);
|
||||
assert!(value["data"]["items"][0].get("member_count").is_none());
|
||||
}
|
||||
|
||||
|
||||
@@ -76,7 +76,7 @@ mod profile_statistics {
|
||||
ProviderProfileActivityInsights, ProviderProfileDailyUsage, ProviderProfileInvocation,
|
||||
ProviderProfileStatistics, ProviderProfileStatisticsSummary,
|
||||
};
|
||||
use gateway_api::admin::accounts::AccountProfileStatisticsData;
|
||||
use gateway_api::admin::accounts::{AccountProfileAvatarQuery, AccountProfileStatisticsData};
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
@@ -134,6 +134,74 @@ mod profile_statistics {
|
||||
assert!(value.get("models").is_none());
|
||||
assert!(value.get("estimatedCost").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn profile_avatar_query_accepts_only_bounded_cache_versions() {
|
||||
let valid: AccountProfileAvatarQuery = serde_json::from_value(json!({
|
||||
"accountId": "acct_openai",
|
||||
"version": "k9m2z1"
|
||||
}))
|
||||
.expect("decode avatar query");
|
||||
valid.validate().expect("validate avatar query");
|
||||
|
||||
for version in ["", "contains-dash", "x23456789012345678901234567890123"] {
|
||||
let query: AccountProfileAvatarQuery = serde_json::from_value(json!({
|
||||
"accountId": "acct_openai",
|
||||
"version": version
|
||||
}))
|
||||
.expect("decode invalid avatar query");
|
||||
assert_eq!(
|
||||
query.validate().expect_err("reject version").field(),
|
||||
"version"
|
||||
);
|
||||
}
|
||||
assert!(
|
||||
serde_json::from_value::<AccountProfileAvatarQuery>(json!({
|
||||
"accountId": "acct_openai",
|
||||
"sourceUrl": "https://example.test/avatar"
|
||||
}))
|
||||
.is_err()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
mod profile_avatar {
|
||||
use axum::body::to_bytes;
|
||||
use bytes::Bytes;
|
||||
use futures::stream;
|
||||
use gateway_admin::model::provider_credentials::ProviderProfileAvatar;
|
||||
use gateway_api::admin::accounts::profile_avatar_response;
|
||||
|
||||
#[tokio::test]
|
||||
async fn profile_avatar_response_preserves_stream_metadata_and_private_cache() {
|
||||
let response = profile_avatar_response(ProviderProfileAvatar {
|
||||
content_type: Some("image/svg+xml".to_owned()),
|
||||
content_length: Some(6),
|
||||
etag: Some("\"v1\"".to_owned()),
|
||||
body: Box::pin(stream::iter([Ok(Bytes::from_static(b"avatar"))])),
|
||||
});
|
||||
|
||||
assert_eq!(response.status(), 200);
|
||||
assert_eq!(response.headers()["content-type"], "image/svg+xml");
|
||||
assert_eq!(response.headers()["content-length"], "6");
|
||||
assert_eq!(response.headers()["etag"], "\"v1\"");
|
||||
assert_eq!(response.headers()["cache-control"], "private, max-age=3600");
|
||||
assert_eq!(response.headers()["x-content-type-options"], "nosniff");
|
||||
assert_eq!(
|
||||
response.headers()["cross-origin-resource-policy"],
|
||||
"same-origin"
|
||||
);
|
||||
assert_eq!(
|
||||
response.headers()["content-security-policy"],
|
||||
"sandbox; default-src 'none'"
|
||||
);
|
||||
assert_eq!(
|
||||
to_bytes(response.into_body(), 16)
|
||||
.await
|
||||
.expect("avatar response body"),
|
||||
Bytes::from_static(b"avatar")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
mod batch_update {
|
||||
|
||||
@@ -15,14 +15,14 @@ use gateway_admin::{
|
||||
MutationContext, Revision,
|
||||
account_groups::{
|
||||
AccountGroupAccountSummary, AccountGroupCapacity, AccountGroupColor,
|
||||
AccountGroupListQuery, AccountGroupMutation, AccountGroupPage, AccountGroupRecord,
|
||||
AccountGroupUsage, DeleteAccountGroup, NewAccountGroup, SetAccountGroupEnabled,
|
||||
UpdateAccountGroup,
|
||||
AccountGroupListQuery, AccountGroupMemberFact, AccountGroupMutation, AccountGroupPage,
|
||||
AccountGroupRecord, AccountGroupUsage, DeleteAccountGroup, NewAccountGroup,
|
||||
SetAccountGroupEnabled, UpdateAccountGroup,
|
||||
},
|
||||
accounts::{
|
||||
AccountListQuery, AccountPage, AccountPageItem, AccountUpdateResult, AccountUsage,
|
||||
AccountUsageWindowQuery, AccountUsageWindowResult, AccountsUpdateResult,
|
||||
BatchUpdateAccounts, DeleteAccounts, UpdateAccount,
|
||||
AccountListQuery, AccountPage, AccountPageItem, AccountRuntimeSnapshot,
|
||||
AccountUpdateResult, AccountUsage, AccountUsageWindowQuery, AccountUsageWindowResult,
|
||||
AccountsUpdateResult, BatchUpdateAccounts, DeleteAccounts, UpdateAccount,
|
||||
},
|
||||
auth::{AdminAuditEvent, AdminSession},
|
||||
client_keys::{
|
||||
@@ -32,7 +32,7 @@ use gateway_admin::{
|
||||
observability::{
|
||||
DashboardObservation, DecimalAmount, DiagnosticDimension, DiagnosticObservation,
|
||||
OpsError, OpsErrorPage, OpsErrorQuery, RequestMetricPoint, TimeRange, UsageDetail,
|
||||
UsageFilter, UsageOverview, UsagePage, UsageQuery, UsageRecord,
|
||||
UsageFilter, UsageListRecord, UsageOverview, UsagePage, UsageQuery,
|
||||
},
|
||||
provider_credentials::{
|
||||
AuthorizationCommit, AuthorizationStarted, CompleteAuthorization, CredentialDetails,
|
||||
@@ -52,8 +52,9 @@ use gateway_admin::{
|
||||
ports::{
|
||||
provider::{ProviderAdmin, ProviderAdminError, ProviderAdminErrorKind},
|
||||
store::{
|
||||
AccountGroupStore, AccountStore, AdminStoreError, AdminStoreErrorKind, AdminStorePorts,
|
||||
AdminStoreResult, AuthStore, ClientKeyStore, ObservabilityStore, SettingsStore,
|
||||
AccountGroupStore, AccountRuntimeStore, AccountStore, AdminAccountStorePorts,
|
||||
AdminStoreError, AdminStoreErrorKind, AdminStorePorts, AdminStoreResult, AuthStore,
|
||||
ClientKeyStore, ObservabilityStore, SettingsStore,
|
||||
},
|
||||
system::{
|
||||
SystemOperationError, SystemOperationErrorKind, SystemOperations,
|
||||
@@ -64,7 +65,7 @@ use gateway_admin::{
|
||||
use gateway_api::admin::AdminSessionState;
|
||||
use gateway_core::{
|
||||
engine::{
|
||||
credential::ProviderAccountId,
|
||||
credential::{AccountStatusFacts, CredentialState, ProviderAccountId, QuotaState},
|
||||
probe::{AccountProbe, AccountProbeError, AccountProbeRequest, AccountProbeResult},
|
||||
},
|
||||
policy::{ClientApiKeyId, RateLimits},
|
||||
@@ -87,7 +88,7 @@ pub(super) struct AdminTestFixture {
|
||||
pub services: AdminServices,
|
||||
pub auth: Arc<MemoryAuthStore>,
|
||||
pub settings: Arc<MemorySettingsStore>,
|
||||
pub usage_records: Arc<Mutex<Vec<UsageRecord>>>,
|
||||
pub usage_records: Arc<Mutex<Vec<UsageListRecord>>>,
|
||||
pub usage_detail: Arc<Mutex<Option<UsageDetail>>>,
|
||||
pub diagnostics: Arc<Mutex<Vec<DiagnosticObservation>>>,
|
||||
pub ops_errors: Arc<Mutex<Vec<OpsError>>>,
|
||||
@@ -121,8 +122,7 @@ impl AdminTestFixture {
|
||||
dashboard_summary_range: Arc::clone(&dashboard_summary_range),
|
||||
});
|
||||
let stores = AdminStorePorts::new(
|
||||
unused.clone(),
|
||||
account_groups.clone(),
|
||||
AdminAccountStorePorts::new(unused.clone(), unused.clone(), account_groups.clone()),
|
||||
auth.clone(),
|
||||
client_keys.clone(),
|
||||
unused,
|
||||
@@ -466,6 +466,45 @@ impl AccountGroupStore for MemoryAccountGroupStore {
|
||||
})
|
||||
}
|
||||
|
||||
async fn load_account_group_members(
|
||||
&self,
|
||||
group_ids: &[gateway_core::routing::AccountGroupId],
|
||||
) -> AdminStoreResult<Vec<AccountGroupMemberFact>> {
|
||||
if !group_ids.iter().any(|id| id.as_str() == PRIMARY_GROUP_ID) {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
Ok(vec![
|
||||
AccountGroupMemberFact {
|
||||
group_id: group_id(PRIMARY_GROUP_ID),
|
||||
account_id: "acct_group_ready".to_owned(),
|
||||
status: AccountStatusFacts {
|
||||
enabled: true,
|
||||
credential_state: CredentialState::Ready,
|
||||
access_token_expires_at: None,
|
||||
quota: QuotaState::default(),
|
||||
rate_limited_until: None,
|
||||
last_error_reason: None,
|
||||
last_error_message: None,
|
||||
},
|
||||
total_slots: 1,
|
||||
},
|
||||
AccountGroupMemberFact {
|
||||
group_id: group_id(PRIMARY_GROUP_ID),
|
||||
account_id: "acct_group_disabled".to_owned(),
|
||||
status: AccountStatusFacts {
|
||||
enabled: false,
|
||||
credential_state: CredentialState::Ready,
|
||||
access_token_expires_at: None,
|
||||
quota: QuotaState::default(),
|
||||
rate_limited_until: None,
|
||||
last_error_reason: None,
|
||||
last_error_message: None,
|
||||
},
|
||||
total_slots: 1,
|
||||
},
|
||||
])
|
||||
}
|
||||
|
||||
async fn create_account_group(
|
||||
&self,
|
||||
command: NewAccountGroup,
|
||||
@@ -622,7 +661,7 @@ impl ClientKeyStore for MemoryClientKeyStore {
|
||||
}
|
||||
|
||||
struct UnusedStore {
|
||||
usage_records: Arc<Mutex<Vec<UsageRecord>>>,
|
||||
usage_records: Arc<Mutex<Vec<UsageListRecord>>>,
|
||||
usage_detail: Arc<Mutex<Option<UsageDetail>>>,
|
||||
diagnostics: Arc<Mutex<Vec<DiagnosticObservation>>>,
|
||||
ops_errors: Arc<Mutex<Vec<OpsError>>>,
|
||||
@@ -632,11 +671,19 @@ struct UnusedStore {
|
||||
|
||||
#[async_trait]
|
||||
impl AccountStore for UnusedStore {
|
||||
async fn list_accounts(&self, _: AccountListQuery) -> AdminStoreResult<AccountPage> {
|
||||
async fn list_accounts(
|
||||
&self,
|
||||
_: AccountListQuery,
|
||||
_: AccountRuntimeSnapshot,
|
||||
) -> AdminStoreResult<AccountPage> {
|
||||
Err(unavailable("account list"))
|
||||
}
|
||||
|
||||
async fn load_account(&self, _: &str) -> AdminStoreResult<Option<AccountPageItem>> {
|
||||
async fn load_account(
|
||||
&self,
|
||||
_: &str,
|
||||
_: AccountRuntimeSnapshot,
|
||||
) -> AdminStoreResult<Option<AccountPageItem>> {
|
||||
Err(unavailable("account"))
|
||||
}
|
||||
|
||||
@@ -752,6 +799,23 @@ impl AccountStore for UnusedStore {
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AccountRuntimeStore for UnusedStore {
|
||||
async fn active_rate_limits(&self) -> AdminStoreResult<AccountRuntimeSnapshot> {
|
||||
Ok(AccountRuntimeSnapshot {
|
||||
rate_limited_until: BTreeMap::new(),
|
||||
in_flight: Some(BTreeMap::new()),
|
||||
})
|
||||
}
|
||||
|
||||
async fn account_runtime(&self, _: &[String]) -> AdminStoreResult<AccountRuntimeSnapshot> {
|
||||
Ok(AccountRuntimeSnapshot {
|
||||
rate_limited_until: BTreeMap::new(),
|
||||
in_flight: Some(BTreeMap::new()),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ObservabilityStore for UnusedStore {
|
||||
async fn dashboard_summary(
|
||||
@@ -791,12 +855,13 @@ impl ObservabilityStore for UnusedStore {
|
||||
Err(unavailable("usage billing facts"))
|
||||
}
|
||||
|
||||
async fn list_usage_records(&self, _: UsageQuery) -> AdminStoreResult<UsagePage> {
|
||||
async fn list_usage_records(&self, query: UsageQuery) -> AdminStoreResult<UsagePage> {
|
||||
let items = self.usage_records.lock().expect("usage records").clone();
|
||||
Ok(UsagePage {
|
||||
current_page: query.current_page,
|
||||
page_size: query.page_size.get(),
|
||||
total: items.len() as u64,
|
||||
items,
|
||||
next_cursor: None,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -821,12 +886,13 @@ impl ObservabilityStore for UnusedStore {
|
||||
Ok(self.diagnostics.lock().expect("diagnostics").clone())
|
||||
}
|
||||
|
||||
async fn list_ops_errors(&self, _: OpsErrorQuery) -> AdminStoreResult<OpsErrorPage> {
|
||||
async fn list_ops_errors(&self, query: OpsErrorQuery) -> AdminStoreResult<OpsErrorPage> {
|
||||
let items = self.ops_errors.lock().expect("ops errors").clone();
|
||||
Ok(OpsErrorPage {
|
||||
current_page: query.current_page,
|
||||
page_size: query.page_size.get(),
|
||||
total: items.len() as u64,
|
||||
items,
|
||||
next_cursor: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -1037,10 +1103,12 @@ fn capacity(used_slots: Option<u64>, total_slots: u64) -> AccountGroupCapacity {
|
||||
}
|
||||
}
|
||||
|
||||
fn usage(today_usd: &str, total_usd: &str) -> AccountGroupUsage {
|
||||
fn usage(today_usd: &str, retained_total_usd: &str) -> AccountGroupUsage {
|
||||
AccountGroupUsage {
|
||||
today_usd: today_usd.parse::<DecimalAmount>().expect("today cost"),
|
||||
total_usd: total_usd.parse::<DecimalAmount>().expect("total cost"),
|
||||
retained_total_usd: retained_total_usd
|
||||
.parse::<DecimalAmount>()
|
||||
.expect("retained total cost"),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -19,27 +19,34 @@ fn dashboard_query_should_reject_unknown_trend_kind() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usage_query_should_bound_page_size_and_cursor() {
|
||||
fn usage_query_should_accept_element_plus_pagination() {
|
||||
let query: UsageQuery = serde_json::from_value(json!({
|
||||
"page": 2,
|
||||
"pageSize": 100,
|
||||
"cursor": "opaque"
|
||||
"currentPage": 129,
|
||||
"pageSize": 10
|
||||
}))
|
||||
.unwrap();
|
||||
assert_eq!(query.validate_page().unwrap(), (2, 100));
|
||||
assert!(query.validate_cursor().is_ok());
|
||||
assert_eq!(query.validate_pagination().unwrap(), (129, 10));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usage_query_should_reject_zero_page() {
|
||||
let query: UsageQuery = serde_json::from_value(json!({"page": 0})).unwrap();
|
||||
assert_eq!(query.validate_page().unwrap_err().field(), "page");
|
||||
fn usage_query_should_reject_removed_total_toggle() {
|
||||
assert!(serde_json::from_value::<UsageQuery>(json!({"includeTotal": true})).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usage_query_should_reject_removed_cursor_contract() {
|
||||
assert!(serde_json::from_value::<UsageQuery>(json!({"cursor": "opaque"})).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usage_query_should_reject_nonstandard_page_name() {
|
||||
assert!(serde_json::from_value::<UsageQuery>(json!({"page": 1})).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ops_query_should_reject_page_size_above_terminal_limit() {
|
||||
let query: OpsQuery = serde_json::from_value(json!({"pageSize": 101})).unwrap();
|
||||
assert_eq!(query.validate_page().unwrap_err().field(), "pageSize");
|
||||
assert_eq!(query.validate_pagination().unwrap_err().field(), "pageSize");
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -2,12 +2,11 @@
|
||||
|
||||
use chrono::{TimeZone, Utc};
|
||||
use gateway_admin::model::observability::DesktopReleaseStatus;
|
||||
use gateway_api::admin::PageMeta;
|
||||
use gateway_api::admin::observability::{
|
||||
BillingView, CostCoverageView, CursorWire, DashboardAccountRequestBucketView,
|
||||
DashboardAccountUsageView, DashboardDesktopReleaseStatusView, DashboardWireAttributeView,
|
||||
DashboardWireProfileView, DashboardWireTargetView, PageData, TokenDetailsView, TrendData,
|
||||
TrendKind, TrendPointView, TrendSummaryView,
|
||||
BillingView, CostCoverageView, DashboardAccountRequestBucketView, DashboardAccountUsageView,
|
||||
DashboardDesktopReleaseStatusView, DashboardWireAttributeView, DashboardWireProfileView,
|
||||
DashboardWireTargetView, PageData, TokenDetailsView, TrendData, TrendKind, TrendPointView,
|
||||
TrendSummaryView,
|
||||
};
|
||||
use serde_json::json;
|
||||
|
||||
@@ -15,12 +14,15 @@ use serde_json::json;
|
||||
fn usage_page_should_keep_terminal_camel_case_shape() {
|
||||
let data = PageData {
|
||||
items: vec![json!({"id": "request_1"})],
|
||||
page: PageMeta::new(1, 50, 1, 1),
|
||||
next_cursor: Some("cursor".to_owned()),
|
||||
current_page: 129,
|
||||
page_size: 10,
|
||||
total: 1287,
|
||||
};
|
||||
let value = serde_json::to_value(data).unwrap();
|
||||
assert_eq!(value["page"]["pageSize"], 50);
|
||||
assert_eq!(value["nextCursor"], "cursor");
|
||||
assert_eq!(value["currentPage"], 129);
|
||||
assert_eq!(value["pageSize"], 10);
|
||||
assert_eq!(value["total"], 1287);
|
||||
assert!(value.get("nextCursor").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -155,11 +157,7 @@ fn sensitive_response_views_do_not_require_debug_or_add_secret_fields() {
|
||||
image_output_tokens_display: "-".to_owned(),
|
||||
total_tokens_display: "3".to_owned(),
|
||||
};
|
||||
let cursor = CursorWire {
|
||||
observed_at: Utc.timestamp_opt(0, 0).single().unwrap(),
|
||||
stable_id: "request_1".to_owned(),
|
||||
};
|
||||
let value = serde_json::to_value((&coverage, &token_details, &cursor)).unwrap();
|
||||
let value = serde_json::to_value((&coverage, &token_details)).unwrap();
|
||||
assert!(value.to_string().contains("known"));
|
||||
assert!(!value.to_string().contains("secret"));
|
||||
}
|
||||
@@ -209,7 +207,7 @@ async fn dashboard_summary_should_include_quota_exhaustion_in_unavailable_headli
|
||||
};
|
||||
use chrono::{Duration, Utc};
|
||||
use gateway_admin::model::observability::{
|
||||
AccountPoolMetrics, AttemptMetrics, DashboardObservation, RequestMetrics, TimeRange,
|
||||
AccountPoolMetrics, DashboardObservation, TimeRange,
|
||||
};
|
||||
use gateway_api::admin::observability;
|
||||
use tower::ServiceExt as _;
|
||||
@@ -224,8 +222,6 @@ async fn dashboard_summary_should_include_quota_exhaustion_in_unavailable_headli
|
||||
.lock()
|
||||
.expect("dashboard observation") = Some(DashboardObservation {
|
||||
range: TimeRange::new(now - Duration::hours(1), now).expect("dashboard range"),
|
||||
requests: RequestMetrics::default(),
|
||||
attempts: AttemptMetrics::default(),
|
||||
totals: Default::default(),
|
||||
provider_accounts: AccountPoolMetrics {
|
||||
total: 1002,
|
||||
@@ -293,8 +289,7 @@ async fn dashboard_summary_should_use_lifetime_totals_for_card_footers() {
|
||||
};
|
||||
use chrono::{Duration, Utc};
|
||||
use gateway_admin::model::observability::{
|
||||
AccountPoolMetrics, AttemptMetrics, DashboardObservation, DashboardTotals, RequestMetrics,
|
||||
TimeRange,
|
||||
AccountPoolMetrics, DashboardObservation, DashboardTotals, TimeRange,
|
||||
};
|
||||
use gateway_api::admin::observability;
|
||||
use tower::ServiceExt as _;
|
||||
@@ -304,20 +299,11 @@ async fn dashboard_summary_should_use_lifetime_totals_for_card_footers() {
|
||||
let fixture = AdminTestFixture::new().await;
|
||||
fixture.auth.insert_session("valid-session");
|
||||
let now = Utc::now();
|
||||
let range_requests = RequestMetrics {
|
||||
request_count: 2,
|
||||
input_tokens: 8,
|
||||
cached_tokens: 2,
|
||||
total_tokens: 12,
|
||||
..RequestMetrics::default()
|
||||
};
|
||||
*fixture
|
||||
.dashboard_observation
|
||||
.lock()
|
||||
.expect("dashboard observation") = Some(DashboardObservation {
|
||||
range: TimeRange::new(now - Duration::hours(1), now).expect("dashboard range"),
|
||||
requests: range_requests,
|
||||
attempts: AttemptMetrics::default(),
|
||||
totals: DashboardTotals {
|
||||
request_count: 42,
|
||||
input_tokens: 800,
|
||||
@@ -374,8 +360,7 @@ async fn dashboard_summary_should_default_to_current_china_day() {
|
||||
};
|
||||
use chrono::{Duration, Utc};
|
||||
use gateway_admin::model::observability::{
|
||||
AccountPoolMetrics, AttemptMetrics, DashboardObservation, RequestMetrics, TimeRange,
|
||||
china_day_start,
|
||||
AccountPoolMetrics, DashboardObservation, TimeRange, china_day_start,
|
||||
};
|
||||
use gateway_api::admin::observability;
|
||||
use tower::ServiceExt as _;
|
||||
@@ -390,8 +375,6 @@ async fn dashboard_summary_should_default_to_current_china_day() {
|
||||
.lock()
|
||||
.expect("dashboard observation") = Some(DashboardObservation {
|
||||
range: TimeRange::new(now - Duration::hours(1), now).expect("dashboard range"),
|
||||
requests: RequestMetrics::default(),
|
||||
attempts: AttemptMetrics::default(),
|
||||
totals: Default::default(),
|
||||
provider_accounts: AccountPoolMetrics::default(),
|
||||
trend: Vec::new(),
|
||||
@@ -556,6 +539,10 @@ async fn ops_errors_should_keep_account_label_and_authentication_contract() {
|
||||
client_api_key_ref: Some("key_err".to_owned()),
|
||||
component: "model_request".to_owned(),
|
||||
operation: "responses".to_owned(),
|
||||
protocol: Some("openai".to_owned()),
|
||||
client_transport: Some("http_sse".to_owned()),
|
||||
requested_model_id: Some("gpt-5.4".to_owned()),
|
||||
service_tier: Some("priority".to_owned()),
|
||||
endpoint: Some("/v1/responses".to_owned()),
|
||||
provider_kind: Some("openai".to_owned()),
|
||||
provider_account_ref: Some("acct_err".to_owned()),
|
||||
@@ -573,6 +560,13 @@ async fn ops_errors_should_keep_account_label_and_authentication_contract() {
|
||||
latency_ms: Some(90),
|
||||
message: "snapshot error".to_owned(),
|
||||
occurrence_count: 1,
|
||||
client_ip: Some("203.0.113.8".to_owned()),
|
||||
user_agent: Some("codex-cli/0.144.0".to_owned()),
|
||||
reasoning_effort: Some("medium".to_owned()),
|
||||
reasoning_preset: None,
|
||||
request_kind: Some("root".to_owned()),
|
||||
subagent_kind: None,
|
||||
compact: Some(false),
|
||||
occurred_at: Utc::now(),
|
||||
stable_sort_id: "model_request:req_err".to_owned(),
|
||||
});
|
||||
@@ -602,6 +596,11 @@ async fn ops_errors_should_keep_account_label_and_authentication_contract() {
|
||||
"accountLabel": value["data"]["items"][0]["metadata"]["accountLabel"],
|
||||
"clientStatusCode": value["data"]["items"][0]["clientStatusCode"],
|
||||
"route": value["data"]["items"][0]["route"],
|
||||
"requestedModel": value["data"]["items"][0]["requestedModel"],
|
||||
"clientTransport": value["data"]["items"][0]["clientTransport"],
|
||||
"clientIp": value["data"]["items"][0]["clientIp"],
|
||||
"userAgent": value["data"]["items"][0]["userAgent"],
|
||||
"providerErrorCode": value["data"]["items"][0]["providerErrorCode"],
|
||||
}),
|
||||
serde_json::json!({
|
||||
"provider": "openai",
|
||||
@@ -611,6 +610,11 @@ async fn ops_errors_should_keep_account_label_and_authentication_contract() {
|
||||
"accountLabel": "err@example.invalid",
|
||||
"clientStatusCode": 502,
|
||||
"route": "/v1/responses",
|
||||
"requestedModel": "gpt-5.4",
|
||||
"clientTransport": "http_sse",
|
||||
"clientIp": "203.0.113.8",
|
||||
"userAgent": "codex-cli/0.144.0",
|
||||
"providerErrorCode": "upstream",
|
||||
})
|
||||
);
|
||||
}
|
||||
@@ -798,7 +802,7 @@ async fn usage_route_should_forward_a_bounded_unknown_outcome_filter() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn usage_route_should_expose_image_and_websocket_facts() {
|
||||
async fn usage_route_should_expose_table_facts_without_detail_payload() {
|
||||
use std::str::FromStr as _;
|
||||
|
||||
use axum::{
|
||||
@@ -807,8 +811,7 @@ async fn usage_route_should_expose_image_and_websocket_facts() {
|
||||
};
|
||||
use chrono::Utc;
|
||||
use gateway_admin::model::observability::{
|
||||
CalculatedBillingBreakdown, CurrencyCost, DecimalAmount, RequestOutcome, UsageBilling,
|
||||
UsageRecord,
|
||||
CalculatedBillingBreakdown, CurrencyCost, DecimalAmount, UsageBilling, UsageListRecord,
|
||||
};
|
||||
use gateway_api::admin::observability;
|
||||
use tower::ServiceExt as _;
|
||||
@@ -825,17 +828,10 @@ async fn usage_route_should_expose_image_and_websocket_facts() {
|
||||
.usage_records
|
||||
.lock()
|
||||
.expect("usage records")
|
||||
.push(UsageRecord {
|
||||
.push(UsageListRecord {
|
||||
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(),
|
||||
client_transport: "http_sse".to_owned(),
|
||||
client_transport: "websocket".to_owned(),
|
||||
requested_model_id: Some("grok-4.5".to_owned()),
|
||||
provider_kind: Some("xai".to_owned()),
|
||||
provider_account_ref: Some("acct_snapshot".to_owned()),
|
||||
@@ -844,30 +840,7 @@ async fn usage_route_should_expose_image_and_websocket_facts() {
|
||||
provider_account_authentication_kind: Some("oauth".to_owned()),
|
||||
upstream_model_id: Some("grok-4.5".to_owned()),
|
||||
upstream_transport: Some("http_sse".to_owned()),
|
||||
http_version: Some("h2".to_owned()),
|
||||
websocket_pool: Some("reuse".to_owned()),
|
||||
service_tier: Some("default".to_owned()),
|
||||
provider_metadata_json: Some(
|
||||
serde_json::json!({
|
||||
"effectiveModel": "grok-4.5",
|
||||
"requestSummary": {"inputItemsCount": 2},
|
||||
"transport": "must-not-overwrite-core",
|
||||
})
|
||||
.to_string(),
|
||||
),
|
||||
attempt_count: 1,
|
||||
upstream_send_state: "sent".to_owned(),
|
||||
downstream_committed_at: None,
|
||||
outcome: RequestOutcome::Succeeded,
|
||||
client_status_code: Some(200),
|
||||
upstream_status_code: Some(200),
|
||||
client_response_id: Some("resp_usage_metadata".to_owned()),
|
||||
upstream_request_id: Some("upstream_usage_metadata".to_owned()),
|
||||
upstream_response_id: None,
|
||||
error_kind: None,
|
||||
provider_error_code: None,
|
||||
error_message: None,
|
||||
retry_after_ms: None,
|
||||
input_tokens: Some(1),
|
||||
output_tokens: Some(1),
|
||||
cached_tokens: Some(0),
|
||||
@@ -895,31 +868,26 @@ async fn usage_route_should_expose_image_and_websocket_facts() {
|
||||
multiplier_percent: 100,
|
||||
},
|
||||
))),
|
||||
transport_decision_wait_ms: Some(7),
|
||||
connect_ms: Some(11),
|
||||
headers_ms: Some(13),
|
||||
transport_decision_wait_ms: Some(4),
|
||||
connect_ms: Some(5),
|
||||
headers_ms: Some(7),
|
||||
first_event_ms: Some(17),
|
||||
first_reasoning_ms: Some(19),
|
||||
first_text_ms: Some(23),
|
||||
first_token_ms: Some(19),
|
||||
provider_processing_ms: Some(29),
|
||||
first_reasoning_ms: Some(18),
|
||||
first_text_ms: Some(19),
|
||||
first_token_ms: Some(20),
|
||||
provider_processing_ms: Some(21),
|
||||
latency_ms: Some(31),
|
||||
admission_decision_ms: Some(3),
|
||||
account_selection_wait_ms: Some(5),
|
||||
capacity_used_slots: Some(7),
|
||||
capacity_total_slots: Some(10),
|
||||
client_ip: Some("127.0.0.1".to_owned()),
|
||||
user_agent: Some("usage-metadata-test".to_owned()),
|
||||
admission_decision_ms: Some(1),
|
||||
account_selection_wait_ms: Some(2),
|
||||
capacity_used_slots: Some(3),
|
||||
capacity_total_slots: Some(8),
|
||||
client_ip: Some("192.0.2.10".to_owned()),
|
||||
user_agent: Some("codex-cli/1.0".to_owned()),
|
||||
reasoning_effort: Some("max".to_owned()),
|
||||
reasoning_preset: Some("ultra".to_owned()),
|
||||
request_kind: Some("review".to_owned()),
|
||||
subagent_kind: Some("worker".to_owned()),
|
||||
compact: true,
|
||||
image_generation_requested: true,
|
||||
image_generation_succeeded: Some(true),
|
||||
started_at: Utc::now(),
|
||||
deadline_at: Utc::now(),
|
||||
completed_at: Some(Utc::now()),
|
||||
});
|
||||
let response = observability::router::<AdminTestState>()
|
||||
.with_state(fixture.state())
|
||||
@@ -959,9 +927,6 @@ 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"],
|
||||
@@ -969,32 +934,22 @@ async fn usage_route_should_expose_image_and_websocket_facts() {
|
||||
"authenticationKind": value["data"]["items"][0]["authenticationKind"],
|
||||
"imageInputTokens": value["data"]["items"][0]["tokenDetails"]["imageInputTokens"],
|
||||
"imageOutputTokens": value["data"]["items"][0]["tokenDetails"]["imageOutputTokens"],
|
||||
"websocketPool": value["data"]["items"][0]["websocketPool"],
|
||||
"imageGenerationRequested": value["data"]["items"][0]["imageGenerationRequested"],
|
||||
"imageGenerationSucceeded": value["data"]["items"][0]["imageGenerationSucceeded"],
|
||||
"requestedModel": value["data"]["items"][0]["requestedModel"],
|
||||
"upstreamModel": value["data"]["items"][0]["upstreamModel"],
|
||||
"clientIp": value["data"]["items"][0]["clientIp"],
|
||||
"userAgent": value["data"]["items"][0]["userAgent"],
|
||||
"reasoningEffort": value["data"]["items"][0]["reasoningEffort"],
|
||||
"reasoningPreset": value["data"]["items"][0]["reasoningPreset"],
|
||||
"requestKind": value["data"]["items"][0]["requestKind"],
|
||||
"subagentKind": value["data"]["items"][0]["subagentKind"],
|
||||
"compact": value["data"]["items"][0]["compact"],
|
||||
"transport": value["data"]["items"][0]["transport"],
|
||||
"transportDecisionWaitMs": value["data"]["items"][0]["latencyDetails"]["transportDecisionWaitMs"],
|
||||
"wsConnectMs": value["data"]["items"][0]["latencyDetails"]["wsConnectMs"],
|
||||
"firstReasoningMs": value["data"]["items"][0]["latencyDetails"]["firstReasoningMs"],
|
||||
"firstTextMs": value["data"]["items"][0]["latencyDetails"]["firstTextMs"],
|
||||
"firstTokenMs": value["data"]["items"][0]["latencyDetails"]["firstTokenMs"],
|
||||
"effectiveModel": value["data"]["items"][0]["metadata"]["effectiveModel"],
|
||||
"requestSummary": value["data"]["items"][0]["metadata"]["requestSummary"],
|
||||
"clientTransport": value["data"]["items"][0]["clientTransport"],
|
||||
"upstreamTransport": value["data"]["items"][0]["upstreamTransport"],
|
||||
"latencyDetails": value["data"]["items"][0]["latencyDetails"],
|
||||
"firstTokenLatencyMs": value["data"]["items"][0]["firstTokenLatencyMs"],
|
||||
"latencyMs": value["data"]["items"][0]["latencyMs"],
|
||||
"clientIp": value["data"]["items"][0]["clientIp"],
|
||||
"userAgent": value["data"]["items"][0]["userAgent"],
|
||||
}),
|
||||
serde_json::json!({
|
||||
"route": "/v1/responses",
|
||||
"routingScope": "groups",
|
||||
"routingGroupRefs": ["grp_usage"],
|
||||
"routingGroupNamesSnapshot": ["Usage group"],
|
||||
"serviceTier": "default",
|
||||
"accountId": "acct_snapshot",
|
||||
"accountName": "Snapshot Alpha",
|
||||
@@ -1002,26 +957,33 @@ async fn usage_route_should_expose_image_and_websocket_facts() {
|
||||
"authenticationKind": "oauth",
|
||||
"imageInputTokens": 31,
|
||||
"imageOutputTokens": 9,
|
||||
"websocketPool": {"kind": "reuse"},
|
||||
"imageGenerationRequested": true,
|
||||
"imageGenerationSucceeded": true,
|
||||
"requestedModel": "grok-4.5",
|
||||
"upstreamModel": "grok-4.5",
|
||||
"clientIp": "127.0.0.1",
|
||||
"userAgent": "usage-metadata-test",
|
||||
"reasoningEffort": "max",
|
||||
"reasoningPreset": "ultra",
|
||||
"requestKind": "review",
|
||||
"subagentKind": "worker",
|
||||
"compact": true,
|
||||
"transport": "http_sse",
|
||||
"transportDecisionWaitMs": 7,
|
||||
"wsConnectMs": 11,
|
||||
"firstReasoningMs": 19,
|
||||
"firstTextMs": 23,
|
||||
"firstTokenMs": 19,
|
||||
"effectiveModel": "grok-4.5",
|
||||
"requestSummary": {"inputItemsCount": 2},
|
||||
"clientTransport": "websocket",
|
||||
"upstreamTransport": "http_sse",
|
||||
"latencyDetails": {
|
||||
"admissionDecisionMs": 1,
|
||||
"accountSelectionWaitMs": 2,
|
||||
"capacityUsedSlots": 3,
|
||||
"capacityTotalSlots": 8,
|
||||
"transportDecisionWaitMs": 4,
|
||||
"wsConnectMs": 5,
|
||||
"upstreamHeadersMs": 7,
|
||||
"firstEventMs": 17,
|
||||
"firstReasoningMs": 18,
|
||||
"firstTextMs": 19,
|
||||
"firstTokenMs": 20,
|
||||
"openaiProcessingMs": 21,
|
||||
},
|
||||
"firstTokenLatencyMs": 20,
|
||||
"latencyMs": 31,
|
||||
"clientIp": "192.0.2.10",
|
||||
"userAgent": "codex-cli/1.0",
|
||||
})
|
||||
);
|
||||
assert!(value["data"]["items"][0].get("metadata").is_none());
|
||||
}
|
||||
|
||||
@@ -62,7 +62,10 @@ fn source_tree_should_match_frozen_machine_manifest() {
|
||||
"src/openai/responses/mod.rs",
|
||||
"src/openai/responses/request.rs",
|
||||
"src/openai/responses/response.rs",
|
||||
"src/openai/responses/websocket.rs",
|
||||
"src/openai/responses/websocket/connection.rs",
|
||||
"src/openai/responses/websocket/forward.rs",
|
||||
"src/openai/responses/websocket/mod.rs",
|
||||
"src/openai/responses/websocket/protocol.rs",
|
||||
"src/openai/router.rs",
|
||||
"src/openai/service.rs",
|
||||
];
|
||||
@@ -99,7 +102,9 @@ fn test_tree_should_match_frozen_rust_mirror() {
|
||||
"tests/openai/models.rs",
|
||||
"tests/openai/responses/http.rs",
|
||||
"tests/openai/responses/mod.rs",
|
||||
"tests/openai/responses/websocket.rs",
|
||||
"tests/openai/responses/websocket/connection.rs",
|
||||
"tests/openai/responses/websocket/mod.rs",
|
||||
"tests/openai/responses/websocket/protocol.rs",
|
||||
"tests/openai/router.rs",
|
||||
];
|
||||
expected.sort_unstable();
|
||||
|
||||
@@ -320,8 +320,8 @@ fn decoder_should_preserve_ordinary_request_headers_as_opaque_multivalues() {
|
||||
vec![b"future-ascii".to_vec(), b"\x80\xff".to_vec()]
|
||||
);
|
||||
assert_eq!(
|
||||
values("version"),
|
||||
vec![b"future-v1".to_vec(), b"future-\x80".to_vec()]
|
||||
openai_protocol_context(&decoded).get("version"),
|
||||
Some(&json!("future-v1"))
|
||||
);
|
||||
assert_eq!(
|
||||
values("openai-beta"),
|
||||
@@ -349,6 +349,7 @@ fn decoder_should_preserve_ordinary_request_headers_as_opaque_multivalues() {
|
||||
// 上游指纹由运行时画像统一生成,客户端不得覆盖。
|
||||
"user-agent",
|
||||
"originator",
|
||||
"version",
|
||||
"x-oai-attestation",
|
||||
"x-oai-is",
|
||||
"x-oai-is-update",
|
||||
|
||||
@@ -0,0 +1,651 @@
|
||||
use std::{
|
||||
fmt,
|
||||
pin::Pin,
|
||||
sync::{
|
||||
Arc,
|
||||
atomic::{AtomicBool, Ordering},
|
||||
},
|
||||
task::{Context, Poll},
|
||||
time::{Duration, SystemTime},
|
||||
};
|
||||
|
||||
use axum::{
|
||||
extract::ws::Message,
|
||||
http::{StatusCode, header::AUTHORIZATION},
|
||||
};
|
||||
use futures::{Sink, SinkExt, Stream, StreamExt, future::BoxFuture};
|
||||
use gateway_core::{
|
||||
engine::{
|
||||
CancellationToken, CommitRequirement, CoordinatedEvent, EngineError,
|
||||
execution::{
|
||||
AuthenticatedClient, ClientAuthenticationError, ExecutionService, ExecutionSession,
|
||||
StartExecution, StartProviderExecution, StartedExecution,
|
||||
},
|
||||
},
|
||||
error::GatewayError,
|
||||
event::{GatewayEvent, ProtocolWireEvent, ProviderEvent, ProviderResponseHeader, ResponseMeta},
|
||||
routing::PublicModelId,
|
||||
};
|
||||
use serde_json::{Value, json};
|
||||
use tokio::sync::mpsc::{
|
||||
UnboundedReceiver, UnboundedSender, error::TryRecvError, unbounded_channel,
|
||||
};
|
||||
use tokio_tungstenite::{
|
||||
connect_async,
|
||||
tungstenite::{
|
||||
Message as ClientMessage, client::IntoClientRequest, protocol::frame::coding::CloseCode,
|
||||
},
|
||||
};
|
||||
|
||||
use super::{AtomicFailureExecution, AtomicFailureTrace};
|
||||
use crate::openai::{api_router, authenticated_client};
|
||||
|
||||
#[derive(Default)]
|
||||
struct CrossingLimitTrace {
|
||||
release_terminal: tokio::sync::Notify,
|
||||
started_emitted: AtomicBool,
|
||||
terminal_emitted: AtomicBool,
|
||||
finalized: AtomicBool,
|
||||
}
|
||||
|
||||
struct CrossingLimitSession {
|
||||
trace: Arc<CrossingLimitTrace>,
|
||||
}
|
||||
|
||||
impl ExecutionSession for CrossingLimitSession {
|
||||
fn next_event(&mut self) -> BoxFuture<'_, Result<Option<CoordinatedEvent>, EngineError>> {
|
||||
Box::pin(async move {
|
||||
if !self.trace.started_emitted.swap(true, Ordering::AcqRel) {
|
||||
return Ok(Some(crossing_limit_started_batch()));
|
||||
}
|
||||
if !self.trace.terminal_emitted.load(Ordering::Acquire) {
|
||||
self.trace.release_terminal.notified().await;
|
||||
if !self.trace.terminal_emitted.swap(true, Ordering::AcqRel) {
|
||||
return Ok(Some(crossing_limit_terminal_batch()));
|
||||
}
|
||||
}
|
||||
self.trace.finalized.store(true, Ordering::Release);
|
||||
Ok(None)
|
||||
})
|
||||
}
|
||||
|
||||
fn collect_uncommitted(&mut self) -> BoxFuture<'_, Result<Vec<ProviderEvent>, EngineError>> {
|
||||
Box::pin(async { Err(EngineError::InvalidDeliveryState) })
|
||||
}
|
||||
|
||||
fn response_headers(&self) -> &[ProviderResponseHeader] {
|
||||
&[]
|
||||
}
|
||||
|
||||
fn commit_downstream(&mut self, _: Option<u16>) -> BoxFuture<'_, Result<(), EngineError>> {
|
||||
Box::pin(async { Ok(()) })
|
||||
}
|
||||
|
||||
fn record_client_status(&mut self, _: u16) -> BoxFuture<'_, Result<(), EngineError>> {
|
||||
Box::pin(async { Ok(()) })
|
||||
}
|
||||
|
||||
fn is_finalized(&self) -> bool {
|
||||
self.trace.finalized.load(Ordering::Acquire)
|
||||
}
|
||||
|
||||
fn cancel(&self) {}
|
||||
|
||||
fn detach_finalize(self: Box<Self>) -> BoxFuture<'static, ()> {
|
||||
Box::pin(async move {
|
||||
self.trace.finalized.store(true, Ordering::Release);
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
struct CrossingLimitExecution {
|
||||
client: AuthenticatedClient,
|
||||
trace: Arc<CrossingLimitTrace>,
|
||||
}
|
||||
|
||||
impl ExecutionService for CrossingLimitExecution {
|
||||
fn authenticate(
|
||||
&self,
|
||||
plaintext: &str,
|
||||
) -> Result<AuthenticatedClient, ClientAuthenticationError> {
|
||||
(plaintext == "sk_ws_crossing_limit")
|
||||
.then(|| self.client.clone())
|
||||
.ok_or(ClientAuthenticationError::InvalidKey)
|
||||
}
|
||||
|
||||
fn public_models(&self, _: &AuthenticatedClient) -> Vec<PublicModelId> {
|
||||
vec![PublicModelId::new("model-a").expect("public model")]
|
||||
}
|
||||
|
||||
fn contains_public_model(&self, _: &AuthenticatedClient, model: &PublicModelId) -> bool {
|
||||
model.as_str() == "model-a"
|
||||
}
|
||||
|
||||
fn start(
|
||||
&self,
|
||||
request: StartExecution,
|
||||
) -> BoxFuture<'_, Result<StartedExecution, GatewayError>> {
|
||||
Box::pin(async move {
|
||||
Ok(StartedExecution {
|
||||
request_id: gateway_core::engine::ModelRequestId::new("req_ws_crossing_limit")
|
||||
.expect("request id"),
|
||||
created_at: SystemTime::now(),
|
||||
stream: request.metadata.stream,
|
||||
session: Box::new(CrossingLimitSession {
|
||||
trace: Arc::clone(&self.trace),
|
||||
}),
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
fn start_provider_endpoint(
|
||||
&self,
|
||||
_: StartProviderExecution,
|
||||
) -> BoxFuture<'_, Result<StartedExecution, GatewayError>> {
|
||||
Box::pin(async { unreachable!("WebSocket test does not execute provider endpoints") })
|
||||
}
|
||||
}
|
||||
|
||||
fn crossing_limit_started_batch() -> CoordinatedEvent {
|
||||
let started = ProviderEvent::canonical_with_wire(
|
||||
vec![GatewayEvent::Started(ResponseMeta::new(
|
||||
"resp_ws_crossing_limit",
|
||||
"model-a",
|
||||
))],
|
||||
ProtocolWireEvent::json(
|
||||
"openai",
|
||||
Some("response.created".to_owned()),
|
||||
json!({
|
||||
"type": "response.created",
|
||||
"response": {
|
||||
"id": "resp_ws_crossing_limit",
|
||||
"model": "model-a",
|
||||
"status": "in_progress"
|
||||
}
|
||||
}),
|
||||
)
|
||||
.expect("created wire"),
|
||||
);
|
||||
CoordinatedEvent::try_batch(vec![started], CommitRequirement::CommitBeforeDelivery)
|
||||
.expect("started WebSocket batch")
|
||||
}
|
||||
|
||||
fn crossing_limit_terminal_batch() -> CoordinatedEvent {
|
||||
let completed = ProviderEvent::wire(
|
||||
ProtocolWireEvent::json(
|
||||
"openai",
|
||||
Some("response.completed".to_owned()),
|
||||
json!({
|
||||
"type": "response.completed",
|
||||
"response": {
|
||||
"id": "resp_ws_crossing_limit",
|
||||
"model": "model-a",
|
||||
"status": "completed"
|
||||
}
|
||||
}),
|
||||
)
|
||||
.expect("completed wire"),
|
||||
);
|
||||
CoordinatedEvent::try_batch(vec![completed], CommitRequirement::AlreadyCommitted)
|
||||
.expect("terminal WebSocket batch")
|
||||
}
|
||||
|
||||
// 连接 runtime 没有正式公共 API;镜像测试直接编译该私有模块,以测试 socket pump,
|
||||
// 不把测试钩子暴露给 gateway-api 的消费者。
|
||||
#[expect(
|
||||
dead_code,
|
||||
reason = "the mirrored test only exercises pump-facing seams"
|
||||
)]
|
||||
#[path = "../../../../src/openai/responses/websocket/connection.rs"]
|
||||
mod connection_under_test;
|
||||
|
||||
use connection_under_test::{
|
||||
ConnectionConfig, ConnectionEvent, ConnectionWriteError, FramePhase,
|
||||
ResponsesWebSocketConnection, WriteContext, spawn_connection,
|
||||
};
|
||||
|
||||
#[derive(Debug)]
|
||||
struct TestSocketError;
|
||||
|
||||
impl fmt::Display for TestSocketError {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter.write_str("test socket error")
|
||||
}
|
||||
}
|
||||
|
||||
struct TestSocket {
|
||||
incoming: UnboundedReceiver<Result<Message, TestSocketError>>,
|
||||
written: UnboundedSender<Message>,
|
||||
stall_writes: bool,
|
||||
dropped: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
impl Stream for TestSocket {
|
||||
type Item = Result<Message, TestSocketError>;
|
||||
|
||||
fn poll_next(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
||||
self.incoming.poll_recv(context)
|
||||
}
|
||||
}
|
||||
|
||||
impl Sink<Message> for TestSocket {
|
||||
type Error = TestSocketError;
|
||||
|
||||
fn poll_ready(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
||||
if self.stall_writes {
|
||||
Poll::Pending
|
||||
} else {
|
||||
Poll::Ready(Ok(()))
|
||||
}
|
||||
}
|
||||
|
||||
fn start_send(self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> {
|
||||
self.written.send(item).map_err(|_| TestSocketError)
|
||||
}
|
||||
|
||||
fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
||||
Poll::Ready(Ok(()))
|
||||
}
|
||||
|
||||
fn poll_close(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
|
||||
Poll::Ready(Ok(()))
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for TestSocket {
|
||||
fn drop(&mut self) {
|
||||
self.dropped.store(true, Ordering::Release);
|
||||
}
|
||||
}
|
||||
|
||||
struct PumpHarness {
|
||||
connection: ResponsesWebSocketConnection,
|
||||
incoming: UnboundedSender<Result<Message, TestSocketError>>,
|
||||
written: UnboundedReceiver<Message>,
|
||||
dropped: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
fn test_connection(stall_writes: bool) -> PumpHarness {
|
||||
let (incoming_tx, incoming_rx) = unbounded_channel();
|
||||
let (written_tx, written_rx) = unbounded_channel();
|
||||
let dropped = Arc::new(AtomicBool::new(false));
|
||||
let socket = TestSocket {
|
||||
incoming: incoming_rx,
|
||||
written: written_tx,
|
||||
stall_writes,
|
||||
dropped: Arc::clone(&dropped),
|
||||
};
|
||||
let connection = spawn_connection(
|
||||
socket,
|
||||
Arc::from("ws_test"),
|
||||
CancellationToken::new(),
|
||||
ConnectionConfig::PRODUCTION,
|
||||
);
|
||||
PumpHarness {
|
||||
connection,
|
||||
incoming: incoming_tx,
|
||||
written: written_rx,
|
||||
dropped,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pump_replies_to_ping_with_the_same_pong_payload() {
|
||||
let PumpHarness {
|
||||
connection,
|
||||
incoming,
|
||||
mut written,
|
||||
..
|
||||
} = test_connection(false);
|
||||
incoming
|
||||
.send(Ok(Message::Ping(vec![1, 2, 3].into())))
|
||||
.expect("send Ping to test socket");
|
||||
|
||||
let message = written.recv().await.expect("pump must write Pong");
|
||||
let Message::Pong(payload) = message else {
|
||||
panic!("pump wrote a non-Pong control frame");
|
||||
};
|
||||
assert_eq!(payload.as_ref(), &[1, 2, 3]);
|
||||
drop(connection);
|
||||
}
|
||||
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn pump_consumes_pong_without_forwarding_a_business_event() {
|
||||
let PumpHarness {
|
||||
mut connection,
|
||||
incoming,
|
||||
..
|
||||
} = test_connection(false);
|
||||
incoming
|
||||
.send(Ok(Message::Pong(vec![4, 5, 6].into())))
|
||||
.expect("send Pong to test socket");
|
||||
|
||||
let result = tokio::time::timeout(Duration::from_secs(1), connection.next_event()).await;
|
||||
assert!(result.is_err(), "Pong unexpectedly reached the coordinator");
|
||||
}
|
||||
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn pump_does_not_emit_an_active_ping() {
|
||||
let PumpHarness {
|
||||
connection,
|
||||
incoming: _incoming,
|
||||
mut written,
|
||||
..
|
||||
} = test_connection(false);
|
||||
tokio::time::advance(Duration::from_secs(5 * 60)).await;
|
||||
tokio::task::yield_now().await;
|
||||
|
||||
assert!(matches!(written.try_recv(), Err(TryRecvError::Empty)));
|
||||
drop(connection);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn outbound_commands_are_written_in_acknowledged_order() {
|
||||
let PumpHarness {
|
||||
mut connection,
|
||||
incoming: _incoming,
|
||||
mut written,
|
||||
..
|
||||
} = test_connection(false);
|
||||
let request_id = Arc::<str>::from("req_test");
|
||||
|
||||
connection
|
||||
.send_text(
|
||||
"first".to_owned(),
|
||||
WriteContext::request(&request_id, FramePhase::First),
|
||||
)
|
||||
.await
|
||||
.expect("first write acknowledgement");
|
||||
connection
|
||||
.send_text(
|
||||
"second".to_owned(),
|
||||
WriteContext::request(&request_id, FramePhase::Terminal),
|
||||
)
|
||||
.await
|
||||
.expect("second write acknowledgement");
|
||||
|
||||
let first = written.recv().await.expect("first wire message");
|
||||
let second = written.recv().await.expect("second wire message");
|
||||
assert!(
|
||||
matches!(first, Message::Text(payload) if payload.as_str() == "first")
|
||||
&& matches!(second, Message::Text(payload) if payload.as_str() == "second")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn stalled_write_times_out_at_three_hundred_seconds_and_aborts_the_pump() {
|
||||
let PumpHarness {
|
||||
mut connection,
|
||||
incoming: _incoming,
|
||||
written: _written,
|
||||
dropped,
|
||||
} = test_connection(true);
|
||||
let request_id = Arc::<str>::from("req_test");
|
||||
|
||||
let error = connection
|
||||
.send_text(
|
||||
"payload".to_owned(),
|
||||
WriteContext::request(&request_id, FramePhase::First),
|
||||
)
|
||||
.await
|
||||
.expect_err("stalled write must time out");
|
||||
tokio::task::yield_now().await;
|
||||
|
||||
assert!(
|
||||
matches!(error, ConnectionWriteError::Timeout { timeout } if timeout == Duration::from_secs(300))
|
||||
&& dropped.load(Ordering::Acquire)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn pump_marks_the_connection_expired_at_sixty_minutes() {
|
||||
let PumpHarness {
|
||||
mut connection,
|
||||
incoming: _incoming,
|
||||
written: _written,
|
||||
..
|
||||
} = test_connection(false);
|
||||
|
||||
let event = connection
|
||||
.next_event()
|
||||
.await
|
||||
.expect("deadline must emit an event");
|
||||
|
||||
assert!(matches!(event, ConnectionEvent::Expired) && connection.is_expired());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dropping_the_connection_aborts_and_drops_the_socket_owner() {
|
||||
let PumpHarness {
|
||||
connection,
|
||||
incoming: _incoming,
|
||||
written: _written,
|
||||
dropped,
|
||||
} = test_connection(false);
|
||||
|
||||
drop(connection);
|
||||
tokio::task::yield_now().await;
|
||||
|
||||
assert!(dropped.load(Ordering::Acquire));
|
||||
}
|
||||
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn idle_connection_reaches_the_official_limit_without_starting_an_execution() {
|
||||
let trace = Arc::new(AtomicFailureTrace::default());
|
||||
let execution = Arc::new(AtomicFailureExecution {
|
||||
client: authenticated_client("sk_ws_atomic"),
|
||||
trace: Arc::clone(&trace),
|
||||
response_headers: Vec::new(),
|
||||
});
|
||||
let app = api_router(execution).await;
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("bind WebSocket test server");
|
||||
let address = listener.local_addr().expect("test server address");
|
||||
let server = tokio::spawn(async move {
|
||||
axum::serve(listener, app)
|
||||
.await
|
||||
.expect("serve WebSocket test router");
|
||||
});
|
||||
let mut request = format!("ws://{address}/v1/responses")
|
||||
.into_client_request()
|
||||
.expect("WebSocket request");
|
||||
request.headers_mut().insert(
|
||||
AUTHORIZATION,
|
||||
"Bearer sk_ws_atomic".parse().expect("authorization"),
|
||||
);
|
||||
let (mut socket, response) = connect_async(request).await.expect("upgrade WebSocket");
|
||||
assert_eq!(response.status(), StatusCode::SWITCHING_PROTOCOLS);
|
||||
|
||||
tokio::time::advance(Duration::from_secs(60 * 60)).await;
|
||||
tokio::task::yield_now().await;
|
||||
let message = tokio::time::timeout(Duration::from_secs(1), socket.next())
|
||||
.await
|
||||
.expect("connection limit response timeout")
|
||||
.expect("connection remains available for the limit error")
|
||||
.expect("valid WebSocket frame");
|
||||
let ClientMessage::Text(text) = message else {
|
||||
panic!("connection limit must be a text error event");
|
||||
};
|
||||
let value = serde_json::from_str::<Value>(&text).expect("connection limit JSON");
|
||||
|
||||
assert_eq!(
|
||||
value,
|
||||
json!({
|
||||
"type": "error",
|
||||
"status": 400,
|
||||
"error": {
|
||||
"type": "invalid_request_error",
|
||||
"code": "websocket_connection_limit_reached",
|
||||
"message": "Responses websocket connection limit reached (60 minutes). Create a new websocket connection to continue."
|
||||
}
|
||||
})
|
||||
);
|
||||
assert_eq!(trace.starts.load(Ordering::Acquire), 0);
|
||||
|
||||
socket.close(None).await.expect("close WebSocket");
|
||||
server.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn active_response_crossing_the_limit_writes_terminal_before_the_limit_error() {
|
||||
let trace = Arc::new(CrossingLimitTrace::default());
|
||||
let execution = Arc::new(CrossingLimitExecution {
|
||||
client: authenticated_client("sk_ws_crossing_limit"),
|
||||
trace: Arc::clone(&trace),
|
||||
});
|
||||
let app = api_router(execution).await;
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("bind WebSocket test server");
|
||||
let address = listener.local_addr().expect("test server address");
|
||||
let server = tokio::spawn(async move {
|
||||
axum::serve(listener, app)
|
||||
.await
|
||||
.expect("serve WebSocket test router");
|
||||
});
|
||||
let mut request = format!("ws://{address}/v1/responses")
|
||||
.into_client_request()
|
||||
.expect("WebSocket request");
|
||||
request.headers_mut().insert(
|
||||
AUTHORIZATION,
|
||||
"Bearer sk_ws_crossing_limit"
|
||||
.parse()
|
||||
.expect("authorization"),
|
||||
);
|
||||
let (mut socket, response) = connect_async(request).await.expect("upgrade WebSocket");
|
||||
assert_eq!(response.status(), StatusCode::SWITCHING_PROTOCOLS);
|
||||
socket
|
||||
.send(ClientMessage::Text(
|
||||
json!({
|
||||
"type": "response.create",
|
||||
"model": "model-a",
|
||||
"input": "hello"
|
||||
})
|
||||
.to_string()
|
||||
.into(),
|
||||
))
|
||||
.await
|
||||
.expect("send response.create");
|
||||
|
||||
let mut initial_types = Vec::new();
|
||||
loop {
|
||||
tokio::task::yield_now().await;
|
||||
let message = tokio::time::timeout(Duration::from_secs(1), socket.next())
|
||||
.await
|
||||
.unwrap_or_else(|_| panic!("initial response timeout after {initial_types:?}"))
|
||||
.expect("WebSocket remains open")
|
||||
.expect("valid WebSocket frame");
|
||||
let ClientMessage::Text(text) = message else {
|
||||
panic!("unexpected initial WebSocket frame after {initial_types:?}: {message:?}");
|
||||
};
|
||||
let value = serde_json::from_str::<Value>(&text).expect("initial response JSON");
|
||||
let event_type = value.get("type").and_then(Value::as_str);
|
||||
initial_types.extend(event_type.map(str::to_owned));
|
||||
if event_type == Some("response.created") {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
tokio::time::pause();
|
||||
tokio::time::advance(Duration::from_secs(60 * 60)).await;
|
||||
tokio::task::yield_now().await;
|
||||
tokio::time::resume();
|
||||
trace.release_terminal.notify_one();
|
||||
let mut event_types = Vec::new();
|
||||
loop {
|
||||
tokio::task::yield_now().await;
|
||||
let message = tokio::time::timeout(Duration::from_secs(1), socket.next())
|
||||
.await
|
||||
.expect("post-limit response timeout")
|
||||
.expect("WebSocket remains open through the limit error")
|
||||
.expect("valid WebSocket frame");
|
||||
let ClientMessage::Text(text) = message else {
|
||||
continue;
|
||||
};
|
||||
let value = serde_json::from_str::<Value>(&text).expect("post-limit response JSON");
|
||||
let Some(event_type) = value.get("type").and_then(Value::as_str) else {
|
||||
continue;
|
||||
};
|
||||
event_types.push(event_type.to_owned());
|
||||
if event_type == "error" {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
assert_eq!(event_types, ["response.completed", "error"]);
|
||||
assert!(trace.finalized.load(Ordering::Acquire));
|
||||
server.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn second_text_frame_during_an_active_response_keeps_the_single_response_contract() {
|
||||
let trace = Arc::new(CrossingLimitTrace::default());
|
||||
let execution = Arc::new(CrossingLimitExecution {
|
||||
client: authenticated_client("sk_ws_crossing_limit"),
|
||||
trace,
|
||||
});
|
||||
let app = api_router(execution).await;
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("bind WebSocket test server");
|
||||
let address = listener.local_addr().expect("test server address");
|
||||
let server = tokio::spawn(async move {
|
||||
axum::serve(listener, app)
|
||||
.await
|
||||
.expect("serve WebSocket test router");
|
||||
});
|
||||
let mut request = format!("ws://{address}/v1/responses")
|
||||
.into_client_request()
|
||||
.expect("WebSocket request");
|
||||
request.headers_mut().insert(
|
||||
AUTHORIZATION,
|
||||
"Bearer sk_ws_crossing_limit"
|
||||
.parse()
|
||||
.expect("authorization"),
|
||||
);
|
||||
let (mut socket, _) = connect_async(request).await.expect("upgrade WebSocket");
|
||||
let response_create = ClientMessage::Text(
|
||||
json!({
|
||||
"type": "response.create",
|
||||
"model": "model-a",
|
||||
"input": "hello"
|
||||
})
|
||||
.to_string()
|
||||
.into(),
|
||||
);
|
||||
socket
|
||||
.send(response_create.clone())
|
||||
.await
|
||||
.expect("send first response.create");
|
||||
loop {
|
||||
let message = tokio::time::timeout(Duration::from_secs(2), socket.next())
|
||||
.await
|
||||
.expect("initial response timeout")
|
||||
.expect("WebSocket remains open")
|
||||
.expect("valid WebSocket frame");
|
||||
let ClientMessage::Text(text) = message else {
|
||||
continue;
|
||||
};
|
||||
let value = serde_json::from_str::<Value>(&text).expect("initial response JSON");
|
||||
if value.get("type").and_then(Value::as_str) == Some("response.created") {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
socket
|
||||
.send(response_create)
|
||||
.await
|
||||
.expect("send overlapping response.create");
|
||||
let message = tokio::time::timeout(Duration::from_secs(2), socket.next())
|
||||
.await
|
||||
.expect("policy close timeout")
|
||||
.expect("WebSocket returns a policy close")
|
||||
.expect("valid WebSocket frame");
|
||||
let ClientMessage::Close(Some(frame)) = message else {
|
||||
panic!("overlapping response.create must close the connection");
|
||||
};
|
||||
|
||||
assert_eq!(frame.code, CloseCode::Policy);
|
||||
server.abort();
|
||||
}
|
||||
+5
@@ -1,3 +1,6 @@
|
||||
mod connection;
|
||||
mod protocol;
|
||||
|
||||
use std::sync::{
|
||||
Arc,
|
||||
atomic::{AtomicBool, AtomicUsize, Ordering},
|
||||
@@ -171,6 +174,7 @@ fn response_create_should_reject_non_boolean_stream_without_disclosing_body_valu
|
||||
|
||||
#[derive(Default)]
|
||||
struct AtomicFailureTrace {
|
||||
starts: AtomicUsize,
|
||||
next_calls: AtomicUsize,
|
||||
committed: AtomicBool,
|
||||
finalized: AtomicBool,
|
||||
@@ -258,6 +262,7 @@ impl ExecutionService for AtomicFailureExecution {
|
||||
&self,
|
||||
request: StartExecution,
|
||||
) -> BoxFuture<'_, Result<StartedExecution, GatewayError>> {
|
||||
self.trace.starts.fetch_add(1, Ordering::AcqRel);
|
||||
Box::pin(async move {
|
||||
Ok(StartedExecution {
|
||||
request_id: gateway_core::engine::ModelRequestId::new("req_ws_atomic")
|
||||
@@ -0,0 +1,33 @@
|
||||
use gateway_core::operation::Operation;
|
||||
use serde_json::json;
|
||||
|
||||
use super::decode_response_create;
|
||||
|
||||
#[test]
|
||||
fn latest_official_codex_response_create_fixture_decodes_unchanged() {
|
||||
// Audited against openai/codex main 94cbbddafc1776d5e377bca1b05932c697e82238.
|
||||
let decoded = decode_response_create(
|
||||
&json!({
|
||||
"type": "response.create",
|
||||
"model": "smart-code",
|
||||
"input": "hello",
|
||||
"store": false
|
||||
})
|
||||
.to_string(),
|
||||
)
|
||||
.expect("latest official Codex response.create fixture must decode");
|
||||
let Operation::Generate(request) = decoded.operation() else {
|
||||
panic!("Responses must map to Generate");
|
||||
};
|
||||
|
||||
assert_eq!(
|
||||
request.protocol_payload().body(),
|
||||
json!({
|
||||
"model": "smart-code",
|
||||
"input": "hello",
|
||||
"store": false
|
||||
})
|
||||
.as_object()
|
||||
.expect("fixture body")
|
||||
);
|
||||
}
|
||||
@@ -12,11 +12,11 @@ use crate::engine::continuation::{
|
||||
};
|
||||
use crate::engine::provider::{Provider, ProviderCallMetadata, ProviderRequest, ProviderStream};
|
||||
use crate::engine::{
|
||||
AccountAttemptContext, AttemptContext, AttemptRecord, AttemptTrigger, CancellationToken,
|
||||
CommitRequirement, ContinuationAttempt, CoordinatedEvent, EngineError, ExecutionOutcome,
|
||||
ExecutionStore, GatewayEngine, IntermediateFailure, ModelRequestFinalization, ModelRequestId,
|
||||
ModelRequestTimings, NewModelRequest, ProviderAccountStateOwner, ProviderAttemptOutcome,
|
||||
RequestAttemptContext, UpstreamSendState,
|
||||
AccountAttemptContext, AttemptContext, AttemptRecord, AttemptTransport, AttemptTrigger,
|
||||
CancellationToken, CommitRequirement, ContinuationAttempt, CoordinatedEvent, EngineError,
|
||||
ExecutionOutcome, ExecutionStore, GatewayEngine, IntermediateFailure, ModelRequestFinalization,
|
||||
ModelRequestId, ModelRequestTimings, NewModelRequest, ProviderAccountStateOwner,
|
||||
ProviderAttemptOutcome, RequestAttemptContext, UpstreamSendState,
|
||||
};
|
||||
use crate::error::{GatewayError, GatewayErrorKind, ProviderError, ProviderErrorKind, StoreError};
|
||||
use crate::event::{
|
||||
@@ -166,6 +166,7 @@ where
|
||||
excluded_accounts: BTreeSet::new(),
|
||||
credential_recovery_attempted_accounts: BTreeSet::new(),
|
||||
recovery_account: None,
|
||||
transport_fallback_account: None,
|
||||
current: None,
|
||||
send_state_watermark: UpstreamSendState::NotSent,
|
||||
downstream_committed_at: None,
|
||||
@@ -245,6 +246,9 @@ pub struct ResponseExecutionSession<S: ?Sized> {
|
||||
/// attempt 建立时即被消费,后续可重试错误仍可换号消耗剩余重试预算。
|
||||
/// 与 `required_account`(外部指定、贯穿整个请求)语义不同,不可合并。
|
||||
recovery_account: Option<crate::engine::credential::ProviderAccountId>,
|
||||
/// 上游传输在提交边界前失败后的一次性同账号备用传输钉选。
|
||||
/// 与凭据恢复分开保存,避免把 transport fallback 误记为 OAuth 恢复。
|
||||
transport_fallback_account: Option<crate::engine::credential::ProviderAccountId>,
|
||||
current: Option<CurrentAttempt>,
|
||||
/// 请求级发送状态水位;跨 attempt 单调不降,终态写回不得低于此档。
|
||||
send_state_watermark: UpstreamSendState,
|
||||
@@ -610,14 +614,26 @@ where
|
||||
.checked_add(1)
|
||||
.and_then(NonZeroU32::new)
|
||||
.ok_or(EngineError::EmptyRoutingPlan)?;
|
||||
// recovery 钉选在此被一次性消费,只绑定本次 replay attempt;
|
||||
// 请求局部恢复钉选在此被一次性消费,只绑定本次 replay attempt;
|
||||
// 外部 required_account 每次 attempt 都重新生效。
|
||||
let pinned_account = match &self.account_selection {
|
||||
AccountSelection::Diagnostic(account) => Some(account.clone()),
|
||||
AccountSelection::Scheduled(_) => self
|
||||
.recovery_account
|
||||
.take()
|
||||
.or_else(|| self.account_selection.required_account().cloned()),
|
||||
let (pinned_account, attempt_transport) = match &self.account_selection {
|
||||
AccountSelection::Diagnostic(account) => {
|
||||
(Some(account.clone()), AttemptTransport::Default)
|
||||
}
|
||||
AccountSelection::Scheduled(_) => {
|
||||
let transport_fallback_account = self.transport_fallback_account.take();
|
||||
let transport = if transport_fallback_account.is_some() {
|
||||
AttemptTransport::Fallback
|
||||
} else {
|
||||
AttemptTransport::Default
|
||||
};
|
||||
(
|
||||
transport_fallback_account
|
||||
.or_else(|| self.recovery_account.take())
|
||||
.or_else(|| self.account_selection.required_account().cloned()),
|
||||
transport,
|
||||
)
|
||||
}
|
||||
};
|
||||
let account_context = match &self.account_selection {
|
||||
AccountSelection::Diagnostic(account) => AccountAttemptContext::diagnostic(
|
||||
@@ -646,7 +662,8 @@ where
|
||||
self.continuation.clone(),
|
||||
self.cancellation.clone(),
|
||||
)
|
||||
.with_continuation_attempt(self.continuation_attempt);
|
||||
.with_continuation_attempt(self.continuation_attempt)
|
||||
.with_transport(attempt_transport);
|
||||
let trigger = if self.attempts == 0 {
|
||||
AttemptTrigger::Initial
|
||||
} else {
|
||||
@@ -1023,6 +1040,11 @@ where
|
||||
&& !self.delivery_pending
|
||||
&& error.allows_pre_delivery_retry()
|
||||
&& self.attempts < self.plan.max_attempts().get();
|
||||
let transport_fallback_retry = pre_delivery_retry
|
||||
&& matches!(
|
||||
error.pre_delivery_retry(),
|
||||
Some(crate::error::PreDeliveryRetry::SameAccountTransportFallback)
|
||||
);
|
||||
let ordinary_retry = self.account_selection.required_account().is_none()
|
||||
&& self.continuation_attempt == ContinuationAttempt::None
|
||||
&& self.downstream_committed_at.is_none()
|
||||
@@ -1053,6 +1075,9 @@ where
|
||||
// 只钉住紧随其后的 replay attempt;replay 再遇可重试错误时,
|
||||
// ordinary/continuation 重试门不受影响,仍可换号。
|
||||
self.recovery_account = Some(account);
|
||||
} else if transport_fallback_retry {
|
||||
self.transport_fallback_account =
|
||||
Some(current.metadata.provider_account_id().clone());
|
||||
} else if !continuation_retry {
|
||||
self.excluded_accounts
|
||||
.insert(current.metadata.provider_account_id().clone());
|
||||
|
||||
@@ -910,6 +910,90 @@ pub struct LoadedCredential {
|
||||
pub credential: PlaintextCredential,
|
||||
}
|
||||
|
||||
/// Provider 已计算好时间边界的有界 OAuth refresh 候选查询。
|
||||
///
|
||||
/// Store 只负责按持久事实筛选和稳定排序,不拥有提前量或恢复窗口语义。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ProviderRefreshQuery {
|
||||
provider: ProviderKind,
|
||||
refresh_due_before: SystemTime,
|
||||
force_due_before: SystemTime,
|
||||
observed_at: SystemTime,
|
||||
excluded_account_ids: Vec<ProviderAccountId>,
|
||||
limit: NonZeroU32,
|
||||
}
|
||||
|
||||
impl ProviderRefreshQuery {
|
||||
#[must_use]
|
||||
pub fn new(
|
||||
provider: ProviderKind,
|
||||
refresh_due_before: SystemTime,
|
||||
force_due_before: SystemTime,
|
||||
observed_at: SystemTime,
|
||||
excluded_account_ids: Vec<ProviderAccountId>,
|
||||
limit: NonZeroU32,
|
||||
) -> Self {
|
||||
Self {
|
||||
provider,
|
||||
refresh_due_before,
|
||||
force_due_before,
|
||||
observed_at,
|
||||
excluded_account_ids,
|
||||
limit,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn provider(&self) -> &ProviderKind {
|
||||
&self.provider
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn refresh_due_before(&self) -> SystemTime {
|
||||
self.refresh_due_before
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn force_due_before(&self) -> SystemTime {
|
||||
self.force_due_before
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn observed_at(&self) -> SystemTime {
|
||||
self.observed_at
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn excluded_account_ids(&self) -> &[ProviderAccountId] {
|
||||
&self.excluded_account_ids
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn limit(&self) -> NonZeroU32 {
|
||||
self.limit
|
||||
}
|
||||
|
||||
/// 判断内存 Store 中的账号是否满足与 PostgreSQL 相同的候选谓词。
|
||||
#[must_use]
|
||||
pub fn contains(&self, account: &ProviderAccount) -> bool {
|
||||
account.provider() == &self.provider
|
||||
&& account.enabled()
|
||||
&& account.has_refresh_token()
|
||||
&& matches!(
|
||||
account.credential_state(),
|
||||
CredentialState::Unknown | CredentialState::Ready
|
||||
)
|
||||
&& !self.excluded_account_ids.contains(account.id())
|
||||
&& account.access_token_expires_at().is_some_and(|expires_at| {
|
||||
expires_at <= self.force_due_before
|
||||
|| (expires_at <= self.refresh_due_before
|
||||
&& account
|
||||
.next_refresh_at()
|
||||
.is_none_or(|retry_at| retry_at <= self.observed_at))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Admin/Provider import 创建账号时的一次性明文输入。
|
||||
#[derive(Clone, PartialEq)]
|
||||
pub struct NewProviderAccount {
|
||||
@@ -1191,6 +1275,12 @@ pub trait ProviderAccountStore: Send + Sync {
|
||||
provider: &ProviderKind,
|
||||
) -> Result<Vec<ProviderAccount>, StoreError>;
|
||||
|
||||
/// 一次有界查询返回账号事实和 revision-fenced 明文 credential。
|
||||
async fn list_refresh_candidates(
|
||||
&self,
|
||||
query: ProviderRefreshQuery,
|
||||
) -> Result<Vec<LoadedCredential>, StoreError>;
|
||||
|
||||
async fn load_credential(
|
||||
&self,
|
||||
account: &ProviderAccountId,
|
||||
|
||||
@@ -106,6 +106,16 @@ pub enum AttemptTrigger {
|
||||
AccountRetry,
|
||||
}
|
||||
|
||||
/// Provider 可为当前 attempt 选择的请求局部传输档位。
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash)]
|
||||
pub enum AttemptTransport {
|
||||
/// 使用 Provider 的默认传输策略。
|
||||
#[default]
|
||||
Default,
|
||||
/// 使用 Provider 定义的备用传输。
|
||||
Fallback,
|
||||
}
|
||||
|
||||
impl AttemptTrigger {
|
||||
#[must_use]
|
||||
pub const fn as_str(self) -> &'static str {
|
||||
@@ -224,8 +234,8 @@ fn lock_unpoisoned<T>(mutex: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
|
||||
|
||||
/// 单次 attempt 的账号选择与账号绑定状态事实。
|
||||
///
|
||||
/// `required_account` 只用于管理端 connection test 等必须命中唯一账号的内部请求;
|
||||
/// 一旦设置,Core 与 Provider 都不得换号或切换 target。
|
||||
/// `required_account` 用于管理端 connection test,或请求局部的同账号恢复 attempt;
|
||||
/// 一旦设置,Provider 在本 attempt 内不得换号或切换 target。
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct AccountAttemptContext {
|
||||
excluded_accounts: BTreeSet<ProviderAccountId>,
|
||||
@@ -425,6 +435,7 @@ pub struct AttemptContext {
|
||||
account: AccountAttemptContext,
|
||||
continuation: Option<ContinuationBinding>,
|
||||
continuation_attempt: ContinuationAttempt,
|
||||
transport: AttemptTransport,
|
||||
cancellation: CancellationToken,
|
||||
}
|
||||
|
||||
@@ -452,6 +463,7 @@ impl AttemptContext {
|
||||
account,
|
||||
continuation,
|
||||
continuation_attempt,
|
||||
transport: AttemptTransport::Default,
|
||||
cancellation,
|
||||
}
|
||||
}
|
||||
@@ -466,6 +478,13 @@ impl AttemptContext {
|
||||
self
|
||||
}
|
||||
|
||||
/// 覆盖本次 attempt 的 Provider 传输档位。
|
||||
#[must_use]
|
||||
pub const fn with_transport(mut self, transport: AttemptTransport) -> Self {
|
||||
self.transport = transport;
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn request_id(&self) -> &ModelRequestId {
|
||||
self.request.request_id()
|
||||
@@ -503,7 +522,7 @@ impl AttemptContext {
|
||||
self.account.excluded_accounts()
|
||||
}
|
||||
|
||||
/// 管理端 connection test 等内部请求强制使用的唯一账号。
|
||||
/// 管理端诊断或请求局部恢复 attempt 强制使用的唯一账号。
|
||||
#[must_use]
|
||||
pub const fn required_account(&self) -> Option<&ProviderAccountId> {
|
||||
self.account.required_account()
|
||||
@@ -542,6 +561,12 @@ impl AttemptContext {
|
||||
self.continuation_attempt
|
||||
}
|
||||
|
||||
/// 返回本次 attempt 的 Provider 传输档位。
|
||||
#[must_use]
|
||||
pub const fn transport(&self) -> AttemptTransport {
|
||||
self.transport
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn cancellation(&self) -> &CancellationToken {
|
||||
&self.cancellation
|
||||
|
||||
@@ -197,6 +197,15 @@ pub enum ContinuationFailure {
|
||||
HistoryUnavailable,
|
||||
}
|
||||
|
||||
/// Provider 在下游尚未提交事件时请求的隐藏恢复方式。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
||||
pub enum PreDeliveryRetry {
|
||||
/// 排除本次账号,按普通调度策略重新选号。
|
||||
AccountRotation,
|
||||
/// 固定本次账号,并要求 Provider 使用备用传输。
|
||||
SameAccountTransportFallback,
|
||||
}
|
||||
|
||||
/// Adapter 已明确分类为非 bearer 的不透明上游值。
|
||||
///
|
||||
/// Core 不解释或校验内容,只通过自定义 [`Debug`] 避免它意外进入日志。协议 adapter
|
||||
@@ -360,7 +369,7 @@ pub struct ProviderError {
|
||||
retry_after: Option<Duration>,
|
||||
continuation_failure: Option<ContinuationFailure>,
|
||||
replay_safe: bool,
|
||||
pre_delivery_retry: bool,
|
||||
pre_delivery_retry: Option<PreDeliveryRetry>,
|
||||
credential_recovery_required: bool,
|
||||
retry_same_account: bool,
|
||||
sensitive_context_redacted: bool,
|
||||
@@ -390,7 +399,7 @@ impl ProviderError {
|
||||
retry_after: None,
|
||||
continuation_failure: None,
|
||||
replay_safe: false,
|
||||
pre_delivery_retry: false,
|
||||
pre_delivery_retry: None,
|
||||
credential_recovery_required: false,
|
||||
retry_same_account: false,
|
||||
sensitive_context_redacted: false,
|
||||
@@ -456,7 +465,17 @@ impl ProviderError {
|
||||
/// 指定账号或下游已经进入提交状态时必须忽略它。
|
||||
#[must_use]
|
||||
pub const fn with_pre_delivery_retry(mut self) -> Self {
|
||||
self.pre_delivery_retry = true;
|
||||
self.pre_delivery_retry = Some(PreDeliveryRetry::AccountRotation);
|
||||
self
|
||||
}
|
||||
|
||||
/// 允许 Core 在首个客户端事件前固定原账号并改用 Provider 备用传输。
|
||||
///
|
||||
/// 与普通换号恢复相同,Core 仍会拒绝 continuation、指定账号、预算耗尽或
|
||||
/// 下游已提交后的隐藏重放。
|
||||
#[must_use]
|
||||
pub const fn with_pre_delivery_transport_fallback(mut self) -> Self {
|
||||
self.pre_delivery_retry = Some(PreDeliveryRetry::SameAccountTransportFallback);
|
||||
self
|
||||
}
|
||||
|
||||
@@ -561,9 +580,15 @@ impl ProviderError {
|
||||
self.replay_safe
|
||||
}
|
||||
|
||||
/// 返回 Provider 是否允许在首个客户端事件前换号恢复传输失败。
|
||||
/// 返回 Provider 是否允许在首个客户端事件前隐藏恢复传输失败。
|
||||
#[must_use]
|
||||
pub const fn allows_pre_delivery_retry(&self) -> bool {
|
||||
self.pre_delivery_retry.is_some()
|
||||
}
|
||||
|
||||
/// 返回下游提交前由 Provider 选择的隐藏恢复方式。
|
||||
#[must_use]
|
||||
pub const fn pre_delivery_retry(&self) -> Option<PreDeliveryRetry> {
|
||||
self.pre_delivery_retry
|
||||
}
|
||||
|
||||
|
||||
@@ -23,9 +23,9 @@ use gateway_core::engine::provider::{
|
||||
ProviderRegistry, ProviderRequest, ProviderStream, UpstreamTransport,
|
||||
};
|
||||
use gateway_core::engine::{
|
||||
AttemptContext, AttemptCoordinator, AttemptRecord, CancellationToken, CommitRequirement,
|
||||
ContinuationAttempt, EngineError, ExecutionOutcome, ExecutionStore, GatewayEngine,
|
||||
IntermediateFailure, ModelRequestFinalization, ModelRequestId, NewModelRequest,
|
||||
AttemptContext, AttemptCoordinator, AttemptRecord, AttemptTransport, CancellationToken,
|
||||
CommitRequirement, ContinuationAttempt, EngineError, ExecutionOutcome, ExecutionStore,
|
||||
GatewayEngine, IntermediateFailure, ModelRequestFinalization, ModelRequestId, NewModelRequest,
|
||||
ProviderAttemptOutcome, RecoveryReport, UpstreamSendState,
|
||||
};
|
||||
use gateway_core::error::{
|
||||
@@ -2754,6 +2754,49 @@ fn pre_delivery_transport_failure_rotates_account_for_non_idempotent_generation(
|
||||
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Succeeded);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pre_delivery_transport_fallback_retries_the_same_account_with_fallback_transport() {
|
||||
let operation = generate_operation();
|
||||
let route_plan = plan(&operation);
|
||||
let (coordinator, store, provider) = coordinator(vec![
|
||||
Script::Stream {
|
||||
account_id: "acct_first",
|
||||
items: vec![Err(ProviderError::new(
|
||||
ProviderErrorKind::Transport,
|
||||
UpstreamSendState::Ambiguous,
|
||||
)
|
||||
.with_pre_delivery_transport_fallback())],
|
||||
},
|
||||
Script::Stream {
|
||||
account_id: "acct_first",
|
||||
items: complete_stream(None),
|
||||
},
|
||||
]);
|
||||
|
||||
let mut session = block_on(coordinator.start(
|
||||
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
|
||||
operation,
|
||||
route_plan,
|
||||
None,
|
||||
None,
|
||||
CancellationToken::new(),
|
||||
))
|
||||
.expect("start execution");
|
||||
block_on(session.collect_uncommitted()).expect("same-account fallback succeeds");
|
||||
block_on(session.commit_downstream(Some(200))).expect("commit winning response");
|
||||
|
||||
let original = ProviderAccountId::new("acct_first").expect("account id");
|
||||
let contexts = provider.contexts.lock().expect("contexts lock");
|
||||
assert_eq!(contexts.len(), 2);
|
||||
assert_eq!(contexts[0].transport(), AttemptTransport::Default);
|
||||
assert_eq!(contexts[1].transport(), AttemptTransport::Fallback);
|
||||
assert_eq!(contexts[1].required_account(), Some(&original));
|
||||
assert!(!contexts[1].excluded_accounts().contains(&original));
|
||||
let state = store.state.lock().expect("store lock");
|
||||
assert_eq!(state.intermediate_failures, 1);
|
||||
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Succeeded);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn final_capacity_exhaustion_returns_the_last_retryable_upstream_failure() {
|
||||
let operation = generate_operation();
|
||||
@@ -2890,7 +2933,7 @@ fn pre_delivery_retry_marker_is_ignored_after_downstream_commit() {
|
||||
))),
|
||||
Err(
|
||||
ProviderError::new(ProviderErrorKind::Transport, UpstreamSendState::Ambiguous)
|
||||
.with_pre_delivery_retry(),
|
||||
.with_pre_delivery_transport_fallback(),
|
||||
),
|
||||
],
|
||||
},
|
||||
|
||||
@@ -52,6 +52,10 @@ pub async fn initialize(mut config: StoreConfig) -> StoreResult<StoreBundle> {
|
||||
|
||||
config.validate_resolved()?;
|
||||
let pool = postgres::connect_and_migrate(&config.database_url()?, config.pool).await?;
|
||||
let observability_query_budget = postgres::ObservabilityQueryBudget::try_new(
|
||||
config.pool.observability_max_connections(),
|
||||
config.pool.acquire_timeout(),
|
||||
)?;
|
||||
let redis_client = ::redis::Client::open(config.redis_url()?)
|
||||
.map_err(|_| redis_unavailable("create Redis client"))?;
|
||||
let redis_connection = redis_client
|
||||
@@ -68,6 +72,10 @@ pub async fn initialize(mut config: StoreConfig) -> StoreResult<StoreBundle> {
|
||||
|
||||
let credential_leases =
|
||||
redis::RedisCredentialLeaseRepository::new(redis_connection.clone(), REDIS_NAMESPACE)?;
|
||||
let admin_account_runtime = Arc::new(redis::RedisAdminAccountRuntimeStore::new(
|
||||
cooldowns.as_ref().clone(),
|
||||
credential_leases.clone(),
|
||||
));
|
||||
let provider_leases = Arc::new(redis::RedisProviderLeaseCoordinator::new(
|
||||
credential_leases.clone(),
|
||||
));
|
||||
@@ -90,15 +98,15 @@ pub async fn initialize(mut config: StoreConfig) -> StoreResult<StoreBundle> {
|
||||
)?);
|
||||
|
||||
let admin_ports = AdminStorePorts::new(
|
||||
Arc::new(postgres::PgAdminAccountStore::new(
|
||||
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>,
|
||||
)),
|
||||
AdminAccountStorePorts::new(
|
||||
Arc::new(postgres::PgAdminAccountStore::new(
|
||||
pool.clone(),
|
||||
Some(Arc::clone(&cooldowns) as Arc<dyn ProviderCooldownPort>),
|
||||
observability_query_budget.clone(),
|
||||
)),
|
||||
admin_account_runtime,
|
||||
Arc::new(postgres::PgAccountGroupRepository::new(pool.clone())),
|
||||
),
|
||||
Arc::new(AdminAuthStoreAdapter {
|
||||
security: postgres::PgAdminSecurityAuditRepository::new(pool.clone()),
|
||||
settings: postgres::PgRuntimeSettingsRepository::new(pool.clone()),
|
||||
@@ -112,6 +120,7 @@ pub async fn initialize(mut config: StoreConfig) -> StoreResult<StoreBundle> {
|
||||
pool.clone(),
|
||||
Some(credential_leases.clone()),
|
||||
Some(Arc::clone(&cooldowns) as Arc<dyn ProviderCooldownPort>),
|
||||
observability_query_budget,
|
||||
)),
|
||||
Arc::new(AdminSettingsStoreAdapter {
|
||||
control_plane: postgres::PgControlPlaneRepository::new(pool.clone()),
|
||||
@@ -123,6 +132,8 @@ pub async fn initialize(mut config: StoreConfig) -> StoreResult<StoreBundle> {
|
||||
let (execution, execution_writer) =
|
||||
postgres::BufferedExecutionStore::new(Arc::clone(&execution_repository));
|
||||
let execution = Arc::new(execution);
|
||||
let (client_key_usage, client_key_usage_writer) =
|
||||
postgres::PgClientApiKeyUsageSink::new(pool.clone());
|
||||
let retention = Arc::new(postgres::PgRetentionRepository::new(pool.clone()));
|
||||
let admissions: Arc<dyn gateway_core::engine::admission::ClientAdmissionPort> = Arc::new(
|
||||
redis::RedisClientAdmissionRepository::new(redis_connection.clone(), REDIS_NAMESPACE)?,
|
||||
@@ -159,7 +170,7 @@ pub async fn initialize(mut config: StoreConfig) -> StoreResult<StoreBundle> {
|
||||
REDIS_NAMESPACE,
|
||||
)?),
|
||||
),
|
||||
Arc::new(postgres::PgClientApiKeyUsageSink::new(pool.clone())),
|
||||
Arc::new(client_key_usage),
|
||||
);
|
||||
|
||||
let provider_ports = ProviderStorePorts::new(
|
||||
@@ -181,7 +192,10 @@ pub async fn initialize(mut config: StoreConfig) -> StoreResult<StoreBundle> {
|
||||
credential_leases,
|
||||
));
|
||||
let health_probes: Vec<Arc<dyn HealthProbe>> = vec![
|
||||
Arc::new(PostgresHealthProbe { pool: pool.clone() }),
|
||||
Arc::new(PostgresHealthProbe::new(
|
||||
pool.clone(),
|
||||
config.pool.max_connections,
|
||||
)),
|
||||
Arc::new(RedisHealthProbe {
|
||||
connection: redis_connection,
|
||||
}),
|
||||
@@ -189,6 +203,7 @@ pub async fn initialize(mut config: StoreConfig) -> StoreResult<StoreBundle> {
|
||||
let worker_contributions = store_worker_contributions(
|
||||
execution_repository,
|
||||
execution_writer,
|
||||
client_key_usage_writer,
|
||||
admission_release_writer,
|
||||
circuit_feedback_writer,
|
||||
retention,
|
||||
|
||||
@@ -1,6 +1,9 @@
|
||||
//! Store 启动配置、环境变量解析与校验。
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::{
|
||||
path::{Path, PathBuf},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use super::*;
|
||||
|
||||
@@ -8,6 +11,11 @@ pub(crate) const DATABASE_URL_ENV: &str = "CPR_DATABASE_URL";
|
||||
pub(crate) const REDIS_URL_ENV: &str = "CPR_REDIS_URL";
|
||||
pub(crate) const DATABASE_PASSWORD_ENV: &str = "CPR_DATABASE_PASSWORD";
|
||||
pub(crate) const REDIS_PASSWORD_ENV: &str = "CPR_REDIS_PASSWORD";
|
||||
pub(crate) const POSTGRES_STATEMENT_TIMEOUT: Duration = Duration::from_secs(30);
|
||||
pub(crate) const POSTGRES_LOCK_TIMEOUT: Duration = Duration::from_secs(5);
|
||||
pub(crate) const POSTGRES_IDLE_TRANSACTION_TIMEOUT: Duration = Duration::from_secs(30);
|
||||
pub(crate) const POSTGRES_HEALTH_ATTEMPT_TIMEOUT: Duration = Duration::from_millis(450);
|
||||
pub(crate) const POSTGRES_HEALTH_RETRY_DELAY: Duration = Duration::from_millis(50);
|
||||
|
||||
/// Store 自己拥有并校验的启动配置。
|
||||
#[derive(Clone, Deserialize)]
|
||||
@@ -32,22 +40,34 @@ pub struct StorePoolConfig {
|
||||
impl Default for StorePoolConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
max_connections: 10,
|
||||
max_connections: 20,
|
||||
acquire_timeout_seconds: 5,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl StorePoolConfig {
|
||||
fn validate(&self) -> StoreResult<()> {
|
||||
if self.max_connections == 0 || self.acquire_timeout_seconds == 0 {
|
||||
pub(crate) fn validate(&self) -> StoreResult<()> {
|
||||
if self.max_connections < 2 || self.acquire_timeout_seconds == 0 {
|
||||
return Err(StoreError::InvalidData {
|
||||
entity: "store config",
|
||||
message: "pool limits must be positive".to_owned(),
|
||||
message:
|
||||
"pool.max_connections must be at least 2 and acquire timeout must be positive"
|
||||
.to_owned(),
|
||||
});
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 管理观测查询可并发占用的连接数;始终为数据面保留约 20% 的池容量。
|
||||
#[must_use]
|
||||
pub const fn observability_max_connections(self) -> u32 {
|
||||
self.max_connections - self.max_connections.div_ceil(5)
|
||||
}
|
||||
|
||||
pub(crate) const fn acquire_timeout(self) -> Duration {
|
||||
Duration::from_secs(self.acquire_timeout_seconds)
|
||||
}
|
||||
}
|
||||
|
||||
impl StoreConfig {
|
||||
@@ -123,6 +143,7 @@ impl fmt::Debug for StoreConfig {
|
||||
.debug_struct("StoreConfig")
|
||||
.field("database", &"[REDACTED]")
|
||||
.field("redis", &"[REDACTED]")
|
||||
.field("pool", &self.pool)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -15,8 +15,8 @@ use gateway_admin::model::settings::{
|
||||
use gateway_admin::model::{MutationActor, MutationContext, Revision as AdminRevision};
|
||||
use gateway_admin::ports::backup::BackupStorePorts;
|
||||
use gateway_admin::ports::store::{
|
||||
AdminStoreError, AdminStoreErrorKind, AdminStorePorts, AdminStoreResult, AuthStore,
|
||||
SettingsStore,
|
||||
AdminAccountStorePorts, AdminStoreError, AdminStoreErrorKind, AdminStorePorts,
|
||||
AdminStoreResult, AuthStore, SettingsStore,
|
||||
};
|
||||
use gateway_core::CoreStorePorts;
|
||||
use gateway_core::health::{HealthProbe, HealthState};
|
||||
@@ -43,4 +43,5 @@ pub(crate) use admin_adapter::*;
|
||||
pub use bundle::*;
|
||||
pub use config::*;
|
||||
pub use value::*;
|
||||
pub use workers::PostgresHealthProbe;
|
||||
pub(crate) use workers::*;
|
||||
|
||||
@@ -1,10 +1,6 @@
|
||||
//! PostgreSQL owner for provider-neutral account groups and memberships.
|
||||
|
||||
use std::{
|
||||
collections::{BTreeMap, BTreeSet},
|
||||
str::FromStr as _,
|
||||
sync::Arc,
|
||||
};
|
||||
use std::{collections::BTreeMap, str::FromStr as _};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use futures::future::BoxFuture;
|
||||
@@ -13,30 +9,26 @@ use gateway_admin::{
|
||||
MutationContext,
|
||||
account_groups::{
|
||||
AccountGroupAccountSummary, AccountGroupCapacity, AccountGroupColor,
|
||||
AccountGroupListQuery, AccountGroupMutation, AccountGroupPage, AccountGroupRecord,
|
||||
AccountGroupUsage, DeleteAccountGroup, NewAccountGroup, SetAccountGroupEnabled,
|
||||
UpdateAccountGroup,
|
||||
AccountGroupListQuery, AccountGroupMemberFact, 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 gateway_core::{engine::credential::AccountStatusFacts, 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,
|
||||
PgRuntimeSettingsRepository, RuntimeSettingsRepository, account_summary_from_row,
|
||||
append_admin_audit_event_in_transaction, bump_config_revision_in_transaction,
|
||||
completed_usage_fact_predicate,
|
||||
};
|
||||
|
||||
const ENTITY: &str = "account group";
|
||||
@@ -45,31 +37,12 @@ const ENTITY: &str = "account group";
|
||||
#[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),
|
||||
}
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
async fn current_revision(&self) -> AdminStoreResult<gateway_admin::model::Revision> {
|
||||
@@ -82,26 +55,17 @@ impl PgAccountGroupRepository {
|
||||
}
|
||||
|
||||
async fn required_record(&self, id: &AccountGroupId) -> AdminStoreResult<AccountGroupRecord> {
|
||||
let record = load_record(&self.pool, id.as_str())
|
||||
let mut 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))
|
||||
let costs = group_costs(&self.pool, &[id.as_str().to_owned()])
|
||||
.await
|
||||
.map_err(|error| admin_store_error(ENTITY, error))?;
|
||||
if let Some(usage) = costs.get(id.as_str()) {
|
||||
record.usage = usage.clone();
|
||||
}
|
||||
Ok(record)
|
||||
}
|
||||
|
||||
async fn mutate<F>(
|
||||
@@ -167,7 +131,18 @@ impl AccountGroupStore for PgAccountGroupRepository {
|
||||
.map(group_record)
|
||||
.collect::<StoreResult<Vec<_>>>()
|
||||
.map_err(|error| admin_store_error(ENTITY, error))?;
|
||||
self.enrich_records(&mut items).await?;
|
||||
let group_ids = items
|
||||
.iter()
|
||||
.map(|record| record.id.as_str().to_owned())
|
||||
.collect::<Vec<_>>();
|
||||
let costs = group_costs(&self.pool, &group_ids)
|
||||
.await
|
||||
.map_err(|error| admin_store_error(ENTITY, error))?;
|
||||
for record in &mut items {
|
||||
if let Some(usage) = costs.get(record.id.as_str()) {
|
||||
record.usage = usage.clone();
|
||||
}
|
||||
}
|
||||
Ok(AccountGroupPage {
|
||||
config_revision: self.current_revision().await?,
|
||||
items,
|
||||
@@ -177,6 +152,75 @@ impl AccountGroupStore for PgAccountGroupRepository {
|
||||
})
|
||||
}
|
||||
|
||||
async fn load_account_group_members(
|
||||
&self,
|
||||
group_ids: &[AccountGroupId],
|
||||
) -> AdminStoreResult<Vec<AccountGroupMemberFact>> {
|
||||
if group_ids.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let group_ids = group_ids
|
||||
.iter()
|
||||
.map(|group_id| group_id.as_str().to_owned())
|
||||
.collect::<Vec<_>>();
|
||||
let rows = sqlx::query(
|
||||
"select membership.account_group_id,
|
||||
account.id, account.provider_kind, account.name, account.email,
|
||||
account.upstream_user_id, account.upstream_account_id, account.plan_type,
|
||||
account.authentication_kind, account.credential_revision,
|
||||
account.has_refresh_token, account.access_token_expires_at,
|
||||
account.next_refresh_at, account.enabled, account.concurrency_limit,
|
||||
account.weight, account.credential_state, account.quota_access_state,
|
||||
account.quota_evidence, account.quota_access_observed_at,
|
||||
account.quota_reset_at, account.last_error_reason,
|
||||
account.last_error_message, account.credential_observed_at,
|
||||
account.created_at, account.updated_at,
|
||||
settings.max_concurrent_per_account
|
||||
from account_group_accounts membership
|
||||
join provider_accounts account on account.id = membership.provider_account_id
|
||||
cross join runtime_settings settings
|
||||
where settings.id = 1
|
||||
and membership.account_group_id = any($1::text[])
|
||||
order by membership.account_group_id, membership.provider_account_id",
|
||||
)
|
||||
.bind(group_ids)
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_err(|_| admin_store_error(ENTITY, unavailable("load account group members")))?;
|
||||
rows.into_iter()
|
||||
.map(|row| {
|
||||
let group_id = AccountGroupId::new(
|
||||
row.try_get::<String, _>("account_group_id")
|
||||
.map_err(|_| invalid("invalid group ID"))?,
|
||||
)
|
||||
.map_err(|_| invalid("invalid group ID"))?;
|
||||
let default_slots = u64::try_from(
|
||||
row.try_get::<i64, _>("max_concurrent_per_account")
|
||||
.map_err(|_| invalid("invalid default account concurrency"))?,
|
||||
)
|
||||
.map_err(|_| invalid("invalid default account concurrency"))?;
|
||||
let account = account_summary_from_row(row)?;
|
||||
Ok(AccountGroupMemberFact {
|
||||
group_id,
|
||||
account_id: account.id,
|
||||
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: None,
|
||||
last_error_reason: account.last_error_reason,
|
||||
last_error_message: account.last_error_message,
|
||||
},
|
||||
total_slots: account
|
||||
.concurrency_limit
|
||||
.map_or(default_slots, |limit| u64::from(limit.get())),
|
||||
})
|
||||
})
|
||||
.collect::<StoreResult<Vec<_>>>()
|
||||
.map_err(|error| admin_store_error(ENTITY, error))
|
||||
}
|
||||
|
||||
async fn create_account_group(
|
||||
&self,
|
||||
command: NewAccountGroup,
|
||||
@@ -450,7 +494,8 @@ fn group_record(row: &sqlx::postgres::PgRow) -> StoreResult<AccountGroupRecord>
|
||||
},
|
||||
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"))?,
|
||||
retained_total_usd: DecimalAmount::from_str("0")
|
||||
.map_err(|_| invalid("invalid zero cost"))?,
|
||||
},
|
||||
created_at: row
|
||||
.try_get("created_at")
|
||||
@@ -461,182 +506,59 @@ fn group_record(row: &sqlx::postgres::PgRow) -> StoreResult<AccountGroupRecord>
|
||||
})
|
||||
}
|
||||
|
||||
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,
|
||||
if group_ids.is_empty() {
|
||||
return Ok(BTreeMap::new());
|
||||
}
|
||||
let completed_usage = completed_usage_fact_predicate("mr");
|
||||
let statement = format!(
|
||||
"with requested_groups(group_id) as (
|
||||
select unnest($1::text[])
|
||||
)
|
||||
select requested_groups.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
|
||||
coalesce(sum(mr.cost_amount), 0)::text as retained_total_usd
|
||||
from requested_groups
|
||||
cross join runtime_settings settings
|
||||
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
|
||||
on mr.routing_group_refs @> array[requested_groups.group_id]::text[]
|
||||
and mr.started_at >= now() - make_interval(days => settings.usage_retention_days::int)
|
||||
and {completed_usage}
|
||||
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"))?;
|
||||
where settings.id = 1
|
||||
group by requested_groups.group_id"
|
||||
);
|
||||
// 动态片段仅为共享的固定 usage-fact predicate;group IDs 仍使用 bind。
|
||||
let rows = sqlx::query(sqlx::AssertSqlSafe(statement))
|
||||
.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")
|
||||
let group_id: String = row
|
||||
.try_get("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"))?;
|
||||
let retained_total = row
|
||||
.try_get::<String, _>("retained_total_usd")
|
||||
.map_err(|_| invalid("invalid retained 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"))?,
|
||||
retained_total_usd: DecimalAmount::from_str(&retained_total)
|
||||
.map_err(|_| invalid("invalid retained total cost"))?,
|
||||
},
|
||||
))
|
||||
})
|
||||
|
||||
@@ -3,15 +3,13 @@
|
||||
use std::{
|
||||
collections::{BTreeMap, BTreeSet},
|
||||
fmt,
|
||||
sync::{
|
||||
Arc, Mutex,
|
||||
atomic::{AtomicBool, Ordering},
|
||||
},
|
||||
sync::{Arc, Mutex},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, Utc};
|
||||
use futures::future::BoxFuture;
|
||||
use gateway_admin::{
|
||||
model::{
|
||||
MutationContext,
|
||||
@@ -29,12 +27,14 @@ use gateway_admin::{
|
||||
ports::store::{AdminStoreResult, ClientKeyStore},
|
||||
};
|
||||
use gateway_core::{
|
||||
engine::execution::ClientApiKeyUsageSink,
|
||||
engine::{CancellationToken, execution::ClientApiKeyUsageSink},
|
||||
policy::{ClientApiKeyId, PlaintextClientApiKey, RateLimits},
|
||||
routing::{AccountGroupId, ProviderKind},
|
||||
task::{DaemonTask, WorkerTaskError},
|
||||
};
|
||||
use serde::Deserialize;
|
||||
use sqlx::{PgPool, Postgres, QueryBuilder, Transaction};
|
||||
use tokio::sync::Notify;
|
||||
|
||||
use crate::{
|
||||
StoreError, StoreResult, admin_revision, admin_store_error, mutation_audit,
|
||||
@@ -46,6 +46,7 @@ use super::{ControlPlaneRepository, PgControlPlaneRepository};
|
||||
const ENTITY: &str = "client API key";
|
||||
const KEY_LENGTH: usize = 46;
|
||||
const CLIENT_API_KEY_LAST_USED_FLUSH_DELAY: Duration = Duration::from_secs(1);
|
||||
const CLIENT_API_KEY_LAST_USED_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(2);
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ClientApiKeySnapshot {
|
||||
@@ -349,35 +350,8 @@ impl ClientApiKeyRepository for PgClientApiKeyRepository {
|
||||
let mut statement = QueryBuilder::<Postgres>::new(
|
||||
"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
|
||||
k.updated_at, '[]'::jsonb as groups, '{}'::text[] 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());
|
||||
@@ -407,6 +381,7 @@ impl ClientApiKeyRepository for PgClientApiKeyRepository {
|
||||
} else {
|
||||
None
|
||||
};
|
||||
load_client_key_memberships(&self.pool, &mut items).await?;
|
||||
Ok(ClientApiKeyPage {
|
||||
items,
|
||||
total,
|
||||
@@ -501,76 +476,102 @@ impl ClientApiKeyRepository for PgClientApiKeyRepository {
|
||||
/// 该 adapter 仅记录稳定 Key ID;认证材料从不进入异步队列或日志。
|
||||
#[derive(Clone)]
|
||||
pub struct PgClientApiKeyUsageSink {
|
||||
state: Arc<ClientApiKeyUsageBuffer>,
|
||||
}
|
||||
|
||||
struct ClientApiKeyUsageBuffer {
|
||||
pending: Mutex<BTreeMap<String, DateTime<Utc>>>,
|
||||
flush_requested: Notify,
|
||||
}
|
||||
|
||||
pub struct PgClientApiKeyUsageWriter {
|
||||
repository: PgClientApiKeyRepository,
|
||||
pending: Arc<Mutex<BTreeMap<String, DateTime<Utc>>>>,
|
||||
flush_scheduled: Arc<AtomicBool>,
|
||||
state: Arc<ClientApiKeyUsageBuffer>,
|
||||
flush_delay: Duration,
|
||||
}
|
||||
|
||||
impl PgClientApiKeyUsageSink {
|
||||
#[must_use]
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
Self {
|
||||
repository: PgClientApiKeyRepository::new(pool),
|
||||
pending: Arc::new(Mutex::new(BTreeMap::new())),
|
||||
flush_scheduled: Arc::new(AtomicBool::new(false)),
|
||||
}
|
||||
pub fn new(pool: PgPool) -> (Self, PgClientApiKeyUsageWriter) {
|
||||
Self::with_flush_delay(pool, CLIENT_API_KEY_LAST_USED_FLUSH_DELAY)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_flush_delay(
|
||||
pool: PgPool,
|
||||
flush_delay: Duration,
|
||||
) -> (Self, PgClientApiKeyUsageWriter) {
|
||||
let state = Arc::new(ClientApiKeyUsageBuffer {
|
||||
pending: Mutex::new(BTreeMap::new()),
|
||||
flush_requested: Notify::new(),
|
||||
});
|
||||
(
|
||||
Self {
|
||||
state: Arc::clone(&state),
|
||||
},
|
||||
PgClientApiKeyUsageWriter {
|
||||
repository: PgClientApiKeyRepository::new(pool),
|
||||
state,
|
||||
flush_delay: flush_delay.max(Duration::from_millis(1)),
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
fn queue(&self, key_id: &ClientApiKeyId) {
|
||||
self.pending
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
.insert(key_id.as_str().to_owned(), Utc::now());
|
||||
self.schedule_flush();
|
||||
let used_at = Utc::now();
|
||||
let mut pending = lock_unpoisoned(&self.state.pending);
|
||||
pending
|
||||
.entry(key_id.as_str().to_owned())
|
||||
.and_modify(|pending_at| *pending_at = (*pending_at).max(used_at))
|
||||
.or_insert(used_at);
|
||||
drop(pending);
|
||||
self.state.flush_requested.notify_one();
|
||||
}
|
||||
}
|
||||
|
||||
fn schedule_flush(&self) {
|
||||
if self
|
||||
.flush_scheduled
|
||||
.compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
|
||||
.is_err()
|
||||
{
|
||||
return;
|
||||
impl PgClientApiKeyUsageWriter {
|
||||
async fn flush_pending(&self) -> StoreResult<u64> {
|
||||
let updates = std::mem::take(&mut *lock_unpoisoned(&self.state.pending));
|
||||
if updates.is_empty() {
|
||||
return Ok(0);
|
||||
}
|
||||
let Ok(runtime) = tokio::runtime::Handle::try_current() else {
|
||||
self.flush_scheduled.store(false, Ordering::Release);
|
||||
return;
|
||||
};
|
||||
let sink = self.clone();
|
||||
drop(runtime.spawn(async move {
|
||||
tokio::time::sleep(CLIENT_API_KEY_LAST_USED_FLUSH_DELAY).await;
|
||||
sink.flush_pending().await;
|
||||
}));
|
||||
}
|
||||
|
||||
async fn flush_pending(&self) {
|
||||
let updates = std::mem::take(
|
||||
&mut *self
|
||||
.pending
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner),
|
||||
);
|
||||
if let Err(error) = self.repository.touch_client_api_keys(&updates).await {
|
||||
tracing::error!(error = %error, "Failed to flush client API key last-used batch");
|
||||
let mut pending = self
|
||||
.pending
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
for (key_id, used_at) in updates {
|
||||
pending
|
||||
.entry(key_id)
|
||||
.and_modify(|pending_at| *pending_at = (*pending_at).max(used_at))
|
||||
.or_insert(used_at);
|
||||
match self.repository.touch_client_api_keys(&updates).await {
|
||||
Ok(updated) => Ok(updated),
|
||||
Err(error) => {
|
||||
let mut pending = lock_unpoisoned(&self.state.pending);
|
||||
for (key_id, used_at) in updates {
|
||||
pending
|
||||
.entry(key_id)
|
||||
.and_modify(|pending_at| *pending_at = (*pending_at).max(used_at))
|
||||
.or_insert(used_at);
|
||||
}
|
||||
drop(pending);
|
||||
self.state.flush_requested.notify_one();
|
||||
Err(error)
|
||||
}
|
||||
}
|
||||
self.flush_scheduled.store(false, Ordering::Release);
|
||||
if !self
|
||||
.pending
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
.is_empty()
|
||||
}
|
||||
|
||||
async fn flush_on_shutdown(&self) {
|
||||
match tokio::time::timeout(
|
||||
CLIENT_API_KEY_LAST_USED_SHUTDOWN_TIMEOUT,
|
||||
self.flush_pending(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
self.schedule_flush();
|
||||
Ok(Ok(updated)) if updated > 0 => {
|
||||
tracing::info!(updated, "Client API Key last-used 已在关闭前写回");
|
||||
}
|
||||
Ok(Ok(_)) => {}
|
||||
Ok(Err(_)) => {
|
||||
tracing::warn!("Client API Key last-used 关闭写回失败");
|
||||
}
|
||||
Err(_) => {
|
||||
tracing::warn!(
|
||||
pending = lock_unpoisoned(&self.state.pending).len(),
|
||||
"Client API Key last-used 关闭写回超时"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -581,6 +582,41 @@ impl ClientApiKeyUsageSink for PgClientApiKeyUsageSink {
|
||||
}
|
||||
}
|
||||
|
||||
impl DaemonTask for PgClientApiKeyUsageWriter {
|
||||
fn run(&self, cancellation: CancellationToken) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
|
||||
Box::pin(async move {
|
||||
loop {
|
||||
tokio::select! {
|
||||
() = cancellation.cancelled() => {
|
||||
self.flush_on_shutdown().await;
|
||||
return Ok(());
|
||||
}
|
||||
() = self.state.flush_requested.notified() => {}
|
||||
}
|
||||
tokio::select! {
|
||||
() = cancellation.cancelled() => {
|
||||
self.flush_on_shutdown().await;
|
||||
return Ok(());
|
||||
}
|
||||
() = tokio::time::sleep(self.flush_delay) => {}
|
||||
}
|
||||
if self.flush_pending().await.is_err() {
|
||||
tracing::warn!("Client API Key last-used 批量写回失败");
|
||||
return Err(WorkerTaskError::safe(
|
||||
"client API key last-used flush failed",
|
||||
));
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn lock_unpoisoned<T>(mutex: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
|
||||
mutex
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
}
|
||||
|
||||
/// Admin 用例所需的 Client Key 事务能力。
|
||||
#[derive(Clone)]
|
||||
pub struct PgAdminClientKeyStore {
|
||||
@@ -1131,16 +1167,124 @@ async fn count_client_api_keys(pool: &PgPool, search: Option<&str>) -> StoreResu
|
||||
|
||||
fn push_client_key_search(statement: &mut QueryBuilder<Postgres>, search: Option<&str>) {
|
||||
if let Some(search) = search {
|
||||
let prefix = literal_prefix_pattern(search);
|
||||
statement.push(" and (lower(name) like ");
|
||||
statement.push_bind(format!("%{}%", search.to_lowercase()));
|
||||
statement.push_bind(prefix.clone());
|
||||
statement.push(" escape '\\'");
|
||||
statement.push(" or lower(coalesce(label, '')) like ");
|
||||
statement.push_bind(format!("%{}%", search.to_lowercase()));
|
||||
statement.push_bind(prefix.clone());
|
||||
statement.push(" escape '\\'");
|
||||
statement.push(" or lower(left(key, 10)) like ");
|
||||
statement.push_bind(format!("%{}%", search.to_lowercase()));
|
||||
statement.push_bind(prefix);
|
||||
statement.push(" escape '\\'");
|
||||
statement.push(")");
|
||||
}
|
||||
}
|
||||
|
||||
async fn load_client_key_memberships(
|
||||
pool: &PgPool,
|
||||
records: &mut [ClientApiKeyRecord],
|
||||
) -> StoreResult<()> {
|
||||
if records.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
let ids = records
|
||||
.iter()
|
||||
.map(|record| record.id.clone())
|
||||
.collect::<Vec<_>>();
|
||||
let rows = sqlx::query(
|
||||
"with requested_keys(key_id) as (
|
||||
select unnest($1::text[])
|
||||
),
|
||||
global_providers as (
|
||||
select coalesce(array_agg(distinct provider_kind order by provider_kind), '{}')
|
||||
as provider_kinds
|
||||
from provider_accounts
|
||||
)
|
||||
select requested_keys.key_id,
|
||||
groups.id as group_id, groups.name as group_name, groups.color as group_color,
|
||||
groups.enabled as group_enabled,
|
||||
global_providers.provider_kinds as global_provider_kinds,
|
||||
coalesce(array_agg(distinct accounts.provider_kind order by accounts.provider_kind)
|
||||
filter (where accounts.provider_kind is not null), '{}')
|
||||
as group_provider_kinds
|
||||
from requested_keys
|
||||
cross join global_providers
|
||||
left join client_api_key_groups bindings
|
||||
on bindings.client_api_key_id = requested_keys.key_id
|
||||
left join account_groups groups on groups.id = bindings.account_group_id
|
||||
left join account_group_accounts memberships
|
||||
on memberships.account_group_id = groups.id and groups.enabled
|
||||
left join provider_accounts accounts
|
||||
on accounts.id = memberships.provider_account_id
|
||||
group by requested_keys.key_id, global_providers.provider_kinds,
|
||||
groups.id, groups.name, groups.color, groups.enabled
|
||||
order by requested_keys.key_id, groups.id",
|
||||
)
|
||||
.bind(ids)
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("load client API key memberships"))?;
|
||||
let mut groups = BTreeMap::<String, Vec<ClientApiKeyGroupRecord>>::new();
|
||||
let mut providers = BTreeMap::<String, BTreeSet<String>>::new();
|
||||
for row in rows {
|
||||
use sqlx::Row as _;
|
||||
let key_id: String = row
|
||||
.try_get("key_id")
|
||||
.map_err(|_| invalid("invalid client API key membership"))?;
|
||||
let group_id: Option<String> = row
|
||||
.try_get("group_id")
|
||||
.map_err(|_| invalid("invalid client API key group"))?;
|
||||
if let Some(group_id) = group_id {
|
||||
groups
|
||||
.entry(key_id.clone())
|
||||
.or_default()
|
||||
.push(ClientApiKeyGroupRecord {
|
||||
id: group_id,
|
||||
name: row
|
||||
.try_get("group_name")
|
||||
.map_err(|_| invalid("invalid client API key group name"))?,
|
||||
color: row
|
||||
.try_get("group_color")
|
||||
.map_err(|_| invalid("invalid client API key group color"))?,
|
||||
enabled: row
|
||||
.try_get("group_enabled")
|
||||
.map_err(|_| invalid("invalid client API key group state"))?,
|
||||
});
|
||||
providers.entry(key_id).or_default().extend(
|
||||
row.try_get::<Vec<String>, _>("group_provider_kinds")
|
||||
.map_err(|_| invalid("invalid client API key provider kinds"))?,
|
||||
);
|
||||
} else {
|
||||
providers.entry(key_id).or_default().extend(
|
||||
row.try_get::<Vec<String>, _>("global_provider_kinds")
|
||||
.map_err(|_| invalid("invalid global provider kinds"))?,
|
||||
);
|
||||
}
|
||||
}
|
||||
for record in records {
|
||||
record.groups = groups.remove(&record.id).unwrap_or_default();
|
||||
record.provider_kinds = providers
|
||||
.remove(&record.id)
|
||||
.unwrap_or_default()
|
||||
.into_iter()
|
||||
.collect();
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn literal_prefix_pattern(value: &str) -> String {
|
||||
let mut escaped = String::with_capacity(value.len().saturating_add(1));
|
||||
for character in value.to_lowercase().chars() {
|
||||
if matches!(character, '\\' | '%' | '_') {
|
||||
escaped.push('\\');
|
||||
}
|
||||
escaped.push(character);
|
||||
}
|
||||
escaped.push('%');
|
||||
escaped
|
||||
}
|
||||
|
||||
fn push_client_key_cursor(statement: &mut QueryBuilder<Postgres>, cursor: &ClientApiKeyCursor) {
|
||||
let comparison = match cursor.sort.direction {
|
||||
ClientApiKeySortDirection::Asc => " > ",
|
||||
|
||||
@@ -451,16 +451,15 @@ impl ModelRequestRepository for PgExecutionStore {
|
||||
upstream_model_id, upstream_transport, http_version,
|
||||
attempt_count, upstream_send_state, account_selection_wait_ms,
|
||||
capacity_used_slots, capacity_total_slots
|
||||
) values (
|
||||
) select
|
||||
$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, $26,
|
||||
(select name from provider_accounts where id = $25),
|
||||
(select email from provider_accounts where id = $25),
|
||||
(select authentication_kind from provider_accounts where id = $25),
|
||||
account.name, account.email, account.authentication_kind,
|
||||
$27, $28, $29, 1, 'not_sent', $30, $31, $32
|
||||
)",
|
||||
from (values (true)) as seed(present)
|
||||
left join provider_accounts account on account.id = $25",
|
||||
)
|
||||
.bind(request.id)
|
||||
.bind(request.client_api_key_id)
|
||||
@@ -525,14 +524,11 @@ impl ModelRequestRepository for PgExecutionStore {
|
||||
set provider_kind = $2,
|
||||
provider_account_id = $3,
|
||||
provider_account_ref = $4,
|
||||
provider_account_name_snapshot = (
|
||||
select name from provider_accounts where id = $3
|
||||
),
|
||||
provider_account_email_snapshot = (
|
||||
select email from provider_accounts where id = $3
|
||||
),
|
||||
provider_account_authentication_kind_snapshot = (
|
||||
select authentication_kind from provider_accounts where id = $3
|
||||
(provider_account_name_snapshot,
|
||||
provider_account_email_snapshot,
|
||||
provider_account_authentication_kind_snapshot) = (
|
||||
select name, email, authentication_kind
|
||||
from provider_accounts where id = $3
|
||||
),
|
||||
upstream_model_id = $5,
|
||||
upstream_transport = $6,
|
||||
|
||||
@@ -1,10 +1,14 @@
|
||||
//! PostgreSQL 业务表的 adapters。
|
||||
|
||||
use async_trait::async_trait;
|
||||
use sqlx::{PgPool, postgres::PgPoolOptions};
|
||||
use sqlx::{
|
||||
PgPool,
|
||||
postgres::{PgConnectOptions, PgPoolOptions},
|
||||
};
|
||||
|
||||
use crate::{
|
||||
Revision, StoreBackend, StoreError, StorePoolConfig, StoreResult, postgres_unavailable,
|
||||
POSTGRES_IDLE_TRANSACTION_TIMEOUT, POSTGRES_LOCK_TIMEOUT, POSTGRES_STATEMENT_TIMEOUT, Revision,
|
||||
StoreBackend, StoreError, StorePoolConfig, StoreResult, postgres_unavailable,
|
||||
};
|
||||
|
||||
mod account_groups;
|
||||
@@ -20,6 +24,7 @@ mod provider_accounts;
|
||||
mod retention;
|
||||
mod runtime_settings;
|
||||
mod snapshot;
|
||||
mod usage_facts;
|
||||
|
||||
pub use account_groups::*;
|
||||
pub use admin_security_audit::*;
|
||||
@@ -34,6 +39,7 @@ pub use provider_accounts::*;
|
||||
pub use retention::*;
|
||||
pub use runtime_settings::*;
|
||||
pub use snapshot::*;
|
||||
pub(crate) use usage_facts::{completed_usage_fact_predicate, push_completed_usage_fact_filter};
|
||||
|
||||
static MIGRATOR: sqlx::migrate::Migrator = sqlx::migrate!("../../migrations");
|
||||
|
||||
@@ -45,24 +51,65 @@ pub async fn connect_and_migrate(
|
||||
if database_url.trim().is_empty() {
|
||||
return Err(postgres_unavailable("connect PostgreSQL"));
|
||||
}
|
||||
let pool = PgPoolOptions::new()
|
||||
.max_connections(pool_config.max_connections)
|
||||
.acquire_timeout(std::time::Duration::from_secs(
|
||||
pool_config.acquire_timeout_seconds,
|
||||
))
|
||||
.connect(database_url)
|
||||
pool_config.validate()?;
|
||||
let connect_options = database_url
|
||||
.parse::<PgConnectOptions>()
|
||||
.map_err(|_| postgres_unavailable("parse PostgreSQL connection options"))?;
|
||||
let migration_pool = PgPoolOptions::new()
|
||||
.max_connections(1)
|
||||
.connect_with(
|
||||
connect_options
|
||||
.clone()
|
||||
.application_name("codex-proxy-rs:migration"),
|
||||
)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("connect PostgreSQL"))?;
|
||||
if let Err(error) = MIGRATOR.run(&pool).await {
|
||||
pool.close().await;
|
||||
.map_err(|_| postgres_unavailable("connect PostgreSQL for migrations"))?;
|
||||
if let Err(error) = MIGRATOR.run(&migration_pool).await {
|
||||
migration_pool.close().await;
|
||||
return Err(StoreError::Unavailable {
|
||||
backend: StoreBackend::PostgreSql,
|
||||
message: format!("apply PostgreSQL migrations: {error}"),
|
||||
});
|
||||
}
|
||||
migration_pool.close().await;
|
||||
|
||||
let statement_timeout = postgres_duration_setting(POSTGRES_STATEMENT_TIMEOUT);
|
||||
let lock_timeout = postgres_duration_setting(POSTGRES_LOCK_TIMEOUT);
|
||||
let idle_in_transaction_session_timeout =
|
||||
postgres_duration_setting(POSTGRES_IDLE_TRANSACTION_TIMEOUT);
|
||||
let pool = PgPoolOptions::new()
|
||||
.max_connections(pool_config.max_connections)
|
||||
.acquire_timeout(std::time::Duration::from_secs(
|
||||
pool_config.acquire_timeout_seconds,
|
||||
))
|
||||
.after_connect(move |connection, _metadata| {
|
||||
let statement_timeout = statement_timeout.clone();
|
||||
let lock_timeout = lock_timeout.clone();
|
||||
let idle_in_transaction_session_timeout = idle_in_transaction_session_timeout.clone();
|
||||
Box::pin(async move {
|
||||
sqlx::query(
|
||||
"select set_config('statement_timeout', $1, false),
|
||||
set_config('lock_timeout', $2, false),
|
||||
set_config('idle_in_transaction_session_timeout', $3, false)",
|
||||
)
|
||||
.bind(statement_timeout)
|
||||
.bind(lock_timeout)
|
||||
.bind(idle_in_transaction_session_timeout)
|
||||
.execute(connection)
|
||||
.await?;
|
||||
Ok(())
|
||||
})
|
||||
})
|
||||
.connect_with(connect_options.application_name("codex-proxy-rs"))
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("connect PostgreSQL"))?;
|
||||
Ok(pool)
|
||||
}
|
||||
|
||||
fn postgres_duration_setting(duration: std::time::Duration) -> String {
|
||||
format!("{}ms", duration.as_millis())
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ControlPlaneSnapshot {
|
||||
pub settings: RuntimeSettings,
|
||||
@@ -131,21 +178,8 @@ impl PgControlPlaneRepository {
|
||||
#[async_trait]
|
||||
impl ControlPlaneRepository for PgControlPlaneRepository {
|
||||
async fn load_control_plane(&self) -> StoreResult<ControlPlaneSnapshot> {
|
||||
let mut transaction = self
|
||||
.pool
|
||||
.begin()
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("begin control plane snapshot"))?;
|
||||
sqlx::query("set transaction isolation level repeatable read read only")
|
||||
.execute(&mut *transaction)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("configure control plane snapshot"))?;
|
||||
let snapshot = load_control_plane_in_transaction(&mut transaction).await?;
|
||||
transaction
|
||||
.commit()
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("commit control plane snapshot"))?;
|
||||
Ok(snapshot)
|
||||
let settings = load_runtime_settings_from_pool(&self.pool).await?;
|
||||
Ok(ControlPlaneSnapshot { settings })
|
||||
}
|
||||
|
||||
async fn replace_control_plane(
|
||||
|
||||
@@ -16,12 +16,21 @@ use crate::redis::{CredentialLeaseRepository as _, RedisCredentialLeaseRepositor
|
||||
pub struct PgObservabilityRepository {
|
||||
pool: PgPool,
|
||||
cooldowns: Option<Arc<dyn ProviderCooldownPort>>,
|
||||
query_budget: ObservabilityQueryBudget,
|
||||
}
|
||||
|
||||
impl PgObservabilityRepository {
|
||||
#[must_use]
|
||||
pub const fn new(pool: PgPool, cooldowns: Option<Arc<dyn ProviderCooldownPort>>) -> Self {
|
||||
Self { pool, cooldowns }
|
||||
pub fn new(
|
||||
pool: PgPool,
|
||||
cooldowns: Option<Arc<dyn ProviderCooldownPort>>,
|
||||
query_budget: ObservabilityQueryBudget,
|
||||
) -> Self {
|
||||
Self {
|
||||
pool,
|
||||
cooldowns,
|
||||
query_budget,
|
||||
}
|
||||
}
|
||||
|
||||
/// 从账号事实和当前冷却一次性派生 Dashboard 五态;不在 SQL 中复制状态机。
|
||||
@@ -29,8 +38,13 @@ impl PgObservabilityRepository {
|
||||
&self,
|
||||
observed_at: DateTime<Utc>,
|
||||
) -> StoreResult<(ProviderAccountMetrics, Vec<ProviderAccountSummary>)> {
|
||||
let accounts = PgProviderAccountRepository::new(self.pool.clone())
|
||||
.list_provider_accounts(None, true)
|
||||
let repository = PgProviderAccountRepository::new(self.pool.clone());
|
||||
let accounts = self
|
||||
.query_budget
|
||||
.run(
|
||||
"load dashboard provider accounts",
|
||||
repository.list_provider_accounts(None, true),
|
||||
)
|
||||
.await?;
|
||||
let now = observed_at.into();
|
||||
let rate_limited_until =
|
||||
@@ -85,9 +99,10 @@ impl PgAdminObservabilityStore {
|
||||
pool: PgPool,
|
||||
runtime_signals: Option<RedisCredentialLeaseRepository>,
|
||||
cooldowns: Option<Arc<dyn ProviderCooldownPort>>,
|
||||
query_budget: ObservabilityQueryBudget,
|
||||
) -> Self {
|
||||
Self {
|
||||
repository: PgObservabilityRepository::new(pool, cooldowns),
|
||||
repository: PgObservabilityRepository::new(pool, cooldowns, query_budget),
|
||||
runtime_signals,
|
||||
}
|
||||
}
|
||||
@@ -110,26 +125,33 @@ impl ObservabilityRepository for PgObservabilityRepository {
|
||||
outcome: Some("succeeded".to_owned()),
|
||||
..UsageRecordFilter::default()
|
||||
},
|
||||
cursor: None,
|
||||
page: ObservabilityPageNumber::new(1)?,
|
||||
current_page: 1,
|
||||
page_size: ObservabilityPageSize::new(10)?,
|
||||
};
|
||||
// 每批最多三路并发,既缩短概览渲染路径,也不与数据面写路径争抢整个连接池。
|
||||
let (requests, attempts, totals) = futures::try_join!(
|
||||
request_metrics(&self.pool, range, &filter),
|
||||
attempt_metrics(&self.pool, range, &filter),
|
||||
dashboard_totals(&self.pool),
|
||||
// 每条 SQL 独立取一个全局观测槽位,避免整包预留造成队头阻塞。
|
||||
let (totals, (provider_accounts, _)) = futures::try_join!(
|
||||
self.query_budget.run(
|
||||
"load dashboard lifetime totals",
|
||||
dashboard_totals(&self.pool)
|
||||
),
|
||||
self.account_status_snapshot(observed_at),
|
||||
)?;
|
||||
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),
|
||||
async { Ok(list_usage_records(&self.pool, recent_query).await?.items) },
|
||||
self.query_budget.run(
|
||||
"load dashboard request trend",
|
||||
dashboard_request_metric_series(&self.pool, range, &filter),
|
||||
),
|
||||
self.query_budget.run(
|
||||
"load dashboard account usage",
|
||||
provider_account_usage(&self.pool, account_usage_query),
|
||||
),
|
||||
self.query_budget
|
||||
.run("load dashboard recent requests", async {
|
||||
list_usage_record_items(&self.pool, &recent_query).await
|
||||
}),
|
||||
)?;
|
||||
Ok(DashboardObservation {
|
||||
range,
|
||||
requests,
|
||||
attempts,
|
||||
totals,
|
||||
provider_accounts,
|
||||
trend,
|
||||
@@ -142,7 +164,12 @@ impl ObservabilityRepository for PgObservabilityRepository {
|
||||
&self,
|
||||
range: ObservabilityRange,
|
||||
) -> StoreResult<Vec<RequestMetricPoint>> {
|
||||
request_metric_series(&self.pool, range, &UsageRecordFilter::default()).await
|
||||
self.query_budget
|
||||
.run(
|
||||
"load dashboard request trend",
|
||||
dashboard_request_metric_series(&self.pool, range, &UsageRecordFilter::default()),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn usage_trend(
|
||||
@@ -150,7 +177,12 @@ impl ObservabilityRepository for PgObservabilityRepository {
|
||||
range: ObservabilityRange,
|
||||
filter: UsageRecordFilter,
|
||||
) -> StoreResult<Vec<RequestMetricPoint>> {
|
||||
request_metric_series(&self.pool, range, &filter).await
|
||||
self.query_budget
|
||||
.run(
|
||||
"load usage request trend",
|
||||
request_metric_series(&self.pool, range, &filter),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn usage_calculated_billing_facts(
|
||||
@@ -158,22 +190,39 @@ impl ObservabilityRepository for PgObservabilityRepository {
|
||||
range: ObservabilityRange,
|
||||
filter: UsageRecordFilter,
|
||||
) -> StoreResult<Vec<CalculatedUsageBillingFact>> {
|
||||
calculated_usage_billing_facts(&self.pool, range, &filter).await
|
||||
self.query_budget
|
||||
.run(
|
||||
"load calculated usage billing facts",
|
||||
calculated_usage_billing_facts(&self.pool, range, &filter),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn provider_account_usage(
|
||||
&self,
|
||||
query: ProviderAccountUsageQuery,
|
||||
) -> StoreResult<Vec<ProviderAccountUsageObservation>> {
|
||||
provider_account_usage(&self.pool, query).await
|
||||
self.query_budget
|
||||
.run(
|
||||
"load provider account usage",
|
||||
provider_account_usage(&self.pool, query),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn list_usage_records(&self, query: UsageRecordQuery) -> StoreResult<UsageRecordPage> {
|
||||
list_usage_records(&self.pool, query).await
|
||||
self.query_budget
|
||||
.run("list usage records", list_usage_records(&self.pool, query))
|
||||
.await
|
||||
}
|
||||
|
||||
async fn usage_record_detail(&self, request_id: &str) -> StoreResult<UsageRecordDetail> {
|
||||
usage_record_detail(&self.pool, request_id).await
|
||||
self.query_budget
|
||||
.run(
|
||||
"load usage record detail",
|
||||
usage_record_detail(&self.pool, request_id),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn usage_summary(
|
||||
@@ -182,9 +231,27 @@ impl ObservabilityRepository for PgObservabilityRepository {
|
||||
filter: UsageRecordFilter,
|
||||
) -> StoreResult<UsageOverview> {
|
||||
filter.validate()?;
|
||||
let requests = request_metrics(&self.pool, range, &filter).await?;
|
||||
let attempts = attempt_metrics(&self.pool, range, &filter).await?;
|
||||
let providers = provider_observations(&self.pool, range, &filter).await?;
|
||||
let requests = self
|
||||
.query_budget
|
||||
.run(
|
||||
"load usage request metrics",
|
||||
request_metrics(&self.pool, range, &filter),
|
||||
)
|
||||
.await?;
|
||||
let attempts = self
|
||||
.query_budget
|
||||
.run(
|
||||
"load usage attempt metrics",
|
||||
attempt_metrics(&self.pool, range, &filter),
|
||||
)
|
||||
.await?;
|
||||
let providers = self
|
||||
.query_budget
|
||||
.run(
|
||||
"load usage provider observations",
|
||||
provider_observations(&self.pool, range, &filter),
|
||||
)
|
||||
.await?;
|
||||
Ok(UsageOverview {
|
||||
range,
|
||||
requests,
|
||||
@@ -199,11 +266,18 @@ impl ObservabilityRepository for PgObservabilityRepository {
|
||||
filter: UsageRecordFilter,
|
||||
dimension: DiagnosticDimension,
|
||||
) -> StoreResult<Vec<DiagnosticObservation>> {
|
||||
usage_diagnostics(&self.pool, range, &filter, dimension).await
|
||||
self.query_budget
|
||||
.run(
|
||||
"load usage diagnostics",
|
||||
usage_diagnostics(&self.pool, range, &filter, dimension),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn list_ops_errors(&self, query: OpsErrorQuery) -> StoreResult<OpsErrorPage> {
|
||||
list_ops_errors(&self.pool, query).await
|
||||
self.query_budget
|
||||
.run("list ops errors", list_ops_errors(&self.pool, query))
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -30,20 +30,13 @@ pub(crate) fn store_request_outcome(outcome: admin_observability::RequestOutcome
|
||||
outcome.as_str().to_owned()
|
||||
}
|
||||
|
||||
pub(crate) fn store_cursor(
|
||||
cursor: admin_observability::ObservabilityCursor,
|
||||
) -> AdminStoreResult<ObservabilityCursor> {
|
||||
ObservabilityCursor::new(cursor.observed_at, cursor.stable_id).map_err(observability_error)
|
||||
}
|
||||
|
||||
pub(crate) fn store_usage_query(
|
||||
query: admin_observability::UsageQuery,
|
||||
) -> AdminStoreResult<UsageRecordQuery> {
|
||||
Ok(UsageRecordQuery {
|
||||
range: store_range(query.range)?,
|
||||
filter: store_usage_filter(query.filter),
|
||||
cursor: query.cursor.map(store_cursor).transpose()?,
|
||||
page: ObservabilityPageNumber::new(query.page.get()).map_err(observability_error)?,
|
||||
current_page: query.current_page,
|
||||
page_size: ObservabilityPageSize::new(query.page_size.get())
|
||||
.map_err(observability_error)?,
|
||||
})
|
||||
@@ -58,6 +51,7 @@ pub(crate) fn store_ops_error_filter(
|
||||
provider_account_ref: filter.provider_account_ref,
|
||||
provider_kind: filter.provider_kind,
|
||||
operation: filter.operation,
|
||||
endpoint: filter.endpoint,
|
||||
model: filter.model,
|
||||
transport: filter.transport,
|
||||
attempt_index: filter.attempt_index,
|
||||
@@ -75,8 +69,7 @@ pub(crate) fn store_ops_error_query(
|
||||
Ok(OpsErrorQuery {
|
||||
range: store_range(query.range)?,
|
||||
filter: store_ops_error_filter(query.filter),
|
||||
cursor: query.cursor.map(store_cursor).transpose()?,
|
||||
page: ObservabilityPageNumber::new(query.page.get()).map_err(observability_error)?,
|
||||
current_page: query.current_page,
|
||||
page_size: ObservabilityPageSize::new(query.page_size.get())
|
||||
.map_err(observability_error)?,
|
||||
})
|
||||
@@ -101,8 +94,6 @@ pub(crate) fn admin_dashboard_observation(
|
||||
) -> AdminStoreResult<admin_observability::DashboardObservation> {
|
||||
let DashboardObservation {
|
||||
range,
|
||||
requests,
|
||||
attempts,
|
||||
totals,
|
||||
provider_accounts,
|
||||
trend,
|
||||
@@ -111,8 +102,6 @@ pub(crate) fn admin_dashboard_observation(
|
||||
} = observation;
|
||||
Ok(admin_observability::DashboardObservation {
|
||||
range: admin_range(range),
|
||||
requests: admin_request_metrics(requests)?,
|
||||
attempts: admin_attempt_metrics(attempts)?,
|
||||
totals: admin_dashboard_totals(totals)?,
|
||||
provider_accounts: admin_account_pool_metrics(provider_accounts),
|
||||
trend: trend
|
||||
@@ -125,7 +114,7 @@ pub(crate) fn admin_dashboard_observation(
|
||||
.collect::<AdminStoreResult<_>>()?,
|
||||
recent_requests: recent_requests
|
||||
.into_iter()
|
||||
.map(admin_usage_record)
|
||||
.map(admin_usage_list_record)
|
||||
.collect::<AdminStoreResult<_>>()?,
|
||||
})
|
||||
}
|
||||
@@ -411,20 +400,79 @@ pub(crate) fn admin_usage_page(
|
||||
items: page
|
||||
.items
|
||||
.into_iter()
|
||||
.map(admin_usage_record)
|
||||
.map(admin_usage_list_record)
|
||||
.collect::<AdminStoreResult<_>>()?,
|
||||
current_page: page.current_page,
|
||||
page_size: page.page_size,
|
||||
total: page.total,
|
||||
next_cursor: page.next_cursor.map(admin_cursor),
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn admin_cursor(
|
||||
cursor: ObservabilityCursor,
|
||||
) -> admin_observability::ObservabilityCursor {
|
||||
admin_observability::ObservabilityCursor {
|
||||
observed_at: cursor.observed_at,
|
||||
stable_id: cursor.stable_id,
|
||||
}
|
||||
pub(crate) fn admin_usage_list_record(
|
||||
record: UsageListRecord,
|
||||
) -> AdminStoreResult<admin_observability::UsageListRecord> {
|
||||
let billing = match (&record.cost_amount, &record.cost_currency) {
|
||||
(Some(amount), Some(currency)) => Some(admin_observability::UsageBilling::Total {
|
||||
source: record.cost_source.clone(),
|
||||
total: admin_observability::CurrencyCost {
|
||||
currency: currency.clone(),
|
||||
amount: admin_decimal_amount(amount.clone())?,
|
||||
},
|
||||
}),
|
||||
(None, None) => None,
|
||||
_ => {
|
||||
return Err(observability_error(StoreError::InvalidData {
|
||||
entity: "observability request billing",
|
||||
message: "cost amount and currency must be present together".to_owned(),
|
||||
}));
|
||||
}
|
||||
};
|
||||
Ok(admin_observability::UsageListRecord {
|
||||
id: record.id,
|
||||
endpoint: record.endpoint,
|
||||
client_transport: record.client_transport,
|
||||
requested_model_id: record.requested_model_id,
|
||||
provider_kind: record.provider_kind,
|
||||
provider_account_ref: record.provider_account_ref,
|
||||
provider_account_name: record.provider_account_name,
|
||||
provider_account_email: record.provider_account_email,
|
||||
provider_account_authentication_kind: record.provider_account_authentication_kind,
|
||||
upstream_model_id: record.upstream_model_id,
|
||||
upstream_transport: record.upstream_transport,
|
||||
service_tier: record.service_tier,
|
||||
input_tokens: record.input_tokens,
|
||||
output_tokens: record.output_tokens,
|
||||
cached_tokens: record.cached_tokens,
|
||||
cache_write_tokens: record.cache_write_tokens,
|
||||
reasoning_tokens: record.reasoning_tokens,
|
||||
image_input_tokens: record.image_input_tokens,
|
||||
image_output_tokens: record.image_output_tokens,
|
||||
total_tokens: record.total_tokens,
|
||||
cost_source: record.cost_source,
|
||||
cost_amount: admin_optional_decimal_amount(record.cost_amount)?,
|
||||
cost_currency: record.cost_currency,
|
||||
billing,
|
||||
transport_decision_wait_ms: record.transport_decision_wait_ms,
|
||||
connect_ms: record.connect_ms,
|
||||
headers_ms: record.headers_ms,
|
||||
first_event_ms: record.first_event_ms,
|
||||
first_reasoning_ms: record.first_reasoning_ms,
|
||||
first_text_ms: record.first_text_ms,
|
||||
first_token_ms: record.first_token_ms,
|
||||
provider_processing_ms: record.provider_processing_ms,
|
||||
latency_ms: record.latency_ms,
|
||||
admission_decision_ms: record.admission_decision_ms,
|
||||
account_selection_wait_ms: record.account_selection_wait_ms,
|
||||
capacity_used_slots: record.capacity_used_slots,
|
||||
capacity_total_slots: record.capacity_total_slots,
|
||||
client_ip: record.client_ip,
|
||||
user_agent: record.user_agent,
|
||||
reasoning_effort: record.reasoning_effort,
|
||||
reasoning_preset: record.reasoning_preset,
|
||||
subagent_kind: record.subagent_kind,
|
||||
compact: record.compact,
|
||||
started_at: record.started_at,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn admin_usage_record(
|
||||
@@ -638,8 +686,9 @@ pub(crate) fn admin_ops_error_page(
|
||||
) -> AdminStoreResult<admin_observability::OpsErrorPage> {
|
||||
Ok(admin_observability::OpsErrorPage {
|
||||
items: page.items.into_iter().map(admin_ops_error).collect(),
|
||||
current_page: page.current_page,
|
||||
page_size: page.page_size,
|
||||
total: page.total,
|
||||
next_cursor: page.next_cursor.map(admin_cursor),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -652,6 +701,10 @@ pub(crate) fn admin_ops_error(error: OpsErrorRecord) -> admin_observability::Ops
|
||||
client_api_key_ref: error.client_api_key_ref,
|
||||
component: error.component,
|
||||
operation: error.operation,
|
||||
protocol: error.protocol,
|
||||
client_transport: error.client_transport,
|
||||
requested_model_id: error.requested_model_id,
|
||||
service_tier: error.service_tier,
|
||||
endpoint: error.endpoint,
|
||||
provider_kind: error.provider_kind,
|
||||
provider_account_ref: error.provider_account_ref,
|
||||
@@ -669,11 +722,68 @@ pub(crate) fn admin_ops_error(error: OpsErrorRecord) -> admin_observability::Ops
|
||||
latency_ms: error.latency_ms,
|
||||
message: error.message,
|
||||
occurrence_count: error.occurrence_count,
|
||||
client_ip: error.client_ip,
|
||||
user_agent: error.user_agent,
|
||||
reasoning_effort: error.reasoning_effort,
|
||||
reasoning_preset: error.reasoning_preset,
|
||||
request_kind: error.request_kind,
|
||||
subagent_kind: error.subagent_kind,
|
||||
compact: error.compact,
|
||||
occurred_at: error.occurred_at,
|
||||
stable_sort_id: error.stable_sort_id,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn usage_list_record_from_row(
|
||||
row: &sqlx::postgres::PgRow,
|
||||
) -> StoreResult<UsageListRecord> {
|
||||
Ok(UsageListRecord {
|
||||
id: get(row, "id")?,
|
||||
endpoint: get(row, "endpoint")?,
|
||||
client_transport: get(row, "client_transport")?,
|
||||
requested_model_id: get(row, "requested_model_id")?,
|
||||
provider_kind: get(row, "provider_kind")?,
|
||||
provider_account_ref: get(row, "provider_account_ref")?,
|
||||
provider_account_name: get(row, "provider_account_name")?,
|
||||
provider_account_email: get(row, "provider_account_email")?,
|
||||
provider_account_authentication_kind: get(row, "provider_account_authentication_kind")?,
|
||||
upstream_model_id: get(row, "upstream_model_id")?,
|
||||
upstream_transport: get(row, "upstream_transport")?,
|
||||
service_tier: get(row, "service_tier")?,
|
||||
input_tokens: optional_unsigned(row, "input_tokens")?,
|
||||
output_tokens: optional_unsigned(row, "output_tokens")?,
|
||||
cached_tokens: optional_unsigned(row, "cached_tokens")?,
|
||||
cache_write_tokens: optional_unsigned(row, "cache_write_tokens")?,
|
||||
reasoning_tokens: optional_unsigned(row, "reasoning_tokens")?,
|
||||
image_input_tokens: optional_unsigned(row, "image_input_tokens")?,
|
||||
image_output_tokens: optional_unsigned(row, "image_output_tokens")?,
|
||||
total_tokens: optional_unsigned(row, "total_tokens")?,
|
||||
cost_source: get(row, "cost_source")?,
|
||||
cost_amount: optional_decimal(row, "cost_amount")?,
|
||||
cost_currency: get(row, "cost_currency")?,
|
||||
transport_decision_wait_ms: optional_unsigned(row, "transport_decision_wait_ms")?,
|
||||
connect_ms: optional_unsigned(row, "connect_ms")?,
|
||||
headers_ms: optional_unsigned(row, "headers_ms")?,
|
||||
first_event_ms: optional_unsigned(row, "first_event_ms")?,
|
||||
first_reasoning_ms: optional_unsigned(row, "first_reasoning_ms")?,
|
||||
first_text_ms: optional_unsigned(row, "first_text_ms")?,
|
||||
first_token_ms: optional_unsigned(row, "first_token_ms")?,
|
||||
provider_processing_ms: optional_unsigned(row, "provider_processing_ms")?,
|
||||
latency_ms: optional_unsigned(row, "latency_ms")?,
|
||||
admission_decision_ms: optional_unsigned(row, "admission_decision_ms")?,
|
||||
account_selection_wait_ms: optional_unsigned(row, "account_selection_wait_ms")?,
|
||||
capacity_used_slots: optional_unsigned(row, "capacity_used_slots")?,
|
||||
capacity_total_slots: optional_unsigned(row, "capacity_total_slots")?,
|
||||
client_ip: get(row, "client_ip")?,
|
||||
user_agent: get(row, "user_agent")?,
|
||||
reasoning_effort: get(row, "reasoning_effort")?,
|
||||
reasoning_preset: get(row, "reasoning_preset")?,
|
||||
subagent_kind: get(row, "subagent_kind")?,
|
||||
compact: get(row, "compact")?,
|
||||
started_at: get(row, "started_at")?,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn usage_record_from_row(row: &sqlx::postgres::PgRow) -> StoreResult<UsageRecord> {
|
||||
Ok(UsageRecord {
|
||||
id: get(row, "id")?,
|
||||
@@ -768,6 +878,10 @@ pub(crate) fn ops_error_from_row(row: &sqlx::postgres::PgRow) -> StoreResult<Ops
|
||||
client_api_key_ref: get(row, "client_api_key_ref")?,
|
||||
component: get(row, "component")?,
|
||||
operation: get(row, "operation")?,
|
||||
protocol: get(row, "protocol")?,
|
||||
client_transport: get(row, "client_transport")?,
|
||||
requested_model_id: get(row, "requested_model_id")?,
|
||||
service_tier: get(row, "service_tier")?,
|
||||
endpoint: get(row, "endpoint")?,
|
||||
provider_kind: get(row, "provider_kind")?,
|
||||
provider_account_ref: get(row, "provider_account_ref")?,
|
||||
@@ -785,6 +899,13 @@ pub(crate) fn ops_error_from_row(row: &sqlx::postgres::PgRow) -> StoreResult<Ops
|
||||
latency_ms: optional_unsigned(row, "latency_ms")?,
|
||||
message: get(row, "message")?,
|
||||
occurrence_count: to_u32(get(row, "occurrence_count")?)?,
|
||||
client_ip: get(row, "client_ip")?,
|
||||
user_agent: get(row, "user_agent")?,
|
||||
reasoning_effort: get(row, "reasoning_effort")?,
|
||||
reasoning_preset: get(row, "reasoning_preset")?,
|
||||
request_kind: get(row, "request_kind")?,
|
||||
subagent_kind: get(row, "subagent_kind")?,
|
||||
compact: get(row, "compact")?,
|
||||
occurred_at: get(row, "occurred_at")?,
|
||||
stable_sort_id: get(row, "stable_sort_id")?,
|
||||
})
|
||||
|
||||
@@ -18,12 +18,16 @@ use crate::{
|
||||
require_nonempty,
|
||||
};
|
||||
|
||||
use super::{completed_usage_fact_predicate, push_completed_usage_fact_filter};
|
||||
|
||||
mod admin_adapter;
|
||||
mod mapping;
|
||||
mod model;
|
||||
mod queries;
|
||||
mod query_budget;
|
||||
|
||||
pub use admin_adapter::*;
|
||||
pub(crate) use mapping::*;
|
||||
pub use model::*;
|
||||
pub(crate) use queries::*;
|
||||
pub use query_budget::*;
|
||||
|
||||
@@ -43,38 +43,17 @@ impl ObservabilityPageSize {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct ObservabilityPageNumber(u32);
|
||||
|
||||
impl ObservabilityPageNumber {
|
||||
pub fn new(value: u32) -> StoreResult<Self> {
|
||||
if value == 0 {
|
||||
return Err(invalid("page number must be positive"));
|
||||
}
|
||||
Ok(Self(value))
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn get(self) -> u32 {
|
||||
self.0
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ObservabilityCursor {
|
||||
pub observed_at: DateTime<Utc>,
|
||||
pub stable_id: String,
|
||||
}
|
||||
|
||||
impl ObservabilityCursor {
|
||||
pub fn new(observed_at: DateTime<Utc>, stable_id: impl Into<String>) -> StoreResult<Self> {
|
||||
let stable_id = stable_id.into();
|
||||
validate_text(&stable_id, MAX_FILTER_BYTES, "cursor ID")?;
|
||||
Ok(Self {
|
||||
observed_at,
|
||||
stable_id,
|
||||
})
|
||||
}
|
||||
pub(crate) fn observability_page_offset(
|
||||
current_page: u32,
|
||||
page_size: ObservabilityPageSize,
|
||||
) -> StoreResult<i64> {
|
||||
let page_index = current_page
|
||||
.checked_sub(1)
|
||||
.ok_or_else(|| invalid("current page must be positive"))?;
|
||||
u64::from(page_index)
|
||||
.checked_mul(u64::from(page_size.get()))
|
||||
.and_then(|offset| i64::try_from(offset).ok())
|
||||
.ok_or_else(|| invalid("page offset is too large"))
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
@@ -136,8 +115,7 @@ impl UsageRecordFilter {
|
||||
pub struct UsageRecordQuery {
|
||||
pub range: ObservabilityRange,
|
||||
pub filter: UsageRecordFilter,
|
||||
pub cursor: Option<ObservabilityCursor>,
|
||||
pub page: ObservabilityPageNumber,
|
||||
pub current_page: u32,
|
||||
pub page_size: ObservabilityPageSize,
|
||||
}
|
||||
|
||||
@@ -148,6 +126,7 @@ pub struct OpsErrorFilter {
|
||||
pub provider_account_ref: Option<String>,
|
||||
pub provider_kind: Option<String>,
|
||||
pub operation: Option<String>,
|
||||
pub endpoint: Option<String>,
|
||||
pub model: Option<String>,
|
||||
pub transport: Option<String>,
|
||||
pub attempt_index: Option<u32>,
|
||||
@@ -169,6 +148,7 @@ impl OpsErrorFilter {
|
||||
),
|
||||
(self.provider_kind.as_deref(), "provider filter"),
|
||||
(self.operation.as_deref(), "operation filter"),
|
||||
(self.endpoint.as_deref(), "endpoint filter"),
|
||||
(self.model.as_deref(), "model filter"),
|
||||
(self.transport.as_deref(), "transport filter"),
|
||||
(
|
||||
@@ -200,8 +180,7 @@ impl OpsErrorFilter {
|
||||
pub struct OpsErrorQuery {
|
||||
pub range: ObservabilityRange,
|
||||
pub filter: OpsErrorFilter,
|
||||
pub cursor: Option<ObservabilityCursor>,
|
||||
pub page: ObservabilityPageNumber,
|
||||
pub current_page: u32,
|
||||
pub page_size: ObservabilityPageSize,
|
||||
}
|
||||
|
||||
@@ -493,13 +472,59 @@ pub struct DashboardTotals {
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct DashboardObservation {
|
||||
pub range: ObservabilityRange,
|
||||
pub requests: RequestMetrics,
|
||||
pub attempts: AttemptMetrics,
|
||||
pub totals: DashboardTotals,
|
||||
pub provider_accounts: ProviderAccountMetrics,
|
||||
pub trend: Vec<RequestMetricPoint>,
|
||||
pub account_usage: Vec<ProviderAccountUsageObservation>,
|
||||
pub recent_requests: Vec<UsageRecord>,
|
||||
pub recent_requests: Vec<UsageListRecord>,
|
||||
}
|
||||
|
||||
/// 使用记录列表所需的窄投影;完整执行、路由和客户端详情按 ID 单独读取。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct UsageListRecord {
|
||||
pub id: String,
|
||||
pub endpoint: String,
|
||||
pub client_transport: String,
|
||||
pub requested_model_id: Option<String>,
|
||||
pub provider_kind: Option<String>,
|
||||
pub provider_account_ref: Option<String>,
|
||||
pub provider_account_name: Option<String>,
|
||||
pub provider_account_email: Option<String>,
|
||||
pub provider_account_authentication_kind: Option<String>,
|
||||
pub upstream_model_id: Option<String>,
|
||||
pub upstream_transport: Option<String>,
|
||||
pub service_tier: Option<String>,
|
||||
pub input_tokens: Option<u64>,
|
||||
pub output_tokens: Option<u64>,
|
||||
pub cached_tokens: Option<u64>,
|
||||
pub cache_write_tokens: Option<u64>,
|
||||
pub reasoning_tokens: Option<u64>,
|
||||
pub image_input_tokens: Option<u64>,
|
||||
pub image_output_tokens: Option<u64>,
|
||||
pub total_tokens: Option<u64>,
|
||||
pub cost_source: String,
|
||||
pub cost_amount: Option<DecimalAmount>,
|
||||
pub cost_currency: Option<String>,
|
||||
pub transport_decision_wait_ms: Option<u64>,
|
||||
pub connect_ms: Option<u64>,
|
||||
pub headers_ms: Option<u64>,
|
||||
pub first_event_ms: Option<u64>,
|
||||
pub first_reasoning_ms: Option<u64>,
|
||||
pub first_text_ms: Option<u64>,
|
||||
pub first_token_ms: Option<u64>,
|
||||
pub provider_processing_ms: Option<u64>,
|
||||
pub latency_ms: Option<u64>,
|
||||
pub admission_decision_ms: Option<u64>,
|
||||
pub account_selection_wait_ms: Option<u64>,
|
||||
pub capacity_used_slots: Option<u64>,
|
||||
pub capacity_total_slots: Option<u64>,
|
||||
pub client_ip: Option<String>,
|
||||
pub user_agent: Option<String>,
|
||||
pub reasoning_effort: Option<String>,
|
||||
pub reasoning_preset: Option<String>,
|
||||
pub subagent_kind: Option<String>,
|
||||
pub compact: bool,
|
||||
pub started_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
@@ -579,9 +604,10 @@ pub struct UsageRecord {
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct UsageRecordPage {
|
||||
pub items: Vec<UsageRecord>,
|
||||
pub items: Vec<UsageListRecord>,
|
||||
pub current_page: u32,
|
||||
pub page_size: u16,
|
||||
pub total: u64,
|
||||
pub next_cursor: Option<ObservabilityCursor>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
@@ -670,6 +696,10 @@ pub struct OpsErrorRecord {
|
||||
pub client_api_key_ref: Option<String>,
|
||||
pub component: String,
|
||||
pub operation: String,
|
||||
pub protocol: Option<String>,
|
||||
pub client_transport: Option<String>,
|
||||
pub requested_model_id: Option<String>,
|
||||
pub service_tier: Option<String>,
|
||||
pub endpoint: Option<String>,
|
||||
pub provider_kind: Option<String>,
|
||||
pub provider_account_ref: Option<String>,
|
||||
@@ -687,6 +717,13 @@ pub struct OpsErrorRecord {
|
||||
pub latency_ms: Option<u64>,
|
||||
pub message: String,
|
||||
pub occurrence_count: u32,
|
||||
pub client_ip: Option<String>,
|
||||
pub user_agent: Option<String>,
|
||||
pub reasoning_effort: Option<String>,
|
||||
pub reasoning_preset: Option<String>,
|
||||
pub request_kind: Option<String>,
|
||||
pub subagent_kind: Option<String>,
|
||||
pub compact: Option<bool>,
|
||||
pub occurred_at: DateTime<Utc>,
|
||||
pub stable_sort_id: String,
|
||||
}
|
||||
@@ -694,8 +731,9 @@ pub struct OpsErrorRecord {
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct OpsErrorPage {
|
||||
pub items: Vec<OpsErrorRecord>,
|
||||
pub current_page: u32,
|
||||
pub page_size: u16,
|
||||
pub total: u64,
|
||||
pub next_cursor: Option<ObservabilityCursor>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
//! 账号用量查询族。
|
||||
//! Provider 账号用量查询族。
|
||||
|
||||
use super::super::*;
|
||||
|
||||
@@ -14,32 +14,18 @@ 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,
|
||||
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,
|
||||
sum(mr.output_tokens)::bigint as output_tokens,
|
||||
sum(mr.cached_tokens)::bigint as cached_tokens,
|
||||
sum(mr.cache_write_tokens)::bigint as cache_write_tokens,
|
||||
sum(mr.reasoning_tokens)::bigint as reasoning_tokens,
|
||||
sum(mr.image_input_tokens)::bigint as image_input_tokens,
|
||||
sum(mr.image_output_tokens)::bigint as image_output_tokens,
|
||||
count(mr.id) filter (where mr.image_generation_succeeded is true)::bigint
|
||||
as image_request_count,
|
||||
count(mr.id) filter (where mr.image_generation_succeeded is false)::bigint
|
||||
as image_request_failed_count,
|
||||
sum(mr.total_tokens)::bigint as total_tokens,
|
||||
count(mr.id) filter (where mr.cost_source = 'provider_reported')::bigint
|
||||
as provider_reported_count,
|
||||
count(mr.id) filter (where mr.cost_source = 'calculated')::bigint
|
||||
as calculated_count,
|
||||
count(mr.id) filter (where mr.cost_source = 'unavailable')::bigint
|
||||
as unavailable_count,
|
||||
max(mr.started_at) as last_used_at
|
||||
from provider_accounts pa
|
||||
left join model_requests mr
|
||||
on mr.provider_account_ref = pa.id and mr.started_at >= ",
|
||||
"with matched as (
|
||||
select pa.id, pa.provider_kind, pa.authentication_kind, pa.name, pa.email,
|
||||
pa.plan_type, mr.id as request_id,
|
||||
coalesce(mr.upstream_model_id, mr.requested_model_id) as model,
|
||||
mr.cost_currency, mr.outcome, mr.input_tokens, mr.output_tokens,
|
||||
mr.cached_tokens, mr.cache_write_tokens, mr.reasoning_tokens,
|
||||
mr.image_input_tokens, mr.image_output_tokens,
|
||||
mr.image_generation_succeeded, mr.total_tokens, mr.cost_source,
|
||||
mr.cost_amount, mr.started_at
|
||||
from provider_accounts pa
|
||||
left join model_requests mr
|
||||
on mr.provider_account_ref = pa.id and mr.started_at >= ",
|
||||
);
|
||||
statement.push_bind(query.range.start);
|
||||
statement.push(" and mr.started_at < ");
|
||||
@@ -51,24 +37,101 @@ pub(crate) async fn provider_account_usage(
|
||||
statement.push("::text[])");
|
||||
}
|
||||
statement.push(
|
||||
" group by pa.id, pa.provider_kind, pa.authentication_kind, pa.name, pa.email,
|
||||
pa.plan_type",
|
||||
"), aggregated as (
|
||||
select id, provider_kind, authentication_kind, name, email, plan_type,
|
||||
model, cost_currency,
|
||||
grouping(model)::integer as model_grouping,
|
||||
grouping(cost_currency)::integer as currency_grouping,
|
||||
count(request_id)::bigint as request_count,
|
||||
count(request_id) filter (where outcome = 'succeeded')::bigint
|
||||
as success_count,
|
||||
sum(input_tokens)::bigint as input_tokens,
|
||||
sum(output_tokens)::bigint as output_tokens,
|
||||
sum(cached_tokens)::bigint as cached_tokens,
|
||||
sum(cache_write_tokens)::bigint as cache_write_tokens,
|
||||
sum(reasoning_tokens)::bigint as reasoning_tokens,
|
||||
sum(image_input_tokens)::bigint as image_input_tokens,
|
||||
sum(image_output_tokens)::bigint as image_output_tokens,
|
||||
count(request_id) filter (where image_generation_succeeded is true)::bigint
|
||||
as image_request_count,
|
||||
count(request_id) filter (where image_generation_succeeded is false)::bigint
|
||||
as image_request_failed_count,
|
||||
sum(total_tokens)::bigint as total_tokens,
|
||||
count(request_id) filter (where cost_source = 'provider_reported')::bigint
|
||||
as provider_reported_count,
|
||||
count(request_id) filter (where cost_source = 'calculated')::bigint
|
||||
as calculated_count,
|
||||
count(request_id) filter (where cost_source = 'unavailable')::bigint
|
||||
as unavailable_count,
|
||||
max(started_at) as last_used_at,
|
||||
sum(cost_amount)::text as amount
|
||||
from matched
|
||||
group by id, provider_kind, authentication_kind, name, email, plan_type,
|
||||
grouping sets ((), (cost_currency), (model), (model, cost_currency))
|
||||
), selected_accounts as (
|
||||
select id,
|
||||
row_number() over (order by last_used_at desc nulls last, name, id)
|
||||
as sort_position
|
||||
from aggregated
|
||||
where model_grouping = 1 and currency_grouping = 1",
|
||||
);
|
||||
if query.account_ids.is_none() {
|
||||
statement.push(" having count(mr.id) > 0");
|
||||
statement.push(" and request_count > 0");
|
||||
}
|
||||
statement.push(" order by max(mr.started_at) desc nulls last, pa.name, pa.id limit ");
|
||||
statement.push(" order by last_used_at desc nulls last, name, id limit ");
|
||||
statement.push_bind(i64::from(query.limit));
|
||||
statement.push(
|
||||
")
|
||||
select aggregated.*
|
||||
from aggregated
|
||||
join selected_accounts selected using (id)
|
||||
order by selected.sort_position, model_grouping desc, currency_grouping desc,
|
||||
model nulls last, cost_currency nulls last",
|
||||
);
|
||||
let rows = statement
|
||||
.build()
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("load provider account usage"))?;
|
||||
|
||||
let mut observations = rows
|
||||
.iter()
|
||||
.map(provider_account_from_row)
|
||||
.collect::<StoreResult<Vec<_>>>()?;
|
||||
let mut observations = Vec::with_capacity(usize::from(query.limit));
|
||||
let mut costs = HashMap::<String, Vec<CurrencyCostTotal>>::new();
|
||||
let mut models = HashMap::<String, Vec<ProviderAccountModelUsageObservation>>::new();
|
||||
let mut model_costs = HashMap::<(String, String), Vec<CurrencyCostTotal>>::new();
|
||||
for row in &rows {
|
||||
let model_grouping: i32 = get(row, "model_grouping")?;
|
||||
let currency_grouping: i32 = get(row, "currency_grouping")?;
|
||||
match (model_grouping, currency_grouping) {
|
||||
(1, 1) => observations.push(provider_account_from_row(row)?),
|
||||
(1, 0) if get::<Option<String>>(row, "cost_currency")?.is_some() => {
|
||||
costs
|
||||
.entry(get(row, "id")?)
|
||||
.or_default()
|
||||
.push(cost_from_row(row)?);
|
||||
}
|
||||
(0, 1) if get::<Option<String>>(row, "model")?.is_some() => {
|
||||
models
|
||||
.entry(get(row, "id")?)
|
||||
.or_default()
|
||||
.push(provider_account_model_from_row(row)?);
|
||||
}
|
||||
(0, 0)
|
||||
if get::<Option<String>>(row, "model")?.is_some()
|
||||
&& get::<Option<String>>(row, "cost_currency")?.is_some() =>
|
||||
{
|
||||
model_costs
|
||||
.entry((get(row, "id")?, get(row, "model")?))
|
||||
.or_default()
|
||||
.push(cost_from_row(row)?);
|
||||
}
|
||||
(0 | 1, 0 | 1) => {}
|
||||
_ => {
|
||||
return Err(postgres_unavailable(
|
||||
"decode provider account usage grouping",
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
if observations.is_empty() {
|
||||
return Ok(observations);
|
||||
}
|
||||
@@ -81,52 +144,23 @@ pub(crate) async fn provider_account_usage(
|
||||
} else {
|
||||
HashMap::new()
|
||||
};
|
||||
let account_costs = account_costs(pool, query.range, &account_ids).await?;
|
||||
let model_rows = account_model_rows(pool, query.range, &account_ids).await?;
|
||||
let model_costs = account_model_costs(pool, query.range, &account_ids).await?;
|
||||
let account_positions = observations
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(index, item)| (item.account_id.clone(), index))
|
||||
.collect::<HashMap<_, _>>();
|
||||
for observation in &mut observations {
|
||||
observation.costs = account_costs
|
||||
.get(&observation.account_id)
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
observation.costs = costs.remove(&observation.account_id).unwrap_or_default();
|
||||
observation.request_buckets = request_buckets
|
||||
.remove(&observation.account_id)
|
||||
.unwrap_or_default();
|
||||
}
|
||||
for row in model_rows {
|
||||
let account_id: String = get(&row, "provider_account_ref")?;
|
||||
let Some(position) = account_positions.get(&account_id).copied() else {
|
||||
continue;
|
||||
};
|
||||
let model: String = get(&row, "model")?;
|
||||
observations[position]
|
||||
.models
|
||||
.push(ProviderAccountModelUsageObservation {
|
||||
costs: model_costs
|
||||
.get(&(account_id, model.clone()))
|
||||
.cloned()
|
||||
.unwrap_or_default(),
|
||||
model,
|
||||
request_count: unsigned(&row, "request_count")?,
|
||||
success_count: unsigned(&row, "success_count")?,
|
||||
input_tokens: optional_unsigned(&row, "input_tokens")?,
|
||||
output_tokens: optional_unsigned(&row, "output_tokens")?,
|
||||
cached_tokens: optional_unsigned(&row, "cached_tokens")?,
|
||||
cache_write_tokens: optional_unsigned(&row, "cache_write_tokens")?,
|
||||
reasoning_tokens: optional_unsigned(&row, "reasoning_tokens")?,
|
||||
image_input_tokens: optional_unsigned(&row, "image_input_tokens")?,
|
||||
image_output_tokens: optional_unsigned(&row, "image_output_tokens")?,
|
||||
image_request_count: unsigned(&row, "image_request_count")?,
|
||||
image_request_failed_count: unsigned(&row, "image_request_failed_count")?,
|
||||
total_tokens: optional_unsigned(&row, "total_tokens")?,
|
||||
cost_coverage: coverage_from_row(&row)?,
|
||||
last_used_at: get(&row, "last_used_at")?,
|
||||
});
|
||||
observation.models = models.remove(&observation.account_id).unwrap_or_default();
|
||||
for model in &mut observation.models {
|
||||
model.costs = model_costs
|
||||
.remove(&(observation.account_id.clone(), model.model.clone()))
|
||||
.unwrap_or_default();
|
||||
}
|
||||
observation.models.sort_by(|left, right| {
|
||||
right
|
||||
.request_count
|
||||
.cmp(&left.request_count)
|
||||
.then_with(|| left.model.cmp(&right.model))
|
||||
});
|
||||
}
|
||||
Ok(observations)
|
||||
}
|
||||
@@ -136,25 +170,25 @@ pub(crate) async fn provider_account_request_buckets(
|
||||
range: ObservabilityRange,
|
||||
account_ids: &[String],
|
||||
) -> StoreResult<HashMap<String, Vec<ProviderAccountRequestBucket>>> {
|
||||
let rows = sqlx::query(
|
||||
let completed_usage = completed_usage_fact_predicate("mr");
|
||||
let statement = format!(
|
||||
"select provider_account_ref,
|
||||
floor(extract(epoch from (started_at - $1)) / 3600)::bigint as bucket_index,
|
||||
count(*)::bigint as request_count
|
||||
from model_requests mr
|
||||
where mr.provider_account_ref = any($2::text[])
|
||||
and mr.started_at >= $1 and mr.started_at < $3
|
||||
and mr.outcome = 'succeeded'
|
||||
and mr.downstream_committed_at is not null
|
||||
and mr.client_status_code between 200 and 399
|
||||
and {completed_usage}
|
||||
group by mr.provider_account_ref, bucket_index
|
||||
order by mr.provider_account_ref, bucket_index",
|
||||
)
|
||||
.bind(range.start)
|
||||
.bind(account_ids)
|
||||
.bind(range.end)
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("load provider account request timeline"))?;
|
||||
order by mr.provider_account_ref, bucket_index"
|
||||
);
|
||||
let rows = sqlx::query(sqlx::AssertSqlSafe(statement))
|
||||
.bind(range.start)
|
||||
.bind(account_ids)
|
||||
.bind(range.end)
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("load provider account request timeline"))?;
|
||||
|
||||
let mut observed = HashMap::<String, BTreeMap<u64, u64>>::new();
|
||||
for row in rows {
|
||||
@@ -191,121 +225,27 @@ pub(crate) async fn provider_account_request_buckets(
|
||||
Ok(timelines)
|
||||
}
|
||||
|
||||
pub(crate) async fn account_costs(
|
||||
pool: &PgPool,
|
||||
range: ObservabilityRange,
|
||||
account_ids: &[String],
|
||||
) -> StoreResult<HashMap<String, Vec<CurrencyCostTotal>>> {
|
||||
let rows = sqlx::query(
|
||||
"select mr.provider_account_ref, mr.cost_currency, sum(mr.cost_amount)::text as amount
|
||||
from model_requests mr
|
||||
where mr.provider_account_ref = any($1::text[])
|
||||
and mr.started_at >= $2 and mr.started_at < $3
|
||||
and mr.outcome = 'succeeded'
|
||||
and mr.downstream_committed_at is not null
|
||||
and mr.client_status_code between 200 and 399
|
||||
and mr.cost_amount is not null and mr.cost_currency is not null
|
||||
group by mr.provider_account_ref, mr.cost_currency
|
||||
order by mr.provider_account_ref, mr.cost_currency",
|
||||
)
|
||||
.bind(account_ids)
|
||||
.bind(range.start)
|
||||
.bind(range.end)
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("load provider account costs"))?;
|
||||
let mut costs = HashMap::<String, Vec<CurrencyCostTotal>>::new();
|
||||
for row in rows {
|
||||
costs
|
||||
.entry(get(&row, "provider_account_ref")?)
|
||||
.or_default()
|
||||
.push(cost_from_row(&row)?);
|
||||
}
|
||||
Ok(costs)
|
||||
}
|
||||
|
||||
pub(crate) async fn account_model_rows(
|
||||
pool: &PgPool,
|
||||
range: ObservabilityRange,
|
||||
account_ids: &[String],
|
||||
) -> StoreResult<Vec<sqlx::postgres::PgRow>> {
|
||||
sqlx::query(
|
||||
"select mr.provider_account_ref,
|
||||
coalesce(mr.upstream_model_id, mr.requested_model_id) as model,
|
||||
count(*)::bigint as request_count,
|
||||
count(*) filter (where mr.outcome = 'succeeded')::bigint as success_count,
|
||||
sum(mr.input_tokens)::bigint as input_tokens,
|
||||
sum(mr.output_tokens)::bigint as output_tokens,
|
||||
sum(mr.cached_tokens)::bigint as cached_tokens,
|
||||
sum(mr.cache_write_tokens)::bigint as cache_write_tokens,
|
||||
sum(mr.reasoning_tokens)::bigint as reasoning_tokens,
|
||||
sum(mr.image_input_tokens)::bigint as image_input_tokens,
|
||||
sum(mr.image_output_tokens)::bigint as image_output_tokens,
|
||||
count(*) filter (where mr.image_generation_succeeded is true)::bigint
|
||||
as image_request_count,
|
||||
count(*) filter (where mr.image_generation_succeeded is false)::bigint
|
||||
as image_request_failed_count,
|
||||
sum(mr.total_tokens)::bigint as total_tokens,
|
||||
count(*) filter (where mr.cost_source = 'provider_reported')::bigint
|
||||
as provider_reported_count,
|
||||
count(*) filter (where mr.cost_source = 'calculated')::bigint
|
||||
as calculated_count,
|
||||
count(*) filter (where mr.cost_source = 'unavailable')::bigint
|
||||
as unavailable_count,
|
||||
max(mr.started_at) as last_used_at
|
||||
from model_requests mr
|
||||
where mr.provider_account_ref = any($1::text[])
|
||||
and mr.started_at >= $2 and mr.started_at < $3
|
||||
and mr.outcome = 'succeeded'
|
||||
and mr.downstream_committed_at is not null
|
||||
and mr.client_status_code between 200 and 399
|
||||
and coalesce(mr.upstream_model_id, mr.requested_model_id) is not null
|
||||
group by mr.provider_account_ref, coalesce(mr.upstream_model_id, mr.requested_model_id)
|
||||
order by mr.provider_account_ref, request_count desc, model",
|
||||
)
|
||||
.bind(account_ids)
|
||||
.bind(range.start)
|
||||
.bind(range.end)
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("load provider account model usage"))
|
||||
}
|
||||
|
||||
pub(crate) async fn account_model_costs(
|
||||
pool: &PgPool,
|
||||
range: ObservabilityRange,
|
||||
account_ids: &[String],
|
||||
) -> StoreResult<HashMap<(String, String), Vec<CurrencyCostTotal>>> {
|
||||
let rows = sqlx::query(
|
||||
"select mr.provider_account_ref,
|
||||
coalesce(mr.upstream_model_id, mr.requested_model_id) as model,
|
||||
mr.cost_currency, sum(mr.cost_amount)::text as amount
|
||||
from model_requests mr
|
||||
where mr.provider_account_ref = any($1::text[])
|
||||
and mr.started_at >= $2 and mr.started_at < $3
|
||||
and mr.outcome = 'succeeded'
|
||||
and mr.downstream_committed_at is not null
|
||||
and mr.client_status_code between 200 and 399
|
||||
and mr.cost_amount is not null and mr.cost_currency is not null
|
||||
and coalesce(mr.upstream_model_id, mr.requested_model_id) is not null
|
||||
group by mr.provider_account_ref, coalesce(mr.upstream_model_id, mr.requested_model_id),
|
||||
mr.cost_currency
|
||||
order by mr.provider_account_ref, model, mr.cost_currency",
|
||||
)
|
||||
.bind(account_ids)
|
||||
.bind(range.start)
|
||||
.bind(range.end)
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("load provider account model costs"))?;
|
||||
let mut costs = HashMap::<(String, String), Vec<CurrencyCostTotal>>::new();
|
||||
for row in rows {
|
||||
costs
|
||||
.entry((get(&row, "provider_account_ref")?, get(&row, "model")?))
|
||||
.or_default()
|
||||
.push(cost_from_row(&row)?);
|
||||
}
|
||||
Ok(costs)
|
||||
fn provider_account_model_from_row(
|
||||
row: &sqlx::postgres::PgRow,
|
||||
) -> StoreResult<ProviderAccountModelUsageObservation> {
|
||||
Ok(ProviderAccountModelUsageObservation {
|
||||
model: get(row, "model")?,
|
||||
request_count: unsigned(row, "request_count")?,
|
||||
success_count: unsigned(row, "success_count")?,
|
||||
input_tokens: optional_unsigned(row, "input_tokens")?,
|
||||
output_tokens: optional_unsigned(row, "output_tokens")?,
|
||||
cached_tokens: optional_unsigned(row, "cached_tokens")?,
|
||||
cache_write_tokens: optional_unsigned(row, "cache_write_tokens")?,
|
||||
reasoning_tokens: optional_unsigned(row, "reasoning_tokens")?,
|
||||
image_input_tokens: optional_unsigned(row, "image_input_tokens")?,
|
||||
image_output_tokens: optional_unsigned(row, "image_output_tokens")?,
|
||||
image_request_count: unsigned(row, "image_request_count")?,
|
||||
image_request_failed_count: unsigned(row, "image_request_failed_count")?,
|
||||
total_tokens: optional_unsigned(row, "total_tokens")?,
|
||||
cost_coverage: coverage_from_row(row)?,
|
||||
costs: Vec::new(),
|
||||
last_used_at: get(row, "last_used_at")?,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn provider_account_from_row(
|
||||
|
||||
@@ -133,6 +133,23 @@ pub(crate) async fn request_metric_series(
|
||||
pool: &PgPool,
|
||||
range: ObservabilityRange,
|
||||
filter: &UsageRecordFilter,
|
||||
) -> StoreResult<Vec<RequestMetricPoint>> {
|
||||
request_metric_series_inner(pool, range, filter, true).await
|
||||
}
|
||||
|
||||
pub(crate) async fn dashboard_request_metric_series(
|
||||
pool: &PgPool,
|
||||
range: ObservabilityRange,
|
||||
filter: &UsageRecordFilter,
|
||||
) -> StoreResult<Vec<RequestMetricPoint>> {
|
||||
request_metric_series_inner(pool, range, filter, false).await
|
||||
}
|
||||
|
||||
async fn request_metric_series_inner(
|
||||
pool: &PgPool,
|
||||
range: ObservabilityRange,
|
||||
filter: &UsageRecordFilter,
|
||||
load_costs: bool,
|
||||
) -> StoreResult<Vec<RequestMetricPoint>> {
|
||||
filter.validate()?;
|
||||
let granularity = granularity_for(range);
|
||||
@@ -252,10 +269,12 @@ pub(crate) async fn request_metric_series(
|
||||
},
|
||||
);
|
||||
}
|
||||
let bucket_costs = request_costs_by_bucket(pool, range, filter, granularity).await?;
|
||||
for (bucket, costs) in bucket_costs {
|
||||
if let Some(point) = points.get_mut(&bucket) {
|
||||
point.costs = costs;
|
||||
if load_costs {
|
||||
let bucket_costs = request_costs_by_bucket(pool, range, filter, granularity).await?;
|
||||
for (bucket, costs) in bucket_costs {
|
||||
if let Some(point) = points.get_mut(&bucket) {
|
||||
point.costs = costs;
|
||||
}
|
||||
}
|
||||
}
|
||||
fill_metric_gaps(range, granularity, points)
|
||||
@@ -341,63 +360,99 @@ pub(crate) async fn attempt_metrics(
|
||||
filter: &UsageRecordFilter,
|
||||
) -> StoreResult<AttemptMetrics> {
|
||||
filter.validate()?;
|
||||
let fact = completed_usage_fact_predicate("sr");
|
||||
// PERF: `model_requests` 包含宽请求元数据。这里只物化两个聚合族会复用的列,
|
||||
// 避免生产基数下 `mr.*` 与多次标量子查询把 CTE 反复写入临时文件。
|
||||
let mut query = QueryBuilder::<Postgres>::new(
|
||||
"with selected_requests as (
|
||||
select mr.* from model_requests mr where mr.started_at >= ",
|
||||
"with selected_requests as materialized (
|
||||
select mr.id, mr.attempt_count, mr.outcome, mr.error_kind,
|
||||
mr.upstream_status_code, mr.client_status_code,
|
||||
mr.downstream_committed_at, mr.client_transport, mr.cost_source
|
||||
from model_requests mr where mr.started_at >= ",
|
||||
);
|
||||
query.push_bind(range.start);
|
||||
query.push(" and mr.started_at < ");
|
||||
query.push_bind(range.end);
|
||||
push_usage_filter(&mut query, filter, "mr");
|
||||
query.push(
|
||||
"), usage_facts as (
|
||||
select * from selected_requests
|
||||
where outcome = 'succeeded'
|
||||
and downstream_committed_at is not null
|
||||
and client_status_code between 200 and 399
|
||||
), failures as (
|
||||
select coalesce(oe.occurrence_count, 1)::bigint as occurrences,
|
||||
oe.failure_kind, oe.status_code
|
||||
query.push(format!(
|
||||
"), request_aggregate as (
|
||||
select coalesce(sum(sr.attempt_count), 0)::bigint as attempt_count,
|
||||
count(*) filter (
|
||||
where sr.outcome = 'succeeded' and sr.attempt_count > 0
|
||||
)::bigint as success_count,
|
||||
count(*) filter (
|
||||
where sr.outcome = 'cancelled' and sr.attempt_count > 0
|
||||
)::bigint as cancelled_count,
|
||||
count(*) filter (
|
||||
where sr.outcome = 'incomplete' and sr.attempt_count > 0
|
||||
)::bigint as incomplete_count,
|
||||
count(*) filter (
|
||||
where sr.outcome = 'failed' and sr.attempt_count > 0
|
||||
)::bigint as request_failure_count,
|
||||
count(*) filter (
|
||||
where sr.outcome = 'failed' and sr.attempt_count > 0
|
||||
and (coalesce(sr.error_kind, 'failed') in
|
||||
('rate_limited', 'quota_exhausted')
|
||||
or coalesce(sr.upstream_status_code,
|
||||
sr.client_status_code) = 429)
|
||||
)::bigint as request_rate_limited_count,
|
||||
count(*) filter (
|
||||
where sr.outcome = 'failed' and sr.attempt_count > 0
|
||||
and (coalesce(sr.error_kind, 'failed') in
|
||||
('authentication', 'authorization', 'invalid_credential')
|
||||
or coalesce(sr.upstream_status_code,
|
||||
sr.client_status_code) in (401, 403))
|
||||
)::bigint as request_auth_failure_count,
|
||||
count(*) filter (
|
||||
where sr.outcome = 'failed' and sr.attempt_count > 0
|
||||
and coalesce(sr.upstream_status_code,
|
||||
sr.client_status_code) between 500 and 599
|
||||
)::bigint as request_provider_5xx_count,
|
||||
count(*) filter (
|
||||
where {fact} and sr.cost_source = 'provider_reported'
|
||||
)::bigint as provider_reported_count,
|
||||
count(*) filter (
|
||||
where {fact} and sr.cost_source = 'calculated'
|
||||
)::bigint as calculated_count,
|
||||
count(*) filter (
|
||||
where {fact} and sr.cost_source = 'unavailable'
|
||||
)::bigint as unavailable_count
|
||||
from selected_requests sr
|
||||
), ops_aggregate as (
|
||||
select coalesce(sum(coalesce(oe.occurrence_count, 1)), 0)::bigint
|
||||
as ops_failure_count,
|
||||
coalesce(sum(coalesce(oe.occurrence_count, 1)) filter (
|
||||
where oe.failure_kind in ('rate_limited', 'quota_exhausted')
|
||||
or oe.status_code = 429
|
||||
), 0)::bigint as ops_rate_limited_count,
|
||||
coalesce(sum(coalesce(oe.occurrence_count, 1)) filter (
|
||||
where oe.failure_kind in
|
||||
('authentication', 'authorization', 'invalid_credential')
|
||||
or oe.status_code in (401, 403)
|
||||
), 0)::bigint as ops_auth_failure_count,
|
||||
coalesce(sum(coalesce(oe.occurrence_count, 1)) filter (
|
||||
where oe.status_code between 500 and 599
|
||||
), 0)::bigint as ops_provider_5xx_count
|
||||
from ops_events oe
|
||||
join selected_requests sr on sr.id = oe.model_request_id
|
||||
union all
|
||||
select 1::bigint, coalesce(sr.error_kind, 'failed'),
|
||||
coalesce(sr.upstream_status_code, sr.client_status_code)
|
||||
from selected_requests sr
|
||||
where sr.outcome = 'failed' and sr.attempt_count > 0
|
||||
)
|
||||
select coalesce((select sum(attempt_count) from selected_requests), 0)::bigint
|
||||
as attempt_count,
|
||||
coalesce((select count(*) from selected_requests
|
||||
where outcome = 'succeeded' and attempt_count > 0), 0)::bigint
|
||||
as success_count,
|
||||
coalesce((select sum(occurrences) from failures), 0)::bigint as failure_count,
|
||||
coalesce((select count(*) from selected_requests
|
||||
where outcome = 'cancelled' and attempt_count > 0), 0)::bigint
|
||||
as cancelled_count,
|
||||
coalesce((select count(*) from selected_requests
|
||||
where outcome = 'incomplete' and attempt_count > 0), 0)::bigint
|
||||
as incomplete_count,
|
||||
coalesce((select sum(occurrences) from failures
|
||||
where failure_kind in ('rate_limited', 'quota_exhausted')
|
||||
or status_code = 429), 0)::bigint as rate_limited_count,
|
||||
coalesce((select sum(occurrences) from failures
|
||||
where failure_kind in ('authentication', 'authorization',
|
||||
'invalid_credential')
|
||||
or status_code in (401, 403)), 0)::bigint as auth_failure_count,
|
||||
coalesce((select sum(occurrences) from failures
|
||||
where status_code between 500 and 599), 0)::bigint
|
||||
select requests.attempt_count,
|
||||
requests.success_count,
|
||||
requests.request_failure_count + ops.ops_failure_count as failure_count,
|
||||
requests.cancelled_count,
|
||||
requests.incomplete_count,
|
||||
requests.request_rate_limited_count + ops.ops_rate_limited_count
|
||||
as rate_limited_count,
|
||||
requests.request_auth_failure_count + ops.ops_auth_failure_count
|
||||
as auth_failure_count,
|
||||
requests.request_provider_5xx_count + ops.ops_provider_5xx_count
|
||||
as provider_5xx_count,
|
||||
coalesce((select count(*) from usage_facts
|
||||
where cost_source = 'provider_reported'), 0)::bigint
|
||||
as provider_reported_count,
|
||||
coalesce((select count(*) from usage_facts
|
||||
where cost_source = 'calculated'), 0)::bigint
|
||||
as calculated_count,
|
||||
coalesce((select count(*) from usage_facts
|
||||
where cost_source = 'unavailable'), 0)::bigint
|
||||
as unavailable_count",
|
||||
);
|
||||
requests.provider_reported_count,
|
||||
requests.calculated_count,
|
||||
requests.unavailable_count
|
||||
from request_aggregate requests
|
||||
cross join ops_aggregate ops"
|
||||
));
|
||||
let row = query
|
||||
.build()
|
||||
.fetch_one(pool)
|
||||
@@ -448,18 +503,20 @@ pub(crate) async fn provider_observations(
|
||||
range: ObservabilityRange,
|
||||
filter: &UsageRecordFilter,
|
||||
) -> StoreResult<Vec<ProviderObservation>> {
|
||||
let mut query = QueryBuilder::<Postgres>::new(
|
||||
let fact = completed_usage_fact_predicate("mr");
|
||||
// 请求、attempt 与失败覆盖执行审计;token 只累计完整交付的用量事实。
|
||||
let mut query = QueryBuilder::<Postgres>::new(format!(
|
||||
"select coalesce(mr.provider_kind, 'unrouted') as provider_kind,
|
||||
count(*)::bigint as request_count,
|
||||
coalesce(sum(mr.attempt_count), 0)::bigint as attempt_count,
|
||||
count(*) filter (where mr.outcome = 'failed')::bigint as failure_count,
|
||||
coalesce(sum(mr.total_tokens), 0)::bigint as total_tokens
|
||||
from model_requests mr where mr.started_at >= ",
|
||||
);
|
||||
coalesce(sum(mr.total_tokens) filter (where {fact}), 0)::bigint
|
||||
as total_tokens
|
||||
from model_requests mr where mr.started_at >= "
|
||||
));
|
||||
query.push_bind(range.start);
|
||||
query.push(" and mr.started_at < ");
|
||||
query.push_bind(range.end);
|
||||
push_completed_usage_fact_filter(&mut query, "mr");
|
||||
push_usage_filter(&mut query, filter, "mr");
|
||||
query.push(" group by coalesce(mr.provider_kind, 'unrouted') order by request_count desc");
|
||||
let rows = query
|
||||
|
||||
@@ -2,43 +2,52 @@
|
||||
|
||||
use super::super::*;
|
||||
|
||||
pub(crate) const OPS_ERRORS_CTE: &str = "with errors as (
|
||||
select 'model_request'::text as source,
|
||||
mr.id as event_id, mr.id as request_id,
|
||||
nullif(mr.attempt_count, 0) as attempt_index,
|
||||
mr.client_api_key_ref, 'model_request'::text as component, mr.operation,
|
||||
mr.endpoint,
|
||||
mr.provider_kind,
|
||||
mr.provider_account_ref,
|
||||
mr.provider_account_name_snapshot as provider_account_name,
|
||||
mr.provider_account_email_snapshot as provider_account_email,
|
||||
mr.provider_account_authentication_kind_snapshot
|
||||
as provider_account_authentication_kind,
|
||||
mr.upstream_model_id, mr.upstream_transport,
|
||||
coalesce(mr.error_kind, 'failed') as failure_kind,
|
||||
mr.client_status_code, mr.upstream_status_code,
|
||||
mr.provider_error_code, mr.client_response_id, mr.upstream_request_id,
|
||||
mr.latency_ms, coalesce(mr.error_message, mr.error_kind, 'request failed') as message,
|
||||
1::integer as occurrence_count,
|
||||
coalesce(mr.completed_at, mr.started_at) as occurred_at,
|
||||
'model_request:' || mr.id as stable_sort_id
|
||||
from model_requests mr where mr.outcome = 'failed'
|
||||
union all
|
||||
select 'ops_event'::text, oe.id, oe.model_request_id, oe.attempt_index,
|
||||
mr.client_api_key_ref, oe.component, oe.operation,
|
||||
null::text as endpoint,
|
||||
oe.provider_kind, oe.provider_account_ref,
|
||||
oe.provider_account_name_snapshot as provider_account_name,
|
||||
oe.provider_account_email_snapshot as provider_account_email,
|
||||
oe.provider_account_authentication_kind_snapshot
|
||||
as provider_account_authentication_kind,
|
||||
oe.upstream_model_id, null::text, oe.failure_kind,
|
||||
null::integer as client_status_code, oe.status_code as upstream_status_code,
|
||||
oe.provider_error_code, mr.client_response_id, oe.upstream_request_id,
|
||||
oe.latency_ms, oe.message, oe.occurrence_count, oe.created_at,
|
||||
'ops_event:' || oe.id
|
||||
from ops_events oe left join model_requests mr on mr.id = oe.model_request_id
|
||||
)";
|
||||
const REQUEST_ERROR_SELECT: &str = "select 'model_request'::text as source,
|
||||
mr.id as event_id, mr.id as request_id,
|
||||
nullif(mr.attempt_count, 0) as attempt_index,
|
||||
mr.client_api_key_ref, 'model_request'::text as component, mr.operation,
|
||||
mr.protocol, mr.client_transport, mr.requested_model_id, mr.service_tier,
|
||||
mr.endpoint, mr.provider_kind, mr.provider_account_ref,
|
||||
mr.provider_account_name_snapshot as provider_account_name,
|
||||
mr.provider_account_email_snapshot as provider_account_email,
|
||||
mr.provider_account_authentication_kind_snapshot
|
||||
as provider_account_authentication_kind,
|
||||
mr.upstream_model_id, mr.upstream_transport,
|
||||
coalesce(mr.error_kind, 'failed') as failure_kind,
|
||||
mr.client_status_code, mr.upstream_status_code,
|
||||
mr.provider_error_code, mr.client_response_id, mr.upstream_request_id,
|
||||
mr.latency_ms,
|
||||
coalesce(mr.error_message, mr.error_kind, 'request failed') as message,
|
||||
1::integer as occurrence_count,
|
||||
host(mr.client_ip) as client_ip, mr.user_agent,
|
||||
mr.reasoning_effort, mr.reasoning_preset, mr.request_kind,
|
||||
mr.subagent_kind, mr.compact,
|
||||
mr.completed_at as occurred_at,
|
||||
'model_request:' || mr.id as stable_sort_id
|
||||
from model_requests mr
|
||||
where mr.outcome = 'failed'";
|
||||
|
||||
const OPS_EVENT_SELECT: &str = "select 'ops_event'::text as source,
|
||||
oe.id as event_id, oe.model_request_id as request_id, oe.attempt_index,
|
||||
mr.client_api_key_ref, oe.component, oe.operation,
|
||||
mr.protocol, mr.client_transport, mr.requested_model_id, mr.service_tier,
|
||||
mr.endpoint, oe.provider_kind, oe.provider_account_ref,
|
||||
oe.provider_account_name_snapshot as provider_account_name,
|
||||
oe.provider_account_email_snapshot as provider_account_email,
|
||||
oe.provider_account_authentication_kind_snapshot
|
||||
as provider_account_authentication_kind,
|
||||
oe.upstream_model_id, null::text as upstream_transport, oe.failure_kind,
|
||||
null::integer as client_status_code, oe.status_code as upstream_status_code,
|
||||
oe.provider_error_code, mr.client_response_id, oe.upstream_request_id,
|
||||
oe.latency_ms, oe.message, oe.occurrence_count,
|
||||
host(mr.client_ip) as client_ip, mr.user_agent,
|
||||
mr.reasoning_effort, mr.reasoning_preset, mr.request_kind,
|
||||
mr.subagent_kind, mr.compact,
|
||||
oe.created_at as occurred_at,
|
||||
'ops_event:' || oe.id as stable_sort_id
|
||||
from ops_events oe
|
||||
left join model_requests mr on mr.id = oe.model_request_id
|
||||
where true";
|
||||
|
||||
pub(crate) async fn list_ops_errors(
|
||||
pool: &PgPool,
|
||||
@@ -46,65 +55,31 @@ pub(crate) async fn list_ops_errors(
|
||||
) -> StoreResult<OpsErrorPage> {
|
||||
query.filter.validate()?;
|
||||
let total = count_ops_errors(pool, query.range, &query.filter).await?;
|
||||
let mut statement = QueryBuilder::<Postgres>::new(OPS_ERRORS_CTE);
|
||||
statement.push(
|
||||
" select source, event_id, request_id, attempt_index, client_api_key_ref,
|
||||
component, operation, endpoint, provider_kind, provider_account_ref,
|
||||
provider_account_name, provider_account_email,
|
||||
provider_account_authentication_kind,
|
||||
upstream_model_id, upstream_transport,
|
||||
failure_kind, client_status_code, upstream_status_code,
|
||||
provider_error_code, client_response_id,
|
||||
upstream_request_id, latency_ms, message, occurrence_count, occurred_at,
|
||||
stable_sort_id
|
||||
from errors e where e.occurred_at >= ",
|
||||
);
|
||||
statement.push_bind(query.range.start);
|
||||
statement.push(" and e.occurred_at < ");
|
||||
statement.push_bind(query.range.end);
|
||||
push_ops_filter(&mut statement, &query.filter);
|
||||
if let Some(cursor) = &query.cursor {
|
||||
statement.push(" and (e.occurred_at, e.stable_sort_id) < (");
|
||||
statement.push_bind(cursor.observed_at);
|
||||
statement.push(", ");
|
||||
statement.push_bind(cursor.stable_id.clone());
|
||||
statement.push(")");
|
||||
}
|
||||
statement.push(" order by e.occurred_at desc, e.stable_sort_id desc limit ");
|
||||
statement.push_bind(i64::from(query.page_size.get()) + 1);
|
||||
if query.cursor.is_none() {
|
||||
let offset = u64::from(query.page.get() - 1)
|
||||
.checked_mul(u64::from(query.page_size.get()))
|
||||
.and_then(|value| i64::try_from(value).ok())
|
||||
.ok_or_else(|| invalid("ops error page offset is too large"))?;
|
||||
statement.push(" offset ");
|
||||
statement.push_bind(offset);
|
||||
}
|
||||
let offset = observability_page_offset(query.current_page, query.page_size)?;
|
||||
let mut statement = QueryBuilder::<Postgres>::new("select * from (");
|
||||
statement.push(REQUEST_ERROR_SELECT);
|
||||
push_request_error_predicates(&mut statement, query.range, &query.filter);
|
||||
statement.push(" union all ");
|
||||
statement.push(OPS_EVENT_SELECT);
|
||||
push_ops_event_predicates(&mut statement, query.range, &query.filter);
|
||||
statement.push(") e order by occurred_at desc, stable_sort_id desc limit ");
|
||||
statement.push_bind(i64::from(query.page_size.get()));
|
||||
statement.push(" offset ");
|
||||
statement.push_bind(offset);
|
||||
let rows = statement
|
||||
.build()
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("list ops errors"))?;
|
||||
let mut items = rows
|
||||
let items = rows
|
||||
.iter()
|
||||
.map(ops_error_from_row)
|
||||
.collect::<StoreResult<Vec<_>>>()?;
|
||||
let has_more = items.len() > usize::from(query.page_size.get());
|
||||
if has_more {
|
||||
items.pop();
|
||||
}
|
||||
let next_cursor = if has_more {
|
||||
items
|
||||
.last()
|
||||
.map(|item| ObservabilityCursor::new(item.occurred_at, item.stable_sort_id.clone()))
|
||||
.transpose()?
|
||||
} else {
|
||||
None
|
||||
};
|
||||
Ok(OpsErrorPage {
|
||||
items,
|
||||
current_page: query.current_page,
|
||||
page_size: query.page_size.get(),
|
||||
total,
|
||||
next_cursor,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -113,12 +88,15 @@ pub(crate) async fn count_ops_errors(
|
||||
range: ObservabilityRange,
|
||||
filter: &OpsErrorFilter,
|
||||
) -> StoreResult<u64> {
|
||||
let mut statement = QueryBuilder::<Postgres>::new(OPS_ERRORS_CTE);
|
||||
statement.push(" select count(*)::bigint from errors e where e.occurred_at >= ");
|
||||
statement.push_bind(range.start);
|
||||
statement.push(" and e.occurred_at < ");
|
||||
statement.push_bind(range.end);
|
||||
push_ops_filter(&mut statement, filter);
|
||||
let mut statement = QueryBuilder::<Postgres>::new(
|
||||
"select coalesce(sum(source_count), 0)::bigint from (select count(*)::bigint as source_count from model_requests mr where mr.outcome = 'failed'",
|
||||
);
|
||||
push_request_error_predicates(&mut statement, range, filter);
|
||||
statement.push(
|
||||
" union all select count(*)::bigint as source_count from ops_events oe left join model_requests mr on mr.id = oe.model_request_id where true",
|
||||
);
|
||||
push_ops_event_predicates(&mut statement, range, filter);
|
||||
statement.push(") counts");
|
||||
let total = statement
|
||||
.build_query_scalar::<i64>()
|
||||
.fetch_one(pool)
|
||||
@@ -127,45 +105,141 @@ pub(crate) async fn count_ops_errors(
|
||||
to_u64(total)
|
||||
}
|
||||
|
||||
pub(crate) fn push_ops_filter(statement: &mut QueryBuilder<Postgres>, filter: &OpsErrorFilter) {
|
||||
fn push_request_error_predicates(
|
||||
statement: &mut QueryBuilder<Postgres>,
|
||||
range: ObservabilityRange,
|
||||
filter: &OpsErrorFilter,
|
||||
) {
|
||||
push_range(statement, "mr.completed_at", range);
|
||||
for (column, value) in [
|
||||
("client_api_key_ref", &filter.client_api_key_ref),
|
||||
("request_id", &filter.request_id),
|
||||
("provider_account_ref", &filter.provider_account_ref),
|
||||
("provider_kind", &filter.provider_kind),
|
||||
("operation", &filter.operation),
|
||||
("upstream_transport", &filter.transport),
|
||||
("upstream_request_id", &filter.upstream_request_id),
|
||||
("failure_kind", &filter.failure_kind),
|
||||
("mr.client_api_key_ref", &filter.client_api_key_ref),
|
||||
("mr.id", &filter.request_id),
|
||||
("mr.provider_account_ref", &filter.provider_account_ref),
|
||||
("mr.provider_kind", &filter.provider_kind),
|
||||
("mr.operation", &filter.operation),
|
||||
("mr.endpoint", &filter.endpoint),
|
||||
("mr.upstream_transport", &filter.transport),
|
||||
("mr.upstream_request_id", &filter.upstream_request_id),
|
||||
] {
|
||||
if let Some(value) = value {
|
||||
statement.push(format!(" and e.{column} = "));
|
||||
statement.push_bind(value.clone());
|
||||
}
|
||||
push_text_equality(statement, column, value);
|
||||
}
|
||||
if let Some(value) = &filter.response_id {
|
||||
statement.push(" and e.client_response_id = ");
|
||||
statement.push_bind(value.as_bytes().to_vec());
|
||||
}
|
||||
if let Some(model) = &filter.model {
|
||||
statement.push(" and e.upstream_model_id = ");
|
||||
statement.push_bind(model.clone());
|
||||
if let Some(value) = &filter.failure_kind {
|
||||
statement.push(" and coalesce(mr.error_kind, 'failed') = ");
|
||||
statement.push_bind(value.clone());
|
||||
}
|
||||
push_response_id_filter(statement, "mr.client_response_id", filter);
|
||||
push_text_equality(statement, "mr.upstream_model_id", &filter.model);
|
||||
if let Some(index) = filter.attempt_index {
|
||||
statement.push(" and e.attempt_index = ");
|
||||
statement.push(" and nullif(mr.attempt_count, 0) = ");
|
||||
statement.push_bind(i32::try_from(index).unwrap_or(i32::MAX));
|
||||
}
|
||||
if let Some(status) = filter.status_code {
|
||||
statement.push(" and e.upstream_status_code = ");
|
||||
statement.push(" and mr.upstream_status_code = ");
|
||||
statement.push_bind(i32::from(status));
|
||||
}
|
||||
if let Some(search) = &filter.search {
|
||||
statement.push(
|
||||
" and lower(concat_ws(' ', e.event_id, e.request_id, e.client_api_key_ref,
|
||||
e.component, e.operation, e.provider_kind, e.provider_account_ref,
|
||||
e.upstream_model_id, e.failure_kind, e.provider_error_code,
|
||||
e.upstream_request_id, e.message)) like ",
|
||||
push_prefix_search(
|
||||
statement,
|
||||
&[
|
||||
"mr.id",
|
||||
"mr.client_api_key_ref",
|
||||
"mr.provider_account_ref",
|
||||
"mr.upstream_request_id",
|
||||
"mr.provider_error_code",
|
||||
],
|
||||
search,
|
||||
);
|
||||
statement.push_bind(format!("%{}%", search.to_lowercase()));
|
||||
}
|
||||
}
|
||||
|
||||
fn push_ops_event_predicates(
|
||||
statement: &mut QueryBuilder<Postgres>,
|
||||
range: ObservabilityRange,
|
||||
filter: &OpsErrorFilter,
|
||||
) {
|
||||
push_range(statement, "oe.created_at", range);
|
||||
for (column, value) in [
|
||||
("mr.client_api_key_ref", &filter.client_api_key_ref),
|
||||
("oe.model_request_id", &filter.request_id),
|
||||
("oe.provider_account_ref", &filter.provider_account_ref),
|
||||
("oe.provider_kind", &filter.provider_kind),
|
||||
("oe.operation", &filter.operation),
|
||||
("mr.endpoint", &filter.endpoint),
|
||||
("oe.upstream_request_id", &filter.upstream_request_id),
|
||||
("oe.failure_kind", &filter.failure_kind),
|
||||
] {
|
||||
push_text_equality(statement, column, value);
|
||||
}
|
||||
// Ops events do not persist upstream transport. A transport filter therefore
|
||||
// intentionally excludes this source instead of matching an unrelated request fact.
|
||||
if filter.transport.is_some() {
|
||||
statement.push(" and false");
|
||||
}
|
||||
push_response_id_filter(statement, "mr.client_response_id", filter);
|
||||
push_text_equality(statement, "oe.upstream_model_id", &filter.model);
|
||||
if let Some(index) = filter.attempt_index {
|
||||
statement.push(" and oe.attempt_index = ");
|
||||
statement.push_bind(i32::try_from(index).unwrap_or(i32::MAX));
|
||||
}
|
||||
if let Some(status) = filter.status_code {
|
||||
statement.push(" and oe.status_code = ");
|
||||
statement.push_bind(i32::from(status));
|
||||
}
|
||||
if let Some(search) = &filter.search {
|
||||
push_prefix_search(
|
||||
statement,
|
||||
&[
|
||||
"oe.id",
|
||||
"oe.model_request_id",
|
||||
"mr.client_api_key_ref",
|
||||
"oe.provider_account_ref",
|
||||
"oe.upstream_request_id",
|
||||
"oe.provider_error_code",
|
||||
],
|
||||
search,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
fn push_range(statement: &mut QueryBuilder<Postgres>, column: &str, range: ObservabilityRange) {
|
||||
statement.push(format!(" and {column} >= "));
|
||||
statement.push_bind(range.start);
|
||||
statement.push(format!(" and {column} < "));
|
||||
statement.push_bind(range.end);
|
||||
}
|
||||
|
||||
fn push_text_equality(
|
||||
statement: &mut QueryBuilder<Postgres>,
|
||||
column: &str,
|
||||
value: &Option<String>,
|
||||
) {
|
||||
if let Some(value) = value {
|
||||
statement.push(format!(" and {column} = "));
|
||||
statement.push_bind(value.clone());
|
||||
}
|
||||
}
|
||||
|
||||
fn push_response_id_filter(
|
||||
statement: &mut QueryBuilder<Postgres>,
|
||||
column: &str,
|
||||
filter: &OpsErrorFilter,
|
||||
) {
|
||||
if let Some(value) = &filter.response_id {
|
||||
statement.push(format!(" and {column} = "));
|
||||
statement.push_bind(value.as_bytes().to_vec());
|
||||
}
|
||||
}
|
||||
|
||||
fn push_prefix_search(statement: &mut QueryBuilder<Postgres>, columns: &[&str], value: &str) {
|
||||
let pattern = literal_prefix_pattern(value);
|
||||
statement.push(" and (");
|
||||
for (index, column) in columns.iter().enumerate() {
|
||||
if index > 0 {
|
||||
statement.push(" or ");
|
||||
}
|
||||
statement.push(format!("{column} like "));
|
||||
statement.push_bind(pattern.clone());
|
||||
statement.push(" escape '\\'");
|
||||
}
|
||||
statement.push(")");
|
||||
}
|
||||
|
||||
@@ -75,34 +75,54 @@ pub(crate) fn push_usage_filter(
|
||||
query.push_bind(value.clone());
|
||||
}
|
||||
if let Some(value) = &filter.search {
|
||||
query.push(format!(
|
||||
" and lower(concat_ws(' ', {alias}.id, {alias}.client_api_key_ref,
|
||||
{alias}.requested_model_id, {alias}.provider_kind,
|
||||
{alias}.provider_account_ref, {alias}.upstream_model_id,
|
||||
encode({alias}.client_response_id, 'escape'), {alias}.upstream_request_id,
|
||||
{alias}.error_kind, {alias}.provider_error_code,
|
||||
{alias}.error_message)) like "
|
||||
));
|
||||
query.push_bind(format!("%{}%", value.to_lowercase()));
|
||||
let pattern = literal_prefix_pattern(value);
|
||||
query.push(format!(" and ({alias}.id like "));
|
||||
query.push_bind(pattern.clone());
|
||||
for column in [
|
||||
"client_api_key_ref",
|
||||
"provider_account_ref",
|
||||
"requested_model_id",
|
||||
"upstream_model_id",
|
||||
"upstream_request_id",
|
||||
] {
|
||||
query.push(format!(" escape '\\' or {alias}.{column} like "));
|
||||
query.push_bind(pattern.clone());
|
||||
}
|
||||
query.push(" escape '\\')");
|
||||
}
|
||||
}
|
||||
|
||||
/// 只把已完整交付给客户端的成功响应投影为使用事实。
|
||||
///
|
||||
/// `model_requests` 同时承担执行审计:包括上游已发送、但尚未收到首个事件就断开的
|
||||
/// WebSocket 请求。这些失败仍必须留给 Ops Errors 和调度健康度分析,不能混入用量、
|
||||
/// 成本、账号使用次数或请求明细。
|
||||
pub(crate) fn completed_usage_fact_predicate(alias: &str) -> String {
|
||||
pub(crate) fn literal_prefix_pattern(value: &str) -> String {
|
||||
format!(
|
||||
"{alias}.outcome = 'succeeded' and {alias}.downstream_committed_at is not null and {alias}.client_status_code between 200 and 399"
|
||||
"{}%",
|
||||
value
|
||||
.replace('\\', "\\\\")
|
||||
.replace('%', "\\%")
|
||||
.replace('_', "\\_")
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn push_completed_usage_fact_filter(query: &mut QueryBuilder<Postgres>, alias: &str) {
|
||||
query.push(format!(" and {}", completed_usage_fact_predicate(alias)));
|
||||
}
|
||||
pub(crate) const USAGE_LIST_RECORD_SELECT: &str =
|
||||
"select mr.id, 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,
|
||||
mr.provider_account_email_snapshot as provider_account_email,
|
||||
mr.provider_account_authentication_kind_snapshot
|
||||
as provider_account_authentication_kind,
|
||||
mr.upstream_model_id, mr.upstream_transport, mr.service_tier,
|
||||
mr.input_tokens, mr.output_tokens, mr.cached_tokens, mr.cache_write_tokens,
|
||||
mr.reasoning_tokens, mr.image_input_tokens, mr.image_output_tokens,
|
||||
mr.total_tokens, mr.cost_source, mr.cost_amount::text, mr.cost_currency,
|
||||
mr.transport_decision_wait_ms, mr.connect_ms, mr.headers_ms,
|
||||
mr.first_event_ms, mr.first_reasoning_ms, mr.first_text_ms, mr.first_token_ms,
|
||||
mr.provider_processing_ms, mr.latency_ms, mr.admission_decision_ms,
|
||||
mr.account_selection_wait_ms, mr.capacity_used_slots, mr.capacity_total_slots,
|
||||
host(mr.client_ip) as client_ip, mr.user_agent,
|
||||
mr.reasoning_effort, mr.reasoning_preset, mr.subagent_kind, mr.compact,
|
||||
mr.started_at
|
||||
from model_requests mr";
|
||||
|
||||
pub(crate) const USAGE_RECORD_SELECT: &str =
|
||||
pub(crate) const USAGE_RECORD_DETAIL_SELECT: &str =
|
||||
"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,
|
||||
@@ -125,7 +145,7 @@ pub(crate) const USAGE_RECORD_SELECT: &str =
|
||||
mr.first_event_ms, mr.first_reasoning_ms, mr.first_text_ms, mr.first_token_ms,
|
||||
mr.provider_processing_ms, mr.latency_ms, mr.admission_decision_ms,
|
||||
mr.account_selection_wait_ms, mr.capacity_used_slots, mr.capacity_total_slots,
|
||||
mr.client_ip::text as client_ip,
|
||||
host(mr.client_ip) as client_ip,
|
||||
mr.user_agent, mr.reasoning_effort, mr.reasoning_preset, mr.request_kind,
|
||||
mr.subagent_kind, mr.compact, mr.image_generation_requested,
|
||||
mr.image_generation_succeeded, mr.started_at, mr.deadline_at, mr.completed_at
|
||||
@@ -137,56 +157,42 @@ pub(crate) async fn list_usage_records(
|
||||
) -> StoreResult<UsageRecordPage> {
|
||||
query.filter.validate()?;
|
||||
let total = count_usage_records(pool, query.range, &query.filter).await?;
|
||||
let mut statement = QueryBuilder::<Postgres>::new(USAGE_RECORD_SELECT);
|
||||
let items = list_usage_record_items(pool, &query).await?;
|
||||
Ok(UsageRecordPage {
|
||||
items,
|
||||
current_page: query.current_page,
|
||||
page_size: query.page_size.get(),
|
||||
total,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) async fn list_usage_record_items(
|
||||
pool: &PgPool,
|
||||
query: &UsageRecordQuery,
|
||||
) -> StoreResult<Vec<UsageListRecord>> {
|
||||
query.filter.validate()?;
|
||||
let offset = observability_page_offset(query.current_page, query.page_size)?;
|
||||
let mut statement = QueryBuilder::<Postgres>::new(USAGE_LIST_RECORD_SELECT);
|
||||
statement.push(" where mr.started_at >= ");
|
||||
statement.push_bind(query.range.start);
|
||||
statement.push(" and mr.started_at < ");
|
||||
statement.push_bind(query.range.end);
|
||||
push_completed_usage_fact_filter(&mut statement, "mr");
|
||||
push_usage_filter(&mut statement, &query.filter, "mr");
|
||||
if let Some(cursor) = &query.cursor {
|
||||
statement.push(" and (mr.started_at, mr.id) < (");
|
||||
statement.push_bind(cursor.observed_at);
|
||||
statement.push(", ");
|
||||
statement.push_bind(cursor.stable_id.clone());
|
||||
statement.push(")");
|
||||
}
|
||||
statement.push(" order by mr.started_at desc, mr.id desc limit ");
|
||||
statement.push_bind(i64::from(query.page_size.get()) + 1);
|
||||
if query.cursor.is_none() {
|
||||
let offset = u64::from(query.page.get() - 1)
|
||||
.checked_mul(u64::from(query.page_size.get()))
|
||||
.and_then(|value| i64::try_from(value).ok())
|
||||
.ok_or_else(|| invalid("usage page offset is too large"))?;
|
||||
statement.push(" offset ");
|
||||
statement.push_bind(offset);
|
||||
}
|
||||
statement.push_bind(i64::from(query.page_size.get()));
|
||||
statement.push(" offset ");
|
||||
statement.push_bind(offset);
|
||||
let rows = statement
|
||||
.build()
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("list usage records"))?;
|
||||
let mut items = rows
|
||||
let items = rows
|
||||
.iter()
|
||||
.map(usage_record_from_row)
|
||||
.map(usage_list_record_from_row)
|
||||
.collect::<StoreResult<Vec<_>>>()?;
|
||||
let has_more = items.len() > usize::from(query.page_size.get());
|
||||
if has_more {
|
||||
items.pop();
|
||||
}
|
||||
let next_cursor = if has_more {
|
||||
items
|
||||
.last()
|
||||
.map(|item| ObservabilityCursor::new(item.started_at, item.id.clone()))
|
||||
.transpose()?
|
||||
} else {
|
||||
None
|
||||
};
|
||||
Ok(UsageRecordPage {
|
||||
items,
|
||||
total,
|
||||
next_cursor,
|
||||
})
|
||||
Ok(items)
|
||||
}
|
||||
|
||||
pub(crate) async fn count_usage_records(
|
||||
@@ -216,7 +222,7 @@ pub(crate) async fn usage_record_detail(
|
||||
) -> StoreResult<UsageRecordDetail> {
|
||||
require_nonempty("model request", "id", request_id)?;
|
||||
validate_text(request_id, MAX_FILTER_BYTES, "request ID")?;
|
||||
let mut statement = QueryBuilder::<Postgres>::new(USAGE_RECORD_SELECT);
|
||||
let mut statement = QueryBuilder::<Postgres>::new(USAGE_RECORD_DETAIL_SELECT);
|
||||
statement.push(" where mr.id = ");
|
||||
statement.push_bind(request_id.to_owned());
|
||||
push_completed_usage_fact_filter(&mut statement, "mr");
|
||||
@@ -341,47 +347,85 @@ pub(crate) async fn usage_diagnostics(
|
||||
) -> StoreResult<Vec<DiagnosticObservation>> {
|
||||
filter.validate()?;
|
||||
let dimension_sql = diagnostic_dimension_sql(dimension);
|
||||
let mut statement = QueryBuilder::<Postgres>::new("select ");
|
||||
let completed_usage = completed_usage_fact_predicate("mr");
|
||||
let mut statement = QueryBuilder::<Postgres>::new("with matched as (select ");
|
||||
statement.push(dimension_sql);
|
||||
statement.push(
|
||||
" as dimension_name,
|
||||
count(*)::bigint as request_count,
|
||||
count(*) filter (where mr.outcome = 'succeeded')::bigint as success_count,
|
||||
count(*) filter (where mr.outcome = 'failed')::bigint as failure_count,
|
||||
coalesce(sum(mr.attempt_count), 0)::bigint as attempt_count,
|
||||
coalesce(sum(mr.total_tokens), 0)::bigint as total_tokens,
|
||||
round(avg(mr.latency_ms))::bigint as average_latency_ms,
|
||||
round(percentile_cont(0.95) within group (order by mr.latency_ms))::bigint
|
||||
as latency_p95_ms,
|
||||
round(percentile_cont(0.95) within group (order by mr.first_token_ms))::bigint
|
||||
as first_token_p95_ms,
|
||||
count(*) filter (where mr.outcome in ('cancelled', 'incomplete'))::bigint
|
||||
as non_completion_count,
|
||||
coalesce(sum(greatest(mr.attempt_count - 1, 0)), 0)::bigint as retry_count,
|
||||
count(*) filter (where mr.cost_source = 'provider_reported')::bigint
|
||||
as provider_reported_count,
|
||||
count(*) filter (where mr.cost_source = 'calculated')::bigint
|
||||
as calculated_count,
|
||||
count(*) filter (where mr.cost_source = 'unavailable')::bigint
|
||||
as unavailable_count
|
||||
statement.push(format!(
|
||||
" as dimension_name, mr.outcome, mr.attempt_count, mr.total_tokens,
|
||||
mr.latency_ms, mr.first_token_ms, mr.cost_source, mr.cost_amount,
|
||||
mr.cost_currency, mr.downstream_committed_at, mr.client_transport,
|
||||
mr.client_status_code, ({completed_usage}) as is_completed_usage
|
||||
from model_requests mr where mr.started_at >= ",
|
||||
);
|
||||
));
|
||||
statement.push_bind(range.start);
|
||||
statement.push(" and mr.started_at < ");
|
||||
statement.push_bind(range.end);
|
||||
push_diagnostic_dimension_filter(&mut statement, dimension);
|
||||
push_usage_filter(&mut statement, filter, "mr");
|
||||
statement.push(" group by dimension_name order by request_count desc, dimension_name limit ");
|
||||
statement.push(
|
||||
"), aggregated as (
|
||||
select dimension_name, cost_currency,
|
||||
grouping(cost_currency)::integer as currency_grouping,
|
||||
count(*)::bigint as request_count,
|
||||
count(*) filter (where outcome = 'succeeded')::bigint as success_count,
|
||||
count(*) filter (where outcome = 'failed')::bigint as failure_count,
|
||||
coalesce(sum(attempt_count), 0)::bigint as attempt_count,
|
||||
coalesce(sum(total_tokens) filter (where is_completed_usage), 0)::bigint
|
||||
as total_tokens,
|
||||
round(avg(latency_ms) filter (where is_completed_usage))::bigint
|
||||
as average_latency_ms,
|
||||
round(percentile_cont(0.95) within group (order by latency_ms)
|
||||
filter (where is_completed_usage))::bigint
|
||||
as latency_p95_ms,
|
||||
round(percentile_cont(0.95) within group (order by first_token_ms)
|
||||
filter (where is_completed_usage))::bigint
|
||||
as first_token_p95_ms,
|
||||
count(*) filter (where outcome in ('cancelled', 'incomplete'))::bigint
|
||||
as non_completion_count,
|
||||
coalesce(sum(greatest(attempt_count - 1, 0)), 0)::bigint as retry_count,
|
||||
count(*) filter (
|
||||
where is_completed_usage and cost_source = 'provider_reported'
|
||||
)::bigint
|
||||
as provider_reported_count,
|
||||
count(*) filter (
|
||||
where is_completed_usage and cost_source = 'calculated'
|
||||
)::bigint
|
||||
as calculated_count,
|
||||
count(*) filter (
|
||||
where is_completed_usage and cost_source = 'unavailable'
|
||||
)::bigint
|
||||
as unavailable_count,
|
||||
sum(cost_amount) filter (
|
||||
where is_completed_usage
|
||||
)::text as amount
|
||||
from matched
|
||||
group by dimension_name, grouping sets ((), (cost_currency))
|
||||
), selected_dimensions as (
|
||||
select dimension_name,
|
||||
row_number() over (order by request_count desc, dimension_name)
|
||||
as sort_position
|
||||
from aggregated
|
||||
where currency_grouping = 1
|
||||
order by request_count desc, dimension_name limit ",
|
||||
);
|
||||
statement.push_bind(DIAGNOSTIC_LIMIT);
|
||||
statement.push(
|
||||
")
|
||||
select aggregated.*
|
||||
from aggregated
|
||||
join selected_dimensions selected using (dimension_name)
|
||||
order by selected.sort_position, currency_grouping desc, cost_currency nulls last",
|
||||
);
|
||||
let rows = statement
|
||||
.build()
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("load usage diagnostics"))?;
|
||||
let mut observations = rows
|
||||
.iter()
|
||||
.map(|row| {
|
||||
Ok(DiagnosticObservation {
|
||||
let mut observations = Vec::with_capacity(DIAGNOSTIC_LIMIT as usize);
|
||||
let mut costs = HashMap::<String, Vec<CurrencyCostTotal>>::new();
|
||||
for row in &rows {
|
||||
match get::<i32>(row, "currency_grouping")? {
|
||||
1 => observations.push(DiagnosticObservation {
|
||||
key: get(row, "dimension_name")?,
|
||||
name: get(row, "dimension_name")?,
|
||||
request_count: unsigned(row, "request_count")?,
|
||||
@@ -396,9 +440,19 @@ pub(crate) async fn usage_diagnostics(
|
||||
retry_count: unsigned(row, "retry_count")?,
|
||||
cost_coverage: coverage_from_row(row)?,
|
||||
costs: Vec::new(),
|
||||
})
|
||||
})
|
||||
.collect::<StoreResult<Vec<_>>>()?;
|
||||
}),
|
||||
0 if get::<Option<String>>(row, "cost_currency")?.is_some()
|
||||
&& get::<Option<String>>(row, "amount")?.is_some() =>
|
||||
{
|
||||
costs
|
||||
.entry(get(row, "dimension_name")?)
|
||||
.or_default()
|
||||
.push(cost_from_row(row)?);
|
||||
}
|
||||
0 => {}
|
||||
_ => return Err(postgres_unavailable("decode usage diagnostic grouping")),
|
||||
}
|
||||
}
|
||||
let mut display_names = match dimension {
|
||||
DiagnosticDimension::Account => {
|
||||
diagnostic_account_display_names(
|
||||
@@ -426,17 +480,7 @@ pub(crate) async fn usage_diagnostics(
|
||||
if let Some(name) = display_names.remove(&observation.key) {
|
||||
observation.name = name;
|
||||
}
|
||||
}
|
||||
let positions = observations
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(index, item)| (item.key.clone(), index))
|
||||
.collect::<HashMap<_, _>>();
|
||||
let costs = diagnostic_costs(pool, range, filter, dimension_sql, dimension).await?;
|
||||
for (key, values) in costs {
|
||||
if let Some(position) = positions.get(&key).copied() {
|
||||
observations[position].costs = values;
|
||||
}
|
||||
observation.costs = costs.remove(&observation.key).unwrap_or_default();
|
||||
}
|
||||
Ok(observations)
|
||||
}
|
||||
@@ -449,17 +493,27 @@ pub(crate) async fn diagnostic_account_display_names(
|
||||
return Ok(HashMap::new());
|
||||
}
|
||||
let rows = sqlx::query(
|
||||
"select provider_account_ref,
|
||||
"select requested.account_id as provider_account_ref,
|
||||
coalesce(
|
||||
(array_agg(provider_account_email_snapshot order by started_at desc)
|
||||
filter (where provider_account_email_snapshot is not null))[1],
|
||||
(array_agg(provider_account_name_snapshot order by started_at desc)
|
||||
filter (where provider_account_name_snapshot is not null))[1],
|
||||
provider_account_ref
|
||||
(
|
||||
select request.provider_account_email_snapshot
|
||||
from model_requests request
|
||||
where request.provider_account_ref = requested.account_id
|
||||
and request.provider_account_email_snapshot is not null
|
||||
order by request.started_at desc, request.id desc
|
||||
limit 1
|
||||
),
|
||||
(
|
||||
select request.provider_account_name_snapshot
|
||||
from model_requests request
|
||||
where request.provider_account_ref = requested.account_id
|
||||
and request.provider_account_name_snapshot is not null
|
||||
order by request.started_at desc, request.id desc
|
||||
limit 1
|
||||
),
|
||||
requested.account_id
|
||||
) as display_name
|
||||
from model_requests
|
||||
where provider_account_ref = any($1::text[])
|
||||
group by provider_account_ref",
|
||||
from unnest($1::text[]) as requested(account_id)",
|
||||
)
|
||||
.bind(account_ids)
|
||||
.fetch_all(pool)
|
||||
@@ -494,44 +548,6 @@ pub(crate) async fn diagnostic_api_key_display_names(
|
||||
Ok(display_names)
|
||||
}
|
||||
|
||||
pub(crate) async fn diagnostic_costs(
|
||||
pool: &PgPool,
|
||||
range: ObservabilityRange,
|
||||
filter: &UsageRecordFilter,
|
||||
dimension_sql: &'static str,
|
||||
dimension: DiagnosticDimension,
|
||||
) -> StoreResult<HashMap<String, Vec<CurrencyCostTotal>>> {
|
||||
let mut statement = QueryBuilder::<Postgres>::new("select ");
|
||||
statement.push(dimension_sql);
|
||||
statement.push(
|
||||
" as dimension_name, mr.cost_currency, sum(mr.cost_amount)::text as amount
|
||||
from model_requests mr where mr.started_at >= ",
|
||||
);
|
||||
statement.push_bind(range.start);
|
||||
statement.push(" and mr.started_at < ");
|
||||
statement.push_bind(range.end);
|
||||
statement.push(" and mr.cost_amount is not null and mr.cost_currency is not null");
|
||||
push_completed_usage_fact_filter(&mut statement, "mr");
|
||||
push_diagnostic_dimension_filter(&mut statement, dimension);
|
||||
push_usage_filter(&mut statement, filter, "mr");
|
||||
statement.push(
|
||||
" group by dimension_name, mr.cost_currency order by dimension_name, mr.cost_currency",
|
||||
);
|
||||
let rows = statement
|
||||
.build()
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("load diagnostic costs"))?;
|
||||
let mut result = HashMap::<String, Vec<CurrencyCostTotal>>::new();
|
||||
for row in rows {
|
||||
result
|
||||
.entry(get(&row, "dimension_name")?)
|
||||
.or_default()
|
||||
.push(cost_from_row(&row)?);
|
||||
}
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
pub(crate) fn diagnostic_dimension_sql(dimension: DiagnosticDimension) -> &'static str {
|
||||
match dimension {
|
||||
DiagnosticDimension::Provider => "coalesce(mr.provider_kind, 'unrouted')",
|
||||
|
||||
@@ -0,0 +1,89 @@
|
||||
//! 管理观测查询的 PostgreSQL 连接槽位预算。
|
||||
|
||||
use std::{
|
||||
future::Future,
|
||||
sync::Arc,
|
||||
time::{Duration, Instant},
|
||||
};
|
||||
|
||||
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
|
||||
|
||||
use crate::{StoreBackend, StoreError, StoreResult};
|
||||
|
||||
/// 所有管理观测入口共享的 PostgreSQL 连接槽位预算。
|
||||
#[derive(Clone)]
|
||||
pub struct ObservabilityQueryBudget {
|
||||
slots: Arc<Semaphore>,
|
||||
wait_timeout: Duration,
|
||||
}
|
||||
|
||||
impl ObservabilityQueryBudget {
|
||||
/// 预算表示观测类 SQL 最多可同时占用的连接数。
|
||||
pub fn try_new(max_connections: u32, wait_timeout: Duration) -> StoreResult<Self> {
|
||||
if max_connections == 0 || wait_timeout.is_zero() {
|
||||
return Err(StoreError::InvalidData {
|
||||
entity: "observability query budget",
|
||||
message: "requires a positive connection count and wait timeout".to_owned(),
|
||||
});
|
||||
}
|
||||
let permits = usize::try_from(max_connections).map_err(|_| StoreError::InvalidData {
|
||||
entity: "observability query budget",
|
||||
message: "connection budget does not fit this platform".to_owned(),
|
||||
})?;
|
||||
Ok(Self {
|
||||
slots: Arc::new(Semaphore::new(permits)),
|
||||
wait_timeout,
|
||||
})
|
||||
}
|
||||
|
||||
/// 在共享管理观测预算内执行一条 PostgreSQL 查询。
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// 预算在等待期限内没有空闲槽位、预算已关闭,或查询本身失败时返回错误。
|
||||
pub async fn run<T, F>(&self, operation: &'static str, query: F) -> StoreResult<T>
|
||||
where
|
||||
F: Future<Output = StoreResult<T>>,
|
||||
{
|
||||
let _slot = self.acquire(operation).await?;
|
||||
query.await
|
||||
}
|
||||
|
||||
async fn acquire(&self, operation: &'static str) -> StoreResult<OwnedSemaphorePermit> {
|
||||
let started_at = Instant::now();
|
||||
let acquisition = Arc::clone(&self.slots).acquire_owned();
|
||||
match tokio::time::timeout(self.wait_timeout, acquisition).await {
|
||||
Ok(Ok(permit)) => {
|
||||
tracing::debug!(
|
||||
query_class = "observability",
|
||||
operation,
|
||||
wait_milliseconds = elapsed_milliseconds(started_at),
|
||||
available_slots = self.slots.available_permits(),
|
||||
"PostgreSQL query budget acquired"
|
||||
);
|
||||
Ok(permit)
|
||||
}
|
||||
Ok(Err(_)) => Err(StoreError::Unavailable {
|
||||
backend: StoreBackend::PostgreSql,
|
||||
message: "observability PostgreSQL connection budget is closed".to_owned(),
|
||||
}),
|
||||
Err(_) => {
|
||||
tracing::warn!(
|
||||
query_class = "observability",
|
||||
operation,
|
||||
wait_milliseconds = elapsed_milliseconds(started_at),
|
||||
available_slots = self.slots.available_permits(),
|
||||
"PostgreSQL query budget exhausted"
|
||||
);
|
||||
Err(StoreError::Unavailable {
|
||||
backend: StoreBackend::PostgreSql,
|
||||
message: "observability PostgreSQL connection budget is exhausted".to_owned(),
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn elapsed_milliseconds(started_at: Instant) -> u64 {
|
||||
u64::try_from(started_at.elapsed().as_millis()).unwrap_or(u64::MAX)
|
||||
}
|
||||
@@ -6,6 +6,7 @@ use std::sync::Arc;
|
||||
use gateway_core::provider_ports::ProviderCooldownPort;
|
||||
|
||||
use super::*;
|
||||
use crate::postgres::ObservabilityQueryBudget;
|
||||
|
||||
/// Admin 账号用例所需的公共账号、留存观测与 revision 事务能力。
|
||||
///
|
||||
@@ -17,17 +18,23 @@ pub struct PgAdminAccountStore {
|
||||
observability: PgObservabilityRepository,
|
||||
control_plane: PgControlPlaneRepository,
|
||||
cooldowns: Option<Arc<dyn ProviderCooldownPort>>,
|
||||
query_budget: ObservabilityQueryBudget,
|
||||
}
|
||||
|
||||
impl PgAdminAccountStore {
|
||||
#[must_use]
|
||||
pub fn new(pool: PgPool, cooldowns: Option<Arc<dyn ProviderCooldownPort>>) -> Self {
|
||||
pub fn new(
|
||||
pool: PgPool,
|
||||
cooldowns: Option<Arc<dyn ProviderCooldownPort>>,
|
||||
query_budget: ObservabilityQueryBudget,
|
||||
) -> Self {
|
||||
Self {
|
||||
pool: pool.clone(),
|
||||
accounts: PgProviderAccountRepository::new(pool.clone()),
|
||||
observability: PgObservabilityRepository::new(pool.clone(), cooldowns.clone()),
|
||||
observability: PgObservabilityRepository::new(pool.clone(), None, query_budget.clone()),
|
||||
control_plane: PgControlPlaneRepository::new(pool),
|
||||
cooldowns,
|
||||
query_budget,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -97,62 +104,61 @@ impl PgAdminAccountStore {
|
||||
.iter()
|
||||
.map(|window| window.account_id.clone())
|
||||
.collect::<Vec<_>>();
|
||||
let (usage_rows, cost_rows, model_rows, model_cost_rows) = futures::try_join!(
|
||||
sqlx::query(ACCOUNT_USAGE_BY_WINDOWS_SQL)
|
||||
.bind(account_ids.clone())
|
||||
.bind(keys.clone())
|
||||
.bind(starts.clone())
|
||||
.bind(ends.clone())
|
||||
.fetch_all(&self.pool),
|
||||
sqlx::query(ACCOUNT_USAGE_COSTS_BY_WINDOWS_SQL)
|
||||
.bind(account_ids.clone())
|
||||
.bind(keys.clone())
|
||||
.bind(starts.clone())
|
||||
.bind(ends.clone())
|
||||
.fetch_all(&self.pool),
|
||||
sqlx::query(ACCOUNT_USAGE_MODELS_BY_WINDOWS_SQL)
|
||||
.bind(account_ids.clone())
|
||||
.bind(keys.clone())
|
||||
.bind(starts.clone())
|
||||
.bind(ends.clone())
|
||||
.fetch_all(&self.pool),
|
||||
sqlx::query(ACCOUNT_USAGE_MODEL_COSTS_BY_WINDOWS_SQL)
|
||||
.bind(account_ids)
|
||||
.bind(keys)
|
||||
.bind(starts)
|
||||
.bind(ends)
|
||||
.fetch_all(&self.pool),
|
||||
)
|
||||
.map_err(|_| {
|
||||
admin_store_error(
|
||||
ENTITY,
|
||||
postgres_unavailable("load provider account quota window usage"),
|
||||
)
|
||||
})?;
|
||||
let rows = self
|
||||
.query_budget
|
||||
.run("load account usage windows", async {
|
||||
sqlx::query(sqlx::AssertSqlSafe(account_usage_by_windows_sql()))
|
||||
.bind(account_ids)
|
||||
.bind(keys)
|
||||
.bind(starts)
|
||||
.bind(ends)
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("load provider account quota window usage"))
|
||||
})
|
||||
.await
|
||||
.map_err(|error| admin_store_error(ENTITY, error))?;
|
||||
let mut usage_rows = Vec::with_capacity(windows.len());
|
||||
let mut costs_by_window = BTreeMap::<(String, String), Vec<AccountCost>>::new();
|
||||
for row in &cost_rows {
|
||||
let (key, cost) = admin_account_usage_window_cost(row)?;
|
||||
costs_by_window.entry(key).or_default().push(cost);
|
||||
}
|
||||
let mut model_costs = BTreeMap::<(String, String, String), Vec<AccountCost>>::new();
|
||||
for row in &model_cost_rows {
|
||||
let (key, cost) = admin_account_usage_window_model_cost(row)?;
|
||||
model_costs.entry(key).or_default().push(cost);
|
||||
}
|
||||
let mut models_by_window = BTreeMap::<(String, String), Vec<AccountModelUsage>>::new();
|
||||
for row in &model_rows {
|
||||
let ((account_id, window_key, model), mut usage) =
|
||||
admin_account_usage_window_model(row)?;
|
||||
usage.costs = model_costs
|
||||
.remove(&(account_id.clone(), window_key.clone(), model))
|
||||
.unwrap_or_default();
|
||||
models_by_window
|
||||
.entry((account_id, window_key))
|
||||
.or_default()
|
||||
.push(usage);
|
||||
for row in &rows {
|
||||
let model_grouping = window_usage_value::<i32>(row, "model_grouping")?;
|
||||
let currency_grouping = window_usage_value::<i32>(row, "currency_grouping")?;
|
||||
match (model_grouping, currency_grouping) {
|
||||
(1, 1) => usage_rows.push(row),
|
||||
(1, 0) if window_usage_value::<Option<String>>(row, "cost_currency")?.is_some() => {
|
||||
let (key, cost) = admin_account_usage_window_cost(row)?;
|
||||
costs_by_window.entry(key).or_default().push(cost);
|
||||
}
|
||||
(0, 1) if window_usage_value::<Option<String>>(row, "model")?.is_some() => {
|
||||
let ((account_id, window_key, _), usage) =
|
||||
admin_account_usage_window_model(row)?;
|
||||
models_by_window
|
||||
.entry((account_id, window_key))
|
||||
.or_default()
|
||||
.push(usage);
|
||||
}
|
||||
(0, 0)
|
||||
if window_usage_value::<Option<String>>(row, "model")?.is_some()
|
||||
&& window_usage_value::<Option<String>>(row, "cost_currency")?
|
||||
.is_some() =>
|
||||
{
|
||||
let (key, cost) = admin_account_usage_window_model_cost(row)?;
|
||||
model_costs.entry(key).or_default().push(cost);
|
||||
}
|
||||
(0 | 1, 0 | 1) => {}
|
||||
_ => {
|
||||
return Err(AdminStoreError::new(
|
||||
AdminStoreErrorKind::Unavailable,
|
||||
ENTITY,
|
||||
"account usage window query returned an invalid grouping marker",
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
let mut results = usage_rows
|
||||
.iter()
|
||||
.into_iter()
|
||||
.map(admin_account_usage_window)
|
||||
.collect::<AdminStoreResult<Vec<_>>>()?;
|
||||
for result in &mut results {
|
||||
@@ -162,6 +168,21 @@ impl PgAdminAccountStore {
|
||||
result.usage.models = models_by_window
|
||||
.remove(&(result.account_id.clone(), result.key.clone()))
|
||||
.unwrap_or_default();
|
||||
for model in &mut result.usage.models {
|
||||
model.costs = model_costs
|
||||
.remove(&(
|
||||
result.account_id.clone(),
|
||||
result.key.clone(),
|
||||
model.model.clone(),
|
||||
))
|
||||
.unwrap_or_default();
|
||||
}
|
||||
result.usage.models.sort_by(|left, right| {
|
||||
right
|
||||
.request_count
|
||||
.cmp(&left.request_count)
|
||||
.then_with(|| left.model.cmp(&right.model))
|
||||
});
|
||||
}
|
||||
Ok(results)
|
||||
}
|
||||
@@ -331,67 +352,15 @@ impl PgAdminAccountStore {
|
||||
}
|
||||
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]
|
||||
impl AccountStore for PgAdminAccountStore {
|
||||
async fn list_accounts(&self, query: AdminAccountListQuery) -> AdminStoreResult<AccountPage> {
|
||||
async fn list_accounts(
|
||||
&self,
|
||||
query: AdminAccountListQuery,
|
||||
runtime: gateway_admin::model::accounts::AccountRuntimeSnapshot,
|
||||
) -> AdminStoreResult<AccountPage> {
|
||||
if query.page == 0 {
|
||||
return Err(AdminStoreError::new(
|
||||
AdminStoreErrorKind::Invalid,
|
||||
@@ -399,93 +368,57 @@ impl AccountStore for PgAdminAccountStore {
|
||||
"page number must be positive",
|
||||
));
|
||||
}
|
||||
let (control_plane, accounts) = futures::try_join!(
|
||||
self.control_plane.load_control_plane(),
|
||||
self.accounts.list_provider_accounts(None, true),
|
||||
)
|
||||
.map_err(|error| admin_store_error(ENTITY, error))?;
|
||||
let now = Utc::now();
|
||||
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
|
||||
let rate_limited_until = runtime
|
||||
.rate_limited_until
|
||||
.into_iter()
|
||||
.filter_map(|account| {
|
||||
let until = rate_limited_until.get(&account.id).copied();
|
||||
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<_>>();
|
||||
|
||||
if query.sort.is_some_and(|sort| {
|
||||
matches!(
|
||||
sort.field,
|
||||
AdminAccountSortField::Usage | AdminAccountSortField::LastUsedAt
|
||||
)
|
||||
}) {
|
||||
let range = retained_usage_range(control_plane.settings.usage_retention_days, now)?;
|
||||
let account_ids = items
|
||||
.iter()
|
||||
.map(|item| item.account.id.clone())
|
||||
.collect::<Vec<_>>();
|
||||
let mut usage_by_account = self
|
||||
.usage_observations(range, &account_ids)
|
||||
.await?
|
||||
.into_iter()
|
||||
.map(|usage| (usage.account_id.clone(), usage))
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
for item in &mut items {
|
||||
item.usage = usage_by_account.remove(&item.account.id);
|
||||
}
|
||||
}
|
||||
sort_admin_account_items(&mut items, query.sort);
|
||||
|
||||
let total = u64::try_from(items.len()).unwrap_or(u64::MAX);
|
||||
let page_size = usize::from(query.page_size.get());
|
||||
let offset = u64::from(query.page - 1).saturating_mul(u64::from(query.page_size.get()));
|
||||
let offset = usize::try_from(offset).unwrap_or(usize::MAX);
|
||||
let items = items
|
||||
.into_iter()
|
||||
.skip(offset)
|
||||
.take(page_size)
|
||||
.collect::<Vec<_>>();
|
||||
let item_ids = items
|
||||
.map(|(account_id, until)| (account_id, until.into()))
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
let now_system_time = std::time::SystemTime::from(now);
|
||||
let active_rate_limited_ids = rate_limited_until
|
||||
.iter()
|
||||
.map(|item| item.account.id.clone())
|
||||
.filter(|(_, until)| **until > now_system_time)
|
||||
.map(|(account_id, _)| account_id.clone())
|
||||
.collect::<Vec<_>>();
|
||||
let page =
|
||||
load_admin_account_page(&self.pool, &query, now, active_rate_limited_ids).await?;
|
||||
let item_ids = page
|
||||
.accounts
|
||||
.iter()
|
||||
.map(|account| account.id.clone())
|
||||
.collect::<Vec<_>>();
|
||||
let mut groups_by_account = self.account_groups_by_account(&item_ids).await?;
|
||||
let items = items
|
||||
let items = page
|
||||
.accounts
|
||||
.into_iter()
|
||||
.map(|item| {
|
||||
let account_id = item.account.id.clone();
|
||||
let mut account = admin_account_record(item.account)?;
|
||||
.map(|summary| {
|
||||
let account_id = summary.id.clone();
|
||||
let projection = account_status_projection(
|
||||
&summary,
|
||||
now.into(),
|
||||
rate_limited_until.get(&account_id).copied(),
|
||||
);
|
||||
let mut account = admin_account_record(summary)?;
|
||||
account.groups = groups_by_account.remove(&account_id).unwrap_or_default();
|
||||
Ok(AccountPageItem {
|
||||
account,
|
||||
projection: item.projection,
|
||||
projection,
|
||||
})
|
||||
})
|
||||
.collect::<AdminStoreResult<Vec<_>>>()?;
|
||||
Ok(AccountPage {
|
||||
config_revision: admin_revision(control_plane.settings.config_revision)?,
|
||||
config_revision: page.config_revision,
|
||||
items,
|
||||
total,
|
||||
summary,
|
||||
total: page.total,
|
||||
summary: page.summary,
|
||||
})
|
||||
}
|
||||
|
||||
async fn load_account(&self, account_id: &str) -> AdminStoreResult<Option<AccountPageItem>> {
|
||||
async fn load_account(
|
||||
&self,
|
||||
account_id: &str,
|
||||
runtime: gateway_admin::model::accounts::AccountRuntimeSnapshot,
|
||||
) -> AdminStoreResult<Option<AccountPageItem>> {
|
||||
let record = self
|
||||
.accounts
|
||||
.load_provider_account(account_id)
|
||||
@@ -495,14 +428,11 @@ impl AccountStore for PgAdminAccountStore {
|
||||
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 rate_limited_until = runtime
|
||||
.rate_limited_until
|
||||
.get(account_id)
|
||||
.copied()
|
||||
.map(Into::into);
|
||||
let projection = account_status_projection(&record.summary, now.into(), rate_limited_until);
|
||||
let account_id = record.summary.id.clone();
|
||||
let mut groups = self
|
||||
|
||||
@@ -0,0 +1,291 @@
|
||||
//! 管理账号目录的数据库分页 read model。
|
||||
|
||||
use super::*;
|
||||
|
||||
pub(crate) struct AdminAccountPageRows {
|
||||
pub(crate) config_revision: AdminRevision,
|
||||
pub(crate) accounts: Vec<ProviderAccountSummary>,
|
||||
pub(crate) total: u64,
|
||||
pub(crate) summary: AccountSummary,
|
||||
}
|
||||
|
||||
pub(crate) async fn load_admin_account_page(
|
||||
pool: &PgPool,
|
||||
query: &AdminAccountListQuery,
|
||||
now: DateTime<Utc>,
|
||||
active_rate_limited_ids: Vec<String>,
|
||||
) -> AdminStoreResult<AdminAccountPageRows> {
|
||||
validate_group_filter(pool, query.group_filter.as_ref()).await?;
|
||||
|
||||
let provider_kind = query
|
||||
.provider_kind
|
||||
.as_ref()
|
||||
.map(|provider| provider.as_str().to_owned());
|
||||
let search_prefix = query.search.as_deref().map(literal_prefix_pattern);
|
||||
let status = query.status.map(|status| status.as_str().to_owned());
|
||||
let (group_mode, group_id) = match query.group_filter.as_ref() {
|
||||
None => (0_i16, None),
|
||||
Some(AccountGroupFilter::Group(group_id)) => (1_i16, Some(group_id.as_str().to_owned())),
|
||||
Some(AccountGroupFilter::Ungrouped) => (2_i16, None),
|
||||
};
|
||||
let page_size = i64::from(query.page_size.get());
|
||||
let offset = i64::from(query.page - 1).saturating_mul(page_size);
|
||||
let order_by = admin_account_order(query.sort);
|
||||
let usage_join = query.sort.is_some_and(|sort| {
|
||||
matches!(
|
||||
sort.field,
|
||||
AdminAccountSortField::Usage | AdminAccountSortField::LastUsedAt
|
||||
)
|
||||
});
|
||||
let completed_usage = completed_usage_fact_predicate("mr");
|
||||
let usage_join = if usage_join {
|
||||
format!(
|
||||
"left join (
|
||||
select mr.provider_account_ref,
|
||||
coalesce(sum(coalesce(
|
||||
mr.total_tokens,
|
||||
coalesce(mr.input_tokens, 0) + coalesce(mr.output_tokens, 0)
|
||||
)), 0)::bigint as total_tokens,
|
||||
max(mr.started_at) as last_used_at
|
||||
from model_requests mr
|
||||
where mr.started_at >= $2 - make_interval(days => settings.usage_retention_days::int)
|
||||
and mr.started_at < $2
|
||||
and {completed_usage}
|
||||
group by mr.provider_account_ref
|
||||
) usage on usage.provider_account_ref = a.id"
|
||||
)
|
||||
} else {
|
||||
String::new()
|
||||
};
|
||||
|
||||
let statement = format!(
|
||||
"with account_statuses as (
|
||||
select a.id,
|
||||
case
|
||||
when not a.enabled then 'disabled'
|
||||
when a.credential_state <> 'ready'
|
||||
or a.access_token_expires_at <= $2 then 'error'
|
||||
when a.quota_access_state = 'exhausted' then 'quota_exhausted'
|
||||
when a.id = any($1::text[]) then 'rate_limited'
|
||||
else 'normal'
|
||||
end as admin_status
|
||||
from provider_accounts a
|
||||
),
|
||||
global_summary as (
|
||||
select count(*)::bigint as summary_total,
|
||||
count(*) filter (where admin_status = 'normal')::bigint as summary_normal,
|
||||
count(*) filter (where admin_status = 'quota_exhausted')::bigint
|
||||
as summary_quota_exhausted,
|
||||
count(*) filter (where admin_status = 'rate_limited')::bigint
|
||||
as summary_rate_limited,
|
||||
count(*) filter (where admin_status = 'disabled')::bigint
|
||||
as summary_disabled,
|
||||
count(*) filter (where admin_status = 'error')::bigint as summary_error
|
||||
from account_statuses
|
||||
),
|
||||
filtered as materialized (
|
||||
select status.id, status.admin_status
|
||||
from account_statuses status
|
||||
join provider_accounts a on a.id = status.id
|
||||
where ($3::text is null or a.provider_kind = $3)
|
||||
and ($4::text is null or
|
||||
lower(a.id) like $4 escape '\\' or
|
||||
lower(a.name) like $4 escape '\\' or
|
||||
lower(coalesce(a.email, '')) like $4 escape '\\' or
|
||||
lower(coalesce(a.upstream_account_id, '')) like $4 escape '\\' or
|
||||
lower(coalesce(a.upstream_user_id, '')) like $4 escape '\\')
|
||||
and ($5::text is null or status.admin_status = $5)
|
||||
and ($6::smallint = 0 or
|
||||
($6 = 1 and exists (
|
||||
select 1
|
||||
from account_group_accounts membership
|
||||
where membership.provider_account_id = a.id
|
||||
and membership.account_group_id = $7
|
||||
)) or
|
||||
($6 = 2 and not exists (
|
||||
select 1
|
||||
from account_group_accounts membership
|
||||
where membership.provider_account_id = a.id
|
||||
)))
|
||||
),
|
||||
filtered_total as (
|
||||
select count(*)::bigint as filtered_total from filtered
|
||||
),
|
||||
settings as (
|
||||
select config_revision, usage_retention_days
|
||||
from runtime_settings
|
||||
where id = 1
|
||||
)
|
||||
select a.id, a.provider_kind, a.name, a.email, a.upstream_user_id,
|
||||
a.upstream_account_id, a.plan_type, a.authentication_kind,
|
||||
a.credential_revision, a.has_refresh_token, a.access_token_expires_at,
|
||||
a.next_refresh_at, a.enabled, a.concurrency_limit, a.weight,
|
||||
a.credential_state, a.quota_access_state, a.quota_evidence,
|
||||
a.quota_access_observed_at, a.quota_reset_at, a.last_error_reason,
|
||||
a.last_error_message, a.credential_observed_at, a.created_at, a.updated_at,
|
||||
filtered_total.filtered_total,
|
||||
global_summary.summary_total, global_summary.summary_normal,
|
||||
global_summary.summary_quota_exhausted, global_summary.summary_rate_limited,
|
||||
global_summary.summary_disabled, global_summary.summary_error,
|
||||
settings.config_revision
|
||||
from filtered_total
|
||||
cross join global_summary
|
||||
cross join settings
|
||||
left join lateral (
|
||||
select ordered.*, row_number() over ()::bigint as page_position
|
||||
from (
|
||||
select filtered.id, filtered.admin_status
|
||||
from filtered
|
||||
join provider_accounts a on a.id = filtered.id
|
||||
{usage_join}
|
||||
order by {order_by}
|
||||
limit $8 offset $9
|
||||
) ordered
|
||||
) page on true
|
||||
left join provider_accounts a on a.id = page.id
|
||||
order by page.page_position"
|
||||
);
|
||||
// 动态片段只来自上面的封闭排序枚举与固定 usage predicate;所有请求值仍使用 bind。
|
||||
let rows = sqlx::query(sqlx::AssertSqlSafe(statement))
|
||||
.bind(active_rate_limited_ids)
|
||||
.bind(now)
|
||||
.bind(provider_kind)
|
||||
.bind(search_prefix)
|
||||
.bind(status)
|
||||
.bind(group_mode)
|
||||
.bind(group_id)
|
||||
.bind(page_size)
|
||||
.bind(offset)
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
.map_err(|_| admin_store_error(ENTITY, postgres_unavailable("load admin account page")))?;
|
||||
let metadata = rows.first().ok_or_else(|| {
|
||||
AdminStoreError::new(
|
||||
AdminStoreErrorKind::Unavailable,
|
||||
ENTITY,
|
||||
"admin account page metadata is unavailable",
|
||||
)
|
||||
})?;
|
||||
let config_revision = revision_from_row(metadata)?;
|
||||
let total = unsigned_metadata(metadata, "filtered_total")?;
|
||||
let summary = AccountSummary {
|
||||
total: unsigned_metadata(metadata, "summary_total")?,
|
||||
normal: unsigned_metadata(metadata, "summary_normal")?,
|
||||
quota_exhausted: unsigned_metadata(metadata, "summary_quota_exhausted")?,
|
||||
rate_limited: unsigned_metadata(metadata, "summary_rate_limited")?,
|
||||
disabled: unsigned_metadata(metadata, "summary_disabled")?,
|
||||
error: unsigned_metadata(metadata, "summary_error")?,
|
||||
};
|
||||
let accounts = rows
|
||||
.into_iter()
|
||||
.filter_map(|row| {
|
||||
row.try_get::<Option<String>, _>("id")
|
||||
.ok()
|
||||
.flatten()
|
||||
.map(|_| row)
|
||||
})
|
||||
.map(account_summary_from_row)
|
||||
.collect::<StoreResult<Vec<_>>>()
|
||||
.map_err(|error| admin_store_error(ENTITY, error))?;
|
||||
Ok(AdminAccountPageRows {
|
||||
config_revision,
|
||||
accounts,
|
||||
total,
|
||||
summary,
|
||||
})
|
||||
}
|
||||
|
||||
async fn validate_group_filter(
|
||||
pool: &PgPool,
|
||||
filter: Option<&AccountGroupFilter>,
|
||||
) -> AdminStoreResult<()> {
|
||||
let Some(AccountGroupFilter::Group(group_id)) = filter else {
|
||||
return Ok(());
|
||||
};
|
||||
let exists =
|
||||
sqlx::query_scalar::<_, bool>("select exists(select 1 from account_groups where id = $1)")
|
||||
.bind(group_id.as_str())
|
||||
.fetch_one(pool)
|
||||
.await
|
||||
.map_err(|_| {
|
||||
admin_store_error(
|
||||
ENTITY,
|
||||
postgres_unavailable("validate admin account group filter"),
|
||||
)
|
||||
})?;
|
||||
if exists {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(admin_store_error(
|
||||
ENTITY,
|
||||
StoreError::NotFound {
|
||||
entity: "account group",
|
||||
id: group_id.as_str().to_owned(),
|
||||
},
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
fn admin_account_order(sort: Option<AdminAccountSort>) -> String {
|
||||
let Some(sort) = sort else {
|
||||
return "a.id asc".to_owned();
|
||||
};
|
||||
let expression = match sort.field {
|
||||
AdminAccountSortField::Email => "a.email",
|
||||
AdminAccountSortField::Status => {
|
||||
"case filtered.admin_status when 'normal' then 0 when 'rate_limited' then 1 when 'quota_exhausted' then 2 when 'error' then 3 else 4 end"
|
||||
}
|
||||
AdminAccountSortField::PlanType => "a.plan_type",
|
||||
AdminAccountSortField::Usage => "usage.total_tokens",
|
||||
AdminAccountSortField::LastUsedAt => "usage.last_used_at",
|
||||
AdminAccountSortField::ExpiresAt => "a.access_token_expires_at",
|
||||
};
|
||||
match sort.direction {
|
||||
AdminSortDirection::Asc => format!("{expression} asc nulls first, a.id asc"),
|
||||
AdminSortDirection::Desc => format!("{expression} desc nulls last, a.id desc"),
|
||||
}
|
||||
}
|
||||
|
||||
fn literal_prefix_pattern(value: &str) -> String {
|
||||
let mut escaped = String::with_capacity(value.len().saturating_add(1));
|
||||
for character in value.to_lowercase().chars() {
|
||||
if matches!(character, '\\' | '%' | '_') {
|
||||
escaped.push('\\');
|
||||
}
|
||||
escaped.push(character);
|
||||
}
|
||||
escaped.push('%');
|
||||
escaped
|
||||
}
|
||||
|
||||
fn revision_from_row(row: &sqlx::postgres::PgRow) -> AdminStoreResult<AdminRevision> {
|
||||
let value = row.try_get::<i64, _>("config_revision").map_err(|_| {
|
||||
AdminStoreError::new(
|
||||
AdminStoreErrorKind::Invalid,
|
||||
ENTITY,
|
||||
"persisted config revision is invalid",
|
||||
)
|
||||
})?;
|
||||
let revision = Revision::new(u64::try_from(value).map_err(|_| {
|
||||
AdminStoreError::new(
|
||||
AdminStoreErrorKind::Invalid,
|
||||
ENTITY,
|
||||
"persisted config revision is invalid",
|
||||
)
|
||||
})?)
|
||||
.map_err(|error| admin_store_error(ENTITY, error))?;
|
||||
admin_revision(revision)
|
||||
}
|
||||
|
||||
fn unsigned_metadata(row: &sqlx::postgres::PgRow, column: &'static str) -> AdminStoreResult<u64> {
|
||||
row.try_get::<i64, _>(column)
|
||||
.ok()
|
||||
.and_then(|value| u64::try_from(value).ok())
|
||||
.ok_or_else(|| {
|
||||
AdminStoreError::new(
|
||||
AdminStoreErrorKind::Invalid,
|
||||
ENTITY,
|
||||
format!("persisted {column} is invalid"),
|
||||
)
|
||||
})
|
||||
}
|
||||
@@ -2,6 +2,15 @@
|
||||
|
||||
use super::*;
|
||||
|
||||
fn loaded_credential_from_record(
|
||||
record: ProviderAccountRecord,
|
||||
) -> Result<LoadedCredential, CoreStoreError> {
|
||||
Ok(LoadedCredential {
|
||||
account: core_account_from_summary(record.summary)?,
|
||||
credential: PlaintextCredential::new(record.provider_credentials_json.fields().clone()),
|
||||
})
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ProviderAccountStore for PgProviderAccountRepository {
|
||||
async fn create_account(&self, account: CoreNewProviderAccount) -> Result<(), CoreStoreError> {
|
||||
@@ -72,6 +81,35 @@ impl ProviderAccountStore for PgProviderAccountRepository {
|
||||
.collect()
|
||||
}
|
||||
|
||||
async fn list_refresh_candidates(
|
||||
&self,
|
||||
query: CoreProviderRefreshQuery,
|
||||
) -> Result<Vec<LoadedCredential>, CoreStoreError> {
|
||||
let excluded_account_ids = query
|
||||
.excluded_account_ids()
|
||||
.iter()
|
||||
.map(|account_id| account_id.as_str())
|
||||
.collect::<Vec<_>>();
|
||||
let rows = sqlx::query(REFRESH_CANDIDATES_SELECT)
|
||||
.bind(query.provider().as_str())
|
||||
.bind(DateTime::<Utc>::from(query.refresh_due_before()))
|
||||
.bind(DateTime::<Utc>::from(query.force_due_before()))
|
||||
.bind(DateTime::<Utc>::from(query.observed_at()))
|
||||
.bind(excluded_account_ids)
|
||||
.bind(i64::from(query.limit().get()))
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_err(|_| CoreStoreError::new(CoreStoreErrorKind::Unavailable))?;
|
||||
rows.into_iter()
|
||||
.map(account_record_from_row)
|
||||
.map(|record| {
|
||||
record
|
||||
.map_err(core_store_error)
|
||||
.and_then(loaded_credential_from_record)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
async fn load_credential(
|
||||
&self,
|
||||
account: &CoreProviderAccountId,
|
||||
@@ -93,10 +131,7 @@ impl ProviderAccountStore for PgProviderAccountRepository {
|
||||
.await
|
||||
.map_err(core_store_error)?
|
||||
.ok_or_else(|| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?;
|
||||
Ok(LoadedCredential {
|
||||
account: core_account_from_summary(record.summary)?,
|
||||
credential: PlaintextCredential::new(record.provider_credentials_json.fields().clone()),
|
||||
})
|
||||
loaded_credential_from_record(record)
|
||||
}
|
||||
|
||||
async fn compare_and_swap_credential(
|
||||
|
||||
@@ -1,139 +1,7 @@
|
||||
//! Admin 账号列表、用量与成本的领域映射。
|
||||
|
||||
use std::time::SystemTime;
|
||||
|
||||
use super::*;
|
||||
|
||||
pub(crate) struct AdminAccountListItem {
|
||||
pub(crate) account: ProviderAccountSummary,
|
||||
pub(crate) projection: gateway_core::engine::credential::AccountStatusProjection,
|
||||
pub(crate) usage: Option<ProviderAccountUsageObservation>,
|
||||
}
|
||||
|
||||
pub(crate) fn retained_usage_range(
|
||||
retention_days: u32,
|
||||
now: DateTime<Utc>,
|
||||
) -> AdminStoreResult<ObservabilityRange> {
|
||||
let duration = TimeDelta::try_days(i64::from(retention_days)).ok_or_else(|| {
|
||||
AdminStoreError::new(
|
||||
AdminStoreErrorKind::Invalid,
|
||||
ENTITY,
|
||||
"usage retention is outside the supported range",
|
||||
)
|
||||
})?;
|
||||
let start = now.checked_sub_signed(duration).ok_or_else(|| {
|
||||
AdminStoreError::new(
|
||||
AdminStoreErrorKind::Invalid,
|
||||
ENTITY,
|
||||
"usage retention start is outside the supported range",
|
||||
)
|
||||
})?;
|
||||
ObservabilityRange::new(start, now).map_err(|error| admin_store_error(ENTITY, error))
|
||||
}
|
||||
|
||||
pub(crate) fn account_matches_admin_query(
|
||||
account: &ProviderAccountSummary,
|
||||
status: AdminAccountStatus,
|
||||
query: &AdminAccountListQuery,
|
||||
) -> bool {
|
||||
let provider_matches = query
|
||||
.provider_kind
|
||||
.as_ref()
|
||||
.is_none_or(|provider| account.provider_kind == provider.as_str());
|
||||
let search_matches = query.search.as_ref().is_none_or(|search| {
|
||||
let search = search.to_lowercase();
|
||||
[
|
||||
Some(account.id.as_str()),
|
||||
Some(account.name.as_str()),
|
||||
account.email.as_deref(),
|
||||
account.upstream_account_id.as_deref(),
|
||||
account.upstream_user_id.as_deref(),
|
||||
]
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.any(|value| value.to_lowercase().contains(&search))
|
||||
});
|
||||
let status_matches = query.status.is_none_or(|expected| expected == status);
|
||||
provider_matches && search_matches && status_matches
|
||||
}
|
||||
|
||||
pub(crate) fn admin_account_summary(
|
||||
accounts: &[ProviderAccountSummary],
|
||||
now: DateTime<Utc>,
|
||||
rate_limited_until: &BTreeMap<String, SystemTime>,
|
||||
) -> AccountSummary {
|
||||
let mut summary = AccountSummary {
|
||||
total: u64::try_from(accounts.len()).unwrap_or(u64::MAX),
|
||||
normal: 0,
|
||||
quota_exhausted: 0,
|
||||
rate_limited: 0,
|
||||
disabled: 0,
|
||||
error: 0,
|
||||
};
|
||||
for account in accounts {
|
||||
let until = rate_limited_until.get(&account.id).copied();
|
||||
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);
|
||||
}
|
||||
AdminAccountStatus::RateLimited => {
|
||||
summary.rate_limited = summary.rate_limited.saturating_add(1);
|
||||
}
|
||||
AdminAccountStatus::Disabled => {
|
||||
summary.disabled = summary.disabled.saturating_add(1);
|
||||
}
|
||||
AdminAccountStatus::Error => summary.error = summary.error.saturating_add(1),
|
||||
}
|
||||
}
|
||||
summary
|
||||
}
|
||||
|
||||
pub(crate) fn sort_admin_account_items(
|
||||
items: &mut [AdminAccountListItem],
|
||||
sort: Option<AdminAccountSort>,
|
||||
) {
|
||||
let Some(sort) = sort else {
|
||||
items.sort_by(|left, right| left.account.id.cmp(&right.account.id));
|
||||
return;
|
||||
};
|
||||
items.sort_by(|left, right| {
|
||||
let ordering = match sort.field {
|
||||
AdminAccountSortField::Email => left.account.email.cmp(&right.account.email),
|
||||
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
|
||||
.as_ref()
|
||||
.and_then(|usage| usage.total_tokens)
|
||||
.cmp(&right.usage.as_ref().and_then(|usage| usage.total_tokens)),
|
||||
AdminAccountSortField::LastUsedAt => left
|
||||
.usage
|
||||
.as_ref()
|
||||
.and_then(|usage| usage.last_used_at.as_ref())
|
||||
.cmp(
|
||||
&right
|
||||
.usage
|
||||
.as_ref()
|
||||
.and_then(|usage| usage.last_used_at.as_ref()),
|
||||
),
|
||||
AdminAccountSortField::ExpiresAt => left
|
||||
.account
|
||||
.access_token_expires_at
|
||||
.cmp(&right.account.access_token_expires_at),
|
||||
}
|
||||
.then_with(|| left.account.id.cmp(&right.account.id));
|
||||
match sort.direction {
|
||||
AdminSortDirection::Asc => ordering,
|
||||
AdminSortDirection::Desc => ordering.reverse(),
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
pub(crate) fn admin_account_record(
|
||||
summary: ProviderAccountSummary,
|
||||
) -> AdminStoreResult<AccountRecord> {
|
||||
|
||||
@@ -1,10 +1,6 @@
|
||||
//! 明文 `provider_accounts` 与凭证 revision CAS 的唯一 PostgreSQL owner。
|
||||
|
||||
use std::{
|
||||
collections::{BTreeMap, BTreeSet},
|
||||
fmt,
|
||||
str::FromStr,
|
||||
};
|
||||
use std::{collections::BTreeSet, fmt, str::FromStr};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, TimeDelta, Utc};
|
||||
@@ -16,10 +12,9 @@ use gateway_admin::{
|
||||
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,
|
||||
AccountSummary, AccountUpdateResult, AccountUsage, AccountUsageWindowQuery,
|
||||
AccountUsageWindowResult, AccountsUpdateResult, BatchUpdateAccounts, DeleteAccounts,
|
||||
SortDirection as AdminSortDirection, UpdateAccount,
|
||||
},
|
||||
observability::{
|
||||
CostCoverage as AdminCostCoverage, DecimalAmount as AdminDecimalAmount, TimeRange,
|
||||
@@ -34,8 +29,6 @@ use gateway_admin::{
|
||||
},
|
||||
ports::store::{AccountStore, AdminStoreError, AdminStoreErrorKind, AdminStoreResult},
|
||||
};
|
||||
use sqlx::{PgPool, Postgres, Row, Transaction};
|
||||
|
||||
use gateway_core::engine::credential::{
|
||||
AccountConcurrencyLimit, AccountErrorReason, AccountStateChange, AccountWeight,
|
||||
CredentialCasOutcome, CredentialCasUpdate, CredentialCasUpdateParts,
|
||||
@@ -43,11 +36,13 @@ use gateway_core::engine::credential::{
|
||||
NewProviderAccount as CoreNewProviderAccount, OpaqueProviderData, PlaintextCredential,
|
||||
ProviderAccount as CoreProviderAccount, ProviderAccountId as CoreProviderAccountId,
|
||||
ProviderAccountIdentity, ProviderAccountStore,
|
||||
ProviderAccountUpdate as CoreProviderAccountUpdate, QuotaAccessChange, QuotaAccessState,
|
||||
ProviderAccountUpdate as CoreProviderAccountUpdate,
|
||||
ProviderRefreshQuery as CoreProviderRefreshQuery, QuotaAccessChange, QuotaAccessState,
|
||||
QuotaEvidence, QuotaObservation, QuotaObservationTouch, QuotaState, QuotaWriteOutcome,
|
||||
};
|
||||
use gateway_core::error::{StoreError as CoreStoreError, StoreErrorKind as CoreStoreErrorKind};
|
||||
use gateway_core::routing::{AccountGroupId, ProviderKind};
|
||||
use sqlx::{PgPool, Postgres, Row, Transaction};
|
||||
|
||||
use crate::{
|
||||
ConflictKind, JsonObject, Revision, StoreError, StoreResult, admin_revision, admin_store_error,
|
||||
@@ -59,10 +54,11 @@ use super::{
|
||||
ObservabilityRepository, PgControlPlaneRepository, PgObservabilityRepository,
|
||||
ProviderAccountModelUsageObservation, ProviderAccountUsageObservation,
|
||||
ProviderAccountUsageQuery, append_admin_audit_event_in_transaction,
|
||||
bump_config_revision_in_transaction,
|
||||
bump_config_revision_in_transaction, completed_usage_fact_predicate,
|
||||
};
|
||||
|
||||
mod admin_adapter;
|
||||
mod admin_queries;
|
||||
mod core_adapter;
|
||||
mod mapping;
|
||||
mod repository;
|
||||
@@ -70,6 +66,7 @@ mod rows;
|
||||
mod runtime;
|
||||
|
||||
pub use admin_adapter::*;
|
||||
pub(crate) use admin_queries::*;
|
||||
pub(crate) use mapping::*;
|
||||
pub use repository::*;
|
||||
pub use rows::*;
|
||||
|
||||
@@ -9,139 +9,74 @@ pub(crate) const QUOTA_MAX_BYTES: usize = 128 * 1024;
|
||||
pub(crate) const MAX_ADMIN_IMPORT_BATCH: usize = 200;
|
||||
pub(crate) const ADMIN_USAGE_CHUNK_SIZE: usize = 200;
|
||||
|
||||
pub(crate) const ACCOUNT_USAGE_BY_WINDOWS_SQL: &str = "with requested_windows as (
|
||||
/// 四种窗口投影通过一组 `GROUPING SETS` 从同一批匹配请求派生。`grouping_* = 1`
|
||||
/// 表示对应维度未参与该行聚合,调用方据此区分窗口、模型与成本行。
|
||||
pub(crate) fn account_usage_by_windows_sql() -> String {
|
||||
let completed_usage = completed_usage_fact_predicate("mr");
|
||||
format!(
|
||||
"with requested_windows as (
|
||||
select *
|
||||
from unnest($1::text[], $2::text[], $3::timestamptz[], $4::timestamptz[])
|
||||
as requested(account_id, window_key, window_start, window_end)
|
||||
), matched as (
|
||||
select requested.account_id,
|
||||
requested.window_key,
|
||||
mr.id as request_id,
|
||||
coalesce(mr.upstream_model_id, mr.requested_model_id) as model,
|
||||
mr.outcome, mr.input_tokens, mr.output_tokens, mr.cached_tokens,
|
||||
mr.cache_write_tokens, mr.reasoning_tokens,
|
||||
mr.image_input_tokens, mr.image_output_tokens,
|
||||
mr.image_generation_succeeded, mr.total_tokens,
|
||||
mr.cost_source, mr.cost_amount, mr.cost_currency, mr.started_at
|
||||
from requested_windows requested
|
||||
left join model_requests mr
|
||||
on mr.provider_account_ref = requested.account_id
|
||||
and mr.started_at >= requested.window_start
|
||||
and mr.started_at < requested.window_end
|
||||
and {completed_usage}
|
||||
)
|
||||
select requested.account_id,
|
||||
requested.window_key,
|
||||
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,
|
||||
sum(mr.output_tokens)::bigint as output_tokens,
|
||||
sum(mr.cached_tokens)::bigint as cached_tokens,
|
||||
sum(mr.cache_write_tokens)::bigint as cache_write_tokens,
|
||||
sum(mr.reasoning_tokens)::bigint as reasoning_tokens,
|
||||
sum(mr.image_input_tokens)::bigint as image_input_tokens,
|
||||
sum(mr.image_output_tokens)::bigint as image_output_tokens,
|
||||
count(mr.id) filter (where mr.image_generation_succeeded is true)::bigint
|
||||
select account_id,
|
||||
window_key,
|
||||
model,
|
||||
cost_currency,
|
||||
grouping(model)::integer as model_grouping,
|
||||
grouping(cost_currency)::integer as currency_grouping,
|
||||
count(request_id)::bigint as request_count,
|
||||
count(request_id) filter (where outcome = 'succeeded')::bigint as success_count,
|
||||
sum(input_tokens)::bigint as input_tokens,
|
||||
sum(output_tokens)::bigint as output_tokens,
|
||||
sum(cached_tokens)::bigint as cached_tokens,
|
||||
sum(cache_write_tokens)::bigint as cache_write_tokens,
|
||||
sum(reasoning_tokens)::bigint as reasoning_tokens,
|
||||
sum(image_input_tokens)::bigint as image_input_tokens,
|
||||
sum(image_output_tokens)::bigint as image_output_tokens,
|
||||
count(request_id) filter (where image_generation_succeeded is true)::bigint
|
||||
as image_request_count,
|
||||
count(mr.id) filter (where mr.image_generation_succeeded is false)::bigint
|
||||
count(request_id) filter (where image_generation_succeeded is false)::bigint
|
||||
as image_request_failed_count,
|
||||
coalesce(sum(coalesce(
|
||||
mr.total_tokens,
|
||||
coalesce(mr.input_tokens, 0) + coalesce(mr.output_tokens, 0)
|
||||
total_tokens,
|
||||
coalesce(input_tokens, 0) + coalesce(output_tokens, 0)
|
||||
)), 0)::bigint as total_tokens,
|
||||
count(mr.id) filter (where mr.cost_source = 'provider_reported')::bigint
|
||||
count(request_id) filter (where cost_source = 'provider_reported')::bigint
|
||||
as provider_reported_count,
|
||||
count(mr.id) filter (where mr.cost_source = 'calculated')::bigint
|
||||
count(request_id) filter (where cost_source = 'calculated')::bigint
|
||||
as calculated_count,
|
||||
count(mr.id) filter (where mr.cost_source = 'unavailable')::bigint
|
||||
count(request_id) filter (where cost_source = 'unavailable')::bigint
|
||||
as unavailable_count,
|
||||
max(mr.started_at) as last_used_at
|
||||
from requested_windows requested
|
||||
left join model_requests mr
|
||||
on mr.provider_account_ref = requested.account_id
|
||||
and mr.started_at >= requested.window_start
|
||||
and mr.started_at < requested.window_end
|
||||
and mr.outcome = 'succeeded'
|
||||
and mr.downstream_committed_at is not null
|
||||
and mr.client_status_code between 200 and 399
|
||||
group by requested.account_id, requested.window_key
|
||||
order by requested.account_id, requested.window_key";
|
||||
|
||||
pub(crate) const ACCOUNT_USAGE_COSTS_BY_WINDOWS_SQL: &str = "with requested_windows as (
|
||||
select *
|
||||
from unnest($1::text[], $2::text[], $3::timestamptz[], $4::timestamptz[])
|
||||
as requested(account_id, window_key, window_start, window_end)
|
||||
)
|
||||
select requested.account_id,
|
||||
requested.window_key,
|
||||
mr.cost_currency,
|
||||
sum(mr.cost_amount)::text as amount
|
||||
from requested_windows requested
|
||||
join model_requests mr
|
||||
on mr.provider_account_ref = requested.account_id
|
||||
and mr.started_at >= requested.window_start
|
||||
and mr.started_at < requested.window_end
|
||||
and mr.outcome = 'succeeded'
|
||||
and mr.downstream_committed_at is not null
|
||||
and mr.client_status_code between 200 and 399
|
||||
and mr.cost_amount is not null
|
||||
and mr.cost_currency is not null
|
||||
group by requested.account_id, requested.window_key, mr.cost_currency
|
||||
order by requested.account_id, requested.window_key, mr.cost_currency";
|
||||
|
||||
pub(crate) const ACCOUNT_USAGE_MODELS_BY_WINDOWS_SQL: &str = "with requested_windows as (
|
||||
select *
|
||||
from unnest($1::text[], $2::text[], $3::timestamptz[], $4::timestamptz[])
|
||||
as requested(account_id, window_key, window_start, window_end)
|
||||
)
|
||||
select requested.account_id,
|
||||
requested.window_key,
|
||||
coalesce(mr.upstream_model_id, mr.requested_model_id) as model,
|
||||
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,
|
||||
sum(mr.output_tokens)::bigint as output_tokens,
|
||||
sum(mr.cached_tokens)::bigint as cached_tokens,
|
||||
sum(mr.cache_write_tokens)::bigint as cache_write_tokens,
|
||||
sum(mr.reasoning_tokens)::bigint as reasoning_tokens,
|
||||
sum(mr.image_input_tokens)::bigint as image_input_tokens,
|
||||
sum(mr.image_output_tokens)::bigint as image_output_tokens,
|
||||
count(mr.id) filter (where mr.image_generation_succeeded is true)::bigint
|
||||
as image_request_count,
|
||||
count(mr.id) filter (where mr.image_generation_succeeded is false)::bigint
|
||||
as image_request_failed_count,
|
||||
coalesce(sum(coalesce(
|
||||
mr.total_tokens,
|
||||
coalesce(mr.input_tokens, 0) + coalesce(mr.output_tokens, 0)
|
||||
)), 0)::bigint as total_tokens,
|
||||
count(mr.id) filter (where mr.cost_source = 'provider_reported')::bigint
|
||||
as provider_reported_count,
|
||||
count(mr.id) filter (where mr.cost_source = 'calculated')::bigint
|
||||
as calculated_count,
|
||||
count(mr.id) filter (where mr.cost_source = 'unavailable')::bigint
|
||||
as unavailable_count,
|
||||
max(mr.started_at) as last_used_at
|
||||
from requested_windows requested
|
||||
join model_requests mr
|
||||
on mr.provider_account_ref = requested.account_id
|
||||
and mr.started_at >= requested.window_start
|
||||
and mr.started_at < requested.window_end
|
||||
and mr.outcome = 'succeeded'
|
||||
and mr.downstream_committed_at is not null
|
||||
and mr.client_status_code between 200 and 399
|
||||
and coalesce(mr.upstream_model_id, mr.requested_model_id) is not null
|
||||
group by requested.account_id, requested.window_key,
|
||||
coalesce(mr.upstream_model_id, mr.requested_model_id)
|
||||
order by requested.account_id, requested.window_key, request_count desc, model";
|
||||
|
||||
pub(crate) const ACCOUNT_USAGE_MODEL_COSTS_BY_WINDOWS_SQL: &str = "with requested_windows as (
|
||||
select *
|
||||
from unnest($1::text[], $2::text[], $3::timestamptz[], $4::timestamptz[])
|
||||
as requested(account_id, window_key, window_start, window_end)
|
||||
)
|
||||
select requested.account_id,
|
||||
requested.window_key,
|
||||
coalesce(mr.upstream_model_id, mr.requested_model_id) as model,
|
||||
mr.cost_currency,
|
||||
sum(mr.cost_amount)::text as amount
|
||||
from requested_windows requested
|
||||
join model_requests mr
|
||||
on mr.provider_account_ref = requested.account_id
|
||||
and mr.started_at >= requested.window_start
|
||||
and mr.started_at < requested.window_end
|
||||
and mr.outcome = 'succeeded'
|
||||
and mr.downstream_committed_at is not null
|
||||
and mr.client_status_code between 200 and 399
|
||||
and mr.cost_amount is not null
|
||||
and mr.cost_currency is not null
|
||||
and coalesce(mr.upstream_model_id, mr.requested_model_id) is not null
|
||||
group by requested.account_id, requested.window_key,
|
||||
coalesce(mr.upstream_model_id, mr.requested_model_id), mr.cost_currency
|
||||
order by requested.account_id, requested.window_key, model, mr.cost_currency";
|
||||
max(started_at) as last_used_at,
|
||||
sum(cost_amount)::text as amount
|
||||
from matched
|
||||
group by grouping sets (
|
||||
(account_id, window_key),
|
||||
(account_id, window_key, cost_currency),
|
||||
(account_id, window_key, model),
|
||||
(account_id, window_key, model, cost_currency)
|
||||
)
|
||||
order by account_id, window_key, model_grouping desc, currency_grouping desc,
|
||||
model nulls last, cost_currency nulls last"
|
||||
)
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ProviderAccountAdminScope {
|
||||
@@ -455,6 +390,29 @@ pub(crate) const ACCOUNT_SELECT_BY_IDS: &str = "select id, provider_kind, name,
|
||||
where id = any($1::text[]) and provider_kind = $2
|
||||
order by id";
|
||||
|
||||
pub(crate) const REFRESH_CANDIDATES_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, 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 provider_kind = $1
|
||||
and enabled
|
||||
and has_refresh_token
|
||||
and credential_state in ('unknown', 'ready')
|
||||
and access_token_expires_at is not null
|
||||
and (
|
||||
access_token_expires_at <= $3
|
||||
or (
|
||||
access_token_expires_at <= $2
|
||||
and (next_refresh_at is null or next_refresh_at <= $4)
|
||||
)
|
||||
)
|
||||
and not (id = any($5::text[]))
|
||||
order by access_token_expires_at, id
|
||||
limit $6";
|
||||
|
||||
pub(crate) fn account_record_from_row(
|
||||
row: sqlx::postgres::PgRow,
|
||||
) -> StoreResult<ProviderAccountRecord> {
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, Utc};
|
||||
use sqlx::PgPool;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use crate::{StoreError, StoreResult, postgres_unavailable};
|
||||
|
||||
@@ -18,6 +19,48 @@ pub struct RetentionReport {
|
||||
pub model_requests: u64,
|
||||
pub ops_events: u64,
|
||||
pub admin_audit_events: u64,
|
||||
pub batches: u32,
|
||||
pub budget_exhausted: bool,
|
||||
}
|
||||
|
||||
/// 一次 retention cycle 可占用的删除预算;不从部署配置暴露低价值开关。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct RetentionCycleBudget {
|
||||
batch_rows: u32,
|
||||
max_batches: u32,
|
||||
max_duration: Duration,
|
||||
batch_pause: Duration,
|
||||
}
|
||||
|
||||
impl RetentionCycleBudget {
|
||||
pub fn try_new(
|
||||
batch_rows: u32,
|
||||
max_batches: u32,
|
||||
max_duration: Duration,
|
||||
batch_pause: Duration,
|
||||
) -> StoreResult<Self> {
|
||||
if batch_rows == 0 || max_batches == 0 || max_duration.is_zero() {
|
||||
return Err(StoreError::InvalidData {
|
||||
entity: "retention cycle budget",
|
||||
message: "row, batch, and duration budgets must be positive".to_owned(),
|
||||
});
|
||||
}
|
||||
Ok(Self {
|
||||
batch_rows,
|
||||
max_batches,
|
||||
max_duration,
|
||||
batch_pause,
|
||||
})
|
||||
}
|
||||
|
||||
const fn production() -> Self {
|
||||
Self {
|
||||
batch_rows: 5_000,
|
||||
max_batches: 12,
|
||||
max_duration: Duration::from_secs(30),
|
||||
batch_pause: Duration::from_millis(50),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -33,12 +76,21 @@ pub trait RetentionRepository: Send + Sync {
|
||||
#[derive(Clone)]
|
||||
pub struct PgRetentionRepository {
|
||||
pool: PgPool,
|
||||
cycle_budget: RetentionCycleBudget,
|
||||
}
|
||||
|
||||
impl PgRetentionRepository {
|
||||
#[must_use]
|
||||
pub const fn new(pool: PgPool) -> Self {
|
||||
Self { pool }
|
||||
Self {
|
||||
pool,
|
||||
cycle_budget: RetentionCycleBudget::production(),
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn with_cycle_budget(pool: PgPool, cycle_budget: RetentionCycleBudget) -> Self {
|
||||
Self { pool, cycle_budget }
|
||||
}
|
||||
}
|
||||
|
||||
@@ -78,83 +130,118 @@ impl RetentionRepository for PgRetentionRepository {
|
||||
});
|
||||
}
|
||||
|
||||
// 分批删除,每批独立提交:清理是幂等的,无需跨表原子性;
|
||||
// 有界批量避免百万行单事务的长锁与 WAL 峰值。
|
||||
let model_requests = purge_in_batches(
|
||||
&self.pool,
|
||||
"delete from model_requests
|
||||
where ctid in (
|
||||
select ctid from model_requests
|
||||
where outcome <> 'running'
|
||||
and completed_at < $1 - ($2 * interval '1 day')
|
||||
limit $3
|
||||
)",
|
||||
now,
|
||||
settings.usage_retention_days,
|
||||
"delete expired model requests",
|
||||
)
|
||||
.await?;
|
||||
let ops_events = purge_in_batches(
|
||||
&self.pool,
|
||||
"delete from ops_events
|
||||
where ctid in (
|
||||
select ctid from ops_events
|
||||
where model_request_id is null
|
||||
and created_at < $1 - ($2 * interval '1 day')
|
||||
limit $3
|
||||
)",
|
||||
now,
|
||||
settings.ops_event_retention_days,
|
||||
"delete expired ops events",
|
||||
)
|
||||
.await?;
|
||||
let admin_audit_events = purge_in_batches(
|
||||
&self.pool,
|
||||
"delete from admin_audit_events
|
||||
where ctid in (
|
||||
select ctid from admin_audit_events
|
||||
where created_at < $1 - ($2 * interval '1 day')
|
||||
limit $3
|
||||
)",
|
||||
now,
|
||||
settings.audit_retention_days,
|
||||
"delete expired admin audit events",
|
||||
)
|
||||
.await?;
|
||||
// 各表轮转执行单批独立事务;行数、批数和 wall-clock 同时有界。
|
||||
let mut targets = [
|
||||
RetentionTarget::new(
|
||||
"delete from model_requests
|
||||
where ctid in (
|
||||
select ctid from model_requests
|
||||
where outcome <> 'running'
|
||||
and completed_at < $1 - ($2 * interval '1 day')
|
||||
limit $3
|
||||
)",
|
||||
settings.usage_retention_days,
|
||||
"delete expired model requests",
|
||||
),
|
||||
RetentionTarget::new(
|
||||
"delete from ops_events
|
||||
where ctid in (
|
||||
select ctid from ops_events
|
||||
where model_request_id is null
|
||||
and created_at < $1 - ($2 * interval '1 day')
|
||||
limit $3
|
||||
)",
|
||||
settings.ops_event_retention_days,
|
||||
"delete expired ops events",
|
||||
),
|
||||
RetentionTarget::new(
|
||||
"delete from admin_audit_events
|
||||
where ctid in (
|
||||
select ctid from admin_audit_events
|
||||
where created_at < $1 - ($2 * interval '1 day')
|
||||
limit $3
|
||||
)",
|
||||
settings.audit_retention_days,
|
||||
"delete expired admin audit events",
|
||||
),
|
||||
];
|
||||
let started_at = Instant::now();
|
||||
let mut batches = 0_u32;
|
||||
'cycles: loop {
|
||||
let mut attempted = false;
|
||||
for target in &mut targets {
|
||||
if target.complete {
|
||||
continue;
|
||||
}
|
||||
if batches >= self.cycle_budget.max_batches
|
||||
|| started_at.elapsed() >= self.cycle_budget.max_duration
|
||||
{
|
||||
break 'cycles;
|
||||
}
|
||||
attempted = true;
|
||||
let deleted =
|
||||
purge_batch(&self.pool, target, now, self.cycle_budget.batch_rows).await?;
|
||||
batches = batches.saturating_add(1);
|
||||
target.deleted = target.deleted.saturating_add(deleted);
|
||||
target.complete = deleted < u64::from(self.cycle_budget.batch_rows);
|
||||
if !target.complete
|
||||
&& batches < self.cycle_budget.max_batches
|
||||
&& started_at.elapsed() < self.cycle_budget.max_duration
|
||||
&& !self.cycle_budget.batch_pause.is_zero()
|
||||
{
|
||||
tokio::time::sleep(self.cycle_budget.batch_pause).await;
|
||||
}
|
||||
}
|
||||
if !attempted || targets.iter().all(|target| target.complete) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Ok(RetentionReport {
|
||||
model_requests,
|
||||
ops_events,
|
||||
admin_audit_events,
|
||||
model_requests: targets[0].deleted,
|
||||
ops_events: targets[1].deleted,
|
||||
admin_audit_events: targets[2].deleted,
|
||||
batches,
|
||||
budget_exhausted: targets.iter().any(|target| !target.complete),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
const RETENTION_DELETE_BATCH_ROWS: i64 = 10_000;
|
||||
|
||||
async fn purge_in_batches(
|
||||
pool: &PgPool,
|
||||
struct RetentionTarget {
|
||||
delete_sql: &'static str,
|
||||
now: DateTime<Utc>,
|
||||
retention_days: u32,
|
||||
label: &'static str,
|
||||
) -> StoreResult<u64> {
|
||||
let mut total = 0u64;
|
||||
loop {
|
||||
let deleted = sqlx::query(delete_sql)
|
||||
.bind(now)
|
||||
.bind(i64::from(retention_days))
|
||||
.bind(RETENTION_DELETE_BATCH_ROWS)
|
||||
.execute(pool)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable(label))?
|
||||
.rows_affected();
|
||||
total += deleted;
|
||||
if deleted < RETENTION_DELETE_BATCH_ROWS.unsigned_abs() {
|
||||
return Ok(total);
|
||||
deleted: u64,
|
||||
complete: bool,
|
||||
}
|
||||
|
||||
impl RetentionTarget {
|
||||
const fn new(delete_sql: &'static str, retention_days: u32, label: &'static str) -> Self {
|
||||
Self {
|
||||
delete_sql,
|
||||
retention_days,
|
||||
label,
|
||||
deleted: 0,
|
||||
complete: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn purge_batch(
|
||||
pool: &PgPool,
|
||||
target: &RetentionTarget,
|
||||
now: DateTime<Utc>,
|
||||
batch_rows: u32,
|
||||
) -> StoreResult<u64> {
|
||||
sqlx::query(target.delete_sql)
|
||||
.bind(now)
|
||||
.bind(i64::from(target.retention_days))
|
||||
.bind(i64::from(batch_rows))
|
||||
.execute(pool)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable(target.label))
|
||||
.map(|result| result.rows_affected())
|
||||
}
|
||||
|
||||
fn to_u32(value: i64) -> StoreResult<u32> {
|
||||
u32::try_from(value).map_err(|_| StoreError::InvalidData {
|
||||
entity: "retention settings",
|
||||
|
||||
@@ -128,38 +128,7 @@ impl PgRuntimeSettingsRepository {
|
||||
#[async_trait]
|
||||
impl RuntimeSettingsRepository for PgRuntimeSettingsRepository {
|
||||
async fn load_runtime_settings(&self) -> StoreResult<RuntimeSettings> {
|
||||
let row = sqlx::query_as::<
|
||||
_,
|
||||
(
|
||||
i64,
|
||||
Option<String>,
|
||||
i64,
|
||||
i64,
|
||||
i64,
|
||||
i64,
|
||||
String,
|
||||
sqlx::types::Json<BTreeMap<String, String>>,
|
||||
i64,
|
||||
i64,
|
||||
i64,
|
||||
DateTime<Utc>,
|
||||
),
|
||||
>(
|
||||
"select config_revision, admin_api_key, refresh_margin_seconds,
|
||||
refresh_concurrency, max_concurrent_per_account, request_interval_ms,
|
||||
rotation_strategy, model_mappings_json, usage_retention_days, ops_event_retention_days,
|
||||
audit_retention_days, updated_at
|
||||
from runtime_settings where id = 1",
|
||||
)
|
||||
.fetch_optional(&self.pool)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("load runtime settings"))?
|
||||
.ok_or_else(|| StoreError::NotFound {
|
||||
entity: "runtime settings",
|
||||
id: "1".to_owned(),
|
||||
})?;
|
||||
|
||||
runtime_settings_from_row(row)
|
||||
load_runtime_settings_from_pool(&self.pool).await
|
||||
}
|
||||
|
||||
async fn update_runtime_settings(
|
||||
@@ -180,6 +149,24 @@ impl RuntimeSettingsRepository for PgRuntimeSettingsRepository {
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn load_runtime_settings_from_pool(pool: &PgPool) -> StoreResult<RuntimeSettings> {
|
||||
let row = sqlx::query_as::<_, RuntimeSettingsRow>(
|
||||
"select config_revision, admin_api_key, refresh_margin_seconds,
|
||||
refresh_concurrency, max_concurrent_per_account, request_interval_ms,
|
||||
rotation_strategy, model_mappings_json, usage_retention_days, ops_event_retention_days,
|
||||
audit_retention_days, updated_at
|
||||
from runtime_settings where id = 1",
|
||||
)
|
||||
.fetch_optional(pool)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("load runtime settings"))?
|
||||
.ok_or_else(|| StoreError::NotFound {
|
||||
entity: "runtime settings",
|
||||
id: "1".to_owned(),
|
||||
})?;
|
||||
runtime_settings_from_row(row)
|
||||
}
|
||||
|
||||
impl ProviderRuntimePolicyPort for PgRuntimeSettingsRepository {
|
||||
fn load_refresh_policy(
|
||||
&self,
|
||||
|
||||
@@ -3,7 +3,6 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, Utc};
|
||||
use gateway_core::engine::credential::ProviderAccountId;
|
||||
use gateway_core::routing::{
|
||||
AccountGroupId, ConfigRevision,
|
||||
@@ -29,14 +28,6 @@ pub struct SnapshotRuntimeSettings {
|
||||
pub model_mappings: BTreeMap<String, String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct CredentialRevisionVectorEntry {
|
||||
pub provider_account_id: String,
|
||||
pub credential_revision: Revision,
|
||||
pub credential_observed_at: DateTime<Utc>,
|
||||
pub quota_observed_at: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct RuntimeSnapshotData {
|
||||
pub config_revision: Revision,
|
||||
@@ -71,7 +62,6 @@ pub struct SnapshotGroupMembershipData {
|
||||
pub trait RuntimeSnapshotRepository: Send + Sync {
|
||||
async fn load_runtime_snapshot(&self) -> StoreResult<RuntimeSnapshotData>;
|
||||
async fn current_config_revision(&self) -> StoreResult<Revision>;
|
||||
async fn credential_revision_vector(&self) -> StoreResult<Vec<CredentialRevisionVectorEntry>>;
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
@@ -135,26 +125,6 @@ impl RuntimeSnapshotRepository for PgRuntimeSnapshotRepository {
|
||||
})?;
|
||||
revision_from_i64(revision)
|
||||
}
|
||||
|
||||
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, credential_observed_at, quota_observed_at
|
||||
from provider_accounts order by id",
|
||||
)
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("read credential revision vector"))?;
|
||||
rows.into_iter()
|
||||
.map(|row| {
|
||||
Ok(CredentialRevisionVectorEntry {
|
||||
provider_account_id: row.0,
|
||||
credential_revision: revision_from_i64(row.1)?,
|
||||
credential_observed_at: row.2,
|
||||
quota_observed_at: row.3,
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
impl SnapshotStorePort for PgRuntimeSnapshotRepository {
|
||||
|
||||
@@ -0,0 +1,20 @@
|
||||
//! `model_requests` 完成用量事实的 PostgreSQL 语义。
|
||||
|
||||
use sqlx::{Postgres, QueryBuilder};
|
||||
|
||||
/// 只把已完整交付给客户端的成功响应投影为用量事实。
|
||||
///
|
||||
/// `model_requests` 同时承担执行审计:包括上游已发送、但尚未收到首个事件就断开的
|
||||
/// WebSocket 请求。这些失败仍必须留给 Ops Errors 和调度健康度分析,不能混入用量、
|
||||
/// 成本、账号使用次数或请求明细。
|
||||
/// 客户端 WebSocket 的每个 `response.create` 没有独立 HTTP 状态码;成功终态与下游
|
||||
/// 提交边界已足以证明交付,连接握手的 101 不能冒充单次请求状态。
|
||||
pub(crate) fn completed_usage_fact_predicate(alias: &str) -> String {
|
||||
format!(
|
||||
"{alias}.outcome = 'succeeded' and {alias}.downstream_committed_at is not null and (({alias}.client_transport = 'websocket' and {alias}.client_status_code is null) or {alias}.client_status_code between 200 and 399)"
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn push_completed_usage_fact_filter(query: &mut QueryBuilder<Postgres>, alias: &str) {
|
||||
query.push(format!(" and {}", completed_usage_fact_predicate(alias)));
|
||||
}
|
||||
@@ -0,0 +1,79 @@
|
||||
//! Admin query service 使用的账号运行态组合 adapter。
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use futures::future::join_all;
|
||||
use gateway_admin::{
|
||||
model::accounts::AccountRuntimeSnapshot,
|
||||
ports::store::{AccountRuntimeStore, AdminStoreResult},
|
||||
};
|
||||
|
||||
use super::{
|
||||
CredentialCooldownRepository as _, CredentialLeaseRepository as _,
|
||||
RedisCredentialCooldownRepository, RedisCredentialLeaseRepository,
|
||||
};
|
||||
|
||||
/// 只组合可丢失 Redis 事实;不持有 PostgreSQL,也不执行状态投影。
|
||||
#[derive(Clone)]
|
||||
pub struct RedisAdminAccountRuntimeStore {
|
||||
cooldowns: RedisCredentialCooldownRepository,
|
||||
leases: RedisCredentialLeaseRepository,
|
||||
}
|
||||
|
||||
impl RedisAdminAccountRuntimeStore {
|
||||
#[must_use]
|
||||
pub const fn new(
|
||||
cooldowns: RedisCredentialCooldownRepository,
|
||||
leases: RedisCredentialLeaseRepository,
|
||||
) -> Self {
|
||||
Self { cooldowns, leases }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AccountRuntimeStore for RedisAdminAccountRuntimeStore {
|
||||
async fn active_rate_limits(&self) -> AdminStoreResult<AccountRuntimeSnapshot> {
|
||||
self.cooldowns
|
||||
.active_cooldowns()
|
||||
.await
|
||||
.map_err(|error| crate::admin_store_error("account runtime", error))
|
||||
}
|
||||
|
||||
async fn account_runtime(
|
||||
&self,
|
||||
account_ids: &[String],
|
||||
) -> AdminStoreResult<AccountRuntimeSnapshot> {
|
||||
let reads = account_ids.iter().map(|account_id| async move {
|
||||
self.cooldowns
|
||||
.read_credential_cooldown(account_id)
|
||||
.await
|
||||
.map(|cooldown| {
|
||||
cooldown.map(|cooldown| (account_id.clone(), cooldown.cooldown_until))
|
||||
})
|
||||
});
|
||||
let mut rate_limited_until = BTreeMap::new();
|
||||
for result in join_all(reads).await {
|
||||
if let Some((account_id, until)) =
|
||||
result.map_err(|error| crate::admin_store_error("account runtime", error))?
|
||||
{
|
||||
rate_limited_until.insert(account_id, until);
|
||||
}
|
||||
}
|
||||
let in_flight = self
|
||||
.leases
|
||||
.credential_runtime_signals(account_ids)
|
||||
.await
|
||||
.ok()
|
||||
.map(|signals| {
|
||||
signals
|
||||
.into_iter()
|
||||
.map(|signal| (signal.resource_id, u64::from(signal.in_flight)))
|
||||
.collect()
|
||||
});
|
||||
Ok(AccountRuntimeSnapshot {
|
||||
rate_limited_until,
|
||||
in_flight,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -2,6 +2,7 @@
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, Utc};
|
||||
use gateway_admin::model::accounts::AccountRuntimeSnapshot;
|
||||
use gateway_core::{
|
||||
engine::credential::{CredentialRevision, ProviderAccountId},
|
||||
provider_ports::{
|
||||
@@ -26,11 +27,13 @@ local clock = redis.call('TIME')
|
||||
local now_ms = (tonumber(clock[1]) * 1000) + math.floor(tonumber(clock[2]) / 1000)
|
||||
if incoming_until <= now_ms then
|
||||
redis.call('DEL', KEYS[1])
|
||||
if #KEYS > 1 then redis.call('ZREM', KEYS[2], ARGV[3]) end
|
||||
return 0
|
||||
end
|
||||
redis.call('HSET', KEYS[1], 'revision', ARGV[1], 'until_ms', ARGV[2])
|
||||
local ttl = incoming_until - now_ms + 60000
|
||||
redis.call('PEXPIRE', KEYS[1], ttl)
|
||||
if #KEYS > 1 then redis.call('ZADD', KEYS[2], incoming_until, ARGV[3]) end
|
||||
return 1
|
||||
"#;
|
||||
|
||||
@@ -39,12 +42,14 @@ local revision = redis.call('HGET', KEYS[1], 'revision')
|
||||
local until_ms = redis.call('HGET', KEYS[1], 'until_ms')
|
||||
if revision == false or until_ms == false then
|
||||
redis.call('DEL', KEYS[1])
|
||||
if #KEYS > 1 then redis.call('ZREM', KEYS[2], ARGV[1]) end
|
||||
return {0, '0', '0'}
|
||||
end
|
||||
local clock = redis.call('TIME')
|
||||
local now_ms = (tonumber(clock[1]) * 1000) + math.floor(tonumber(clock[2]) / 1000)
|
||||
if tonumber(until_ms) <= now_ms then
|
||||
redis.call('DEL', KEYS[1])
|
||||
if #KEYS > 1 then redis.call('ZREM', KEYS[2], ARGV[1]) end
|
||||
return {0, '0', '0'}
|
||||
end
|
||||
return {1, revision, until_ms}
|
||||
@@ -54,9 +59,17 @@ const INVALIDATE_SCRIPT: &str = r#"
|
||||
local current = tonumber(redis.call('HGET', KEYS[1], 'revision') or '0')
|
||||
if current > tonumber(ARGV[1]) then return 0 end
|
||||
redis.call('DEL', KEYS[1])
|
||||
if #KEYS > 1 then redis.call('ZREM', KEYS[2], ARGV[2]) end
|
||||
return 1
|
||||
"#;
|
||||
|
||||
const ACTIVE_COOLDOWNS_SCRIPT: &str = r#"
|
||||
local clock = redis.call('TIME')
|
||||
local now_ms = (tonumber(clock[1]) * 1000) + math.floor(tonumber(clock[2]) / 1000)
|
||||
redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', now_ms)
|
||||
return redis.call('ZRANGEBYSCORE', KEYS[1], '(' .. now_ms, '+inf', 'WITHSCORES')
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct CredentialCooldown {
|
||||
pub provider_account_id: String,
|
||||
@@ -99,6 +112,10 @@ impl RedisCredentialCooldownRepository {
|
||||
Ok(format!("{}:account:{fingerprint}:cooldown", self.namespace))
|
||||
}
|
||||
|
||||
fn active_index_key(&self) -> String {
|
||||
format!("{}:account:active-cooldowns", self.namespace)
|
||||
}
|
||||
|
||||
fn scoped_key(
|
||||
&self,
|
||||
provider_account_id: &str,
|
||||
@@ -118,29 +135,57 @@ impl RedisCredentialCooldownRepository {
|
||||
key: String,
|
||||
credential_revision: Revision,
|
||||
cooldown_until: DateTime<Utc>,
|
||||
index_member: Option<&str>,
|
||||
) -> StoreResult<bool> {
|
||||
let until_ms = cooldown_until.timestamp_millis();
|
||||
if until_ms <= 0 {
|
||||
return Err(invalid("cooldown expiry must be positive"));
|
||||
}
|
||||
let mut connection = self.connection.clone();
|
||||
let written = Script::new(WRITE_SCRIPT)
|
||||
.key(key)
|
||||
.arg(credential_revision.get())
|
||||
.arg(until_ms)
|
||||
.invoke_async::<i64>(&mut connection)
|
||||
.await
|
||||
.map_err(|_| redis_unavailable("cache credential cooldown"))?;
|
||||
let written = if let Some(index_member) = index_member {
|
||||
Script::new(WRITE_SCRIPT)
|
||||
.key(key)
|
||||
.key(self.active_index_key())
|
||||
.arg(credential_revision.get())
|
||||
.arg(until_ms)
|
||||
.arg(index_member)
|
||||
.invoke_async::<i64>(&mut connection)
|
||||
.await
|
||||
} else {
|
||||
Script::new(WRITE_SCRIPT)
|
||||
.key(key)
|
||||
.arg(credential_revision.get())
|
||||
.arg(until_ms)
|
||||
.arg("")
|
||||
.invoke_async::<i64>(&mut connection)
|
||||
.await
|
||||
}
|
||||
.map_err(|_| redis_unavailable("cache credential cooldown"))?;
|
||||
Ok(written == 1)
|
||||
}
|
||||
|
||||
async fn read_at_key(&self, key: String) -> StoreResult<Option<(Revision, DateTime<Utc>)>> {
|
||||
async fn read_at_key(
|
||||
&self,
|
||||
key: String,
|
||||
index_member: Option<&str>,
|
||||
) -> StoreResult<Option<(Revision, DateTime<Utc>)>> {
|
||||
let mut connection = self.connection.clone();
|
||||
let (present, revision, until_ms): (i64, String, String) = Script::new(READ_SCRIPT)
|
||||
.key(key)
|
||||
.invoke_async(&mut connection)
|
||||
.await
|
||||
.map_err(|_| redis_unavailable("read credential cooldown"))?;
|
||||
let result = if let Some(index_member) = index_member {
|
||||
Script::new(READ_SCRIPT)
|
||||
.key(key)
|
||||
.key(self.active_index_key())
|
||||
.arg(index_member)
|
||||
.invoke_async(&mut connection)
|
||||
.await
|
||||
} else {
|
||||
Script::new(READ_SCRIPT)
|
||||
.key(key)
|
||||
.arg("")
|
||||
.invoke_async(&mut connection)
|
||||
.await
|
||||
};
|
||||
let (present, revision, until_ms): (i64, String, String) =
|
||||
result.map_err(|_| redis_unavailable("read credential cooldown"))?;
|
||||
if present == 0 {
|
||||
return Ok(None);
|
||||
}
|
||||
@@ -159,16 +204,53 @@ impl RedisCredentialCooldownRepository {
|
||||
&self,
|
||||
key: String,
|
||||
through_revision: Revision,
|
||||
index_member: Option<&str>,
|
||||
) -> StoreResult<bool> {
|
||||
let mut connection = self.connection.clone();
|
||||
let removed = Script::new(INVALIDATE_SCRIPT)
|
||||
.key(key)
|
||||
.arg(through_revision.get())
|
||||
.invoke_async::<i64>(&mut connection)
|
||||
.await
|
||||
.map_err(|_| redis_unavailable("invalidate credential cooldown"))?;
|
||||
let removed = if let Some(index_member) = index_member {
|
||||
Script::new(INVALIDATE_SCRIPT)
|
||||
.key(key)
|
||||
.key(self.active_index_key())
|
||||
.arg(through_revision.get())
|
||||
.arg(index_member)
|
||||
.invoke_async::<i64>(&mut connection)
|
||||
.await
|
||||
} else {
|
||||
Script::new(INVALIDATE_SCRIPT)
|
||||
.key(key)
|
||||
.arg(through_revision.get())
|
||||
.arg("")
|
||||
.invoke_async::<i64>(&mut connection)
|
||||
.await
|
||||
}
|
||||
.map_err(|_| redis_unavailable("invalidate credential cooldown"))?;
|
||||
Ok(removed == 1)
|
||||
}
|
||||
|
||||
pub(crate) async fn active_cooldowns(&self) -> StoreResult<AccountRuntimeSnapshot> {
|
||||
let mut connection = self.connection.clone();
|
||||
let values = Script::new(ACTIVE_COOLDOWNS_SCRIPT)
|
||||
.key(self.active_index_key())
|
||||
.invoke_async::<Vec<String>>(&mut connection)
|
||||
.await
|
||||
.map_err(|_| redis_unavailable("list active credential cooldowns"))?;
|
||||
if values.len() % 2 != 0 {
|
||||
return Err(invalid("active cooldown index is invalid"));
|
||||
}
|
||||
let mut rate_limited_until = std::collections::BTreeMap::new();
|
||||
for pair in values.chunks_exact(2) {
|
||||
let until_ms = pair[1]
|
||||
.parse::<i64>()
|
||||
.map_err(|_| invalid("active cooldown expiry is invalid"))?;
|
||||
let until = DateTime::from_timestamp_millis(until_ms)
|
||||
.ok_or_else(|| invalid("active cooldown expiry is invalid"))?;
|
||||
rate_limited_until.insert(pair[0].clone(), until);
|
||||
}
|
||||
Ok(AccountRuntimeSnapshot {
|
||||
rate_limited_until,
|
||||
in_flight: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -183,6 +265,7 @@ impl CredentialCooldownRepository for RedisCredentialCooldownRepository {
|
||||
self.key(&cooldown.provider_account_id)?,
|
||||
cooldown.credential_revision,
|
||||
cooldown.cooldown_until,
|
||||
Some(&cooldown.provider_account_id),
|
||||
)
|
||||
.await
|
||||
}
|
||||
@@ -196,7 +279,7 @@ impl CredentialCooldownRepository for RedisCredentialCooldownRepository {
|
||||
"provider_account_id",
|
||||
provider_account_id,
|
||||
)?;
|
||||
self.read_at_key(self.key(provider_account_id)?)
|
||||
self.read_at_key(self.key(provider_account_id)?, Some(provider_account_id))
|
||||
.await
|
||||
.map(|value| {
|
||||
value.map(|(credential_revision, cooldown_until)| CredentialCooldown {
|
||||
@@ -217,8 +300,12 @@ impl CredentialCooldownRepository for RedisCredentialCooldownRepository {
|
||||
"provider_account_id",
|
||||
provider_account_id,
|
||||
)?;
|
||||
self.invalidate_at_key(self.key(provider_account_id)?, through_revision)
|
||||
.await
|
||||
self.invalidate_at_key(
|
||||
self.key(provider_account_id)?,
|
||||
through_revision,
|
||||
Some(provider_account_id),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn delete_account_cooldowns(&self, provider_account_id: &str) -> StoreResult<bool> {
|
||||
@@ -262,7 +349,13 @@ impl CredentialCooldownRepository for RedisCredentialCooldownRepository {
|
||||
.query_async(&mut connection)
|
||||
.await
|
||||
.map_err(|_| redis_unavailable("delete account cooldown keys"))?;
|
||||
Ok(removed > 0)
|
||||
let index_removed: i64 = redis::cmd("ZREM")
|
||||
.arg(self.active_index_key())
|
||||
.arg(provider_account_id)
|
||||
.query_async(&mut connection)
|
||||
.await
|
||||
.map_err(|_| redis_unavailable("remove active account cooldown"))?;
|
||||
Ok(removed > 0 || index_removed > 0)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -337,6 +430,7 @@ impl ProviderCooldownPort for RedisCredentialCooldownRepository {
|
||||
.map_err(|_| provider_invalid("encode scoped credential cooldown"))?,
|
||||
revision,
|
||||
cooldown.until().into(),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.map_err(|_| provider_unavailable("cache scoped credential cooldown"))
|
||||
@@ -353,6 +447,7 @@ impl ProviderCooldownPort for RedisCredentialCooldownRepository {
|
||||
self.read_at_key(
|
||||
self.scoped_key(account_id.as_str(), scope)
|
||||
.map_err(|_| provider_invalid("encode scoped credential cooldown"))?,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.map_err(|_| provider_unavailable("read scoped credential cooldown"))?
|
||||
@@ -382,6 +477,7 @@ impl ProviderCooldownPort for RedisCredentialCooldownRepository {
|
||||
self.scoped_key(account_id.as_str(), scope)
|
||||
.map_err(|_| provider_invalid("encode scoped credential cooldown"))?,
|
||||
revision,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.map_err(|_| provider_unavailable("clear scoped credential cooldown"))
|
||||
|
||||
@@ -7,6 +7,7 @@ use sha2::{Digest, Sha256};
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
mod admin_account_runtime;
|
||||
mod artifact_profile;
|
||||
mod client_admission;
|
||||
mod coordination_buffer;
|
||||
@@ -21,6 +22,7 @@ mod provider_session_exclusion;
|
||||
mod runtime_change;
|
||||
pub(crate) mod worker_lease;
|
||||
|
||||
pub use admin_account_runtime::*;
|
||||
pub use artifact_profile::*;
|
||||
pub use client_admission::*;
|
||||
pub use coordination_buffer::*;
|
||||
|
||||
@@ -5,6 +5,7 @@ use super::*;
|
||||
pub(crate) fn store_worker_contributions(
|
||||
execution: Arc<postgres::PgExecutionStore>,
|
||||
execution_writer: postgres::ExecutionObservationWriter<postgres::PgExecutionStore>,
|
||||
client_key_usage_writer: postgres::PgClientApiKeyUsageWriter,
|
||||
admission_release_writer: redis::ClientAdmissionReleaseWriter,
|
||||
circuit_feedback_writer: redis::ProviderCircuitFeedbackWriter,
|
||||
retention: Arc<postgres::PgRetentionRepository>,
|
||||
@@ -15,6 +16,9 @@ pub(crate) fn store_worker_contributions(
|
||||
WorkerId::try_new(WorkerKind::Retention, "postgres").map_err(worker_definition_error)?;
|
||||
let ops_flush_id =
|
||||
WorkerId::try_new(WorkerKind::OpsFlush, "postgres").map_err(worker_definition_error)?;
|
||||
let client_key_usage_flush_id =
|
||||
WorkerId::try_new(WorkerKind::OpsFlush, "postgres_client_key_usage")
|
||||
.map_err(worker_definition_error)?;
|
||||
let admission_flush_id = WorkerId::try_new(WorkerKind::OpsFlush, "redis_admission")
|
||||
.map_err(worker_definition_error)?;
|
||||
let circuit_flush_id = WorkerId::try_new(WorkerKind::OpsFlush, "redis_circuit")
|
||||
@@ -43,6 +47,16 @@ pub(crate) fn store_worker_contributions(
|
||||
)
|
||||
.map_err(worker_definition_error)?,
|
||||
),
|
||||
WorkerContribution::Registration(
|
||||
WorkerRegistration::try_new(
|
||||
client_key_usage_flush_id,
|
||||
WorkerRunnable::Daemon {
|
||||
restart: ops_flush_restart,
|
||||
task: Box::new(client_key_usage_writer),
|
||||
},
|
||||
)
|
||||
.map_err(worker_definition_error)?,
|
||||
),
|
||||
WorkerContribution::Registration(
|
||||
WorkerRegistration::try_new(
|
||||
admission_flush_id,
|
||||
@@ -129,27 +143,80 @@ pub(crate) struct RetentionTask {
|
||||
impl ScheduledTask for RetentionTask {
|
||||
fn run_cycle(
|
||||
&self,
|
||||
_context: WorkerCycleContext,
|
||||
context: WorkerCycleContext,
|
||||
) -> futures::future::BoxFuture<'_, Result<(), WorkerTaskError>> {
|
||||
Box::pin(async move {
|
||||
let settings =
|
||||
postgres::RetentionRepository::load_retention_settings(self.retention.as_ref())
|
||||
.await
|
||||
.map_err(|_| WorkerTaskError::safe("retention settings read failed"))?;
|
||||
postgres::RetentionRepository::apply_retention(
|
||||
let started_at = std::time::Instant::now();
|
||||
let cleanup = postgres::RetentionRepository::apply_retention(
|
||||
self.retention.as_ref(),
|
||||
chrono::Utc::now(),
|
||||
settings,
|
||||
)
|
||||
.await
|
||||
.map(|_| ())
|
||||
.map_err(|_| WorkerTaskError::safe("retention cleanup failed"))
|
||||
);
|
||||
let report = tokio::select! {
|
||||
() = context.cancellation().cancelled() => return Ok(()),
|
||||
result = cleanup => result
|
||||
.map_err(|_| WorkerTaskError::safe("retention cleanup failed"))?,
|
||||
};
|
||||
tracing::info!(
|
||||
model_requests = report.model_requests,
|
||||
ops_events = report.ops_events,
|
||||
admin_audit_events = report.admin_audit_events,
|
||||
batches = report.batches,
|
||||
budget_exhausted = report.budget_exhausted,
|
||||
elapsed_milliseconds =
|
||||
u64::try_from(started_at.elapsed().as_millis()).unwrap_or(u64::MAX),
|
||||
"PostgreSQL retention cycle completed"
|
||||
);
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) struct PostgresHealthProbe {
|
||||
pub(crate) pool: sqlx::PgPool,
|
||||
pub struct PostgresHealthProbe {
|
||||
pool: sqlx::PgPool,
|
||||
max_connections: u32,
|
||||
}
|
||||
|
||||
impl PostgresHealthProbe {
|
||||
#[must_use]
|
||||
pub const fn new(pool: sqlx::PgPool, max_connections: u32) -> Self {
|
||||
Self {
|
||||
pool,
|
||||
max_connections,
|
||||
}
|
||||
}
|
||||
|
||||
async fn check_once(&self) -> HealthState {
|
||||
let deadline = tokio::time::Instant::now() + POSTGRES_HEALTH_ATTEMPT_TIMEOUT;
|
||||
let mut connection = match tokio::time::timeout_at(deadline, self.pool.acquire()).await {
|
||||
Ok(Ok(connection)) => connection,
|
||||
Ok(Err(_)) => {
|
||||
return HealthState::Unhealthy("PostgreSQL is unavailable".to_owned());
|
||||
}
|
||||
Err(_) => {
|
||||
return postgres_acquire_timeout_state(
|
||||
self.pool.size(),
|
||||
self.pool.num_idle(),
|
||||
self.max_connections,
|
||||
);
|
||||
}
|
||||
};
|
||||
match tokio::time::timeout_at(
|
||||
deadline,
|
||||
sqlx::query_scalar::<_, i32>("select 1").fetch_one(&mut *connection),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Ok(1)) => HealthState::Healthy,
|
||||
Ok(Ok(_)) => HealthState::Unhealthy("PostgreSQL health result is invalid".to_owned()),
|
||||
Ok(Err(_)) => HealthState::Unhealthy("PostgreSQL is unavailable".to_owned()),
|
||||
Err(_) => HealthState::Unhealthy("PostgreSQL health query timed out".to_owned()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl HealthProbe for PostgresHealthProbe {
|
||||
@@ -158,16 +225,37 @@ impl HealthProbe for PostgresHealthProbe {
|
||||
}
|
||||
|
||||
fn check(&self) -> futures::future::BoxFuture<'_, HealthState> {
|
||||
Box::pin(async move {
|
||||
match sqlx::query_scalar::<_, i32>("select 1")
|
||||
.fetch_one(&self.pool)
|
||||
.await
|
||||
{
|
||||
Ok(1) => HealthState::Healthy,
|
||||
Ok(_) => HealthState::Unhealthy("PostgreSQL health result is invalid".to_owned()),
|
||||
Err(_) => HealthState::Unhealthy("PostgreSQL is unavailable".to_owned()),
|
||||
}
|
||||
})
|
||||
Box::pin(health_state_with_one_retry(
|
||||
|| self.check_once(),
|
||||
POSTGRES_HEALTH_RETRY_DELAY,
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
async fn health_state_with_one_retry<F, Fut>(mut check: F, retry_delay: Duration) -> HealthState
|
||||
where
|
||||
F: FnMut() -> Fut,
|
||||
Fut: std::future::Future<Output = HealthState>,
|
||||
{
|
||||
let first = check().await;
|
||||
if first == HealthState::Healthy {
|
||||
return first;
|
||||
}
|
||||
tokio::time::sleep(retry_delay).await;
|
||||
check().await
|
||||
}
|
||||
|
||||
fn postgres_acquire_timeout_state(
|
||||
pool_size: u32,
|
||||
idle_connections: usize,
|
||||
max_connections: u32,
|
||||
) -> HealthState {
|
||||
if pool_size >= max_connections && idle_connections == 0 {
|
||||
HealthState::Degraded(format!(
|
||||
"PostgreSQL pool is saturated ({pool_size}/{max_connections} connections in use)"
|
||||
))
|
||||
} else {
|
||||
HealthState::Unhealthy("PostgreSQL connection acquisition timed out".to_owned())
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
use gateway_store::StoreConfig;
|
||||
use gateway_store::{StoreConfig, StoreError, StorePoolConfig};
|
||||
|
||||
const PASSWORD: &str = "111111111111111111111111111111111111111111111111";
|
||||
|
||||
@@ -17,6 +17,60 @@ fn store_config_should_derive_backup_staging_from_runtime_data_dir() {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn store_pool_config_keeps_the_public_surface_to_connection_pool_limits() {
|
||||
let pool: StorePoolConfig =
|
||||
serde_json::from_value(serde_json::json!({})).expect("default pool configuration");
|
||||
|
||||
assert_eq!(pool, StorePoolConfig::default());
|
||||
assert_eq!(pool.max_connections, 20);
|
||||
assert_eq!(pool.acquire_timeout_seconds, 5);
|
||||
assert_eq!(pool.observability_max_connections(), 16);
|
||||
assert_eq!(
|
||||
StorePoolConfig {
|
||||
max_connections: 50,
|
||||
..pool
|
||||
}
|
||||
.observability_max_connections(),
|
||||
40
|
||||
);
|
||||
assert!(
|
||||
serde_json::from_value::<StorePoolConfig>(serde_json::json!({
|
||||
"statement_timeout_seconds": 60,
|
||||
}))
|
||||
.is_err()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn store_config_rejects_a_pool_that_cannot_reserve_both_traffic_classes() {
|
||||
let root = tempfile::tempdir().expect("runtime data root");
|
||||
let mut config: StoreConfig = serde_json::from_value(serde_json::json!({
|
||||
"database": {
|
||||
"url": "postgres://codex_proxy@127.0.0.1:5432/codex_proxy",
|
||||
"password": PASSWORD,
|
||||
},
|
||||
"redis": {
|
||||
"url": "redis://127.0.0.1:6379/",
|
||||
"password": PASSWORD,
|
||||
},
|
||||
"pool": {
|
||||
"max_connections": 1,
|
||||
},
|
||||
}))
|
||||
.expect("syntactically valid Store configuration");
|
||||
|
||||
let error = config
|
||||
.resolve_and_validate(root.path())
|
||||
.expect_err("one connection cannot serve both traffic classes");
|
||||
|
||||
assert!(matches!(
|
||||
error,
|
||||
StoreError::InvalidData { message, .. }
|
||||
if message.contains("max_connections must be at least 2")
|
||||
));
|
||||
}
|
||||
|
||||
fn valid_config() -> StoreConfig {
|
||||
serde_json::from_value(serde_json::json!({
|
||||
"database": {
|
||||
|
||||
@@ -87,6 +87,14 @@ async fn groups_aggregate_cross_provider_members_and_key_bindings_without_multip
|
||||
.await
|
||||
.expect("create scoped client key");
|
||||
}
|
||||
seed_group_cost_snapshot(
|
||||
&database.pool,
|
||||
"req_historical_empty_group",
|
||||
"acct_group_openai",
|
||||
EMPTY_GROUP,
|
||||
"1.5",
|
||||
)
|
||||
.await;
|
||||
|
||||
let page = groups
|
||||
.list_account_groups(AccountGroupListQuery {
|
||||
@@ -97,6 +105,15 @@ async fn groups_aggregate_cross_provider_members_and_key_bindings_without_multip
|
||||
})
|
||||
.await
|
||||
.expect("list account groups");
|
||||
let members = groups
|
||||
.load_account_group_members(std::slice::from_ref(&mixed_group))
|
||||
.await
|
||||
.expect("load current-page group members");
|
||||
assert_eq!(members.len(), 2);
|
||||
assert_eq!(
|
||||
members.iter().map(|member| member.total_slots).sum::<u64>(),
|
||||
7
|
||||
);
|
||||
assert_eq!(page.total, 2);
|
||||
let by_id = page
|
||||
.items
|
||||
@@ -110,17 +127,20 @@ async fn groups_aggregate_cross_provider_members_and_key_bindings_without_multip
|
||||
BTreeMap::from([("openai".to_owned(), 1), ("xai".to_owned(), 1)])
|
||||
);
|
||||
assert_eq!(mixed.client_key_count, 2);
|
||||
assert_eq!(mixed.account_summary.available, 2);
|
||||
// PostgreSQL 返回持久页与 member facts;实时状态/容量由 Admin query service 投影。
|
||||
assert_eq!(mixed.account_summary.available, 0);
|
||||
assert_eq!(mixed.account_summary.limited, 0);
|
||||
assert_eq!(mixed.account_summary.total, 2);
|
||||
assert_eq!(mixed.account_summary.total, 0);
|
||||
assert_eq!(mixed.capacity.used_slots, None);
|
||||
assert_eq!(mixed.capacity.total_slots, 7);
|
||||
assert_eq!(mixed.capacity.total_slots, 0);
|
||||
assert_eq!(mixed.usage.today_usd.as_str(), "0");
|
||||
assert_eq!(mixed.usage.total_usd.as_str(), "0");
|
||||
assert_eq!(mixed.usage.retained_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);
|
||||
assert_eq!(empty.usage.today_usd.as_str(), "1.5");
|
||||
assert_eq!(empty.usage.retained_total_usd.as_str(), "1.5");
|
||||
|
||||
let all_key = keys
|
||||
.reveal_client_key(&client_key_id("key_all_accounts"))
|
||||
@@ -214,6 +234,80 @@ async fn groups_aggregate_cross_provider_members_and_key_bindings_without_multip
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn group_costs_should_include_statusless_websocket_but_reject_statusless_http() {
|
||||
let Some(database) = TestDatabase::create("account_group_statusless_websocket_cost").await
|
||||
else {
|
||||
return;
|
||||
};
|
||||
seed_account(
|
||||
&database.pool,
|
||||
"acct_group_statusless",
|
||||
"openai",
|
||||
"Statusless Account",
|
||||
)
|
||||
.await;
|
||||
let groups = PgAccountGroupRepository::new(database.pool.clone());
|
||||
groups
|
||||
.create_account_group(
|
||||
NewAccountGroup {
|
||||
id: group_id(EMPTY_GROUP),
|
||||
name: "Statusless Costs".to_owned(),
|
||||
description: None,
|
||||
color: group_color("#06B6D4CC"),
|
||||
},
|
||||
&context("create-statusless-cost-group"),
|
||||
)
|
||||
.await
|
||||
.expect("create statusless cost group");
|
||||
for (request_id, cost_amount) in [
|
||||
("req_group_http_success", "1.5"),
|
||||
("req_group_statusless_websocket", "2"),
|
||||
("req_group_statusless_http", "4"),
|
||||
] {
|
||||
seed_group_cost_snapshot(
|
||||
&database.pool,
|
||||
request_id,
|
||||
"acct_group_statusless",
|
||||
EMPTY_GROUP,
|
||||
cost_amount,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
sqlx::query(
|
||||
"update model_requests
|
||||
set client_transport = case id
|
||||
when 'req_group_statusless_websocket' then 'websocket'
|
||||
else 'http_sse'
|
||||
end,
|
||||
client_status_code = null
|
||||
where id in ('req_group_statusless_websocket', 'req_group_statusless_http')",
|
||||
)
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.expect("make account group requests statusless");
|
||||
|
||||
let page = groups
|
||||
.list_account_groups(AccountGroupListQuery {
|
||||
page: 1,
|
||||
page_size: PageSize::new(20).expect("page size"),
|
||||
search: None,
|
||||
enabled: None,
|
||||
})
|
||||
.await
|
||||
.expect("list statusless cost group");
|
||||
|
||||
assert_eq!(
|
||||
(
|
||||
page.items[0].usage.today_usd.as_str(),
|
||||
page.items[0].usage.retained_total_usd.as_str(),
|
||||
),
|
||||
("3.5", "3.5"),
|
||||
);
|
||||
|
||||
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 {
|
||||
@@ -281,6 +375,40 @@ async fn assign_accounts(pool: &sqlx::PgPool, group_id: &str, account_ids: &[&st
|
||||
}
|
||||
}
|
||||
|
||||
async fn seed_group_cost_snapshot(
|
||||
pool: &sqlx::PgPool,
|
||||
request_id: &str,
|
||||
account_id: &str,
|
||||
historical_group_id: &str,
|
||||
cost_amount: &str,
|
||||
) {
|
||||
sqlx::query(
|
||||
"insert into model_requests (
|
||||
id, client_api_key_ref, config_revision, protocol, operation, endpoint,
|
||||
client_transport, requested_model_id, provider_kind, provider_account_id,
|
||||
provider_account_ref, upstream_model_id, upstream_transport, attempt_count,
|
||||
upstream_send_state, downstream_committed_at, outcome, client_status_code,
|
||||
upstream_status_code, total_tokens, cost_source, cost_amount, cost_currency,
|
||||
started_at, deadline_at, completed_at,
|
||||
routing_scope, routing_group_refs, routing_group_names_snapshot
|
||||
) values (
|
||||
$1, 'key-group-history', 1, 'openai', 'responses', '/v1/responses',
|
||||
'http_sse', 'gpt-group', 'openai', $2, $2, 'gpt-group', 'http_sse', 1,
|
||||
'sent', now(), 'succeeded', 200, 200, 10,
|
||||
'provider_reported', $4::numeric, 'USD', now() - interval '1 minute',
|
||||
now() + interval '5 minutes', now(),
|
||||
'groups', array[$3]::text[], jsonb_build_array($3::text)
|
||||
)",
|
||||
)
|
||||
.bind(request_id)
|
||||
.bind(account_id)
|
||||
.bind(historical_group_id)
|
||||
.bind(cost_amount)
|
||||
.execute(pool)
|
||||
.await
|
||||
.expect("seed historical group cost snapshot");
|
||||
}
|
||||
|
||||
async fn current_revision(pool: &sqlx::PgPool) -> u64 {
|
||||
let value =
|
||||
sqlx::query_scalar::<_, i64>("select config_revision from runtime_settings where id = 1")
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::{collections::BTreeMap, sync::Arc, time::Duration};
|
||||
|
||||
use chrono::{TimeZone as _, Utc};
|
||||
use gateway_admin::{
|
||||
@@ -9,10 +9,15 @@ use gateway_admin::{
|
||||
},
|
||||
ports::store::ClientKeyStore as _,
|
||||
};
|
||||
use gateway_core::{
|
||||
engine::{CancellationToken, execution::ClientApiKeyUsageSink as _},
|
||||
policy::ClientApiKeyId,
|
||||
task::DaemonTask as _,
|
||||
};
|
||||
use gateway_store::postgres::{
|
||||
ClientApiKeyCursor, ClientApiKeyCursorValue, ClientApiKeyListQuery, ClientApiKeyRepository,
|
||||
ClientApiKeySort, ClientApiKeySortDirection, ClientApiKeySortField, NewClientApiKey,
|
||||
PgAdminClientKeyStore, PgClientApiKeyRepository,
|
||||
PgAdminClientKeyStore, PgClientApiKeyRepository, PgClientApiKeyUsageSink,
|
||||
};
|
||||
|
||||
use super::TestDatabase;
|
||||
@@ -307,6 +312,47 @@ async fn client_key_usage_touch_should_preserve_the_latest_observed_timestamp()
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn client_key_usage_daemon_should_flush_pending_touch_on_shutdown() {
|
||||
let Some(database) = TestDatabase::create("client_key_usage_daemon").await else {
|
||||
return;
|
||||
};
|
||||
sqlx::query(
|
||||
"insert into client_api_keys (
|
||||
id, name, key, enabled, max_concurrency, requests_per_minute,
|
||||
created_at, updated_at
|
||||
) values ('key_usage_daemon', 'usage', $1, true, 0, 0, now(), now())",
|
||||
)
|
||||
.bind(format!("sk_{}", "u".repeat(43)))
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.expect("seed client key");
|
||||
|
||||
let (sink, writer) =
|
||||
PgClientApiKeyUsageSink::with_flush_delay(database.pool.clone(), Duration::from_secs(60));
|
||||
let writer = Arc::new(writer);
|
||||
let cancellation = CancellationToken::new();
|
||||
let task = {
|
||||
let writer = Arc::clone(&writer);
|
||||
let cancellation = cancellation.clone();
|
||||
tokio::spawn(async move { writer.run(cancellation).await })
|
||||
};
|
||||
sink.record_used(&ClientApiKeyId::new("key_usage_daemon").expect("valid client API Key ID"));
|
||||
cancellation.cancel();
|
||||
task.await
|
||||
.expect("join client key usage daemon")
|
||||
.expect("stop client key usage daemon");
|
||||
|
||||
let last_used_at: Option<chrono::DateTime<Utc>> = sqlx::query_scalar(
|
||||
"select last_used_at from client_api_keys where id = 'key_usage_daemon'",
|
||||
)
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.expect("load client key last-used timestamp");
|
||||
assert!(last_used_at.is_some());
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dedicated_reveal_returns_plaintext_without_debug_exposure() {
|
||||
let Some(database) = TestDatabase::create("client_key_reveal").await else {
|
||||
|
||||
@@ -0,0 +1,47 @@
|
||||
use std::time::Duration;
|
||||
|
||||
use gateway_core::health::{HealthProbe as _, HealthState};
|
||||
use gateway_store::PostgresHealthProbe;
|
||||
|
||||
use super::TestDatabase;
|
||||
|
||||
#[tokio::test]
|
||||
async fn postgres_health_reports_a_saturated_pool_as_degraded() {
|
||||
let Some(database) = TestDatabase::create("health_saturated").await else {
|
||||
return;
|
||||
};
|
||||
let connection_1 = database.pool.acquire().await.expect("first connection");
|
||||
let connection_2 = database.pool.acquire().await.expect("second connection");
|
||||
let probe = PostgresHealthProbe::new(database.pool.clone(), 2);
|
||||
|
||||
let state = probe.check().await;
|
||||
|
||||
assert!(matches!(
|
||||
state,
|
||||
HealthState::Degraded(message)
|
||||
if message == "PostgreSQL pool is saturated (2/2 connections in use)"
|
||||
));
|
||||
drop((connection_1, connection_2));
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn postgres_health_tolerates_one_transient_saturation() {
|
||||
let Some(database) = TestDatabase::create("health_transient").await else {
|
||||
return;
|
||||
};
|
||||
let connection_1 = database.pool.acquire().await.expect("first connection");
|
||||
let connection_2 = database.pool.acquire().await.expect("second connection");
|
||||
let probe = PostgresHealthProbe::new(database.pool.clone(), 2);
|
||||
let release = tokio::spawn(async move {
|
||||
tokio::time::sleep(Duration::from_millis(475)).await;
|
||||
drop(connection_1);
|
||||
});
|
||||
|
||||
let state = probe.check().await;
|
||||
|
||||
assert_eq!(state, HealthState::Healthy);
|
||||
release.await.expect("release task must not panic");
|
||||
drop(connection_2);
|
||||
database.close().await;
|
||||
}
|
||||
@@ -1,6 +1,9 @@
|
||||
use std::str::FromStr;
|
||||
use std::{str::FromStr, time::Duration};
|
||||
|
||||
use gateway_store::postgres::connect_and_migrate;
|
||||
use gateway_store::postgres::{
|
||||
ObservabilityQueryBudget, PgAdminAccountStore, PgAdminObservabilityStore,
|
||||
PgObservabilityRepository, connect_and_migrate,
|
||||
};
|
||||
use sqlx::{
|
||||
ConnectOptions as _, PgPool,
|
||||
postgres::{PgConnectOptions, PgPoolOptions},
|
||||
@@ -14,9 +17,11 @@ mod backup;
|
||||
mod client_keys;
|
||||
mod execution;
|
||||
mod execution_buffer;
|
||||
mod health;
|
||||
mod observability;
|
||||
mod ops_events;
|
||||
mod provider_accounts;
|
||||
mod query_budget;
|
||||
mod retention;
|
||||
mod runtime_settings;
|
||||
mod snapshot;
|
||||
@@ -30,6 +35,23 @@ pub(super) struct TestDatabase {
|
||||
schema: String,
|
||||
}
|
||||
|
||||
pub(super) fn observability_query_budget() -> ObservabilityQueryBudget {
|
||||
ObservabilityQueryBudget::try_new(4, Duration::from_secs(1))
|
||||
.expect("valid test observability query budget")
|
||||
}
|
||||
|
||||
pub(super) fn observability_repository(pool: &PgPool) -> PgObservabilityRepository {
|
||||
PgObservabilityRepository::new(pool.clone(), None, observability_query_budget())
|
||||
}
|
||||
|
||||
pub(super) fn admin_observability_store(pool: &PgPool) -> PgAdminObservabilityStore {
|
||||
PgAdminObservabilityStore::new(pool.clone(), None, None, observability_query_budget())
|
||||
}
|
||||
|
||||
pub(super) fn admin_account_store(pool: &PgPool) -> PgAdminAccountStore {
|
||||
PgAdminAccountStore::new(pool.clone(), None, observability_query_budget())
|
||||
}
|
||||
|
||||
impl TestDatabase {
|
||||
pub(super) async fn create(label: &str) -> Option<Self> {
|
||||
let database_url = crate::support::test_env("CPR_TEST_DATABASE_URL")?;
|
||||
@@ -106,18 +128,33 @@ async fn connect_and_migrate_should_apply_all_migrations_once_and_reopen_cleanly
|
||||
.database(&database)
|
||||
.to_url_lossy()
|
||||
.to_string();
|
||||
let first = connect_and_migrate(&isolated_url, gateway_store::StorePoolConfig::default())
|
||||
let pool_config = gateway_store::StorePoolConfig::default();
|
||||
let first = connect_and_migrate(&isolated_url, pool_config)
|
||||
.await
|
||||
.expect("apply migrations through production migrator");
|
||||
let first_table_count = sqlx::query_scalar::<_, i64>(
|
||||
"select count(*) from information_schema.tables where table_schema = 'public'",
|
||||
let session_settings = sqlx::query_as::<_, (String, i64, i64, i64)>(
|
||||
"select current_setting('application_name'),
|
||||
extract(epoch from current_setting('statement_timeout')::interval)::bigint,
|
||||
extract(epoch from current_setting('lock_timeout')::interval)::bigint,
|
||||
extract(epoch from current_setting(
|
||||
'idle_in_transaction_session_timeout'
|
||||
)::interval)::bigint",
|
||||
)
|
||||
.fetch_one(&first)
|
||||
.await
|
||||
.expect("count migrated tables");
|
||||
.expect("load runtime PostgreSQL session settings");
|
||||
let first_tables = sqlx::query_scalar::<_, String>(
|
||||
"select table_name
|
||||
from information_schema.tables
|
||||
where table_schema = 'public'
|
||||
order by table_name",
|
||||
)
|
||||
.fetch_all(&first)
|
||||
.await
|
||||
.expect("load migrated tables");
|
||||
first.close().await;
|
||||
|
||||
let second = connect_and_migrate(&isolated_url, gateway_store::StorePoolConfig::default())
|
||||
let second = connect_and_migrate(&isolated_url, pool_config)
|
||||
.await
|
||||
.expect("reopen database through production migrator");
|
||||
let migration_count =
|
||||
@@ -180,7 +217,25 @@ 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, 13);
|
||||
assert_eq!(
|
||||
first_tables,
|
||||
[
|
||||
"_sqlx_migrations",
|
||||
"account_group_accounts",
|
||||
"account_groups",
|
||||
"admin_audit_events",
|
||||
"admin_users",
|
||||
"backup_records",
|
||||
"backup_settings",
|
||||
"client_api_key_groups",
|
||||
"client_api_keys",
|
||||
"model_requests",
|
||||
"ops_events",
|
||||
"provider_accounts",
|
||||
"runtime_settings",
|
||||
]
|
||||
);
|
||||
assert_eq!(session_settings, ("codex-proxy-rs".to_owned(), 30, 5, 30));
|
||||
assert_eq!(
|
||||
migration_count,
|
||||
i64::try_from(TEST_MIGRATOR.iter().count())
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
use std::{
|
||||
collections::BTreeMap,
|
||||
num::NonZeroU32,
|
||||
sync::Arc,
|
||||
time::{Duration, SystemTime},
|
||||
};
|
||||
@@ -19,13 +18,15 @@ use gateway_core::{
|
||||
},
|
||||
};
|
||||
use gateway_store::postgres::{
|
||||
DiagnosticDimension, ObservabilityPageNumber, ObservabilityPageSize, ObservabilityRange,
|
||||
ObservabilityRepository, OpsErrorFilter, OpsErrorQuery, PgAdminObservabilityStore,
|
||||
PgObservabilityRepository, ProviderAccountUsageQuery, UsageRecordFilter, UsageRecordQuery,
|
||||
DiagnosticDimension, ObservabilityPageSize, ObservabilityRange, ObservabilityRepository,
|
||||
OpsErrorFilter, OpsErrorQuery, PgAdminObservabilityStore, PgObservabilityRepository,
|
||||
ProviderAccountUsageQuery, UsageRecordFilter, UsageRecordQuery,
|
||||
};
|
||||
use sqlx::PgPool;
|
||||
|
||||
use super::TestDatabase;
|
||||
use super::{
|
||||
TestDatabase, admin_observability_store, observability_query_budget, observability_repository,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn observability_range_rejects_empty_window() {
|
||||
@@ -96,7 +97,7 @@ async fn observability_preserves_and_filters_opaque_response_ids() {
|
||||
.expect("persist opaque response IDs");
|
||||
let range = ObservabilityRange::new(now - TimeDelta::hours(1), now + TimeDelta::hours(1))
|
||||
.expect("observability range");
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone(), None);
|
||||
let repository = observability_repository(&database.pool);
|
||||
|
||||
let records = repository
|
||||
.list_usage_records(UsageRecordQuery {
|
||||
@@ -105,25 +106,115 @@ async fn observability_preserves_and_filters_opaque_response_ids() {
|
||||
response_id: Some(response_id.clone()),
|
||||
..UsageRecordFilter::default()
|
||||
},
|
||||
cursor: None,
|
||||
page: ObservabilityPageNumber::new(1).expect("page"),
|
||||
current_page: 1,
|
||||
page_size: ObservabilityPageSize::new(10).expect("page size"),
|
||||
})
|
||||
.await
|
||||
.expect("filter opaque response ID");
|
||||
assert_eq!(records.total, 1);
|
||||
assert_eq!(records.items[0].id, "req_observe_success");
|
||||
let detail = repository
|
||||
.usage_record_detail("req_observe_success")
|
||||
.await
|
||||
.expect("usage detail with opaque response IDs");
|
||||
assert_eq!(
|
||||
records.items[0].client_response_id.as_deref(),
|
||||
detail.request.client_response_id.as_deref(),
|
||||
Some(response_id.as_str())
|
||||
);
|
||||
assert_eq!(
|
||||
records.items[0].upstream_response_id.as_deref(),
|
||||
detail.request.upstream_response_id.as_deref(),
|
||||
Some(response_id.as_str())
|
||||
);
|
||||
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn usage_page_should_always_return_total() {
|
||||
let Some(database) = TestDatabase::create("usage_page_with_total").await else {
|
||||
return;
|
||||
};
|
||||
let now = Utc::now();
|
||||
seed_observability_facts(&database.pool, now)
|
||||
.await
|
||||
.expect("seed observability facts");
|
||||
let range = ObservabilityRange::new(now - TimeDelta::hours(1), now + TimeDelta::hours(1))
|
||||
.expect("observability range");
|
||||
|
||||
let page = observability_repository(&database.pool)
|
||||
.list_usage_records(UsageRecordQuery {
|
||||
range,
|
||||
filter: UsageRecordFilter::default(),
|
||||
current_page: 1,
|
||||
page_size: ObservabilityPageSize::new(10).expect("page size"),
|
||||
})
|
||||
.await
|
||||
.expect("usage page with total");
|
||||
|
||||
assert_eq!(page.total, 1);
|
||||
assert_eq!(page.current_page, 1);
|
||||
assert_eq!(page.page_size, 10);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn usage_search_should_match_literal_prefix_instead_of_substring() {
|
||||
let Some(database) = TestDatabase::create("usage_literal_prefix_search").await else {
|
||||
return;
|
||||
};
|
||||
let now = Utc::now();
|
||||
seed_observability_facts(&database.pool, now)
|
||||
.await
|
||||
.expect("seed observability facts");
|
||||
let range = ObservabilityRange::new(now - TimeDelta::hours(1), now + TimeDelta::hours(1))
|
||||
.expect("observability range");
|
||||
|
||||
let page = observability_repository(&database.pool)
|
||||
.list_usage_records(UsageRecordQuery {
|
||||
range,
|
||||
filter: UsageRecordFilter {
|
||||
search: Some("observe_success".to_owned()),
|
||||
..UsageRecordFilter::default()
|
||||
},
|
||||
current_page: 1,
|
||||
page_size: ObservabilityPageSize::new(10).expect("page size"),
|
||||
})
|
||||
.await
|
||||
.expect("usage substring search");
|
||||
|
||||
assert_eq!(page.total, 0);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ops_search_should_treat_sql_wildcards_as_literals() {
|
||||
let Some(database) = TestDatabase::create("ops_literal_wildcard_search").await else {
|
||||
return;
|
||||
};
|
||||
let now = Utc::now();
|
||||
seed_observability_facts(&database.pool, now)
|
||||
.await
|
||||
.expect("seed observability facts");
|
||||
let range = ObservabilityRange::new(now - TimeDelta::hours(1), now + TimeDelta::hours(1))
|
||||
.expect("observability range");
|
||||
|
||||
let page = observability_repository(&database.pool)
|
||||
.list_ops_errors(OpsErrorQuery {
|
||||
range,
|
||||
filter: OpsErrorFilter {
|
||||
search: Some("req%".to_owned()),
|
||||
..OpsErrorFilter::default()
|
||||
},
|
||||
current_page: 1,
|
||||
page_size: ObservabilityPageSize::new(10).expect("page size"),
|
||||
})
|
||||
.await
|
||||
.expect("ops literal wildcard search");
|
||||
|
||||
assert_eq!(page.total, 0);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dashboard_account_metrics_should_partition_account_statuses() {
|
||||
let Some(database) = TestDatabase::create("dashboard_account_metrics").await else {
|
||||
@@ -161,7 +252,7 @@ async fn dashboard_account_metrics_should_partition_account_statuses() {
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.expect("seed account metric states");
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone(), None);
|
||||
let repository = observability_repository(&database.pool);
|
||||
let range = ObservabilityRange::new(now - TimeDelta::hours(1), now + TimeDelta::hours(1))
|
||||
.expect("dashboard range");
|
||||
|
||||
@@ -246,8 +337,11 @@ async fn dashboard_account_metrics_with_cooldowns_should_only_reclassify_eligibl
|
||||
test_cooldown("acct_metrics_unknown", 6),
|
||||
test_cooldown("acct_metrics_stale_cooldown", 3),
|
||||
]);
|
||||
let repository =
|
||||
PgObservabilityRepository::new(database.pool.clone(), Some(Arc::new(cooldowns)));
|
||||
let repository = PgObservabilityRepository::new(
|
||||
database.pool.clone(),
|
||||
Some(Arc::new(cooldowns)),
|
||||
observability_query_budget(),
|
||||
);
|
||||
let range = ObservabilityRange::new(now - TimeDelta::hours(1), now + TimeDelta::hours(1))
|
||||
.expect("dashboard range");
|
||||
|
||||
@@ -371,7 +465,7 @@ async fn calculated_usage_billing_facts_keep_only_completed_calculated_costs() {
|
||||
.expect("seed calculated billing facts");
|
||||
let range = ObservabilityRange::new(now - TimeDelta::hours(1), now + TimeDelta::hours(1))
|
||||
.expect("observability range");
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone(), None);
|
||||
let repository = observability_repository(&database.pool);
|
||||
|
||||
let facts = repository
|
||||
.usage_calculated_billing_facts(range, UsageRecordFilter::default())
|
||||
@@ -385,7 +479,7 @@ async fn calculated_usage_billing_facts_keep_only_completed_calculated_costs() {
|
||||
assert_eq!(facts[0].service_tier.as_deref(), Some("priority"));
|
||||
assert_eq!(facts[0].total.amount.as_str(), "1.25");
|
||||
|
||||
let store = PgAdminObservabilityStore::new(database.pool.clone(), None, None);
|
||||
let store = admin_observability_store(&database.pool);
|
||||
let facts = store
|
||||
.usage_calculated_billing_facts(
|
||||
admin_observability::TimeRange::new(range.start, range.end)
|
||||
@@ -413,24 +507,37 @@ async fn admin_observability_adapter_preserves_utc_queries_metrics_costs_and_det
|
||||
let range =
|
||||
admin_observability::TimeRange::new(now - TimeDelta::hours(1), now + TimeDelta::hours(1))
|
||||
.expect("admin observability range");
|
||||
let store = PgAdminObservabilityStore::new(database.pool.clone(), None, None);
|
||||
let store = admin_observability_store(&database.pool);
|
||||
|
||||
let dashboard = store
|
||||
.dashboard_summary(range, now)
|
||||
.await
|
||||
.expect("admin dashboard summary");
|
||||
assert_eq!(dashboard.range, range);
|
||||
assert_eq!(dashboard.requests.request_count, 3);
|
||||
assert_eq!(dashboard.requests.success_count, 2);
|
||||
assert_eq!(dashboard.requests.failure_count, 1);
|
||||
assert_eq!(dashboard.requests.first_token_latency_sum_ms, 120);
|
||||
assert_eq!(dashboard.requests.latency_sum_ms, 900);
|
||||
assert_eq!(dashboard.requests.min_latency_ms, Some(900));
|
||||
assert_eq!(dashboard.requests.max_latency_ms, Some(900));
|
||||
assert_eq!(dashboard.attempts.attempt_count, 4);
|
||||
assert_eq!(dashboard.attempts.cost_coverage.provider_reported_count, 1);
|
||||
assert_eq!(dashboard.attempts.cost_coverage.unavailable_count, 0);
|
||||
assert_eq!(dashboard.attempts.costs[0].amount.as_str(), "1.25");
|
||||
assert_eq!(
|
||||
dashboard
|
||||
.trend
|
||||
.iter()
|
||||
.map(|point| point.metrics.request_count)
|
||||
.sum::<u64>(),
|
||||
3
|
||||
);
|
||||
assert_eq!(
|
||||
dashboard
|
||||
.trend
|
||||
.iter()
|
||||
.map(|point| point.metrics.first_token_latency_sum_ms)
|
||||
.sum::<u64>(),
|
||||
120
|
||||
);
|
||||
assert_eq!(
|
||||
dashboard
|
||||
.trend
|
||||
.iter()
|
||||
.map(|point| point.metrics.latency_sum_ms)
|
||||
.sum::<u64>(),
|
||||
900
|
||||
);
|
||||
assert_eq!(dashboard.provider_accounts.total, 1);
|
||||
assert_eq!(dashboard.account_usage[0].request_count, 1);
|
||||
assert_eq!(dashboard.account_usage[0].request_buckets.len(), 2);
|
||||
@@ -447,10 +554,7 @@ async fn admin_observability_adapter_preserves_utc_queries_metrics_costs_and_det
|
||||
dashboard.recent_requests[0].service_tier.as_deref(),
|
||||
Some("priority")
|
||||
);
|
||||
assert_eq!(
|
||||
dashboard.recent_requests[0].outcome,
|
||||
admin_observability::RequestOutcome::Succeeded,
|
||||
);
|
||||
assert_eq!(dashboard.recent_requests[0].id, "req_observe_success");
|
||||
assert_eq!(
|
||||
dashboard.recent_requests[0]
|
||||
.cost_amount
|
||||
@@ -478,6 +582,14 @@ async fn admin_observability_adapter_preserves_utc_queries_metrics_costs_and_det
|
||||
.sum::<u64>(),
|
||||
1,
|
||||
);
|
||||
assert!(
|
||||
dashboard
|
||||
.trend
|
||||
.iter()
|
||||
.chain(&dashboard_trend)
|
||||
.all(|point| point.costs.is_empty()),
|
||||
"Dashboard 趋势不应触发未展示的成本聚合",
|
||||
);
|
||||
|
||||
let trend = store
|
||||
.usage_trend(
|
||||
@@ -511,20 +623,34 @@ async fn admin_observability_adapter_preserves_utc_queries_metrics_costs_and_det
|
||||
.list_usage_records(admin_observability::UsageQuery {
|
||||
range,
|
||||
filter: admin_observability::UsageFilter::default(),
|
||||
cursor: None,
|
||||
page: admin_page(1),
|
||||
current_page: 1,
|
||||
page_size: PageSize::new(1).expect("page size"),
|
||||
})
|
||||
.await
|
||||
.expect("first usage page");
|
||||
assert_eq!(first_page.total, 1);
|
||||
assert_eq!(first_page.items.len(), 1);
|
||||
assert!(first_page.next_cursor.is_none());
|
||||
assert_eq!(first_page.current_page, 1);
|
||||
assert_eq!(first_page.page_size, 1);
|
||||
assert_eq!(
|
||||
first_page.items[0].service_tier.as_deref(),
|
||||
Some("priority")
|
||||
);
|
||||
|
||||
let deep_page = store
|
||||
.list_usage_records(admin_observability::UsageQuery {
|
||||
range,
|
||||
filter: admin_observability::UsageFilter::default(),
|
||||
current_page: 129,
|
||||
page_size: PageSize::new(1).expect("page size"),
|
||||
})
|
||||
.await
|
||||
.expect("direct deep usage page");
|
||||
assert_eq!(deep_page.current_page, 129);
|
||||
assert_eq!(deep_page.page_size, 1);
|
||||
assert_eq!(deep_page.total, 1);
|
||||
assert!(deep_page.items.is_empty());
|
||||
|
||||
let filtered = store
|
||||
.list_usage_records(admin_observability::UsageQuery {
|
||||
range,
|
||||
@@ -543,24 +669,14 @@ async fn admin_observability_adapter_preserves_utc_queries_metrics_costs_and_det
|
||||
upstream_request_id: Some("upstream_req_success".to_owned()),
|
||||
search: Some("req_observe_success".to_owned()),
|
||||
},
|
||||
cursor: None,
|
||||
page: admin_page(1),
|
||||
current_page: 1,
|
||||
page_size: PageSize::new(10).expect("page size"),
|
||||
})
|
||||
.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].id, "req_observe_success");
|
||||
assert_eq!(filtered.items[0].service_tier.as_deref(), Some("priority"));
|
||||
assert_eq!(
|
||||
filtered.items[0].upstream_request_id.as_deref(),
|
||||
Some("upstream_req_success")
|
||||
);
|
||||
|
||||
let other_outcome = store
|
||||
.list_usage_records(admin_observability::UsageQuery {
|
||||
@@ -572,8 +688,7 @@ async fn admin_observability_adapter_preserves_utc_queries_metrics_costs_and_det
|
||||
),
|
||||
..admin_observability::UsageFilter::default()
|
||||
},
|
||||
cursor: None,
|
||||
page: admin_page(1),
|
||||
current_page: 1,
|
||||
page_size: PageSize::new(10).expect("page size"),
|
||||
})
|
||||
.await
|
||||
@@ -596,6 +711,10 @@ async fn admin_observability_adapter_preserves_utc_queries_metrics_costs_and_det
|
||||
detail.request.routing_group_names_snapshot,
|
||||
["Historical group"]
|
||||
);
|
||||
assert_eq!(
|
||||
detail.request.upstream_request_id.as_deref(),
|
||||
Some("upstream_req_success")
|
||||
);
|
||||
assert_eq!(detail.attempts.len(), 1);
|
||||
assert_eq!(
|
||||
detail.attempts[0].outcome,
|
||||
@@ -642,14 +761,15 @@ async fn admin_observability_adapter_preserves_utc_queries_metrics_costs_and_det
|
||||
request_id: Some("req_observe_failed".to_owned()),
|
||||
provider_kind: Some("openai".to_owned()),
|
||||
provider_account_ref: Some("acct_observe".to_owned()),
|
||||
operation: Some("responses".to_owned()),
|
||||
endpoint: Some("/v1/responses".to_owned()),
|
||||
model: Some("upstream-model".to_owned()),
|
||||
failure_kind: Some("rate_limited".to_owned()),
|
||||
status_code: Some(429),
|
||||
search: Some("limited".to_owned()),
|
||||
search: Some("req_observe_failed".to_owned()),
|
||||
..admin_observability::OpsErrorFilter::default()
|
||||
},
|
||||
cursor: None,
|
||||
page: admin_page(1),
|
||||
current_page: 1,
|
||||
page_size: PageSize::new(10).expect("page size"),
|
||||
})
|
||||
.await
|
||||
@@ -657,6 +777,20 @@ async fn admin_observability_adapter_preserves_utc_queries_metrics_costs_and_det
|
||||
assert_eq!(errors.total, 2);
|
||||
assert!(errors.items.iter().all(|item| item.occurred_at <= now));
|
||||
|
||||
let deep_errors = store
|
||||
.list_ops_errors(admin_observability::OpsErrorQuery {
|
||||
range,
|
||||
filter: admin_observability::OpsErrorFilter::default(),
|
||||
current_page: 129,
|
||||
page_size: PageSize::new(1).expect("page size"),
|
||||
})
|
||||
.await
|
||||
.expect("direct deep ops page");
|
||||
assert_eq!(deep_errors.current_page, 129);
|
||||
assert_eq!(deep_errors.page_size, 1);
|
||||
assert_eq!(deep_errors.total, 2);
|
||||
assert!(deep_errors.items.is_empty());
|
||||
|
||||
for filter in [
|
||||
admin_observability::OpsErrorFilter {
|
||||
client_api_key_ref: Some("missing-key".to_owned()),
|
||||
@@ -687,8 +821,7 @@ async fn admin_observability_adapter_preserves_utc_queries_metrics_costs_and_det
|
||||
.list_ops_errors(admin_observability::OpsErrorQuery {
|
||||
range,
|
||||
filter,
|
||||
cursor: None,
|
||||
page: admin_page(1),
|
||||
current_page: 1,
|
||||
page_size: PageSize::new(10).expect("page size"),
|
||||
})
|
||||
.await
|
||||
@@ -699,10 +832,6 @@ async fn admin_observability_adapter_preserves_utc_queries_metrics_costs_and_det
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
fn admin_page(value: u32) -> admin_observability::PageNumber {
|
||||
admin_observability::PageNumber::new(NonZeroU32::new(value).expect("positive page"))
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dashboard_summary_totals_include_history_outside_selected_range() {
|
||||
let Some(database) = TestDatabase::create("observability_dashboard_totals").await else {
|
||||
@@ -723,13 +852,20 @@ async fn dashboard_summary_totals_include_history_outside_selected_range() {
|
||||
.expect("move historical request outside selected range");
|
||||
let range = ObservabilityRange::new(now - TimeDelta::hours(1), now + TimeDelta::hours(1))
|
||||
.expect("observability range");
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone(), None);
|
||||
let repository = observability_repository(&database.pool);
|
||||
|
||||
let dashboard = repository
|
||||
.dashboard_summary(range, now)
|
||||
.await
|
||||
.expect("dashboard summary");
|
||||
assert_eq!(dashboard.requests.request_count, 2);
|
||||
assert_eq!(
|
||||
dashboard
|
||||
.trend
|
||||
.iter()
|
||||
.map(|point| point.metrics.request_count)
|
||||
.sum::<u64>(),
|
||||
2
|
||||
);
|
||||
assert!(dashboard.account_usage.is_empty());
|
||||
|
||||
let explicit_account_usage = repository
|
||||
@@ -774,51 +910,92 @@ async fn observability_queries_preserve_request_account_cost_and_diagnostic_fact
|
||||
.expect("seed observability facts");
|
||||
let range = ObservabilityRange::new(now - TimeDelta::hours(1), now + TimeDelta::hours(1))
|
||||
.expect("observability range");
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone(), None);
|
||||
let repository = observability_repository(&database.pool);
|
||||
|
||||
let dashboard = repository
|
||||
.dashboard_summary(range, now)
|
||||
.await
|
||||
.expect("dashboard summary");
|
||||
assert_eq!(dashboard.requests.request_count, 3);
|
||||
assert_eq!(dashboard.requests.success_count, 2);
|
||||
assert_eq!(dashboard.requests.failure_count, 1);
|
||||
assert_eq!(dashboard.requests.caller_error_count, 0);
|
||||
assert_eq!(dashboard.requests.cache_eligible_request_count, 1);
|
||||
assert_eq!(dashboard.requests.cache_hit_request_count, 1);
|
||||
assert_eq!(dashboard.requests.cache_hit_request_rate(), Some(1.0));
|
||||
assert_eq!(
|
||||
dashboard
|
||||
.requests
|
||||
.latency_percentiles
|
||||
.p50_ms
|
||||
.trend
|
||||
.iter()
|
||||
.map(|point| point.metrics.request_count)
|
||||
.sum::<u64>(),
|
||||
3
|
||||
);
|
||||
assert_eq!(
|
||||
dashboard
|
||||
.trend
|
||||
.iter()
|
||||
.map(|point| point.metrics.success_count)
|
||||
.sum::<u64>(),
|
||||
2
|
||||
);
|
||||
assert_eq!(
|
||||
dashboard
|
||||
.trend
|
||||
.iter()
|
||||
.map(|point| point.metrics.failure_count)
|
||||
.sum::<u64>(),
|
||||
1
|
||||
);
|
||||
assert_eq!(
|
||||
dashboard
|
||||
.trend
|
||||
.iter()
|
||||
.map(|point| point.metrics.caller_error_count)
|
||||
.sum::<u64>(),
|
||||
0
|
||||
);
|
||||
assert_eq!(
|
||||
dashboard
|
||||
.trend
|
||||
.iter()
|
||||
.map(|point| point.metrics.cache_eligible_request_count)
|
||||
.sum::<u64>(),
|
||||
1
|
||||
);
|
||||
assert_eq!(
|
||||
dashboard
|
||||
.trend
|
||||
.iter()
|
||||
.map(|point| point.metrics.cache_hit_request_count)
|
||||
.sum::<u64>(),
|
||||
1
|
||||
);
|
||||
assert_eq!(
|
||||
dashboard
|
||||
.trend
|
||||
.iter()
|
||||
.find_map(|point| point.metrics.latency_percentiles.p50_ms)
|
||||
.expect("latency p50")
|
||||
.as_f64(),
|
||||
900.0
|
||||
);
|
||||
assert_eq!(
|
||||
dashboard
|
||||
.requests
|
||||
.latency_percentiles
|
||||
.p95_ms
|
||||
.trend
|
||||
.iter()
|
||||
.find_map(|point| point.metrics.latency_percentiles.p95_ms)
|
||||
.expect("latency p95")
|
||||
.as_f64(),
|
||||
900.0
|
||||
);
|
||||
assert_eq!(
|
||||
dashboard
|
||||
.requests
|
||||
.latency_percentiles
|
||||
.p99_ms
|
||||
.trend
|
||||
.iter()
|
||||
.find_map(|point| point.metrics.latency_percentiles.p99_ms)
|
||||
.expect("latency p99")
|
||||
.as_f64(),
|
||||
900.0
|
||||
);
|
||||
assert_eq!(
|
||||
dashboard
|
||||
.requests
|
||||
.first_token_latency_percentiles
|
||||
.p50_ms
|
||||
.trend
|
||||
.iter()
|
||||
.find_map(|point| point.metrics.first_token_latency_percentiles.p50_ms)
|
||||
.expect("first token p50")
|
||||
.as_f64(),
|
||||
120.0
|
||||
@@ -899,8 +1076,7 @@ async fn observability_queries_preserve_request_account_cost_and_diagnostic_fact
|
||||
provider_account_ref: Some("acct_observe".to_owned()),
|
||||
..UsageRecordFilter::default()
|
||||
},
|
||||
cursor: None,
|
||||
page: ObservabilityPageNumber::new(1).expect("page"),
|
||||
current_page: 1,
|
||||
page_size: ObservabilityPageSize::new(10).expect("page size"),
|
||||
})
|
||||
.await
|
||||
@@ -913,25 +1089,27 @@ async fn observability_queries_preserve_request_account_cost_and_diagnostic_fact
|
||||
.expect("successful image usage record");
|
||||
assert_eq!(
|
||||
(
|
||||
successful_image.websocket_pool.as_deref(),
|
||||
successful_image.service_tier.as_deref(),
|
||||
successful_image
|
||||
.provider_account_authentication_kind
|
||||
.as_deref(),
|
||||
successful_image.image_input_tokens,
|
||||
successful_image.image_output_tokens,
|
||||
successful_image.image_generation_requested,
|
||||
successful_image.image_generation_succeeded,
|
||||
),
|
||||
(Some("priority"), Some("oauth"), Some(31), Some(9))
|
||||
);
|
||||
|
||||
let successful_detail = repository
|
||||
.usage_record_detail("req_observe_success")
|
||||
.await
|
||||
.expect("successful usage detail");
|
||||
assert_eq!(
|
||||
(
|
||||
Some("reuse"),
|
||||
Some("priority"),
|
||||
Some("oauth"),
|
||||
Some(31),
|
||||
Some(9),
|
||||
true,
|
||||
Some(true)
|
||||
)
|
||||
successful_detail.request.websocket_pool.as_deref(),
|
||||
successful_detail.request.image_generation_requested,
|
||||
successful_detail.request.image_generation_succeeded,
|
||||
),
|
||||
(Some("reuse"), true, Some(true))
|
||||
);
|
||||
|
||||
assert!(
|
||||
@@ -949,6 +1127,15 @@ async fn observability_queries_preserve_request_account_cost_and_diagnostic_fact
|
||||
assert_eq!(overview.attempts.failure_count, 2);
|
||||
assert_eq!(overview.requests.request_count, 3);
|
||||
assert_eq!(overview.requests.failure_count, 1);
|
||||
assert_eq!(
|
||||
(
|
||||
overview.providers[0].request_count,
|
||||
overview.providers[0].attempt_count,
|
||||
overview.providers[0].failure_count,
|
||||
overview.providers[0].total_tokens,
|
||||
),
|
||||
(3, 4, 1, 120),
|
||||
);
|
||||
|
||||
let succeeded = repository
|
||||
.usage_summary(
|
||||
@@ -993,8 +1180,7 @@ async fn observability_queries_preserve_request_account_cost_and_diagnostic_fact
|
||||
.list_ops_errors(OpsErrorQuery {
|
||||
range,
|
||||
filter: OpsErrorFilter::default(),
|
||||
cursor: None,
|
||||
page: ObservabilityPageNumber::new(1).expect("page"),
|
||||
current_page: 1,
|
||||
page_size: ObservabilityPageSize::new(10).expect("page size"),
|
||||
})
|
||||
.await
|
||||
@@ -1007,6 +1193,15 @@ async fn observability_queries_preserve_request_account_cost_and_diagnostic_fact
|
||||
.expect("request error");
|
||||
assert_eq!(request_error.client_status_code, Some(502));
|
||||
assert_eq!(request_error.upstream_status_code, Some(429));
|
||||
assert_eq!(request_error.client_ip.as_deref(), Some("203.0.113.9"));
|
||||
assert_eq!(
|
||||
request_error.user_agent.as_deref(),
|
||||
Some("codex-cli/0.144.0")
|
||||
);
|
||||
assert_eq!(
|
||||
request_error.requested_model_id.as_deref(),
|
||||
Some("public-model")
|
||||
);
|
||||
let attempt_error = errors
|
||||
.items
|
||||
.iter()
|
||||
@@ -1014,6 +1209,12 @@ async fn observability_queries_preserve_request_account_cost_and_diagnostic_fact
|
||||
.expect("attempt error");
|
||||
assert_eq!(attempt_error.client_status_code, None);
|
||||
assert_eq!(attempt_error.upstream_status_code, Some(429));
|
||||
assert_eq!(attempt_error.endpoint.as_deref(), Some("/v1/responses"));
|
||||
assert_eq!(attempt_error.client_ip.as_deref(), Some("203.0.113.9"));
|
||||
assert_eq!(
|
||||
attempt_error.user_agent.as_deref(),
|
||||
Some("codex-cli/0.144.0")
|
||||
);
|
||||
|
||||
database.close().await;
|
||||
}
|
||||
@@ -1102,7 +1303,7 @@ async fn seed_observability_facts(
|
||||
sqlx::query(
|
||||
"insert into model_requests (
|
||||
id, client_api_key_ref, config_revision, protocol, operation, endpoint,
|
||||
client_transport, requested_model_id,
|
||||
client_transport, requested_model_id, service_tier,
|
||||
provider_kind, provider_account_id,
|
||||
provider_account_ref, upstream_model_id, upstream_transport, attempt_count,
|
||||
provider_account_name_snapshot, provider_account_email_snapshot,
|
||||
@@ -1111,16 +1312,20 @@ 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,
|
||||
client_ip, user_agent, reasoning_effort, reasoning_preset,
|
||||
request_kind, subagent_kind, compact,
|
||||
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',
|
||||
'http_sse', 'public-model', 'priority', 'openai', 'acct_observe',
|
||||
'acct_observe', 'upstream-model',
|
||||
'http_sse', 2,
|
||||
'primary', 'account@example.invalid', 'oauth',
|
||||
'sent', 'failed', 502, 429, 'rate_limited', 'rate_limit',
|
||||
'upstream limited', 1000, 0, 0, true, false, 'unavailable', 700,
|
||||
'203.0.113.9', 'codex-cli/0.144.0', 'medium', null,
|
||||
'root', null, false,
|
||||
$1 - interval '10 minutes', $1 + interval '20 minutes', $1 - interval '9 minutes',
|
||||
'all', '{}'::text[], '[]'::jsonb
|
||||
)",
|
||||
|
||||
@@ -1,11 +1,15 @@
|
||||
use std::time::{Duration, SystemTime};
|
||||
use std::{
|
||||
collections::BTreeMap,
|
||||
num::NonZeroU32,
|
||||
time::{Duration, SystemTime},
|
||||
};
|
||||
|
||||
use chrono::{TimeDelta, Utc};
|
||||
use gateway_admin::{
|
||||
model::{
|
||||
MutationActor, MutationContext, PageSize,
|
||||
accounts::{
|
||||
AccountListQuery, AccountSort, AccountSortField, AccountStatus,
|
||||
AccountListQuery, AccountRuntimeSnapshot, AccountSort, AccountSortField, AccountStatus,
|
||||
AccountUsageWindowQuery, BatchUpdateAccounts, DeleteAccounts, SortDirection,
|
||||
UpdateAccount,
|
||||
},
|
||||
@@ -23,8 +27,8 @@ use gateway_core::engine::credential::{
|
||||
AccountErrorReason, AccountStateChange, CredentialCasOutcome, CredentialCasUpdate,
|
||||
CredentialRevision, CredentialState, OpaqueProviderData, PlaintextCredential,
|
||||
ProviderAccountId, ProviderAccountIdentity, ProviderAccountStore, ProviderAccountUpdate,
|
||||
QuotaAccessChange, QuotaAccessState, QuotaEvidence, QuotaObservation, QuotaObservationTouch,
|
||||
QuotaState, QuotaWriteOutcome,
|
||||
ProviderRefreshQuery, QuotaAccessChange, QuotaAccessState, QuotaEvidence, QuotaObservation,
|
||||
QuotaObservationTouch, QuotaState, QuotaWriteOutcome,
|
||||
};
|
||||
use gateway_core::routing::{AccountGroupId, ProviderKind};
|
||||
use gateway_store::{
|
||||
@@ -39,7 +43,7 @@ use gateway_store::{
|
||||
};
|
||||
use serde_json::json;
|
||||
|
||||
use super::TestDatabase;
|
||||
use super::{TestDatabase, admin_account_store};
|
||||
|
||||
#[derive(sqlx::FromRow)]
|
||||
struct RecoveredAccountRow {
|
||||
@@ -84,6 +88,71 @@ fn postgres_provider_account_adapter_implements_core_port() {
|
||||
assert_terminal_admin_port::<PgAdminAccountStore>();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn refresh_candidates_should_filter_order_and_bound_in_one_query() {
|
||||
let Some(database) = TestDatabase::create("provider_refresh_candidates").await else {
|
||||
return;
|
||||
};
|
||||
let repository = PgProviderAccountRepository::new(database.pool.clone());
|
||||
let now = Utc::now();
|
||||
let mut forced = account("acct_refresh_forced", "user-forced");
|
||||
forced.has_refresh_token = true;
|
||||
forced.access_token_expires_at = Some(now - TimeDelta::hours(3));
|
||||
forced.next_refresh_at = Some(now + TimeDelta::hours(1));
|
||||
let mut due = account("acct_refresh_due", "user-due");
|
||||
due.has_refresh_token = true;
|
||||
due.access_token_expires_at = Some(now);
|
||||
let mut retry_later = account("acct_refresh_retry_later", "user-retry-later");
|
||||
retry_later.has_refresh_token = true;
|
||||
retry_later.access_token_expires_at = Some(now);
|
||||
retry_later.next_refresh_at = Some(now + TimeDelta::hours(1));
|
||||
let mut disabled = account("acct_refresh_disabled", "user-disabled");
|
||||
disabled.has_refresh_token = true;
|
||||
disabled.access_token_expires_at = Some(now);
|
||||
disabled.enabled = false;
|
||||
let mut other_provider = account("acct_refresh_other", "user-other");
|
||||
other_provider.provider_kind = "xai".to_owned();
|
||||
other_provider.has_refresh_token = true;
|
||||
other_provider.access_token_expires_at = Some(now);
|
||||
for candidate in [forced, due, retry_later, disabled, other_provider] {
|
||||
repository
|
||||
.insert_provider_account(candidate)
|
||||
.await
|
||||
.expect("insert refresh candidate fixture");
|
||||
}
|
||||
|
||||
let observed_at = SystemTime::from(now);
|
||||
let query = ProviderRefreshQuery::new(
|
||||
ProviderKind::new("openai").expect("Provider kind"),
|
||||
observed_at
|
||||
.checked_add(Duration::from_secs(5 * 60))
|
||||
.expect("refresh boundary"),
|
||||
observed_at
|
||||
.checked_sub(Duration::from_secs(2 * 60 * 60))
|
||||
.expect("force boundary"),
|
||||
observed_at,
|
||||
Vec::new(),
|
||||
NonZeroU32::new(2).expect("positive limit"),
|
||||
);
|
||||
let candidates = repository
|
||||
.list_refresh_candidates(query)
|
||||
.await
|
||||
.expect("load bounded refresh candidates");
|
||||
|
||||
assert_eq!(
|
||||
candidates
|
||||
.iter()
|
||||
.map(|candidate| candidate.account.id().as_str())
|
||||
.collect::<Vec<_>>(),
|
||||
vec!["acct_refresh_forced", "acct_refresh_due"]
|
||||
);
|
||||
assert!(candidates.iter().all(|candidate| {
|
||||
candidate.credential.expose_to_provider()["access_token"] == "initial-secret"
|
||||
}));
|
||||
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn provider_account_should_allow_missing_upstream_user_id() {
|
||||
let Some(database) = TestDatabase::create("provider_account_missing_upstream_identity").await
|
||||
@@ -552,20 +621,23 @@ async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_u
|
||||
.await
|
||||
.expect("seed charlie usage");
|
||||
|
||||
let store = PgAdminAccountStore::new(database.pool.clone(), None);
|
||||
let store = admin_account_store(&database.pool);
|
||||
let usage_page = store
|
||||
.list_accounts(AccountListQuery {
|
||||
page: 1,
|
||||
page_size: PageSize::new(2).expect("page size"),
|
||||
provider_kind: None,
|
||||
group_filter: None,
|
||||
search: None,
|
||||
status: None,
|
||||
sort: Some(AccountSort {
|
||||
field: AccountSortField::Usage,
|
||||
direction: SortDirection::Desc,
|
||||
}),
|
||||
})
|
||||
.list_accounts(
|
||||
AccountListQuery {
|
||||
page: 1,
|
||||
page_size: PageSize::new(2).expect("page size"),
|
||||
provider_kind: None,
|
||||
group_filter: None,
|
||||
search: None,
|
||||
status: None,
|
||||
sort: Some(AccountSort {
|
||||
field: AccountSortField::Usage,
|
||||
direction: SortDirection::Desc,
|
||||
}),
|
||||
},
|
||||
Default::default(),
|
||||
)
|
||||
.await
|
||||
.expect("sort accounts by retained usage");
|
||||
assert_eq!(usage_page.config_revision.get(), 1);
|
||||
@@ -594,18 +666,21 @@ async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_u
|
||||
);
|
||||
|
||||
let last_used_page = store
|
||||
.list_accounts(AccountListQuery {
|
||||
page: 1,
|
||||
page_size: PageSize::new(2).expect("page size"),
|
||||
provider_kind: None,
|
||||
group_filter: None,
|
||||
search: None,
|
||||
status: None,
|
||||
sort: Some(AccountSort {
|
||||
field: AccountSortField::LastUsedAt,
|
||||
direction: SortDirection::Desc,
|
||||
}),
|
||||
})
|
||||
.list_accounts(
|
||||
AccountListQuery {
|
||||
page: 1,
|
||||
page_size: PageSize::new(2).expect("page size"),
|
||||
provider_kind: None,
|
||||
group_filter: None,
|
||||
search: None,
|
||||
status: None,
|
||||
sort: Some(AccountSort {
|
||||
field: AccountSortField::LastUsedAt,
|
||||
direction: SortDirection::Desc,
|
||||
}),
|
||||
},
|
||||
Default::default(),
|
||||
)
|
||||
.await
|
||||
.expect("sort accounts by retained last use");
|
||||
assert_eq!(
|
||||
@@ -618,15 +693,18 @@ async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_u
|
||||
);
|
||||
|
||||
let filtered = store
|
||||
.list_accounts(AccountListQuery {
|
||||
page: 1,
|
||||
page_size: PageSize::new(10).expect("page size"),
|
||||
provider_kind: Some(ProviderKind::new("openai").expect("Provider kind")),
|
||||
group_filter: None,
|
||||
search: Some("ALPHA@EXAMPLE".to_owned()),
|
||||
status: Some(AccountStatus::Normal),
|
||||
sort: None,
|
||||
})
|
||||
.list_accounts(
|
||||
AccountListQuery {
|
||||
page: 1,
|
||||
page_size: PageSize::new(10).expect("page size"),
|
||||
provider_kind: Some(ProviderKind::new("openai").expect("Provider kind")),
|
||||
group_filter: None,
|
||||
search: Some("ALPHA@EXAMPLE".to_owned()),
|
||||
status: Some(AccountStatus::Normal),
|
||||
sort: None,
|
||||
},
|
||||
Default::default(),
|
||||
)
|
||||
.await
|
||||
.expect("filter account directory");
|
||||
assert_eq!(filtered.total, 1);
|
||||
@@ -634,16 +712,62 @@ async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_u
|
||||
assert_eq!(filtered.items[0].account.id, "acct_alpha");
|
||||
assert_eq!(filtered.items[0].account.provider_kind.as_str(), "openai");
|
||||
|
||||
let rate_limited = store
|
||||
.list_accounts(
|
||||
AccountListQuery {
|
||||
page: 1,
|
||||
page_size: PageSize::new(10).expect("page size"),
|
||||
provider_kind: None,
|
||||
group_filter: None,
|
||||
search: None,
|
||||
status: Some(AccountStatus::RateLimited),
|
||||
sort: None,
|
||||
},
|
||||
AccountRuntimeSnapshot {
|
||||
rate_limited_until: BTreeMap::from([(
|
||||
"acct_alpha".to_owned(),
|
||||
now + TimeDelta::minutes(5),
|
||||
)]),
|
||||
in_flight: None,
|
||||
},
|
||||
)
|
||||
.await
|
||||
.expect("filter active runtime status in PostgreSQL page query");
|
||||
assert_eq!(rate_limited.total, 1);
|
||||
assert_eq!(rate_limited.items[0].account.id, "acct_alpha");
|
||||
assert_eq!(rate_limited.summary.rate_limited, 1);
|
||||
assert_eq!(rate_limited.summary.normal, 0);
|
||||
|
||||
let no_contains_compatibility = store
|
||||
.list_accounts(
|
||||
AccountListQuery {
|
||||
page: 1,
|
||||
page_size: PageSize::new(10).expect("page size"),
|
||||
provider_kind: None,
|
||||
group_filter: None,
|
||||
search: Some("example.invalid".to_owned()),
|
||||
status: None,
|
||||
sort: None,
|
||||
},
|
||||
Default::default(),
|
||||
)
|
||||
.await
|
||||
.expect("use structured account prefix search");
|
||||
assert_eq!(no_contains_compatibility.total, 0);
|
||||
|
||||
let error_accounts = store
|
||||
.list_accounts(AccountListQuery {
|
||||
page: 1,
|
||||
page_size: PageSize::new(10).expect("page size"),
|
||||
provider_kind: None,
|
||||
group_filter: None,
|
||||
search: None,
|
||||
status: Some(AccountStatus::Error),
|
||||
sort: None,
|
||||
})
|
||||
.list_accounts(
|
||||
AccountListQuery {
|
||||
page: 1,
|
||||
page_size: PageSize::new(10).expect("page size"),
|
||||
provider_kind: None,
|
||||
group_filter: None,
|
||||
search: None,
|
||||
status: Some(AccountStatus::Error),
|
||||
sort: None,
|
||||
},
|
||||
Default::default(),
|
||||
)
|
||||
.await
|
||||
.expect("filter error accounts");
|
||||
assert_eq!(error_accounts.items.len(), 2);
|
||||
@@ -676,7 +800,7 @@ async fn terminal_credential_list_preserves_grouped_filters_and_unpaged_collecti
|
||||
.expect("insert xAI credential fixture");
|
||||
}
|
||||
|
||||
let store = PgAdminAccountStore::new(database.pool.clone(), None);
|
||||
let store = admin_account_store(&database.pool);
|
||||
let provider = ProviderKind::new("xai").expect("xAI Provider kind");
|
||||
let complete = store
|
||||
.list_credentials(
|
||||
@@ -757,7 +881,7 @@ async fn terminal_admin_usage_chunks_large_selections_and_preserves_exact_costs(
|
||||
let mut account_ids = vec!["acct_usage_exact".to_owned()];
|
||||
account_ids.extend((0..200).map(|index| format!("missing_account_{index}")));
|
||||
|
||||
let usage = PgAdminAccountStore::new(database.pool.clone(), None)
|
||||
let usage = admin_account_store(&database.pool)
|
||||
.load_account_usage(
|
||||
TimeRange {
|
||||
start: now - TimeDelta::hours(1),
|
||||
@@ -849,7 +973,7 @@ async fn terminal_admin_quota_window_usage_prefers_provider_total_and_falls_back
|
||||
.await
|
||||
.expect("separate provider total from fallback components");
|
||||
|
||||
let mut usage = PgAdminAccountStore::new(database.pool.clone(), None)
|
||||
let mut usage = admin_account_store(&database.pool)
|
||||
.load_account_usage_by_windows(&[
|
||||
AccountUsageWindowQuery {
|
||||
account_id: "acct_quota_window".to_owned(),
|
||||
@@ -904,6 +1028,84 @@ async fn terminal_admin_quota_window_usage_prefers_provider_total_and_falls_back
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn terminal_admin_quota_window_usage_should_include_only_statusless_websocket() {
|
||||
let Some(database) = TestDatabase::create("provider_account_quota_window_websocket").await
|
||||
else {
|
||||
return;
|
||||
};
|
||||
let repository = PgProviderAccountRepository::new(database.pool.clone());
|
||||
repository
|
||||
.insert_provider_account(account(
|
||||
"acct_quota_window_websocket",
|
||||
"user-quota-window-websocket",
|
||||
))
|
||||
.await
|
||||
.expect("insert WebSocket quota window account");
|
||||
let now = Utc::now();
|
||||
for seed in [
|
||||
ModelRequestSeed {
|
||||
request_id: "req_quota_window_websocket",
|
||||
account_id: "acct_quota_window_websocket",
|
||||
provider_kind: "openai",
|
||||
model: "gpt-window-websocket",
|
||||
total_tokens: 13,
|
||||
cost_amount: "1.50",
|
||||
started_at: now - TimeDelta::minutes(2),
|
||||
},
|
||||
ModelRequestSeed {
|
||||
request_id: "req_quota_window_statusless_http",
|
||||
account_id: "acct_quota_window_websocket",
|
||||
provider_kind: "openai",
|
||||
model: "gpt-window-statusless-http",
|
||||
total_tokens: 29,
|
||||
cost_amount: "2.50",
|
||||
started_at: now - TimeDelta::minutes(1),
|
||||
},
|
||||
] {
|
||||
seed_model_request(&database.pool, seed)
|
||||
.await
|
||||
.expect("seed statusless quota window request");
|
||||
}
|
||||
sqlx::query(
|
||||
"update model_requests
|
||||
set client_transport = case id
|
||||
when 'req_quota_window_websocket' then 'websocket'
|
||||
else 'http_sse'
|
||||
end,
|
||||
client_status_code = null
|
||||
where id in ('req_quota_window_websocket', 'req_quota_window_statusless_http')",
|
||||
)
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.expect("make quota window requests statusless");
|
||||
|
||||
let usage = admin_account_store(&database.pool)
|
||||
.load_account_usage_by_windows(&[AccountUsageWindowQuery {
|
||||
account_id: "acct_quota_window_websocket".to_owned(),
|
||||
key: "statusless".to_owned(),
|
||||
range: TimeRange {
|
||||
start: now - TimeDelta::hours(1),
|
||||
end: now + TimeDelta::minutes(1),
|
||||
},
|
||||
}])
|
||||
.await
|
||||
.expect("load statusless quota window usage");
|
||||
|
||||
assert_eq!(
|
||||
(
|
||||
usage[0].usage.request_count,
|
||||
usage[0].usage.success_count,
|
||||
usage[0].usage.total_tokens,
|
||||
usage[0].usage.costs[0].amount.as_str(),
|
||||
usage[0].usage.models[0].model.as_str(),
|
||||
),
|
||||
(1, 1, Some(13), "1.5", "gpt-window-websocket"),
|
||||
);
|
||||
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn terminal_admin_mutations_keep_revision_account_and_audit_atomic() {
|
||||
let Some(database) = TestDatabase::create("provider_account_terminal_mutation").await else {
|
||||
@@ -913,7 +1115,7 @@ async fn terminal_admin_mutations_keep_revision_account_and_audit_atomic() {
|
||||
.insert_provider_account(account("acct_terminal_mutation", "user-terminal-mutation"))
|
||||
.await
|
||||
.expect("insert mutation account");
|
||||
let store = PgAdminAccountStore::new(database.pool.clone(), None);
|
||||
let store = admin_account_store(&database.pool);
|
||||
let context = MutationContext {
|
||||
actor: MutationActor::System,
|
||||
request_id: "request_terminal_mutation".to_owned(),
|
||||
@@ -1056,7 +1258,7 @@ async fn account_recovery_resets_status_facts_without_changing_credentials_or_sc
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.expect("load account before recovery");
|
||||
let store = PgAdminAccountStore::new(database.pool.clone(), None);
|
||||
let store = admin_account_store(&database.pool);
|
||||
|
||||
let result = store
|
||||
.recover_account(
|
||||
@@ -1130,7 +1332,7 @@ async fn terminal_batch_update_replaces_state_and_groups_once_or_rolls_back_ever
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.expect("insert batch account group");
|
||||
let store = PgAdminAccountStore::new(database.pool.clone(), None);
|
||||
let store = admin_account_store(&database.pool);
|
||||
let account_ids = vec!["acct_batch_a".to_owned(), "acct_batch_b".to_owned()];
|
||||
let context = MutationContext {
|
||||
actor: MutationActor::System,
|
||||
@@ -1234,7 +1436,7 @@ async fn terminal_admin_delete_removes_enabled_accounts_in_one_transaction() {
|
||||
.expect("insert enabled account");
|
||||
}
|
||||
|
||||
let revision = PgAdminAccountStore::new(database.pool.clone(), None)
|
||||
let revision = admin_account_store(&database.pool)
|
||||
.delete_accounts(
|
||||
DeleteAccounts {
|
||||
account_ids: vec![
|
||||
@@ -1360,7 +1562,7 @@ async fn authorization_create_returns_existing_account_id_when_identity_is_upser
|
||||
unreachable!("credential fixture must be an object");
|
||||
};
|
||||
|
||||
let result = PgAdminAccountStore::new(database.pool.clone(), None)
|
||||
let result = admin_account_store(&database.pool)
|
||||
.commit_authorization(
|
||||
AuthorizationCommit {
|
||||
pending: PendingAuthorizationMutation::new(
|
||||
|
||||
@@ -0,0 +1,108 @@
|
||||
use std::time::Duration;
|
||||
|
||||
use gateway_store::{StoreError, postgres::ObservabilityQueryBudget};
|
||||
use tokio::sync::{mpsc, oneshot};
|
||||
|
||||
#[test]
|
||||
fn query_budget_rejects_an_empty_budget() {
|
||||
assert!(ObservabilityQueryBudget::try_new(0, Duration::from_secs(1)).is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn query_budget_times_out_without_exceeding_its_connection_limit() {
|
||||
let budget = ObservabilityQueryBudget::try_new(4, Duration::from_millis(10))
|
||||
.expect("valid query budget");
|
||||
let (started_tx, mut started_rx) = mpsc::unbounded_channel();
|
||||
let mut releases = Vec::new();
|
||||
let mut occupied = Vec::new();
|
||||
for _ in 0..4 {
|
||||
let occupied_budget = budget.clone();
|
||||
let started_tx = started_tx.clone();
|
||||
let (release_tx, release_rx) = oneshot::channel();
|
||||
releases.push(release_tx);
|
||||
occupied.push(tokio::spawn(async move {
|
||||
occupied_budget
|
||||
.run("occupied test query", async move {
|
||||
started_tx.send(()).expect("report occupied slot");
|
||||
release_rx.await.expect("release occupied slot");
|
||||
Ok::<_, StoreError>(())
|
||||
})
|
||||
.await
|
||||
}));
|
||||
}
|
||||
for _ in 0..4 {
|
||||
started_rx.recv().await.expect("occupied slot started");
|
||||
}
|
||||
|
||||
let error = budget
|
||||
.run("queued test query", async { Ok::<_, StoreError>(()) })
|
||||
.await
|
||||
.expect_err("exhausted query budget must time out");
|
||||
|
||||
assert!(matches!(
|
||||
error,
|
||||
StoreError::Unavailable { message, .. }
|
||||
if message == "observability PostgreSQL connection budget is exhausted"
|
||||
));
|
||||
for release in releases {
|
||||
release.send(()).expect("release occupied slot");
|
||||
}
|
||||
for task in occupied {
|
||||
task.await
|
||||
.expect("occupied task must not panic")
|
||||
.expect("occupied query must succeed");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn queued_query_runs_after_one_slot_is_released() {
|
||||
let budget =
|
||||
ObservabilityQueryBudget::try_new(4, Duration::from_secs(1)).expect("valid query budget");
|
||||
let (started_tx, mut started_rx) = mpsc::unbounded_channel();
|
||||
let mut releases = Vec::new();
|
||||
let mut occupied = Vec::new();
|
||||
for _ in 0..4 {
|
||||
let occupied_budget = budget.clone();
|
||||
let started_tx = started_tx.clone();
|
||||
let (release_tx, release_rx) = oneshot::channel();
|
||||
releases.push(release_tx);
|
||||
occupied.push(tokio::spawn(async move {
|
||||
occupied_budget
|
||||
.run("occupied test query", async move {
|
||||
started_tx.send(()).expect("report occupied slot");
|
||||
release_rx.await.expect("release occupied slot");
|
||||
Ok::<_, StoreError>(())
|
||||
})
|
||||
.await
|
||||
}));
|
||||
}
|
||||
for _ in 0..4 {
|
||||
started_rx.recv().await.expect("occupied slot started");
|
||||
}
|
||||
let queued_budget = budget.clone();
|
||||
let queued = tokio::spawn(async move {
|
||||
queued_budget
|
||||
.run("queued test query", async { Ok::<_, StoreError>(()) })
|
||||
.await
|
||||
});
|
||||
|
||||
releases
|
||||
.pop()
|
||||
.expect("one occupied slot")
|
||||
.send(())
|
||||
.expect("release one occupied slot");
|
||||
tokio::time::timeout(Duration::from_millis(100), queued)
|
||||
.await
|
||||
.expect("one released slot must admit one queued query")
|
||||
.expect("queued task must not panic")
|
||||
.expect("queued query must succeed");
|
||||
|
||||
for release in releases {
|
||||
release.send(()).expect("release occupied slot");
|
||||
}
|
||||
for task in occupied {
|
||||
task.await
|
||||
.expect("occupied task must not panic")
|
||||
.expect("occupied query must succeed");
|
||||
}
|
||||
}
|
||||
@@ -1,4 +1,11 @@
|
||||
use gateway_store::postgres::RuntimeRetentionSettings;
|
||||
use std::time::Duration;
|
||||
|
||||
use chrono::{Duration as ChronoDuration, Utc};
|
||||
use gateway_store::postgres::{
|
||||
PgRetentionRepository, RetentionCycleBudget, RetentionRepository as _, RuntimeRetentionSettings,
|
||||
};
|
||||
|
||||
use super::TestDatabase;
|
||||
|
||||
#[test]
|
||||
fn retention_settings_preserve_independent_windows() {
|
||||
@@ -9,3 +16,52 @@ fn retention_settings_preserve_independent_windows() {
|
||||
};
|
||||
assert_eq!(settings.audit_retention_days, 90);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn retention_cycle_should_stop_at_the_batch_budget() {
|
||||
let Some(database) = TestDatabase::create("retention_cycle_budget").await else {
|
||||
return;
|
||||
};
|
||||
let expired_at = Utc::now() - ChronoDuration::days(100);
|
||||
sqlx::query(
|
||||
"insert into admin_audit_events (
|
||||
id, actor_kind, actor_ref, action, entity_kind, entity_ref,
|
||||
changed_fields, created_at
|
||||
)
|
||||
select 'retention-' || value::text, 'system', 'retention', 'cleanup',
|
||||
'fixture', 'fixture-' || value::text, array[]::text[], $1
|
||||
from generate_series(1, 5) as value",
|
||||
)
|
||||
.bind(expired_at)
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.expect("seed expired audit events");
|
||||
|
||||
let budget = RetentionCycleBudget::try_new(2, 3, Duration::from_secs(1), Duration::ZERO)
|
||||
.expect("retention cycle budget");
|
||||
let repository = PgRetentionRepository::with_cycle_budget(database.pool.clone(), budget);
|
||||
let report = repository
|
||||
.apply_retention(
|
||||
Utc::now(),
|
||||
RuntimeRetentionSettings {
|
||||
usage_retention_days: 31,
|
||||
ops_event_retention_days: 30,
|
||||
audit_retention_days: 90,
|
||||
},
|
||||
)
|
||||
.await
|
||||
.expect("bounded retention cycle");
|
||||
|
||||
assert_eq!(report.model_requests, 0);
|
||||
assert_eq!(report.ops_events, 0);
|
||||
assert_eq!(report.admin_audit_events, 2);
|
||||
assert_eq!(report.batches, 3);
|
||||
assert!(report.budget_exhausted);
|
||||
let remaining: i64 =
|
||||
sqlx::query_scalar("select count(*) from admin_audit_events where id like 'retention-%'")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.expect("count remaining audit events");
|
||||
assert_eq!(remaining, 3);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
@@ -1,13 +1,131 @@
|
||||
use chrono::{DateTime, TimeDelta, Utc};
|
||||
use gateway_store::postgres::{
|
||||
DiagnosticDimension, ModelRequestAttemptStart, ModelRequestRepository, NewModelRequest,
|
||||
ObservabilityPageNumber, ObservabilityPageSize, ObservabilityRange, ObservabilityRepository,
|
||||
OpsErrorFilter, OpsErrorQuery, OpsEvent, OpsEventLevel, OpsEventRepository, PgExecutionStore,
|
||||
PgObservabilityRepository, PgOpsEventRepository, UsageRecordFilter, UsageRecordQuery,
|
||||
ObservabilityPageSize, ObservabilityRange, ObservabilityRepository, OpsErrorFilter,
|
||||
OpsErrorQuery, OpsEvent, OpsEventLevel, OpsEventRepository, PgExecutionStore,
|
||||
PgOpsEventRepository, ProviderAccountUsageQuery, UsageRecordFilter, UsageRecordQuery,
|
||||
};
|
||||
use sqlx::PgPool;
|
||||
|
||||
use super::TestDatabase;
|
||||
use super::{TestDatabase, observability_repository};
|
||||
|
||||
#[tokio::test]
|
||||
async fn completed_usage_projections_should_accept_statusless_websocket_but_reject_statusless_http()
|
||||
{
|
||||
let Some(database) = TestDatabase::create("snapshot_statusless_websocket_usage").await else {
|
||||
return;
|
||||
};
|
||||
let started_at = Utc::now();
|
||||
seed_account(
|
||||
&database.pool,
|
||||
"acct_statusless_transport",
|
||||
"Statusless transport",
|
||||
None,
|
||||
"oauth",
|
||||
started_at,
|
||||
)
|
||||
.await;
|
||||
let store = PgExecutionStore::new(database.pool.clone());
|
||||
|
||||
let mut websocket_request = new_request("req_statusless_websocket", started_at);
|
||||
websocket_request.client_transport = "websocket".to_owned();
|
||||
store
|
||||
.insert_model_request_with_first_attempt(
|
||||
websocket_request,
|
||||
attempt("req_statusless_websocket", 1, "acct_statusless_transport"),
|
||||
)
|
||||
.await
|
||||
.expect("insert statusless WebSocket request");
|
||||
store
|
||||
.insert_model_request_with_first_attempt(
|
||||
new_request("req_statusless_http", started_at),
|
||||
attempt("req_statusless_http", 1, "acct_statusless_transport"),
|
||||
)
|
||||
.await
|
||||
.expect("insert statusless HTTP request");
|
||||
finalize_request_without_client_status(&database.pool, "req_statusless_websocket", started_at)
|
||||
.await;
|
||||
finalize_request_without_client_status(&database.pool, "req_statusless_http", started_at).await;
|
||||
|
||||
let repository = observability_repository(&database.pool);
|
||||
let page = repository
|
||||
.list_usage_records(usage_query(started_at, UsageRecordFilter::default()))
|
||||
.await
|
||||
.expect("list statusless transport usage records");
|
||||
let request_ids = page
|
||||
.items
|
||||
.iter()
|
||||
.map(|record| record.id.as_str())
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
assert_eq!(
|
||||
(page.total, request_ids),
|
||||
(1, vec!["req_statusless_websocket"])
|
||||
);
|
||||
|
||||
let overview = repository
|
||||
.usage_summary(range_around(started_at), UsageRecordFilter::default())
|
||||
.await
|
||||
.expect("summarize statusless transport usage");
|
||||
assert_eq!(
|
||||
(
|
||||
overview.requests.request_count,
|
||||
overview.requests.success_count,
|
||||
overview.requests.total_tokens,
|
||||
overview.attempts.cost_coverage.provider_reported_count,
|
||||
),
|
||||
(2, 2, 10, 1),
|
||||
);
|
||||
|
||||
let account_usage = repository
|
||||
.provider_account_usage(
|
||||
ProviderAccountUsageQuery::for_accounts(
|
||||
range_around(started_at),
|
||||
vec!["acct_statusless_transport".to_owned()],
|
||||
)
|
||||
.expect("statusless account usage query")
|
||||
.with_hourly_request_buckets()
|
||||
.expect("statusless account request timeline"),
|
||||
)
|
||||
.await
|
||||
.expect("load statusless account usage");
|
||||
assert_eq!(
|
||||
(
|
||||
account_usage[0].request_count,
|
||||
account_usage[0].total_tokens,
|
||||
account_usage[0].models[0].request_count,
|
||||
account_usage[0]
|
||||
.request_buckets
|
||||
.iter()
|
||||
.map(|bucket| bucket.request_count)
|
||||
.sum::<u64>(),
|
||||
account_usage[0].costs[0].amount.as_str(),
|
||||
),
|
||||
(1, Some(10), 1, 1, "1.25"),
|
||||
);
|
||||
|
||||
let diagnostics = repository
|
||||
.usage_diagnostics(
|
||||
range_around(started_at),
|
||||
UsageRecordFilter::default(),
|
||||
DiagnosticDimension::Account,
|
||||
)
|
||||
.await
|
||||
.expect("load statusless usage diagnostics");
|
||||
assert_eq!(
|
||||
(
|
||||
diagnostics[0].request_count,
|
||||
diagnostics[0].success_count,
|
||||
diagnostics[0].total_tokens,
|
||||
diagnostics[0].average_latency_ms,
|
||||
diagnostics[0].cost_coverage.provider_reported_count,
|
||||
diagnostics[0].costs[0].amount.as_str(),
|
||||
),
|
||||
(2, 2, 10, Some(500), 1, "1.25"),
|
||||
);
|
||||
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn request_snapshots_should_survive_account_deletion() {
|
||||
@@ -39,7 +157,7 @@ async fn request_snapshots_should_survive_account_deletion() {
|
||||
.await
|
||||
.expect("delete provider account");
|
||||
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone(), None);
|
||||
let repository = observability_repository(&database.pool);
|
||||
let page = repository
|
||||
.list_usage_records(usage_query(started_at, UsageRecordFilter::default()))
|
||||
.await
|
||||
@@ -154,7 +272,7 @@ async fn attempts_should_keep_their_own_account_snapshots() {
|
||||
.await
|
||||
.expect("delete provider accounts");
|
||||
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone(), None);
|
||||
let repository = observability_repository(&database.pool);
|
||||
let detail = repository
|
||||
.usage_record_detail("req_snap_ab")
|
||||
.await
|
||||
@@ -276,13 +394,12 @@ async fn ops_errors_should_keep_request_and_event_snapshots_after_account_deleti
|
||||
.await
|
||||
.expect("delete provider accounts");
|
||||
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone(), None);
|
||||
let repository = observability_repository(&database.pool);
|
||||
let errors = repository
|
||||
.list_ops_errors(OpsErrorQuery {
|
||||
range: range_around(started_at),
|
||||
filter: OpsErrorFilter::default(),
|
||||
cursor: None,
|
||||
page: ObservabilityPageNumber::new(1).expect("page"),
|
||||
current_page: 1,
|
||||
page_size: ObservabilityPageSize::new(10).expect("page size"),
|
||||
})
|
||||
.await
|
||||
@@ -376,7 +493,7 @@ async fn diagnostics_should_group_same_email_accounts_by_stable_ref() {
|
||||
.await
|
||||
.expect("delete provider accounts");
|
||||
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone(), None);
|
||||
let repository = observability_repository(&database.pool);
|
||||
let diagnostics = repository
|
||||
.usage_diagnostics(
|
||||
range_around(started_at),
|
||||
@@ -448,7 +565,7 @@ async fn diagnostics_should_fallback_to_name_then_ref_for_missing_snapshots() {
|
||||
.expect("insert legacy request without live account");
|
||||
finalize_request(&database.pool, "req_snap_legacy", started_at).await;
|
||||
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone(), None);
|
||||
let repository = observability_repository(&database.pool);
|
||||
let diagnostics = repository
|
||||
.usage_diagnostics(
|
||||
range_around(started_at),
|
||||
@@ -470,6 +587,66 @@ async fn diagnostics_should_fallback_to_name_then_ref_for_missing_snapshots() {
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn diagnostics_should_prefer_the_latest_non_null_email_snapshot() {
|
||||
let Some(database) = TestDatabase::create("snapshot_diagnostics_latest_email").await else {
|
||||
return;
|
||||
};
|
||||
let started_at = Utc::now();
|
||||
seed_account(
|
||||
&database.pool,
|
||||
"acct_snapshot_history",
|
||||
"Older Name",
|
||||
Some("history@example.invalid"),
|
||||
"oauth",
|
||||
started_at,
|
||||
)
|
||||
.await;
|
||||
let store = PgExecutionStore::new(database.pool.clone());
|
||||
store
|
||||
.insert_model_request_with_first_attempt(
|
||||
new_request("req_snapshot_history_old", started_at),
|
||||
attempt("req_snapshot_history_old", 1, "acct_snapshot_history"),
|
||||
)
|
||||
.await
|
||||
.expect("insert older account snapshot");
|
||||
finalize_request(&database.pool, "req_snapshot_history_old", started_at).await;
|
||||
sqlx::query(
|
||||
"update provider_accounts
|
||||
set name = 'Newer Name', email = null, updated_at = now()
|
||||
where id = 'acct_snapshot_history'",
|
||||
)
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.expect("update account snapshot source");
|
||||
store
|
||||
.insert_model_request_with_first_attempt(
|
||||
new_request(
|
||||
"req_snapshot_history_new",
|
||||
started_at + chrono::Duration::seconds(1),
|
||||
),
|
||||
attempt("req_snapshot_history_new", 1, "acct_snapshot_history"),
|
||||
)
|
||||
.await
|
||||
.expect("insert newer account snapshot");
|
||||
finalize_request(&database.pool, "req_snapshot_history_new", started_at).await;
|
||||
|
||||
let diagnostics = observability_repository(&database.pool)
|
||||
.usage_diagnostics(
|
||||
range_around(started_at),
|
||||
UsageRecordFilter::default(),
|
||||
DiagnosticDimension::Account,
|
||||
)
|
||||
.await
|
||||
.expect("account diagnostics");
|
||||
|
||||
assert_eq!(diagnostics.len(), 1);
|
||||
assert_eq!(diagnostics[0].key, "acct_snapshot_history");
|
||||
assert_eq!(diagnostics[0].name, "history@example.invalid");
|
||||
assert_eq!(diagnostics[0].request_count, 2);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn failure_diagnostics_should_only_include_errored_requests() {
|
||||
let Some(database) = TestDatabase::create("snapshot_failure_dimension").await else {
|
||||
@@ -514,7 +691,7 @@ async fn failure_diagnostics_should_only_include_errored_requests() {
|
||||
.await
|
||||
.expect("mark request failed");
|
||||
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone(), None);
|
||||
let repository = observability_repository(&database.pool);
|
||||
let diagnostics = repository
|
||||
.usage_diagnostics(
|
||||
range_around(started_at),
|
||||
@@ -572,7 +749,7 @@ async fn api_key_diagnostics_should_display_key_name_and_fallback_to_ref() {
|
||||
.await
|
||||
.expect("set request b latency");
|
||||
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone(), None);
|
||||
let repository = observability_repository(&database.pool);
|
||||
let diagnostics = repository
|
||||
.usage_diagnostics(
|
||||
range_around(started_at),
|
||||
@@ -639,7 +816,7 @@ async fn renaming_account_should_not_rewrite_historical_snapshots() {
|
||||
.await
|
||||
.expect("rename provider account");
|
||||
|
||||
let repository = PgObservabilityRepository::new(database.pool.clone(), None);
|
||||
let repository = observability_repository(&database.pool);
|
||||
let page = repository
|
||||
.list_usage_records(usage_query(started_at, UsageRecordFilter::default()))
|
||||
.await
|
||||
@@ -708,8 +885,7 @@ fn usage_query(started_at: DateTime<Utc>, filter: UsageRecordFilter) -> UsageRec
|
||||
UsageRecordQuery {
|
||||
range: range_around(started_at),
|
||||
filter,
|
||||
cursor: None,
|
||||
page: ObservabilityPageNumber::new(1).expect("page"),
|
||||
current_page: 1,
|
||||
page_size: ObservabilityPageSize::new(10).expect("page size"),
|
||||
}
|
||||
}
|
||||
@@ -767,6 +943,29 @@ async fn finalize_request(pool: &PgPool, id: &str, started_at: DateTime<Utc>) {
|
||||
.expect("finalize request");
|
||||
}
|
||||
|
||||
async fn finalize_request_without_client_status(
|
||||
pool: &PgPool,
|
||||
id: &str,
|
||||
started_at: DateTime<Utc>,
|
||||
) {
|
||||
sqlx::query(
|
||||
"update model_requests
|
||||
set outcome = 'succeeded', upstream_send_state = 'sent',
|
||||
client_status_code = null, upstream_status_code = 200,
|
||||
input_tokens = 7, output_tokens = 3, total_tokens = 10,
|
||||
first_token_ms = 100, latency_ms = 500,
|
||||
cost_source = 'provider_reported', cost_amount = 1.25,
|
||||
cost_currency = 'USD',
|
||||
downstream_committed_at = $2, completed_at = $2
|
||||
where id = $1",
|
||||
)
|
||||
.bind(id)
|
||||
.bind(started_at + chrono::Duration::seconds(5))
|
||||
.execute(pool)
|
||||
.await
|
||||
.expect("finalize request without client status");
|
||||
}
|
||||
|
||||
async fn seed_api_key(pool: &PgPool, id: &str, name: &str, now: DateTime<Utc>) {
|
||||
sqlx::query(
|
||||
"insert into client_api_keys (id, name, key, enabled, created_at, updated_at)
|
||||
|
||||
@@ -23,7 +23,7 @@ fn credential_cooldown_is_revision_fenced() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn credential_cooldown_round_trips_without_raw_account_id_in_key() {
|
||||
async fn credential_cooldown_round_trips_and_indexes_active_account_without_raw_id_in_keys() {
|
||||
let Some((repository, mut connection, namespace)) = repository().await else {
|
||||
return;
|
||||
};
|
||||
@@ -42,9 +42,26 @@ async fn credential_cooldown_round_trips_without_raw_account_id_in_key() {
|
||||
.expect("read cooldown"),
|
||||
Some(cooldown.clone())
|
||||
);
|
||||
let keys = namespace_keys(&mut connection, &namespace).await;
|
||||
assert_eq!(keys.len(), 1);
|
||||
assert!(!keys[0].contains(&cooldown.provider_account_id));
|
||||
let mut keys = namespace_keys(&mut connection, &namespace).await;
|
||||
keys.sort();
|
||||
assert_eq!(keys.len(), 2);
|
||||
assert!(
|
||||
keys.iter()
|
||||
.all(|key| !key.contains(&cooldown.provider_account_id))
|
||||
);
|
||||
|
||||
let index_key = keys
|
||||
.iter()
|
||||
.find(|key| key.ends_with(":account:active-cooldowns"))
|
||||
.expect("active cooldown index key");
|
||||
let indexed_accounts: Vec<String> = redis::cmd("ZRANGE")
|
||||
.arg(index_key)
|
||||
.arg(0)
|
||||
.arg(-1)
|
||||
.query_async(&mut connection)
|
||||
.await
|
||||
.expect("read active cooldown index");
|
||||
assert_eq!(indexed_accounts, [cooldown.provider_account_id]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
@@ -5,6 +5,7 @@ use std::time::SystemTime;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, Utc};
|
||||
use futures::TryStreamExt as _;
|
||||
use gateway_admin::model::accounts::AccountRecord;
|
||||
use gateway_admin::model::observability::{
|
||||
CalculatedBillingBreakdown, CurrencyCost, DashboardDesktopRelease, DashboardWireAttribute,
|
||||
@@ -19,9 +20,10 @@ use gateway_admin::model::provider_credentials::{
|
||||
PreparedCredentialCreate, PreparedCredentialImport, PreparedCredentialRotation,
|
||||
PreparedCredentialRotationFacts, ProviderDocument, ProviderExport,
|
||||
ProviderExportCredentialInput, ProviderModel, ProviderModels, ProviderProfileActivityInsights,
|
||||
ProviderProfileDailyUsage, ProviderProfileInvocation, ProviderProfileStatistics,
|
||||
ProviderProfileStatisticsSummary, ProviderQuota, ProviderQuotaRequest, ProviderQuotaWindow,
|
||||
ProviderQuotaWindowRole, ProviderResetCredit, ProviderResetCreditResult, ProviderResetCredits,
|
||||
ProviderProfileAvatar, ProviderProfileAvatarStreamError, ProviderProfileDailyUsage,
|
||||
ProviderProfileInvocation, ProviderProfileStatistics, ProviderProfileStatisticsSummary,
|
||||
ProviderQuota, ProviderQuotaRequest, ProviderQuotaWindow, ProviderQuotaWindowRole,
|
||||
ProviderResetCredit, ProviderResetCreditResult, ProviderResetCredits,
|
||||
QuotaLocalUsageAttribution,
|
||||
};
|
||||
use gateway_admin::model::{MutationActor, MutationContext, Revision};
|
||||
@@ -49,10 +51,10 @@ use crate::credential::{
|
||||
CodexCredentialAdminService, CodexCredentialCatalogError, CodexCredentialCatalogService,
|
||||
CodexCredentialProfileService, CodexCredentialQuotaError, CodexCredentialQuotaService,
|
||||
CodexOAuthAdmin, CodexOAuthAdminError, CodexOAuthPendingClaimOutcome, CodexOAuthPendingStore,
|
||||
CodexOAuthPendingStoreError, CodexPendingAuthorization, CodexProfileStatisticsError,
|
||||
CodexQuotaWindow, CodexQuotaWindowKind, CodexQuotaWindowRole, CodexResetCreditsError,
|
||||
CompleteCodexOAuthAuthorization, CompletedCodexOAuthCredential, ExportManagedCodexCredential,
|
||||
StartCodexOAuthAuthorization, StoredCodexPendingAuthorization,
|
||||
CodexOAuthPendingStoreError, CodexPendingAuthorization, CodexProfileAvatarError,
|
||||
CodexProfileStatisticsError, CodexQuotaWindow, CodexQuotaWindowKind, CodexQuotaWindowRole,
|
||||
CodexResetCreditsError, CompleteCodexOAuthAuthorization, CompletedCodexOAuthCredential,
|
||||
ExportManagedCodexCredential, StartCodexOAuthAuthorization, StoredCodexPendingAuthorization,
|
||||
};
|
||||
use crate::credential::{
|
||||
CodexCredentialCodec, CodexOAuthSecret, oauth_owner_ref, parse_access_token_expiration,
|
||||
@@ -61,7 +63,9 @@ use crate::transport::CodexWebSocketPool;
|
||||
use crate::transport::profile::{
|
||||
CodexDesktopReleaseSnapshot, CodexDesktopReleaseStatus, CodexWireProfile, CodexWireProfileState,
|
||||
};
|
||||
use crate::transport::{CodexProfileStatistics, OpenAiBillingUsage, openai_billing_breakdown};
|
||||
use crate::transport::{
|
||||
CodexProfileAvatar, CodexProfileStatistics, OpenAiBillingUsage, openai_billing_breakdown,
|
||||
};
|
||||
|
||||
const PROVIDER_NAME: &str = "openai";
|
||||
const PENDING_DOCUMENT_SCHEMA_VERSION: u64 = 3;
|
||||
@@ -462,6 +466,18 @@ impl ProviderAdmin for OpenAiAdminProvider {
|
||||
Ok(project_profile_statistics(statistics))
|
||||
}
|
||||
|
||||
async fn profile_avatar(
|
||||
&self,
|
||||
account_id: &ProviderAccountId,
|
||||
) -> Result<ProviderProfileAvatar, ProviderAdminError> {
|
||||
self.account(account_id).await?;
|
||||
self.profile_statistics
|
||||
.profile_avatar(account_id)
|
||||
.await
|
||||
.map(project_profile_avatar)
|
||||
.map_err(map_profile_avatar_error)
|
||||
}
|
||||
|
||||
async fn reset_credits(
|
||||
&self,
|
||||
account_id: &ProviderAccountId,
|
||||
@@ -865,6 +881,15 @@ fn project_profile_statistics(statistics: CodexProfileStatistics) -> ProviderPro
|
||||
}
|
||||
}
|
||||
|
||||
fn project_profile_avatar(avatar: CodexProfileAvatar) -> ProviderProfileAvatar {
|
||||
ProviderProfileAvatar {
|
||||
content_type: avatar.content_type,
|
||||
content_length: avatar.content_length,
|
||||
etag: avatar.etag,
|
||||
body: Box::pin(avatar.body.map_err(|_| ProviderProfileAvatarStreamError)),
|
||||
}
|
||||
}
|
||||
|
||||
fn project_reset_credits(
|
||||
credits: crate::transport::CodexRateLimitResetCredits,
|
||||
) -> ProviderResetCredits {
|
||||
@@ -1498,6 +1523,19 @@ fn map_profile_statistics_error(error: CodexProfileStatisticsError) -> ProviderA
|
||||
}
|
||||
}
|
||||
|
||||
fn map_profile_avatar_error(error: CodexProfileAvatarError) -> ProviderAdminError {
|
||||
match error {
|
||||
CodexProfileAvatarError::ProfileStatistics(error) => map_profile_statistics_error(error),
|
||||
CodexProfileAvatarError::Missing => provider_admin_error(ProviderAdminErrorKind::NotFound),
|
||||
CodexProfileAvatarError::InvalidSource | CodexProfileAvatarError::Upstream { .. } => {
|
||||
provider_admin_error(ProviderAdminErrorKind::BadGateway)
|
||||
}
|
||||
CodexProfileAvatarError::TransportUnavailable => {
|
||||
provider_admin_error(ProviderAdminErrorKind::Unavailable)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn map_reset_credits_error(error: CodexResetCreditsError) -> ProviderAdminError {
|
||||
use CodexResetCreditsError as Error;
|
||||
|
||||
|
||||
@@ -38,7 +38,9 @@ pub use oauth::{
|
||||
CodexPendingAuthorization, CompleteCodexOAuthAuthorization, CompletedCodexOAuthAuthorization,
|
||||
CompletedCodexOAuthCredential, StartCodexOAuthAuthorization, StoredCodexPendingAuthorization,
|
||||
};
|
||||
pub use profile_statistics::{CodexCredentialProfileService, CodexProfileStatisticsError};
|
||||
pub use profile_statistics::{
|
||||
CodexCredentialProfileService, CodexProfileAvatarError, CodexProfileStatisticsError,
|
||||
};
|
||||
pub use quota::{
|
||||
CodexAccountQuotaSnapshot, CodexCredentialQuotaError, CodexCredentialQuotaService,
|
||||
CodexQuotaFact, CodexQuotaRefreshPolicy, CodexQuotaSyncSummary, CodexQuotaWindow,
|
||||
|
||||
@@ -1,6 +1,10 @@
|
||||
//! OpenAI 当前账号个人资料统计查询编排。
|
||||
//! OpenAI 当前账号个人资料统计与官方头像查询编排。
|
||||
|
||||
use std::time::SystemTime;
|
||||
use std::{
|
||||
collections::HashMap,
|
||||
sync::Mutex,
|
||||
time::{Duration, Instant, SystemTime},
|
||||
};
|
||||
|
||||
use gateway_core::engine::credential::ProviderAccountId;
|
||||
use reqwest::Client;
|
||||
@@ -10,11 +14,14 @@ use uuid::Uuid;
|
||||
|
||||
use crate::transport::profile::CodexWireProfileState;
|
||||
use crate::transport::{
|
||||
CodexBackendClient, CodexClientError, CodexProfileStatistics, CodexRequestContext,
|
||||
CodexBackendClient, CodexClientError, CodexProfileAvatar, CodexProfileAvatarFetchError,
|
||||
CodexProfileStatistics, CodexRequestContext, fetch_profile_avatar,
|
||||
};
|
||||
|
||||
use super::CodexCredentialRepository;
|
||||
|
||||
const PROFILE_AVATAR_SOURCE_TTL: Duration = Duration::from_secs(10 * 60);
|
||||
|
||||
#[derive(Error)]
|
||||
pub enum CodexProfileStatisticsError {
|
||||
#[error("Codex profile-statistics credential data is invalid")]
|
||||
@@ -59,16 +66,55 @@ impl std::fmt::Debug for CodexProfileStatisticsError {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Error)]
|
||||
pub enum CodexProfileAvatarError {
|
||||
#[error(transparent)]
|
||||
ProfileStatistics(#[from] CodexProfileStatisticsError),
|
||||
#[error("Codex profile does not provide an avatar")]
|
||||
Missing,
|
||||
#[error("Codex profile avatar source is invalid")]
|
||||
InvalidSource,
|
||||
#[error("Codex profile avatar upstream returned HTTP {status}")]
|
||||
Upstream { status: u16 },
|
||||
#[error("Codex profile avatar transport is unavailable")]
|
||||
TransportUnavailable,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for CodexProfileAvatarError {
|
||||
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::ProfileStatistics(error) => formatter
|
||||
.debug_tuple("ProfileStatistics")
|
||||
.field(error)
|
||||
.finish(),
|
||||
Self::Missing => formatter.write_str("Missing"),
|
||||
Self::InvalidSource => formatter.write_str("InvalidSource"),
|
||||
Self::Upstream { status } => formatter
|
||||
.debug_struct("Upstream")
|
||||
.field("status", status)
|
||||
.finish(),
|
||||
Self::TransportUnavailable => formatter.write_str("TransportUnavailable"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct CachedProfileAvatarSource {
|
||||
source: String,
|
||||
expires_at: Instant,
|
||||
}
|
||||
|
||||
pub struct CodexCredentialProfileService {
|
||||
repository: CodexCredentialRepository,
|
||||
profile: CodexWireProfileState,
|
||||
http: Client,
|
||||
base_url: String,
|
||||
avatar_sources: Mutex<HashMap<ProviderAccountId, CachedProfileAvatarSource>>,
|
||||
}
|
||||
|
||||
impl CodexCredentialProfileService {
|
||||
#[must_use]
|
||||
pub const fn new(
|
||||
pub fn new(
|
||||
repository: CodexCredentialRepository,
|
||||
profile: CodexWireProfileState,
|
||||
http: Client,
|
||||
@@ -79,6 +125,7 @@ impl CodexCredentialProfileService {
|
||||
profile,
|
||||
http,
|
||||
base_url,
|
||||
avatar_sources: Mutex::new(HashMap::new()),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -114,7 +161,7 @@ impl CodexCredentialProfileService {
|
||||
.authorization_header()
|
||||
.map_err(|_| CodexProfileStatisticsError::InvalidCredentialData)?;
|
||||
let request_id = format!("profile_statistics_{}", Uuid::now_v7().simple());
|
||||
CodexBackendClient::new(
|
||||
let statistics = CodexBackendClient::new(
|
||||
self.http.clone(),
|
||||
self.base_url.clone(),
|
||||
self.profile.clone(),
|
||||
@@ -126,7 +173,72 @@ impl CodexCredentialProfileService {
|
||||
None,
|
||||
))
|
||||
.await
|
||||
.map_err(map_client_error)
|
||||
.map_err(map_client_error)?;
|
||||
self.remember_avatar_source(account_id, statistics.image_url.as_deref());
|
||||
Ok(statistics)
|
||||
}
|
||||
|
||||
pub async fn profile_avatar(
|
||||
&self,
|
||||
account_id: &ProviderAccountId,
|
||||
) -> Result<CodexProfileAvatar, CodexProfileAvatarError> {
|
||||
let source = match self.cached_avatar_source(account_id) {
|
||||
Some(source) => source,
|
||||
None => self
|
||||
.profile_statistics(account_id)
|
||||
.await?
|
||||
.image_url
|
||||
.ok_or(CodexProfileAvatarError::Missing)?,
|
||||
};
|
||||
let desktop_user_agent = self.profile.snapshot().desktop_user_agent();
|
||||
fetch_profile_avatar(&self.http, &self.base_url, &desktop_user_agent, &source)
|
||||
.await
|
||||
.map_err(map_avatar_fetch_error)
|
||||
}
|
||||
|
||||
fn cached_avatar_source(&self, account_id: &ProviderAccountId) -> Option<String> {
|
||||
let now = Instant::now();
|
||||
let mut sources = self
|
||||
.avatar_sources
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
sources.retain(|_, source| source.expires_at > now);
|
||||
sources.get(account_id).map(|source| source.source.clone())
|
||||
}
|
||||
|
||||
fn remember_avatar_source(&self, account_id: &ProviderAccountId, source: Option<&str>) {
|
||||
let now = Instant::now();
|
||||
let mut sources = self
|
||||
.avatar_sources
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
sources.retain(|_, source| source.expires_at > now);
|
||||
match source {
|
||||
Some(source) => {
|
||||
sources.insert(
|
||||
account_id.clone(),
|
||||
CachedProfileAvatarSource {
|
||||
source: source.to_owned(),
|
||||
expires_at: now + PROFILE_AVATAR_SOURCE_TTL,
|
||||
},
|
||||
);
|
||||
}
|
||||
None => {
|
||||
sources.remove(account_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn map_avatar_fetch_error(error: CodexProfileAvatarFetchError) -> CodexProfileAvatarError {
|
||||
match error {
|
||||
CodexProfileAvatarFetchError::InvalidSource => CodexProfileAvatarError::InvalidSource,
|
||||
CodexProfileAvatarFetchError::Upstream { status } => {
|
||||
CodexProfileAvatarError::Upstream { status }
|
||||
}
|
||||
CodexProfileAvatarFetchError::TransportUnavailable => {
|
||||
CodexProfileAvatarError::TransportUnavailable
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -1,13 +1,14 @@
|
||||
//! Codex AT/RT 刷新状态机;Redis lease + ProviderAccountStore CAS,无 SQL。
|
||||
|
||||
use std::collections::BTreeSet;
|
||||
use std::num::NonZeroU32;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, SystemTime};
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use futures::{StreamExt as _, stream};
|
||||
use gateway_core::engine::credential::{
|
||||
AccountErrorReason, CredentialState, ProviderAccount, ProviderAccountId,
|
||||
AccountErrorReason, CredentialState, ProviderAccount, ProviderAccountId, ProviderRefreshQuery,
|
||||
};
|
||||
use gateway_core::provider_ports::{
|
||||
ProviderCredentialStatePort, ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest,
|
||||
@@ -15,6 +16,7 @@ use gateway_core::provider_ports::{
|
||||
ProviderRuntimePolicyPort, ProviderStoreError, provider_refresh_backoff_at,
|
||||
provider_refresh_retry_at,
|
||||
};
|
||||
use gateway_core::routing::ProviderKind;
|
||||
use secrecy::{ExposeSecret, SecretString};
|
||||
use thiserror::Error;
|
||||
|
||||
@@ -132,28 +134,6 @@ impl<'a> RefreshFailureContext<'a> {
|
||||
}
|
||||
}
|
||||
|
||||
fn refresh_due_at(
|
||||
account: &ProviderAccount,
|
||||
policy: ProviderRefreshPolicy,
|
||||
now: SystemTime,
|
||||
) -> Option<SystemTime> {
|
||||
if refresh_recovery_window_exhausted(account.access_token_expires_at(), now) {
|
||||
// 即使 retry-not-before 异常地落在 deadline 之后,也必须进入终态化路径。
|
||||
return refresh_recovery_deadline(account.access_token_expires_at());
|
||||
}
|
||||
let access_token_expires_at = account.access_token_expires_at()?;
|
||||
if !policy.is_refresh_due(access_token_expires_at, now) {
|
||||
return None;
|
||||
}
|
||||
if account
|
||||
.next_refresh_at()
|
||||
.is_some_and(|retry_at| retry_at > now)
|
||||
{
|
||||
return None;
|
||||
}
|
||||
Some(access_token_expires_at)
|
||||
}
|
||||
|
||||
pub struct DueCodexCredential {
|
||||
pub account: ProviderAccount,
|
||||
pub secret: CodexOAuthSecret,
|
||||
@@ -410,31 +390,25 @@ impl CodexCredentialRefreshService {
|
||||
CodexCredentialRefreshError,
|
||||
> {
|
||||
let now = SystemTime::now();
|
||||
let mut accounts = self
|
||||
.repository
|
||||
.store()
|
||||
.list_accounts()
|
||||
.await
|
||||
.map_err(CredentialRepositoryError::from)?;
|
||||
accounts.retain(|account| {
|
||||
account.provider().as_str() == PROVIDER_NAME
|
||||
&& account.enabled()
|
||||
&& account.has_refresh_token()
|
||||
&& !excluded.contains(account.id())
|
||||
&& matches!(
|
||||
account.credential_state(),
|
||||
CredentialState::Unknown | CredentialState::Ready
|
||||
)
|
||||
&& refresh_due_at(account, policy, now).is_some()
|
||||
});
|
||||
accounts
|
||||
.sort_by_key(|account| (refresh_due_at(account, policy, now), account.id().clone()));
|
||||
accounts.truncate(MAX_REFRESH_BATCH as usize);
|
||||
let mut due = Vec::with_capacity(accounts.len());
|
||||
let provider = ProviderKind::new(PROVIDER_NAME)
|
||||
.map_err(|_| CredentialRepositoryError::InvalidCredentialData)?;
|
||||
let limit = NonZeroU32::new(MAX_REFRESH_BATCH)
|
||||
.ok_or(CredentialRepositoryError::InvalidCredentialData)?;
|
||||
let query = ProviderRefreshQuery::new(
|
||||
provider,
|
||||
now.checked_add(policy.margin()).unwrap_or(now),
|
||||
now.checked_sub(REFRESH_RECOVERY_WINDOW)
|
||||
.unwrap_or(SystemTime::UNIX_EPOCH),
|
||||
now,
|
||||
excluded.iter().cloned().collect(),
|
||||
limit,
|
||||
);
|
||||
let candidates = self.repository.list_refresh_candidates(query).await?;
|
||||
let mut due = Vec::with_capacity(candidates.len());
|
||||
let mut failures = Vec::new();
|
||||
for account in accounts {
|
||||
let account_id = account.id().to_string();
|
||||
match self.repository.load_runtime_credential(&account).await {
|
||||
for loaded in candidates {
|
||||
let account_id = loaded.account.id().to_string();
|
||||
match self.repository.decode_runtime_credential(&loaded) {
|
||||
Ok(runtime)
|
||||
if runtime
|
||||
.authentication
|
||||
@@ -445,7 +419,7 @@ impl CodexCredentialRefreshService {
|
||||
unreachable!("OAuth runtime authentication was checked above")
|
||||
};
|
||||
due.push(DueCodexCredential {
|
||||
account,
|
||||
account: loaded.account,
|
||||
secret: secret.clone(),
|
||||
});
|
||||
}
|
||||
|
||||
@@ -5,8 +5,8 @@ use std::time::SystemTime;
|
||||
|
||||
use gateway_core::engine::credential::{
|
||||
AccountErrorReason, AccountStateChange, CredentialCasOutcome, CredentialCasUpdate,
|
||||
CredentialRevision, CredentialState, ProviderAccount, ProviderAccountStore,
|
||||
ProviderAccountUpdate,
|
||||
CredentialRevision, CredentialState, LoadedCredential, ProviderAccount, ProviderAccountStore,
|
||||
ProviderAccountUpdate, ProviderRefreshQuery,
|
||||
};
|
||||
use gateway_core::routing::ProviderKind;
|
||||
use secrecy::ExposeSecret;
|
||||
@@ -129,6 +129,26 @@ impl CodexCredentialRepository {
|
||||
.map_err(Into::into)
|
||||
}
|
||||
|
||||
pub async fn list_refresh_candidates(
|
||||
&self,
|
||||
query: ProviderRefreshQuery,
|
||||
) -> Result<Vec<LoadedCredential>, CredentialRepositoryError> {
|
||||
self.store
|
||||
.list_refresh_candidates(query)
|
||||
.await
|
||||
.map_err(Into::into)
|
||||
}
|
||||
|
||||
pub fn decode_runtime_credential(
|
||||
&self,
|
||||
loaded: &LoadedCredential,
|
||||
) -> Result<CodexRuntimeCredential, CredentialRepositoryError> {
|
||||
if loaded.account.provider().as_str() != PROVIDER_NAME {
|
||||
return Err(CredentialRepositoryError::InvalidCredentialData);
|
||||
}
|
||||
CodexCredentialCodec::decode(&loaded.credential).map_err(Into::into)
|
||||
}
|
||||
|
||||
pub async fn load_runtime_credential(
|
||||
&self,
|
||||
account: &ProviderAccount,
|
||||
@@ -143,7 +163,7 @@ impl CodexCredentialRepository {
|
||||
if loaded.account != *account {
|
||||
return Err(CredentialRepositoryError::RevisionConflict);
|
||||
}
|
||||
CodexCredentialCodec::decode(&loaded.credential).map_err(Into::into)
|
||||
self.decode_runtime_credential(&loaded)
|
||||
}
|
||||
|
||||
pub async fn load_complete_data(
|
||||
|
||||
@@ -16,7 +16,7 @@ use gateway_core::engine::provider::{
|
||||
ProviderSelectionObservation, ProviderStream, UpstreamTransport,
|
||||
};
|
||||
use gateway_core::engine::{
|
||||
AttemptContext, CancellationToken, ContinuationAttempt, UpstreamSendState,
|
||||
AttemptContext, AttemptTransport, CancellationToken, ContinuationAttempt, UpstreamSendState,
|
||||
};
|
||||
use gateway_core::error::{
|
||||
ClientVisibleUpstreamError, ClientVisibleUpstreamResponse, ContinuationFailure,
|
||||
@@ -47,6 +47,7 @@ use secrecy::{ExposeSecret, SecretString};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Map, Value, json};
|
||||
use url::Url;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::credential::{
|
||||
CodexAccountFailure, CodexCredentialCatalogError, CodexCredentialCatalogService,
|
||||
@@ -69,10 +70,8 @@ use crate::transport::profile::{
|
||||
APPCAST_POLL_INTERVAL, CodexDesktopReleaseService, CodexWireProfileState,
|
||||
};
|
||||
use crate::transport::protocol::responses::{
|
||||
CodexResponsesRequest, PreviousResponseScope, ResponseEventSignals,
|
||||
};
|
||||
use crate::transport::protocol::websocket::{
|
||||
PREVIOUS_RESPONSE_NOT_FOUND_CODE, PREVIOUS_RESPONSE_NOT_FOUND_MESSAGE,
|
||||
CodexResponsesRequest, PREVIOUS_RESPONSE_NOT_FOUND_CODE, PREVIOUS_RESPONSE_NOT_FOUND_MESSAGE,
|
||||
PreviousResponseScope, ResponseEventSignals, transport_requirement,
|
||||
};
|
||||
use crate::transport::request::{
|
||||
CodexRequestEncodeError, RequestAccountScope, encode_generate_request, scope_request_to_account,
|
||||
@@ -289,7 +288,10 @@ impl Provider for CodexProvider {
|
||||
derive_codex_session_affinity(&upstream_request, context.client_api_key_ref());
|
||||
let cyber_policy_session_key =
|
||||
derive_codex_cyber_policy_session_key(&upstream_request, context.client_api_key_ref());
|
||||
let transport = selected_transport(&upstream_request);
|
||||
let transport = match context.transport() {
|
||||
AttemptTransport::Default => selected_transport(&upstream_request),
|
||||
AttemptTransport::Fallback => CodexProviderTransport::HttpOnly,
|
||||
};
|
||||
apply_transport(&mut upstream_request, transport);
|
||||
|
||||
let selection_started_at = Instant::now();
|
||||
@@ -753,6 +755,12 @@ fn cold_json_response_stream(request: ColdJsonResponse) -> EventStream {
|
||||
account_selection,
|
||||
)
|
||||
.await;
|
||||
if let Err(CodexHandshakeAttemptError::Client(error)) = &response {
|
||||
log_client_upstream_error(
|
||||
UpstreamErrorLogContext::new(&request.context, &active_account, None),
|
||||
error,
|
||||
);
|
||||
}
|
||||
let response = match response.map_err(map_handshake_attempt_error) {
|
||||
Ok(response) => response,
|
||||
Err(mut failure) => {
|
||||
@@ -879,6 +887,12 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
&cancellation,
|
||||
)
|
||||
.await;
|
||||
if let Err(CodexHandshakeAttemptError::Client(error)) = &response {
|
||||
log_client_upstream_error(
|
||||
UpstreamErrorLogContext::new(&context, &active_account, None),
|
||||
error,
|
||||
);
|
||||
}
|
||||
let response = response.map_err(map_handshake_attempt_error);
|
||||
let response = match response {
|
||||
Ok(response) => response,
|
||||
@@ -946,6 +960,8 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
{
|
||||
active_account = current;
|
||||
}
|
||||
let response_transport = response.transport;
|
||||
let websocket_connection_id = response.websocket_connection_id;
|
||||
let mut body = response.body;
|
||||
let failure_diagnostics = response.diagnostics.clone();
|
||||
let failure_set_cookie_headers = response.set_cookie_headers.clone();
|
||||
@@ -990,7 +1006,17 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
_ = wait_for_replay_grace(replay_grace_deadline) => Ok(PreCommitPoll::GraceElapsed),
|
||||
chunk = body.next() => match chunk {
|
||||
Some(Ok(chunk)) => Ok(PreCommitPoll::Upstream(Some(chunk))),
|
||||
Some(Err(error)) => Err(map_stream_error(error)),
|
||||
Some(Err(error)) => {
|
||||
log_client_upstream_error(
|
||||
UpstreamErrorLogContext::new(
|
||||
&context,
|
||||
&active_account,
|
||||
websocket_connection_id,
|
||||
),
|
||||
&error,
|
||||
);
|
||||
Err(map_stream_error(error))
|
||||
}
|
||||
None => Ok(PreCommitPoll::Upstream(None)),
|
||||
},
|
||||
};
|
||||
@@ -1002,7 +1028,14 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
}
|
||||
continue;
|
||||
}
|
||||
Err(failure) => {
|
||||
Err(mut failure) => {
|
||||
if failure.error.allows_pre_delivery_retry()
|
||||
&& response_transport == CodexBackendTransport::WebSocket
|
||||
&& transport_requirement(&request).allows_pre_delivery_http_fallback()
|
||||
&& !pre_commit_events.is_committed()
|
||||
{
|
||||
failure.error = failure.error.with_pre_delivery_transport_fallback();
|
||||
}
|
||||
let updates = take_rate_limit_updates(rate_limit_updates.as_ref()).await;
|
||||
if !updates.is_empty() {
|
||||
passive_quota_observation.observe(&updates);
|
||||
@@ -1059,6 +1092,15 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
let service_tier_changed = observation_state
|
||||
.observe_upstream_service_tier(decoder.response_service_tier());
|
||||
let terminal_failure = canonical_failure.map(|(error, semantic_output_seen)| {
|
||||
log_canonical_upstream_error(
|
||||
UpstreamErrorLogContext::new(
|
||||
&context,
|
||||
&active_account,
|
||||
websocket_connection_id,
|
||||
),
|
||||
response_transport,
|
||||
&error,
|
||||
);
|
||||
let atomic_upstream_failure = matches!(&error, CodexCanonicalError::Upstream(_));
|
||||
(
|
||||
map_canonical_error(
|
||||
@@ -1150,6 +1192,15 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
|
||||
}
|
||||
};
|
||||
let terminal_failure = canonical_failure.map(|(error, semantic_output_seen)| {
|
||||
log_canonical_upstream_error(
|
||||
UpstreamErrorLogContext::new(
|
||||
&context,
|
||||
&active_account,
|
||||
websocket_connection_id,
|
||||
),
|
||||
response_transport,
|
||||
&error,
|
||||
);
|
||||
let atomic_upstream_failure = matches!(&error, CodexCanonicalError::Upstream(_));
|
||||
(
|
||||
map_canonical_error(
|
||||
@@ -2196,6 +2247,141 @@ struct OpenAiFailureContext<'a> {
|
||||
allows_account_state_mutation: bool,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
struct UpstreamErrorLogContext<'a> {
|
||||
request_id: &'a str,
|
||||
account_id: &'a str,
|
||||
attempt_index: u32,
|
||||
websocket_connection_id: Option<Uuid>,
|
||||
}
|
||||
|
||||
impl<'a> UpstreamErrorLogContext<'a> {
|
||||
fn new(
|
||||
context: &'a AttemptContext,
|
||||
account: &'a ProviderAccount,
|
||||
websocket_connection_id: Option<Uuid>,
|
||||
) -> Self {
|
||||
Self {
|
||||
request_id: context.request_id().as_str(),
|
||||
account_id: account.id().as_str(),
|
||||
attempt_index: context.attempt_index().get(),
|
||||
websocket_connection_id,
|
||||
}
|
||||
}
|
||||
|
||||
fn with_websocket_connection_id(mut self, connection_id: Option<Uuid>) -> Self {
|
||||
self.websocket_connection_id = self.websocket_connection_id.or(connection_id);
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
fn log_client_upstream_error(context: UpstreamErrorLogContext<'_>, error: &CodexClientError) {
|
||||
match error {
|
||||
CodexClientError::Upstream {
|
||||
status,
|
||||
body,
|
||||
transport,
|
||||
..
|
||||
} => log_raw_upstream_body(
|
||||
context,
|
||||
*transport,
|
||||
"http_error_response",
|
||||
Some(status.as_u16()),
|
||||
None,
|
||||
None,
|
||||
body,
|
||||
),
|
||||
CodexClientError::WebSocket(CodexWebSocketExchangeError::Upstream(upstream)) => {
|
||||
log_raw_upstream_body(
|
||||
context,
|
||||
CodexBackendTransport::WebSocket,
|
||||
"websocket_opening_response",
|
||||
Some(upstream.status_code),
|
||||
None,
|
||||
None,
|
||||
&upstream.body,
|
||||
);
|
||||
}
|
||||
CodexClientError::WebSocket(error) => {
|
||||
let Some(close) = error.close_before_terminal() else {
|
||||
return;
|
||||
};
|
||||
let context = context.with_websocket_connection_id(close.connection_id());
|
||||
let websocket_connection_id = context
|
||||
.websocket_connection_id
|
||||
.map(|connection_id| connection_id.to_string())
|
||||
.unwrap_or_default();
|
||||
let upstream_error_raw = close.reason().unwrap_or_default();
|
||||
tracing::warn!(
|
||||
request_id = %context.request_id,
|
||||
account_id = %context.account_id,
|
||||
attempt_index = context.attempt_index,
|
||||
websocket_connection_id,
|
||||
upstream_transport = WEBSOCKET_TRANSPORT,
|
||||
upstream_error_kind = "websocket_close",
|
||||
upstream_close_code = close.code().unwrap_or_default(),
|
||||
upstream_close_code_present = close.code().is_some(),
|
||||
upstream_error_raw,
|
||||
upstream_error_raw_bytes = upstream_error_raw.len(),
|
||||
upstream_error_raw_present = close.reason().is_some(),
|
||||
"OpenAI upstream WebSocket closed before terminal event"
|
||||
);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
fn log_canonical_upstream_error(
|
||||
context: UpstreamErrorLogContext<'_>,
|
||||
transport: CodexBackendTransport,
|
||||
error: &CodexCanonicalError,
|
||||
) {
|
||||
let CodexCanonicalError::Upstream(failure) = error else {
|
||||
return;
|
||||
};
|
||||
log_raw_upstream_body(
|
||||
context,
|
||||
transport,
|
||||
"responses_error_event",
|
||||
failure.explicit_status_code,
|
||||
failure.upstream_code.as_deref(),
|
||||
failure.upstream_type.as_deref(),
|
||||
failure.raw_body(),
|
||||
);
|
||||
}
|
||||
|
||||
fn log_raw_upstream_body(
|
||||
context: UpstreamErrorLogContext<'_>,
|
||||
transport: CodexBackendTransport,
|
||||
error_kind: &'static str,
|
||||
status_code: Option<u16>,
|
||||
upstream_code: Option<&str>,
|
||||
upstream_type: Option<&str>,
|
||||
upstream_error_raw: &str,
|
||||
) {
|
||||
let websocket_connection_id = context
|
||||
.websocket_connection_id
|
||||
.map(|connection_id| connection_id.to_string())
|
||||
.unwrap_or_default();
|
||||
tracing::warn!(
|
||||
request_id = %context.request_id,
|
||||
account_id = %context.account_id,
|
||||
attempt_index = context.attempt_index,
|
||||
websocket_connection_id,
|
||||
upstream_transport = actual_transport_name(transport),
|
||||
upstream_error_kind = error_kind,
|
||||
upstream_status_code = status_code.unwrap_or_default(),
|
||||
upstream_status_code_present = status_code.is_some(),
|
||||
upstream_error_code = upstream_code.unwrap_or_default(),
|
||||
upstream_error_code_present = upstream_code.is_some(),
|
||||
upstream_error_type = upstream_type.unwrap_or_default(),
|
||||
upstream_error_type_present = upstream_type.is_some(),
|
||||
upstream_error_raw,
|
||||
upstream_error_raw_bytes = upstream_error_raw.len(),
|
||||
"OpenAI upstream returned an error payload"
|
||||
);
|
||||
}
|
||||
|
||||
async fn apply_failure(
|
||||
context: &OpenAiFailureContext<'_>,
|
||||
account: &ProviderAccount,
|
||||
@@ -2417,11 +2603,20 @@ fn map_client_error(
|
||||
))
|
||||
}
|
||||
CodexClientError::WebSocket(error) => {
|
||||
let close_code = error.close_before_terminal().and_then(|close| close.code());
|
||||
let client_visible_error = websocket_client_visible_error(&error);
|
||||
let mut failure = MappedProviderFailure::plain(provider_error(
|
||||
websocket_error_kind(&error),
|
||||
websocket_send_state(&error),
|
||||
));
|
||||
if let Some(close_code) = close_code {
|
||||
failure.error =
|
||||
failure
|
||||
.error
|
||||
.with_upstream_code(OpaqueUpstreamValue::new(format!(
|
||||
"websocket_close_{close_code}"
|
||||
)));
|
||||
}
|
||||
if let Some(client_visible_error) = client_visible_error {
|
||||
failure.error = failure
|
||||
.error
|
||||
|
||||
@@ -59,7 +59,7 @@ pub enum CodexCanonicalError {
|
||||
Protocol(#[source] ProviderError),
|
||||
/// 上游在成功建立流后发送了明确的失败事件。
|
||||
#[error("Codex upstream reported a failed response")]
|
||||
Upstream(ResponsesSseFailure),
|
||||
Upstream(Box<ResponsesSseFailure>),
|
||||
}
|
||||
|
||||
impl fmt::Debug for CodexCanonicalError {
|
||||
@@ -291,7 +291,11 @@ impl CodexCanonicalDecoder {
|
||||
let signals = response_event_signals(event_type, &value);
|
||||
self.merge_timing_signals(signals);
|
||||
if matches!(event_type, Some("response.failed" | "error")) {
|
||||
let failure = ResponsesSseFailure::from_event(event_type.unwrap_or_default(), &value);
|
||||
let failure = ResponsesSseFailure::from_raw_event(
|
||||
event_type.unwrap_or_default(),
|
||||
&event.data,
|
||||
&value,
|
||||
);
|
||||
let mut canonical = Vec::new();
|
||||
if !self.started {
|
||||
// 失败首帧可能是唯一携带 response_id 的上游事实。身份投影只做
|
||||
@@ -307,7 +311,7 @@ impl CodexCanonicalDecoder {
|
||||
} else {
|
||||
output.extend(canonical.into_iter().map(ProviderEvent::canonical));
|
||||
}
|
||||
return Err(CodexCanonicalError::Upstream(failure));
|
||||
return Err(CodexCanonicalError::Upstream(Box::new(failure)));
|
||||
}
|
||||
let mut canonical = Vec::new();
|
||||
if let Some(event_type) = event_type {
|
||||
|
||||
@@ -22,6 +22,7 @@ use reqwest::{
|
||||
};
|
||||
use serde_json::{Value, map::Map};
|
||||
use thiserror::Error;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::transport::protocol::responses::{CodexResponsesRequest, TransportRequirement};
|
||||
|
||||
@@ -401,7 +402,7 @@ pub struct CodexRequestContext<'a> {
|
||||
pub beta_features: Option<&'a str>,
|
||||
/// x-responsesapi-include-timing-metrics 头值。
|
||||
pub include_timing_metrics: Option<&'a str>,
|
||||
/// 客户端版本号。
|
||||
/// 下游 `version` 扩展头;存在时上游值由 Desktop 版本画像统一生成。
|
||||
pub version: Option<&'a str>,
|
||||
/// 客户端 window ID。
|
||||
pub codex_window_id: Option<&'a str>,
|
||||
@@ -538,6 +539,8 @@ pub struct CodexBackendStreamingResponse {
|
||||
pub body: CodexBackendSseStream,
|
||||
/// 实际使用的上游传输。
|
||||
pub transport: CodexBackendTransport,
|
||||
/// WebSocket 响应所绑定的连接;HTTP transport 为 `None`。
|
||||
pub websocket_connection_id: Option<Uuid>,
|
||||
/// 响应头里的最新 turn state。
|
||||
pub turn_state: Option<String>,
|
||||
/// 上游透传的 `set-cookie` 列表。
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user