Compare commits

...
28 Commits
Author SHA1 Message Date
ZyyandCodex 05ba116eef chore(release): v3.25.0
Co-authored-by: Codex <noreply@openai.com>
2026-10-12 17:16:50 +08:00
ZyyandCodex 3ef4743e1a fix(store): persist retries after transport commit
Let Core decide replay eligibility after WebSocket preludes or HTTP
keepalives. Persist subsequent attempts while still enforcing the running
lifecycle and consecutive attempt numbers, preserving delivery facts and
the upstream send watermark.

Cover retry metadata, duplicate and skipped attempts, and terminal request
immutability with a PostgreSQL regression test.

Co-authored-by: Codex <noreply@openai.com>
2026-10-12 17:06:03 +08:00
ZyyandCodex 9e02e52fec fix(responses): keep SSE alive before first output
Start downstream keepalives while admitted streaming requests await their first output, including middleware processing. Separate transport commitment from the replay barrier so heartbeat frames preserve safe retries and execution accounting.

Co-authored-by: Codex <noreply@openai.com>
2026-10-12 16:57:28 +08:00
ZyyandCodex 42ff5b74d0 fix(observability): count incomplete requests in availability
Include incomplete final outcomes in service health metrics and lower the
stable threshold to 95%. Keep the health popover at four outcomes by
folding incomplete requests into the service failure display.

Use consistent service success rates across dashboard trends, health
timelines and usage insights. Cover the reported counts, threshold edges,
caller exclusions and final success after failed attempts in regressions.

Refs #406

Co-authored-by: Codex <noreply@openai.com>
2026-10-12 16:42:33 +08:00
ZyyandCodex eb2eefd51e feat(openai): add ultrafast pricing
Add published GPT-6 Astra and GPT-6.1 Sol rates for short and long context. Expose Ultrafast override bands in the pricing API and editor, with billing and validation coverage.

Co-authored-by: Codex <noreply@openai.com>
2026-10-12 15:46:36 +08:00
think-backandClaude Opus 5.5 abd60ce60a fix(store): acknowledge analytics versions after projection commits (#404)
The analytics worker acknowledged projected queue versions inside its
repeatable-read projection transaction. Request updates keep rewriting
pending queue rows through the source triggers, and once such a change
commits after the worker snapshot, the acknowledgement row lock fails with
a serialization error instead of skipping the stale version. Under steady
traffic every cycle rolled back the whole projection batch, the queue kept
growing and the worker stayed in backoff.

Commit the projection and publication first, then acknowledge the exact
versions in a separate read-committed transaction. Row locks there
re-check the latest committed version, so rewritten versions no longer
match and stay queued for replay, while locked rows are still skipped.

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-12 13:50:02 +08:00
ZyyandCodex 2f96c02ea0 fix(routing): bind known models to their providers
Preserve known model ownership when provider catalog refreshes fail.
Reject ambiguous unknown-provider routes and exclude ChatGPT OAuth
accounts from Grok requests while preserving API key passthrough.

Refs #391

Co-authored-by: Codex <noreply@openai.com>
2026-10-12 13:04:32 +08:00
think-backandClaude Opus 5.5 0f575e279a fix(openai): keep account selection through quota observation updates (#400)
Successful responses keep refreshing the quota observation time of the
account they used. The post-selection credential reload compared the
whole account snapshot, so a busy account could hit three observation-only
conflicts in a row and fail with account_snapshot_conflict before any
upstream request was sent.

Accept the reloaded snapshot when only the observation time of a
non-exhausted quota moved. Credential, state, scheduling, scope and quota
conclusion changes still invalidate the selection and retry as before.

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-12 12:50:19 +08:00
dependabot[bot] 634b8fc674 chore(deps): bump node from 662933c to 86f07bc in /deploy (#402)
Bumps node from `662933c` to `86f07bc`.

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

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2026-10-12 12:45:36 +08:00
dependabot[bot] 00bfb1699d chore(deps): bump debian from 3783cc0 to 7c7b2c9 in /deploy (#401)
Bumps debian from `3783cc0` to `7c7b2c9`.

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

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2026-10-12 12:45:27 +08:00
dependabot[bot] 956ca70a14 chore(deps): bump the minor-and-patch group in /backend with 5 updates (#399)
Bumps the minor-and-patch group in /backend with 5 updates:

| Package | From | To |
| --- | --- | --- |
| [uuid](https://github.com/uuid-rs/uuid) | `1.26.1` | `1.27.0` |
| [jsonschema](https://github.com/Stranger6667/jsonschema) | `0.58.4` | `0.58.6` |
| [h2](https://github.com/hyperium/h2) | `0.4.19` | `0.4.20` |
| [regex](https://github.com/rust-lang/regex) | `1.13.0` | `1.13.1` |
| [zeroize](https://github.com/RustCrypto/utils) | `1.9.0` | `1.9.1` |


Updates `uuid` from 1.26.1 to 1.27.0
- [Release notes](https://github.com/uuid-rs/uuid/releases)
- [Commits](https://github.com/uuid-rs/uuid/compare/v1.26.1...v1.27.0)

Updates `jsonschema` from 0.58.4 to 0.58.6
- [Release notes](https://github.com/Stranger6667/jsonschema/releases)
- [Changelog](https://github.com/Stranger6667/jsonschema/blob/master/CHANGELOG.md)
- [Commits](https://github.com/Stranger6667/jsonschema/compare/cli-v0.58.4...ruby-v0.58.6)

Updates `h2` from 0.4.19 to 0.4.20
- [Release notes](https://github.com/hyperium/h2/releases)
- [Changelog](https://github.com/hyperium/h2/blob/master/CHANGELOG.md)
- [Commits](https://github.com/hyperium/h2/compare/v0.4.19...v0.4.20)

Updates `regex` from 1.13.0 to 1.13.1
- [Release notes](https://github.com/rust-lang/regex/releases)
- [Changelog](https://github.com/rust-lang/regex/blob/master/CHANGELOG.md)
- [Commits](https://github.com/rust-lang/regex/compare/1.13.0...1.13.1)

Updates `zeroize` from 1.9.0 to 1.9.1
- [Commits](https://github.com/RustCrypto/utils/compare/zeroize-v1.9.0...zeroize-v1.9.1)

---
updated-dependencies:
- dependency-name: uuid
  dependency-version: 1.27.0
  dependency-type: direct:production
  update-type: version-update:semver-minor
  dependency-group: minor-and-patch
- dependency-name: jsonschema
  dependency-version: 0.58.6
  dependency-type: direct:production
  update-type: version-update:semver-patch
  dependency-group: minor-and-patch
- dependency-name: h2
  dependency-version: 0.4.20
  dependency-type: direct:production
  update-type: version-update:semver-patch
  dependency-group: minor-and-patch
- dependency-name: regex
  dependency-version: 1.13.1
  dependency-type: direct:production
  update-type: version-update:semver-patch
  dependency-group: minor-and-patch
- dependency-name: zeroize
  dependency-version: 1.9.1
  dependency-type: direct:production
  update-type: version-update:semver-patch
  dependency-group: minor-and-patch
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
2026-10-12 12:05:19 +08:00
ZyyandCodex 24dd7b94f4 fix(openai): limit gateway retries to explicit recovery scenarios (#398)
Preserve ordinary upstream failures for client retry decisions and allow
unsent native TLS handshake EOFs through bounded connection recovery.

Co-authored-by: Codex <noreply@openai.com>
2026-10-12 11:57:49 +08:00
ZyyandCodex 94c769fc9b fix(plugins): exclude managed next waits from RPC budgets (#397)
Pause Runtime and SDK processing budgets while middleware delegates to host next, then resume the remaining budget. Preserve request deadlines and cancellation, share callback deadline state, and add timeout diagnostics and regression coverage.

Fixes #392

Co-authored-by: Codex <noreply@openai.com>
2026-10-12 11:40:01 +08:00
ZyyandCodex eea16f5275 chore(release): v3.24.1
Co-authored-by: Codex <noreply@openai.com>
2026-10-12 00:16:21 +08:00
ZyyandCodex 8a406ff594 fix(observability): preserve query budgets and usage history
Hold PostgreSQL observation slots through connection cleanup and optimize exact usage pagination. Keep usage panels visible when analytics queries fail and backfill missing historical identity snapshots in one migration.

Co-authored-by: Codex <noreply@openai.com>
2026-10-12 00:15:40 +08:00
ZyyandCodex 676044c43f chore(release): v3.24.0
Co-authored-by: Codex <noreply@openai.com>
2026-10-11 23:19:45 +08:00
ZyyandCodex e7161da277 refactor(migrations): consolidate unreleased request analytics schema
Fold migrations 0028 through 0035 into one migration while preserving every released migration through v3.23.1. Define the final observation validators, history schema, recovery leases, token columns, indexes, and analytics triggers once.

Refresh the migration manifest, remove superseded intermediate upgrade tests, and retain final-schema business validation. Verify clean installation, released-baseline upgrades, schema equivalence, and historical data using temporary PostgreSQL checks.

Co-authored-by: Codex <noreply@openai.com>
2026-10-11 23:14:54 +08:00
ZyyandCodex 97faca0c45 fix(routing): prioritize known models over discovery fallbacks
Prefer providers whose catalogs contain the mapped model before providers accepting unlisted models. Preserve account scope, route constraints, provider order within each tier, and unknown-model passthrough.

Add compiled mixed-catalog regressions and document candidate ordering. Refs #391.

Co-authored-by: Codex <noreply@openai.com>
2026-10-11 23:04:51 +08:00
45b473c7e2 fix(openai): preserve Lite interrupts across prelude retries (#390)
* fix(openai): preserve Lite interrupts across prelude retries

Deliver known Responses Lite WebSocket preludes before semantic output while retaining existing recovery guards. Keep interrupt intent across retries and retarget it only after the replacement response is delivered.

Co-authored-by: Codex <noreply@openai.com>

* fix(openai): forward Lite events without waiting for output

Deliver unclassified Lite WebSocket events in order and reuse the existing commit boundary to stop replay after their delivery. Preserve error classification, retry budgets, account recovery, and interrupt handling.

Co-authored-by: Codex <noreply@openai.com>

---------

Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
2026-10-11 22:53:59 +08:00
ZyyandCodex cc69368f77 perf(store): scale request analytics and isolate request history (#389)
Build durable request facts and bounded aggregate queries, freeze historical
identities, and move recovery lease updates out of request history rows.
Align token fields across storage, APIs and pages.

Invalidate analytics only after actual request deletions and track embedded
migration inputs during builds.

Co-authored-by: Codex <noreply@openai.com>
2026-10-11 20:57:36 +08:00
ZyyandCodex 67691d18d3 refactor(openai): clarify adapted model catalog generation
Co-authored-by: Codex <noreply@openai.com>
2026-10-11 11:48:00 +08:00
hi_yuebanandCodex 85c64d35d5 fix(store): avoid revalidating unchanged request observations (#387)
Validate observation shape on insert and changed JSON only, preserving atomic account deletion and historical snapshots. Add upgrade, validation-cost, bulk deletion, and rollback coverage.

Co-authored-by: Codex <noreply@openai.com>
2026-10-11 10:30:04 +08:00
ZyyandCodex 9bf86f7a7c fix(frontend): smooth health timeline popover tracking
Co-authored-by: Codex <noreply@openai.com>
2026-10-10 18:13:25 +08:00
ZyyandCodex 6c56a23508 chore(release): v3.23.1
Co-authored-by: Codex <noreply@openai.com>
2026-10-10 16:51:52 +08:00
b5f3cb402c fix(openai): reject sends on closed WebSocket connections (#381)
Record observed closure before terminal delivery can block, reject queued
unsent commands, and skip commands cancelled before socket writes.
Preserve connection diagnostics and test account egress across quota writes.

Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
2026-10-10 16:36:40 +08:00
57e284d7d1 fix(openai): preserve stream recovery before semantic output (#380)
Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
2026-10-10 16:30:26 +08:00
21597f9e98 fix(store): restore observation query optimization (#378)
Scope timing projection barriers to filtered metric and diagnostic queries, preserving ordinary aggregate index and parallel plans.

Refs #377

Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
2026-10-10 16:30:05 +08:00
ZyyandCodex fa0b87dd0a refactor(frontend): adopt UI v0.4.0 component APIs
Migrate the management UI to Z components and managed notifications, register the library loading directive, and align fixed dependencies and submodules with the published UI version.

Co-authored-by: Codex <noreply@openai.com>
2026-10-10 14:54:56 +08:00
433 changed files with 16365 additions and 5852 deletions
+17 -17
View File
@@ -1980,9 +1980,9 @@ dependencies = [
[[package]]
name = "h2"
version = "0.4.19"
version = "0.4.20"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ef8e5e5a340588f4452631496976cf8636d4a7ecf600239fdc27615d2530bc16"
checksum = "7d29020232d6aa3fb1daca64c1127cf662cf97f254ae16c18c05b8ab635fc118"
dependencies = [
"atomic-waker",
"bytes",
@@ -2446,9 +2446,9 @@ dependencies = [
[[package]]
name = "jsonschema"
version = "0.58.4"
version = "0.58.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5761b0ca0a31464030d7c2a98087fde3124ebb5e09b5a112cdb815f03b8ee150"
checksum = "f622c574a29fb3034a0a90645afee194dfedc75fc124228801b6e77041075bd1"
dependencies = [
"ahash",
"bytecount",
@@ -2474,18 +2474,18 @@ dependencies = [
[[package]]
name = "jsonschema-regex"
version = "0.58.5"
version = "0.58.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "854e9e22c420535035e95eed240cbfb4a5f4b7da74d122d63144e0c31926dbfd"
checksum = "a968fffd387b2ab005bf00c1851f5d26e90df05b6f7ad609ed0d620033cca5e7"
dependencies = [
"regex-syntax",
]
[[package]]
name = "jsonschema-value"
version = "0.58.5"
version = "0.58.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4e3bfac11ac357ec620880025757d38c67da8fd5b9ef5024d19f0f360002b4ca"
checksum = "39b7670345e5ce87371339d13705dbb9276b7e3e934bbe1efa2c6734a1f65a96"
dependencies = [
"ahash",
"bytecount",
@@ -3587,9 +3587,9 @@ dependencies = [
[[package]]
name = "referencing"
version = "0.58.5"
version = "0.58.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "590efadb0a669f1712c1e3ab810a55d4c3d20d0127c455f7236dd0710aba0346"
checksum = "558f3855529c0a7e78a48e409fa653456e55f1b167eab4109c16a255b7727e3f"
dependencies = [
"ahash",
"fluent-uri",
@@ -3604,9 +3604,9 @@ dependencies = [
[[package]]
name = "regex"
version = "1.13.0"
version = "1.13.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2a0e75113e14dc5acb068cd0786884f214f1312650a3d36d269f5c4f3cdee8a2"
checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d"
dependencies = [
"aho-corasick",
"memchr",
@@ -4504,7 +4504,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd"
dependencies = [
"fastrand",
"getrandom 0.3.4",
"getrandom 0.4.3",
"once_cell",
"rustix",
"windows-sys 0.61.2",
@@ -5006,9 +5006,9 @@ checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821"
[[package]]
name = "uuid"
version = "1.26.1"
version = "1.27.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2ef6dac1e96601b4fb3acccccff2139741fcb757cb9a36089bf5be91cfb285ce"
checksum = "97277d36b9c3ace13e58fa6e753f8b0bbbf302a18dd193240d70f9e29681059a"
dependencies = [
"getrandom 0.4.3",
"js-sys",
@@ -5546,9 +5546,9 @@ dependencies = [
[[package]]
name = "zeroize"
version = "1.9.0"
version = "1.9.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e13c156562582aa81c60cb29407084cdb54c4164760106ab78e6c5b0858cf64e"
checksum = "e13084392c5e4bc371903e2935a5eaeed24905a7511356b883835e18a78f6879"
dependencies = [
"zeroize_derive",
]
@@ -56,6 +56,12 @@ pub async fn run() -> Result<(), BootstrapError> {
plugin_runtime,
gateway_host::plugin_distribution::HttpPluginDistribution::new(plugin_http)
));
let analytics_worker = or_shutdown!(
plugin_runtime,
store.request_analytics_worker(std::sync::Arc::new(
gateway_admin::ProviderUsageBillingProjector::new(providers.admin.clone()),
))
);
let mut admin = or_shutdown!(
plugin_runtime,
gateway_admin::initialize_with_plugin_accounts(
@@ -142,6 +148,7 @@ pub async fn run() -> Result<(), BootstrapError> {
);
let mut plan = store.take_worker_contributions();
plan.push(analytics_worker);
plan.push(or_shutdown!(
plugin_runtime,
gateway_host::retention::worker(store.retention())
+6 -3
View File
@@ -35,7 +35,7 @@ pub use use_case::{
client_keys::ClientKeyService,
credentials::{CredentialsService, ProviderCredentials},
import_tasks::ImportTasksService,
observability::ObservabilityService,
observability::{ObservabilityService, ProviderUsageBillingProjector},
proxies::ProxiesService,
settings::SettingsService,
system::SystemService,
@@ -441,11 +441,14 @@ async fn initialize_inner(
plugin_inspector.clone(),
));
let system = Arc::new(DefaultSystemService::new(system, system_preflight));
let observations: Arc<dyn ports::store::ObservabilityStore> = Arc::new(
use_case::observability::CachedObservabilityStore::new(store.observability()),
);
let key_usage = Arc::new(use_case::key_usage::DefaultKeyUsageService::new(
auth.clone(),
client_key_verifier,
store.client_keys(),
store.observability(),
observations.clone(),
system.clone(),
timezone,
));
@@ -506,7 +509,7 @@ async fn initialize_inner(
)),
client_distribution: Arc::new(DefaultClientDistributionService::new(client_distribution)),
observability: Arc::new(DefaultObservabilityService::new(
store.observability(),
observations.clone(),
store.accounts(),
store.settings(),
registry.clone(),
@@ -146,9 +146,9 @@ pub struct AccountUsage {
pub success_count: u64,
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 cached_input_tokens: Option<u64>,
pub cache_write_input_tokens: Option<u64>,
pub reasoning_output_tokens: Option<u64>,
pub image_input_tokens: Option<u64>,
pub image_output_tokens: Option<u64>,
pub image_request_count: u64,
@@ -0,0 +1,65 @@
//! 统计投影的新鲜度、覆盖范围与预计算费用合同
use chrono::{DateTime, Utc};
use super::observability::{DecimalAmount, RequestMetricPoint, UsageOverview};
/// 数据可见性状态,不以缺失投影冒充零统计
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub enum AnalyticsState {
Ready,
#[default]
Building,
Lagging,
Stale,
}
impl AnalyticsState {
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::Ready => "ready",
Self::Building => "building",
Self::Lagging => "lagging",
Self::Stale => "stale",
}
}
}
/// 与统计结果在同一读取快照内取得的覆盖和版本事实
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct AnalyticsMetadata {
pub coverage_start: Option<DateTime<Utc>>,
pub coverage_end: Option<DateTime<Utc>>,
pub generated_at: DateTime<Utc>,
pub lag_seconds: u64,
pub state: AnalyticsState,
pub percentiles_approximate: bool,
pub schema_version: u32,
pub histogram_version: u32,
/// 派生数据版本供缓存判断失效,不作为源事务提交顺序
pub generation: u64,
}
/// 单次请求经 Provider 费用合同确认后可累加的 USD 情景金额
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct UsageBillingProjection {
pub standard_amount: Option<DecimalAmount>,
pub no_cache_amount: Option<DecimalAmount>,
}
/// 时间桶内的情景费用,未知金额不改写为零
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UsageProjectedCostPoint {
pub bucket_start: DateTime<Utc>,
pub standard_amount: Option<DecimalAmount>,
pub no_cache_amount: Option<DecimalAmount>,
}
/// 同一数据库快照内的完整洞察输入
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UsageInsightsObservation {
pub overview: UsageOverview,
pub trend: Vec<RequestMetricPoint>,
pub costs: Vec<UsageProjectedCostPoint>,
}
@@ -4,6 +4,7 @@ use std::num::{NonZeroU16, NonZeroU64};
pub mod account_groups;
pub mod accounts;
pub mod analytics;
pub mod audit;
pub mod auth;
pub mod backup;
@@ -5,6 +5,10 @@ use std::{num::NonZeroU16, str::FromStr};
use chrono::{DateTime, TimeDelta, Utc};
use super::AdminModelError;
pub use super::analytics::{
AnalyticsMetadata, AnalyticsState, UsageBillingProjection, UsageInsightsObservation,
UsageProjectedCostPoint,
};
/// 观测记录查询的页大小,由 HTTP、管理用例和存储端口共同使用
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -48,7 +52,7 @@ impl CalendarPeriod {
}
/// 外部观测查询的 UTC 时间范围
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct TimeRange {
pub start: DateTime<Utc>,
pub end: DateTime<Utc>,
@@ -89,7 +93,7 @@ impl TimeRange {
}
/// 请求结果状态
#[derive(Debug, Clone, PartialEq, Eq)]
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum RequestOutcome {
Running,
Succeeded,
@@ -100,7 +104,7 @@ pub enum RequestOutcome {
}
/// 未知但有界的请求结果值;私有字段禁止绕过 [`RequestOutcome::new`] 构造
#[derive(Debug, Clone, PartialEq, Eq)]
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct OtherRequestOutcome(String);
impl RequestOutcome {
@@ -151,8 +155,9 @@ impl OtherRequestOutcome {
}
/// 用量记录过滤条件
#[derive(Debug, Clone, Default, PartialEq, Eq)]
#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)]
pub struct UsageFilter {
pub routing_group_refs: Vec<String>,
pub client_api_key_ref: Option<String>,
pub request_id: Option<String>,
pub provider_account_ref: Option<String>,
@@ -178,7 +183,7 @@ pub struct UsageQuery {
}
/// Admin 确定概览的数量、趋势粒度及最近请求筛选,Store 只执行该选择
#[derive(Debug, Clone, PartialEq, Eq)]
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct DashboardQuery {
pub range: TimeRange,
pub granularity: Granularity,
@@ -206,6 +211,7 @@ impl DashboardQuery {
/// 运维错误过滤条件
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct OpsErrorFilter {
pub routing_group_refs: Vec<String>,
pub client_api_key_ref: Option<String>,
pub request_id: Option<String>,
pub provider_kind: Option<String>,
@@ -230,7 +236,7 @@ pub struct OpsErrorQuery {
}
/// 用量诊断维度
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum DiagnosticDimension {
Provider,
Model,
@@ -363,8 +369,8 @@ pub struct ProviderBillingInput {
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 cached_input_tokens: Option<u64>,
pub cache_write_input_tokens: Option<u64>,
pub total: CurrencyCost,
}
@@ -380,8 +386,8 @@ pub struct UsageCalculatedBillingFact {
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 cached_input_tokens: Option<u64>,
pub cache_write_input_tokens: Option<u64>,
pub total: CurrencyCost,
}
@@ -408,7 +414,7 @@ pub struct CalculatedBillingBreakdown {
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ImageBillingBreakdown {
pub input_tokens: u64,
pub cached_tokens: u64,
pub cached_input_tokens: u64,
pub input_amount: CurrencyCost,
pub cache_read_amount: CurrencyCost,
pub input_price_per_million: CurrencyCost,
@@ -454,9 +460,9 @@ pub struct RequestMetrics {
pub caller_error_count: u64,
pub input_tokens: u64,
pub output_tokens: u64,
pub cached_tokens: u64,
pub cache_write_tokens: u64,
pub reasoning_tokens: u64,
pub cached_input_tokens: u64,
pub cache_write_input_tokens: u64,
pub reasoning_output_tokens: u64,
pub total_tokens: u64,
pub first_token_latency_sum_ms: u64,
pub first_token_latency_count: u64,
@@ -496,7 +502,7 @@ pub struct AttemptMetrics {
}
/// 时间序列粒度
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum Granularity {
FifteenMinutes,
Hour,
@@ -556,9 +562,9 @@ pub struct AccountModelUsage {
pub success_count: u64,
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 cached_input_tokens: Option<u64>,
pub cache_write_input_tokens: Option<u64>,
pub reasoning_output_tokens: Option<u64>,
pub image_input_tokens: Option<u64>,
pub image_output_tokens: Option<u64>,
pub image_request_count: u64,
@@ -591,9 +597,9 @@ pub struct DashboardAccountUsage {
pub success_count: u64,
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 cached_input_tokens: Option<u64>,
pub cache_write_input_tokens: Option<u64>,
pub reasoning_output_tokens: Option<u64>,
pub image_input_tokens: Option<u64>,
pub image_output_tokens: Option<u64>,
pub image_request_count: u64,
@@ -628,7 +634,7 @@ pub struct DashboardRuntimeSlots {
pub struct DashboardTotals {
pub request_count: u64,
pub input_tokens: u64,
pub cached_tokens: u64,
pub cached_input_tokens: u64,
pub total_tokens: u64,
pub billing_usd: Option<DecimalAmount>,
}
@@ -636,6 +642,7 @@ pub struct DashboardTotals {
/// 仪表盘所需的当前区间事实及全历史累计
#[derive(Debug, Clone, PartialEq)]
pub struct DashboardObservation {
pub analytics: AnalyticsMetadata,
pub range: TimeRange,
pub totals: DashboardTotals,
pub provider_accounts: AccountPoolMetrics,
@@ -657,9 +664,9 @@ pub struct UsageListRecord {
pub provider_account_ref: Option<String>,
pub provider_account_name: Option<String>,
pub provider_account_email: Option<String>,
/// 账号当前备注,按内部账号 ID 关联,不属于请求历史快照
/// 请求选号时冻结的账号备注,未采集时保持未知
pub provider_account_notes: Option<String>,
/// 账号当前套餐,不属于请求历史快照
/// 请求选号时冻结的账号套餐
pub provider_account_plan_type: Option<String>,
pub provider_account_plan_type_display: Option<String>,
pub provider_account_authentication_kind: Option<String>,
@@ -669,9 +676,9 @@ pub struct UsageListRecord {
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 cached_input_tokens: Option<u64>,
pub cache_write_input_tokens: Option<u64>,
pub reasoning_output_tokens: Option<u64>,
pub image_input_tokens: Option<u64>,
pub image_output_tokens: Option<u64>,
pub total_tokens: Option<u64>,
@@ -750,9 +757,9 @@ pub struct UsageRecord {
pub retry_after_ms: Option<u64>,
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 cached_input_tokens: Option<u64>,
pub cache_write_input_tokens: Option<u64>,
pub reasoning_output_tokens: Option<u64>,
pub image_input_tokens: Option<u64>,
pub image_output_tokens: Option<u64>,
pub total_tokens: Option<u64>,
@@ -830,9 +837,9 @@ pub struct UsageAttempt {
pub message: 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 cached_input_tokens: Option<u64>,
pub cache_write_input_tokens: Option<u64>,
pub reasoning_output_tokens: Option<u64>,
pub total_tokens: Option<u64>,
pub cost_source: Option<String>,
pub cost_amount: Option<DecimalAmount>,
@@ -862,6 +869,7 @@ pub struct ProviderObservation {
/// 用量总览
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UsageOverview {
pub analytics: AnalyticsMetadata,
pub range: TimeRange,
pub requests: RequestMetrics,
pub attempts: AttemptMetrics,
@@ -878,6 +886,7 @@ pub struct UsageSummary {
/// 诊断聚合结果,分母包含截取展示项之前的全部匹配请求
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct DiagnosticsObservation {
pub analytics: AnalyticsMetadata,
pub total_request_count: u64,
pub items: Vec<DiagnosticObservation>,
}
@@ -1001,7 +1010,7 @@ pub struct TrendSummary {
pub request_count: u64,
pub input_tokens: u64,
pub output_tokens: u64,
pub cached_tokens: u64,
pub cached_input_tokens: u64,
pub total_tokens: u64,
pub service_failure_count: u64,
pub average_latency_ms: Option<u64>,
@@ -1019,6 +1028,7 @@ pub struct TrendSummary {
/// 趋势结果;API 再决定图表标签与数字格式
#[derive(Debug, Clone, PartialEq)]
pub struct Trend {
pub analytics: AnalyticsMetadata,
pub kind: TrendKind,
pub points: Vec<TrendPoint>,
pub summary: TrendSummary,
@@ -1134,7 +1144,7 @@ pub struct DashboardPeriodMetrics {
pub request_count: u64,
pub input_tokens: u64,
pub output_tokens: u64,
pub cached_tokens: u64,
pub cached_input_tokens: u64,
pub total_tokens: u64,
pub cached_token_rate: f64,
pub observed_cached_token_rate: Option<f64>,
@@ -1224,7 +1234,7 @@ pub struct UsageInsightsCostPoint {
pub bucket_start: DateTime<Utc>,
pub input_tokens: u64,
pub output_tokens: u64,
pub cached_tokens: u64,
pub cached_input_tokens: u64,
pub total_tokens: u64,
pub estimated_cost: Option<DecimalAmount>,
pub standard_cost: Option<DecimalAmount>,
@@ -1249,7 +1259,7 @@ pub struct UsageInsightsCost {
pub cache_hit_request_rate: Option<f64>,
pub input_tokens: u64,
pub output_tokens: u64,
pub cached_tokens: u64,
pub cached_input_tokens: u64,
pub total_tokens: u64,
pub points: Vec<UsageInsightsCostPoint>,
pub costs: Vec<CurrencyCost>,
@@ -1259,6 +1269,7 @@ pub struct UsageInsightsCost {
/// 用量洞察的完整用例结果;API 只负责标签、时区与字符串格式化
#[derive(Debug, Clone, PartialEq)]
pub struct UsageInsights {
pub analytics: AnalyticsMetadata,
pub granularity: Granularity,
pub health: UsageInsightsHealth,
pub performance: UsageInsightsPerformance,
@@ -1294,6 +1305,7 @@ pub struct DiagnosticsItem {
/// 诊断结果
#[derive(Debug, Clone, PartialEq)]
pub struct DiagnosticsResult {
pub analytics: AnalyticsMetadata,
pub dimension: DiagnosticDimension,
pub items: Vec<DiagnosticsItem>,
}
@@ -15,7 +15,7 @@ pub struct QuotaForecastUsage {
pub tokens: u64,
pub input_tokens: u64,
pub output_tokens: u64,
pub cached_tokens: u64,
pub cached_input_tokens: u64,
pub missing_token_count: u64,
pub known_cost_count: u64,
pub unavailable_cost_count: u64,
@@ -35,7 +35,9 @@ impl QuotaForecastUsage {
tokens: self.tokens.checked_sub(baseline.tokens)?,
input_tokens: self.input_tokens.checked_sub(baseline.input_tokens)?,
output_tokens: self.output_tokens.checked_sub(baseline.output_tokens)?,
cached_tokens: self.cached_tokens.checked_sub(baseline.cached_tokens)?,
cached_input_tokens: self
.cached_input_tokens
.checked_sub(baseline.cached_input_tokens)?,
missing_token_count: self
.missing_token_count
.checked_sub(baseline.missing_token_count)?,
@@ -6,7 +6,6 @@ use std::{collections::BTreeMap, net::IpAddr, sync::Arc, time::Duration};
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use futures::stream::BoxStream;
use super::backup::BackupStorePorts;
use crate::model::{
@@ -30,7 +29,8 @@ use crate::model::{
observability::{
DashboardObservation, DashboardQuery, DiagnosticDimension, DiagnosticsObservation,
Granularity, OpsErrorPage, OpsErrorQuery, RequestMetricPoint, TimeRange,
UsageCalculatedBillingFact, UsageDetail, UsageFilter, UsageOverview, UsagePage, UsageQuery,
UsageBillingProjection, UsageCalculatedBillingFact, UsageDetail, UsageFilter,
UsageInsightsObservation, UsageOverview, UsagePage, UsageQuery,
},
provider_credentials::{
AuthorizationCommit, CredentialDetails, CredentialImportCommit, CredentialImportResult,
@@ -421,13 +421,20 @@ pub trait AccountGroupStore: Send + Sync {
) -> AdminStoreResult<AccountGroupMutation>;
}
/// 逐条读取的已计算费用事实;消费结束或丢弃时释放查询资源
pub type UsageCalculatedBillingStream<'a> =
BoxStream<'a, AdminStoreResult<UsageCalculatedBillingFact>>;
/// 后台投影复用 Admin 与 Provider 的费用规则,Store 不重新解释价格
pub trait UsageBillingProjector: Send + Sync {
fn project(
&self,
fact: &UsageCalculatedBillingFact,
) -> Result<UsageBillingProjection, crate::model::AdminError>;
}
/// 用量、趋势、诊断与运维错误的只读能力
#[async_trait]
pub trait ObservabilityStore: Send + Sync {
/// 缓存必须与删除和后台投影提交的版本一致
async fn analytics_cache_generation(&self) -> AdminStoreResult<u64>;
/// 返回历史统计区间和指定观测时刻下的实时账号状态
async fn dashboard_summary(
&self,
@@ -448,14 +455,13 @@ pub trait ObservabilityStore: Send + Sync {
granularity: Granularity,
) -> AdminStoreResult<Vec<RequestMetricPoint>>;
/// 流式返回可由 Provider 重新校验的已计算费用事实,不保证顺序
/// 查询及解码错误由流返回;调用方应逐条聚合,避免收集整个区间
fn usage_calculated_billing_facts(
/// 同一读取快照中的汇总、趋势和已投影情景费用
async fn usage_insights_observation(
&self,
range: TimeRange,
filter: UsageFilter,
granularity: Granularity,
) -> UsageCalculatedBillingStream<'_>;
) -> AdminStoreResult<UsageInsightsObservation>;
async fn list_usage_records(&self, query: UsageQuery) -> AdminStoreResult<UsagePage>;
@@ -192,10 +192,12 @@ impl KeyUsageService for DefaultKeyUsageService {
.ok_or_else(|| AdminError::internal("时间超出支持范围"))?,
end: now,
};
let (overview, trend, health_points) = futures::try_join!(
self.observations.usage_summary(query.range, filter.clone()),
self.observations
.usage_trend(query.range, filter, Granularity::for_range(query.range)),
let (observation, health_points) = futures::try_join!(
self.observations.usage_insights_observation(
query.range,
filter,
Granularity::for_range(query.range)
),
self.observations.usage_trend(
today,
usage_filter(&id, None),
@@ -205,8 +207,8 @@ impl KeyUsageService for DefaultKeyUsageService {
.map_err(|error| map_store_error(error, "key usage overview"))?;
Ok(Some(KeyUsageOverview {
key,
overview,
trend,
overview: observation.overview,
trend: observation.trend,
health_timeline: health_timeline_at(&health_points, now, self.timezone)?,
}))
}
@@ -4,12 +4,10 @@ use std::{
collections::{BTreeMap, HashMap, btree_map::Entry},
str::FromStr,
sync::Arc,
time::{Duration as StdDuration, Instant},
};
use async_trait::async_trait;
use chrono::{DateTime, Duration, Utc};
use futures::TryStreamExt;
use gateway_core::{account::ProviderAccountId, routing::ProviderKind};
use crate::{
@@ -21,16 +19,16 @@ use crate::{
DiagnosticDimension, DiagnosticsItem, DiagnosticsResult, Granularity, HealthStatus,
HealthTimeline, HealthTimelinePoint, OpsErrorPage, OpsErrorQuery, ProviderBillingInput,
RequestMetricPoint, RequestMetrics, TimeRange, Trend, TrendKind, TrendPoint,
TrendSummary, UsageBilling, UsageCalculatedBillingFact, UsageDetail, UsageFilter,
UsageInsights, UsageInsightsCost, UsageInsightsCostPoint, UsageInsightsHealth,
UsageInsightsHealthPoint, UsageInsightsPerformance, UsageInsightsPerformancePoint,
UsageOverview, UsagePage, UsageQuery, UsageSummary,
TrendSummary, UsageBilling, UsageBillingProjection, UsageCalculatedBillingFact,
UsageDetail, UsageFilter, UsageInsights, UsageInsightsCost, UsageInsightsCostPoint,
UsageInsightsHealth, UsageInsightsHealthPoint, UsageInsightsPerformance,
UsageInsightsPerformancePoint, UsageOverview, UsagePage, UsageQuery, UsageSummary,
},
provider_credentials::ProviderQuotaRequest,
},
ports::{
provider::{ProviderAdminErrorKind, ProviderAdminRegistry},
store::{AccountStore, ObservabilityStore, SettingsStore, UsageCalculatedBillingStream},
store::{AccountStore, ObservabilityStore, SettingsStore, UsageBillingProjector},
},
};
@@ -39,59 +37,9 @@ use super::{map_provider_error, map_store_error};
const HEALTH_TIMELINE_SLOT_MINUTES: i64 = 15;
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,
)
}
}
const HEALTH_TIMELINE_STABLE_RELIABILITY: f64 = 95.0;
mod store_cache;
pub(crate) use store_cache::CachedObservabilityStore;
/// API 消费的观测控制面服务
#[async_trait]
@@ -130,7 +78,6 @@ pub(crate) struct DefaultObservabilityService {
accounts: Arc<dyn AccountStore>,
settings: Arc<dyn SettingsStore>,
providers: ProviderAdminRegistry,
dashboard_queries: DashboardQueryCache,
}
impl DefaultObservabilityService {
@@ -148,7 +95,6 @@ impl DefaultObservabilityService {
accounts,
settings,
providers,
dashboard_queries: DashboardQueryCache::default(),
}
}
@@ -185,7 +131,7 @@ impl DefaultObservabilityService {
let yesterday = dashboard_period_metrics(&observation.trend, yesterday_start, today_start);
let total_billing_usd = observation.totals.billing_usd.clone();
let total_cached_token_rate = rate_or_zero(
observation.totals.cached_tokens,
observation.totals.cached_input_tokens,
observation.totals.input_tokens,
);
let (first_token_latency_sum_ms, first_token_latency_count) = observation
@@ -199,7 +145,11 @@ impl DefaultObservabilityService {
});
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 trend = trend(
TrendKind::Usage,
observation.trend.clone(),
observation.analytics.clone(),
)?;
let health_timeline =
health_timeline_at(&observation.trend, observation.range.end, self.timezone)?;
let wire_profiles = self
@@ -247,15 +197,12 @@ impl ObservabilityService for DefaultObservabilityService {
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())?;
let mut result = self.load_dashboard_summary(range).await?;
result.trend = trend(
kind,
result.observation.trend.clone(),
result.observation.analytics.clone(),
)?;
Ok(result)
}
@@ -264,12 +211,16 @@ impl ObservabilityService for DefaultObservabilityService {
range: TimeRange,
kind: TrendKind,
) -> Result<Trend, AdminError> {
let points = self
let observation = self
.store
.dashboard_trend(range, Granularity::for_range(range))
.usage_insights_observation(
range,
UsageFilter::default(),
Granularity::for_range(range),
)
.await
.map_err(|error| map_store_error(error, "dashboard trend"))?;
trend(kind, points)
trend(kind, observation.trend, observation.overview.analytics)
}
async fn usage_records(&self, query: UsageQuery) -> Result<UsagePage, AdminError> {
@@ -321,21 +272,22 @@ impl ObservabilityService for DefaultObservabilityService {
filter: UsageFilter,
) -> Result<UsageInsights, AdminError> {
let granularity = Granularity::for_range(range);
let aggregates = async {
futures::try_join!(
self.store.usage_summary(range, filter.clone()),
self.store.usage_trend(range, filter.clone(), granularity),
)
.map_err(|error| map_store_error(error, "usage insights"))
};
let ((overview, trend), standard_costs) = futures::try_join!(
aggregates,
recover_standard_costs(
&self.providers,
self.store
.usage_calculated_billing_facts(range, filter.clone(), granularity),
),
)?;
let observation = self
.store
.usage_insights_observation(range, filter, granularity)
.await
.map_err(|error| map_store_error(error, "usage insights"))?;
let mut standard_costs = UsageCostScenarios::default();
for point in observation.costs {
if let Some(amount) = point.standard_amount {
standard_costs.standard.add(point.bucket_start, amount)?;
}
if let Some(amount) = point.no_cache_amount {
standard_costs.no_cache.add(point.bucket_start, amount)?;
}
}
let overview = observation.overview;
let trend = observation.trend;
build_usage_insights(overview, trend, standard_costs)
}
@@ -389,7 +341,11 @@ impl ObservabilityService for DefaultObservabilityService {
}
})
.collect::<Vec<_>>();
Ok(DiagnosticsResult { dimension, items })
Ok(DiagnosticsResult {
analytics: observation.analytics,
dimension,
items,
})
}
async fn ops_errors(&self, query: OpsErrorQuery) -> Result<OpsErrorPage, AdminError> {
@@ -423,8 +379,9 @@ fn build_usage_insights(
}
let requests = &overview.requests;
let failed_requests = service_failure_count(requests);
let health_requests = requests.success_count.saturating_add(failed_requests);
let health_window = HealthWindow::from_metrics(requests);
let failed_requests = health_window.failed_requests;
let health_requests = health_window.eligible_requests();
let health = UsageInsightsHealth {
total_requests: requests.request_count,
success_requests: requests.success_count,
@@ -442,17 +399,19 @@ fn build_usage_insights(
points: trend
.iter()
.map(|point| {
let failed_requests = service_failure_count(&point.metrics);
let health_requests = point.metrics.success_count.saturating_add(failed_requests);
let health_window = HealthWindow::from_metrics(&point.metrics);
UsageInsightsHealthPoint {
total_requests: point.metrics.request_count,
bucket_start: point.bucket_start,
success_requests: point.metrics.success_count,
failed_requests,
failed_requests: health_window.failed_requests,
cancelled_requests: point.metrics.cancelled_count,
incomplete_requests: point.metrics.incomplete_count,
caller_error_requests: point.metrics.caller_error_count,
error_rate: rate_or_zero(failed_requests, health_requests),
error_rate: rate_or_zero(
health_window.non_success_requests(),
health_window.eligible_requests(),
),
}
})
.collect(),
@@ -533,14 +492,14 @@ fn build_usage_insights(
tier_premium,
tokens_per_request: rate_or_zero(requests.total_tokens, requests.request_count),
cost_per_successful_request,
cached_token_rate: rate_or_zero(requests.cached_tokens, requests.input_tokens),
cached_token_rate: rate_or_zero(requests.cached_input_tokens, requests.input_tokens),
cache_hit_request_rate: ratio(
requests.cache_hit_request_count,
requests.cache_eligible_request_count,
),
input_tokens: requests.input_tokens,
output_tokens: requests.output_tokens,
cached_tokens: requests.cached_tokens,
cached_input_tokens: requests.cached_input_tokens,
total_tokens: requests.total_tokens,
points: trend
.iter()
@@ -548,7 +507,7 @@ fn build_usage_insights(
bucket_start: point.bucket_start,
input_tokens: point.metrics.input_tokens,
output_tokens: point.metrics.output_tokens,
cached_tokens: point.metrics.cached_tokens,
cached_input_tokens: point.metrics.cached_input_tokens,
total_tokens: point.metrics.total_tokens,
estimated_cost: usd_cost(&point.costs),
standard_cost: standard_costs_by_bucket.get(&point.bucket_start).cloned(),
@@ -558,7 +517,7 @@ fn build_usage_insights(
usd_cost(&point.costs).as_ref(),
),
cached_token_rate: rate_or_zero(
point.metrics.cached_tokens,
point.metrics.cached_input_tokens,
point.metrics.input_tokens,
),
cache_hit_request_rate: ratio(
@@ -571,6 +530,7 @@ fn build_usage_insights(
coverage: overview.attempts.cost_coverage.clone(),
};
Ok(UsageInsights {
analytics: overview.analytics,
granularity,
health,
performance,
@@ -620,58 +580,65 @@ struct UsageCostScenarios {
no_cache: ScenarioCosts,
}
async fn recover_standard_costs(
providers: &ProviderAdminRegistry,
mut facts: UsageCalculatedBillingStream<'_>,
) -> Result<UsageCostScenarios, AdminError> {
let mut scenarios = UsageCostScenarios::default();
while let Some(mut fact) = facts
.try_next()
.await
.map_err(|error| map_store_error(error, "usage billing facts"))?
{
/// 在后台投影时恢复费用场景,查询不再逐请求读取计费快照
pub struct ProviderUsageBillingProjector {
providers: ProviderAdminRegistry,
}
impl ProviderUsageBillingProjector {
#[must_use]
pub fn new(providers: ProviderAdminRegistry) -> Self {
Self { providers }
}
}
impl UsageBillingProjector for ProviderUsageBillingProjector {
fn project(
&self,
fact: &UsageCalculatedBillingFact,
) -> Result<UsageBillingProjection, AdminError> {
let Ok(provider_kind) = ProviderKind::new(fact.provider_kind.clone()) else {
continue;
return Ok(UsageBillingProjection::default());
};
let input = ProviderBillingInput {
upstream_model_id: fact.upstream_model_id.clone(),
service_tier: fact.service_tier.clone(),
input_tokens: fact.input_tokens,
output_tokens: fact.output_tokens,
cached_tokens: fact.cached_tokens,
cache_write_tokens: fact.cache_write_tokens,
cached_input_tokens: fact.cached_input_tokens,
cache_write_input_tokens: fact.cache_write_input_tokens,
total: fact.total.clone(),
};
let breakdown = match fact
.breakdown
.take()
.filter(|b| b.total_amount == input.total)
.as_ref()
.filter(|breakdown| breakdown.total_amount == input.total)
.cloned()
.map_or_else(
|| providers.calculated_billing(&provider_kind, &input),
|b| Ok(Some(b)),
|| self.providers.calculated_billing(&provider_kind, &input),
|breakdown| Ok(Some(breakdown)),
) {
Ok(breakdown) => breakdown,
Err(error) if error.kind() == ProviderAdminErrorKind::Unsupported => continue,
Err(error) if error.kind() == ProviderAdminErrorKind::Unsupported => {
return Ok(UsageBillingProjection::default());
}
Err(error) => return Err(map_provider_error(error, "usage billing")),
};
let Some(breakdown) = breakdown else {
continue;
return Ok(UsageBillingProjection::default());
};
if !breakdown
.standard_amount
.currency
.eq_ignore_ascii_case("USD")
{
continue;
}
scenarios
.standard
.add(fact.bucket_start, breakdown.standard_amount.amount.clone())?;
if let Some(no_cache_cost) = no_cache_cost(&fact, &breakdown) {
scenarios.no_cache.add(fact.bucket_start, no_cache_cost)?;
return Ok(UsageBillingProjection::default());
}
Ok(UsageBillingProjection {
standard_amount: Some(breakdown.standard_amount.amount.clone()),
no_cache_amount: no_cache_cost(fact, &breakdown),
})
}
Ok(scenarios)
}
fn no_cache_cost(
@@ -682,19 +649,21 @@ fn no_cache_cost(
let input_tokens = fact
.input_tokens?
.checked_sub(image.map_or(0, |i| i.input_tokens))?;
let cached_tokens = fact
.cached_tokens
let cached_input_tokens = fact
.cached_input_tokens
.unwrap_or_default()
.checked_sub(image.map_or(0, |i| i.cached_tokens))?
.checked_sub(image.map_or(0, |i| i.cached_input_tokens))?
.min(input_tokens);
let cache_write_tokens = fact
.cache_write_tokens
let cache_write_input_tokens = fact
.cache_write_input_tokens
.unwrap_or_default()
.min(input_tokens.saturating_sub(cached_tokens));
.min(input_tokens.saturating_sub(cached_input_tokens));
let input_rate = decimal(&breakdown.input_price_per_million.amount)?;
let replaced_input_amount = input_rate
.scaled()
.checked_mul(u128::from(cached_tokens.saturating_add(cache_write_tokens)))?
.checked_mul(u128::from(
cached_input_tokens.saturating_add(cache_write_input_tokens),
))?
.checked_div(1_000_000)?;
let mut total = decimal(&breakdown.total_amount.amount)?
.scaled()
@@ -704,7 +673,7 @@ fn no_cache_cost(
if let Some(image) = image {
let image_no_cache = decimal(&image.input_price_per_million.amount)?
.scaled()
.checked_mul(u128::from(image.cached_tokens))?
.checked_mul(u128::from(image.cached_input_tokens))?
.checked_div(1_000_000)?;
total = total
.checked_sub(decimal(&image.cache_read_amount.amount)?.scaled())?
@@ -829,8 +798,8 @@ impl DefaultObservabilityService {
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,
cached_input_tokens: record.cached_input_tokens,
cache_write_input_tokens: record.cache_write_input_tokens,
total: total.clone(),
};
if let Ok(Some(breakdown)) = self.providers.calculated_billing(&provider_kind, &input) {
@@ -861,8 +830,8 @@ impl DefaultObservabilityService {
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,
cached_input_tokens: record.cached_input_tokens,
cache_write_input_tokens: record.cache_write_input_tokens,
total: total.clone(),
};
if let Ok(Some(breakdown)) = self.providers.calculated_billing(&provider_kind, &input) {
@@ -951,7 +920,11 @@ pub(super) fn health_timeline_at(
})
}
fn trend(kind: TrendKind, points: Vec<RequestMetricPoint>) -> Result<Trend, AdminError> {
fn trend(
kind: TrendKind,
points: Vec<RequestMetricPoint>,
analytics: crate::model::analytics::AnalyticsMetadata,
) -> Result<Trend, AdminError> {
let mut metrics = RequestMetrics::default();
let mut coverage = CostCoverage::default();
let mut costs = BTreeMap::<String, DecimalAmount>::new();
@@ -986,7 +959,7 @@ fn trend(kind: TrendKind, points: Vec<RequestMetricPoint>) -> Result<Trend, Admi
let service_failure_count = service_failure_count(&metrics);
let success_rate = ratio(
metrics.success_count,
metrics.success_count.saturating_add(service_failure_count),
HealthWindow::from_metrics(&metrics).eligible_requests(),
);
let cache_hit_request_rate = ratio(
metrics.cache_hit_request_count,
@@ -1007,13 +980,14 @@ fn trend(kind: TrendKind, points: Vec<RequestMetricPoint>) -> Result<Trend, Admi
.filter_map(|point| point.metrics.output_throughput_p50)
.min();
Ok(Trend {
analytics,
kind,
points: points.into_iter().map(trend_point).collect(),
summary: TrendSummary {
request_count: metrics.request_count,
input_tokens: metrics.input_tokens,
output_tokens: metrics.output_tokens,
cached_tokens: metrics.cached_tokens,
cached_input_tokens: metrics.cached_input_tokens,
total_tokens: metrics.total_tokens,
service_failure_count,
average_latency_ms,
@@ -1044,12 +1018,18 @@ fn trend_point(point: RequestMetricPoint) -> TrendPoint {
point.metrics.first_token_latency_sum_ms,
point.metrics.first_token_latency_count,
),
cached_token_rate: rate_or_zero(point.metrics.cached_tokens, point.metrics.input_tokens),
cached_token_rate: rate_or_zero(
point.metrics.cached_input_tokens,
point.metrics.input_tokens,
),
cache_hit_request_rate: ratio(
point.metrics.cache_hit_request_count,
point.metrics.cache_eligible_request_count,
),
success_rate: ratio(point.metrics.success_count, point.metrics.request_count),
success_rate: ratio(
point.metrics.success_count,
HealthWindow::from_metrics(&point.metrics).eligible_requests(),
),
metrics: point.metrics,
cost_coverage: point.cost_coverage,
costs: point.costs,
@@ -1070,11 +1050,13 @@ fn dashboard_period_metrics(
result.request_count = result.request_count.saturating_add(metrics.request_count);
result.input_tokens = result.input_tokens.saturating_add(metrics.input_tokens);
result.output_tokens = result.output_tokens.saturating_add(metrics.output_tokens);
result.cached_tokens = result.cached_tokens.saturating_add(metrics.cached_tokens);
result.cached_input_tokens = result
.cached_input_tokens
.saturating_add(metrics.cached_input_tokens);
result.total_tokens = result.total_tokens.saturating_add(metrics.total_tokens);
}
result.cached_token_rate = rate_or_zero(result.cached_tokens, result.input_tokens);
result.observed_cached_token_rate = ratio(result.cached_tokens, result.input_tokens);
result.cached_token_rate = rate_or_zero(result.cached_input_tokens, result.input_tokens);
result.observed_cached_token_rate = ratio(result.cached_input_tokens, result.input_tokens);
result
}
@@ -1091,13 +1073,15 @@ fn add_request_metrics(total: &mut RequestMetrics, value: &RequestMetrics) {
.saturating_add(value.caller_error_count);
total.input_tokens = total.input_tokens.saturating_add(value.input_tokens);
total.output_tokens = total.output_tokens.saturating_add(value.output_tokens);
total.cached_tokens = total.cached_tokens.saturating_add(value.cached_tokens);
total.cache_write_tokens = total
.cache_write_tokens
.saturating_add(value.cache_write_tokens);
total.reasoning_tokens = total
.reasoning_tokens
.saturating_add(value.reasoning_tokens);
total.cached_input_tokens = total
.cached_input_tokens
.saturating_add(value.cached_input_tokens);
total.cache_write_input_tokens = total
.cache_write_input_tokens
.saturating_add(value.cache_write_input_tokens);
total.reasoning_output_tokens = total
.reasoning_output_tokens
.saturating_add(value.reasoning_output_tokens);
total.total_tokens = total.total_tokens.saturating_add(value.total_tokens);
total.first_token_latency_sum_ms = total
.first_token_latency_sum_ms
@@ -1145,20 +1129,29 @@ struct HealthWindow {
}
impl HealthWindow {
fn from_metrics(metrics: &RequestMetrics) -> Self {
Self {
success_requests: metrics.success_count,
failed_requests: service_failure_count(metrics),
cancelled_requests: metrics.cancelled_count,
incomplete_requests: metrics.incomplete_count,
caller_error_requests: metrics.caller_error_count,
}
}
// 未完成保留交付终态,但上游失败、超时和进程中断都不能视为服务成功
fn non_success_requests(self) -> u64 {
self.failed_requests
.saturating_add(self.incomplete_requests)
}
fn eligible_requests(self) -> u64 {
self.success_requests
.saturating_add(self.non_success_requests())
}
fn add_metrics(&mut self, metrics: &RequestMetrics) {
self.success_requests = self.success_requests.saturating_add(metrics.success_count);
self.failed_requests = self
.failed_requests
.saturating_add(service_failure_count(metrics));
self.cancelled_requests = self
.cancelled_requests
.saturating_add(metrics.cancelled_count);
self.incomplete_requests = self
.incomplete_requests
.saturating_add(metrics.incomplete_count);
self.caller_error_requests = self
.caller_error_requests
.saturating_add(metrics.caller_error_count);
self.add_window(Self::from_metrics(metrics));
}
fn add_window(&mut self, other: Self) {
@@ -1177,15 +1170,13 @@ impl HealthWindow {
}
fn health_status(bucket: HealthWindow, is_future: bool) -> HealthStatus {
let eligible_requests = bucket
.success_requests
.saturating_add(bucket.failed_requests);
let eligible_requests = bucket.eligible_requests();
if is_future {
HealthStatus::Future
} else if eligible_requests == 0 {
HealthStatus::NoData
} else if bucket.success_requests == 0
&& bucket.failed_requests >= HEALTH_TIMELINE_UNAVAILABLE_FAILURE_THRESHOLD
&& bucket.non_success_requests() >= HEALTH_TIMELINE_UNAVAILABLE_FAILURE_THRESHOLD
{
HealthStatus::Unavailable
} else if eligible_requests < HEALTH_TIMELINE_MIN_SAMPLE_SIZE {
@@ -1200,9 +1191,7 @@ fn health_status(bucket: HealthWindow, is_future: bool) -> HealthStatus {
}
fn health_reliability(bucket: HealthWindow) -> Option<f64> {
let eligible_requests = bucket
.success_requests
.saturating_add(bucket.failed_requests);
let eligible_requests = bucket.eligible_requests();
(eligible_requests > 0)
.then(|| bucket.success_requests as f64 / eligible_requests as f64 * 100.0)
}
@@ -0,0 +1,298 @@
//! 有界统计缓存:身份和完整筛选参与键,删除版本阻止旧数据复活
use crate::model::observability::*;
use crate::ports::store::{
AdminStoreError, AdminStoreErrorKind, AdminStoreResult, ObservabilityStore,
};
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use std::{collections::HashMap, future::Future, hash::Hash, sync::Arc, time::Duration};
use tokio::{sync::Mutex, time::Instant};
const CAPACITY: usize = 64;
const TTL: Duration = Duration::from_secs(15);
const MAX_STALE: Duration = Duration::from_secs(60);
const REQUEST_TIMEOUT: Duration = Duration::from_secs(10);
type Slot<T> = Arc<Mutex<Option<Cached<T>>>>;
struct Cached<T> {
value: T,
created: Instant,
}
struct Entry<T> {
slot: Slot<T>,
touched: Instant,
}
struct Cache<K, T> {
entries: Mutex<HashMap<K, Entry<T>>>,
}
impl<K, T> Default for Cache<K, T> {
fn default() -> Self {
Self {
entries: Mutex::new(HashMap::new()),
}
}
}
trait AnalyticsValue: Clone {
fn metadata(&mut self) -> &mut AnalyticsMetadata;
}
impl AnalyticsValue for DashboardObservation {
fn metadata(&mut self) -> &mut AnalyticsMetadata {
&mut self.analytics
}
}
impl AnalyticsValue for UsageOverview {
fn metadata(&mut self) -> &mut AnalyticsMetadata {
&mut self.analytics
}
}
impl AnalyticsValue for UsageInsightsObservation {
fn metadata(&mut self) -> &mut AnalyticsMetadata {
&mut self.overview.analytics
}
}
impl AnalyticsValue for DiagnosticsObservation {
fn metadata(&mut self) -> &mut AnalyticsMetadata {
&mut self.analytics
}
}
impl<K: Eq + Hash, T: AnalyticsValue> Cache<K, T> {
async fn get<F, Fut>(
&self,
key: K,
store: &dyn ObservabilityStore,
ttl: Duration,
load: F,
) -> AdminStoreResult<T>
where
F: FnOnce() -> Fut,
Fut: Future<Output = AdminStoreResult<T>>,
{
// 刷新失败后的同键等待也计入预算,避免串行重试把延迟放大为队列长度倍数
tokio::time::timeout(
REQUEST_TIMEOUT,
self.get_within_budget(key, store, ttl, load),
)
.await
.map_err(|source| {
AdminStoreError::new(
AdminStoreErrorKind::Unavailable,
"observability cache",
"observability cache execution budget is exhausted",
)
.with_source(source)
})?
}
async fn get_within_budget<F, Fut>(
&self,
key: K,
store: &dyn ObservabilityStore,
ttl: Duration,
load: F,
) -> AdminStoreResult<T>
where
F: FnOnce() -> Fut,
Fut: Future<Output = AdminStoreResult<T>>,
{
let slot = {
let mut entries = self.entries.lock().await;
if !entries.contains_key(&key) && entries.len() >= CAPACITY {
let oldest = entries
.iter()
.filter(|(_, entry)| Arc::strong_count(&entry.slot) == 1)
.min_by_key(|(_, entry)| entry.touched)
.map(|(_, entry)| entry.touched);
if let Some(oldest) = oldest {
entries.retain(|_, entry| entry.touched != oldest);
}
}
if !entries.contains_key(&key) && entries.len() >= CAPACITY {
drop(entries);
return load().await;
}
let entry = entries.entry(key).or_insert_with(|| Entry {
slot: Arc::new(Mutex::new(None)),
touched: Instant::now(),
});
entry.touched = Instant::now();
entry.slot.clone()
};
// 同键刷新只执行一次;失败不会缓存成成功,也不延长旧结果寿命
let mut cached = slot.lock().await;
// 等锁期间代次可能推进;取得同键锁后再核对,失败时也不能回用已删除事实
let generation = store.analytics_cache_generation().await?;
if let Some(entry) = cached.as_mut()
&& entry.value.metadata().generation == generation
&& entry.created.elapsed() < ttl
{
return Ok(with_age(entry.value.clone(), false));
}
match load().await {
Ok(value) => {
*cached = Some(Cached {
value: value.clone(),
created: Instant::now(),
});
Ok(value)
}
Err(error) => {
if ttl == TTL
&& error.kind() == AdminStoreErrorKind::Unavailable
&& let Some(entry) = cached.as_mut()
&& entry.created.elapsed() <= MAX_STALE
&& Utc::now()
.signed_duration_since(entry.value.metadata().generated_at)
.to_std()
.unwrap_or_default()
<= MAX_STALE
&& entry.value.metadata().state != AnalyticsState::Building
&& entry.value.metadata().generation
== store.analytics_cache_generation().await?
{
return Ok(with_age(entry.value.clone(), true));
}
Err(error)
}
}
}
}
fn with_age<T: AnalyticsValue>(mut value: T, stale: bool) -> T {
let metadata = value.metadata();
metadata.lag_seconds = metadata.lag_seconds.saturating_add(
Utc::now()
.signed_duration_since(metadata.generated_at)
.num_seconds()
.max(0) as u64,
);
if stale {
metadata.state = AnalyticsState::Stale;
}
value
}
fn normalized(mut filter: UsageFilter) -> UsageFilter {
filter.routing_group_refs.sort_unstable();
filter.routing_group_refs.dedup();
filter
}
type UsageKey = (TimeRange, UsageFilter);
pub(crate) struct CachedObservabilityStore {
inner: Arc<dyn ObservabilityStore>,
dashboard: Cache<DashboardQuery, DashboardObservation>,
summaries: Cache<UsageKey, UsageOverview>,
insights: Cache<(UsageKey, Granularity), UsageInsightsObservation>,
diagnostics: Cache<(UsageKey, DiagnosticDimension, u16), DiagnosticsObservation>,
}
impl CachedObservabilityStore {
pub(crate) fn new(inner: Arc<dyn ObservabilityStore>) -> Self {
Self {
inner,
dashboard: Cache::default(),
summaries: Cache::default(),
insights: Cache::default(),
diagnostics: Cache::default(),
}
}
}
#[async_trait]
impl ObservabilityStore for CachedObservabilityStore {
async fn analytics_cache_generation(&self) -> AdminStoreResult<u64> {
self.inner.analytics_cache_generation().await
}
async fn dashboard_summary(
&self,
query: DashboardQuery,
observed_at: DateTime<Utc>,
) -> AdminStoreResult<DashboardObservation> {
// 复合概览还包含账号容量和最近请求,保持原有两秒实时刷新上限
self.dashboard
.get(
query.clone(),
self.inner.as_ref(),
Duration::from_secs(2),
|| self.inner.dashboard_summary(query, observed_at),
)
.await
}
async fn dashboard_trend(
&self,
range: TimeRange,
granularity: Granularity,
) -> AdminStoreResult<Vec<RequestMetricPoint>> {
self.inner.dashboard_trend(range, granularity).await
}
async fn usage_trend(
&self,
range: TimeRange,
filter: UsageFilter,
granularity: Granularity,
) -> AdminStoreResult<Vec<RequestMetricPoint>> {
self.inner.usage_trend(range, filter, granularity).await
}
async fn usage_insights_observation(
&self,
range: TimeRange,
filter: UsageFilter,
granularity: Granularity,
) -> AdminStoreResult<UsageInsightsObservation> {
let filter = normalized(filter);
self.insights
.get(
((range, filter.clone()), granularity),
self.inner.as_ref(),
TTL,
|| {
self.inner
.usage_insights_observation(range, filter, granularity)
},
)
.await
}
async fn usage_summary(
&self,
range: TimeRange,
filter: UsageFilter,
) -> AdminStoreResult<UsageOverview> {
let filter = normalized(filter);
self.summaries
.get((range, filter.clone()), self.inner.as_ref(), TTL, || {
self.inner.usage_summary(range, filter)
})
.await
}
async fn usage_diagnostics(
&self,
range: TimeRange,
filter: UsageFilter,
dimension: DiagnosticDimension,
limit: u16,
) -> AdminStoreResult<DiagnosticsObservation> {
let filter = normalized(filter);
self.diagnostics
.get(
((range, filter.clone()), dimension, limit),
self.inner.as_ref(),
TTL,
|| {
self.inner
.usage_diagnostics(range, filter, dimension, limit)
},
)
.await
}
async fn list_usage_records(&self, query: UsageQuery) -> AdminStoreResult<UsagePage> {
self.inner.list_usage_records(query).await
}
async fn usage_record_detail(&self, request_id: &str) -> AdminStoreResult<UsageDetail> {
self.inner.usage_record_detail(request_id).await
}
async fn list_ops_errors(&self, query: OpsErrorQuery) -> AdminStoreResult<OpsErrorPage> {
self.inner.list_ops_errors(query).await
}
}
@@ -260,10 +260,12 @@ impl SettingsService for DefaultSettingsService {
PricingChange::Replace(pricing) => {
pricing.validate().map_err(AdminError::invalid)?;
if command.provider == "xai"
&& pricing
.bands
.keys()
.any(|band| matches!(band.as_str(), "image" | "flex" | "long_flex"))
&& pricing.bands.keys().any(|band| {
matches!(
band.as_str(),
"image" | "flex" | "long_flex" | "ultrafast" | "long_ultrafast"
)
})
{
return Err(AdminError::invalid("xAI 不支持此价格档位"));
}
+3 -3
View File
@@ -170,9 +170,9 @@ fn usage_window_should_not_use_daily_rolling_usage_for_weekly_statistics() {
success_count: 5,
input_tokens: Some(100),
output_tokens: Some(20),
cached_tokens: Some(80),
cache_write_tokens: Some(0),
reasoning_tokens: Some(0),
cached_input_tokens: Some(80),
cache_write_input_tokens: Some(0),
reasoning_output_tokens: Some(0),
image_input_tokens: Some(0),
image_output_tokens: Some(0),
image_request_count: 0,
@@ -115,9 +115,9 @@ fn usage_window(key: &str, group: &str, seconds: u64) -> ProviderQuotaWindow {
success_count: 0,
input_tokens: None,
output_tokens: None,
cached_tokens: None,
cache_write_tokens: None,
reasoning_tokens: None,
cached_input_tokens: None,
cache_write_input_tokens: None,
reasoning_output_tokens: None,
image_input_tokens: None,
image_output_tokens: None,
image_request_count: 0,
@@ -32,9 +32,9 @@ fn window(key: &str, days: u64) -> ProviderQuotaWindow {
success_count: 10,
input_tokens: Some(900),
output_tokens: Some(100),
cached_tokens: Some(800),
cache_write_tokens: None,
reasoning_tokens: None,
cached_input_tokens: Some(800),
cache_write_input_tokens: None,
reasoning_output_tokens: None,
image_input_tokens: None,
image_output_tokens: None,
image_request_count: 0,
@@ -84,7 +84,7 @@ fn samples(quota: &ProviderQuota) -> Vec<QuotaForecastSample> {
tokens: usage.total_tokens.unwrap_or(0),
input_tokens: usage.input_tokens.unwrap_or(0),
output_tokens: usage.output_tokens.unwrap_or(0),
cached_tokens: usage.cached_tokens.unwrap_or(0),
cached_input_tokens: usage.cached_input_tokens.unwrap_or(0),
missing_token_count: u64::from(usage.total_tokens.is_none()),
known_cost_count: if usd.is_some() {
usage.cost_coverage.known_count()
@@ -3032,9 +3032,9 @@ fn quota_local_usage(account_id: &str, total_tokens: u64) -> AccountUsage {
success_count: 1,
input_tokens: Some(total_tokens),
output_tokens: Some(0),
cached_tokens: Some(0),
cache_write_tokens: Some(0),
reasoning_tokens: Some(0),
cached_input_tokens: Some(0),
cache_write_input_tokens: Some(0),
reasoning_output_tokens: Some(0),
image_input_tokens: Some(0),
image_output_tokens: Some(0),
image_request_count: 0,
@@ -772,15 +772,17 @@ impl ObservabilityStore for UnavailableStore {
Err(unavailable("usage trend"))
}
fn usage_calculated_billing_facts(
async fn analytics_cache_generation(&self) -> AdminStoreResult<u64> {
Err(unavailable("analytics generation"))
}
async fn usage_insights_observation(
&self,
_: TimeRange,
_: UsageFilter,
_: Granularity,
) -> gateway_admin::ports::store::UsageCalculatedBillingStream<'_> {
Box::pin(futures::stream::once(async {
Err(unavailable("usage billing facts"))
}))
) -> AdminStoreResult<gateway_admin::model::observability::UsageInsightsObservation> {
Err(unavailable("usage insights"))
}
async fn list_usage_records(&self, _: UsageQuery) -> AdminStoreResult<UsagePage> {
@@ -17,13 +17,13 @@ use gateway_admin::{
model::{
MutationContext, Revision,
observability::{
AccountPoolMetrics, AttemptMetrics, CostCoverage, CurrencyCost, DashboardAccountUsage,
DashboardObservation, DashboardQuery, DashboardRuntimeSlots, DiagnosticDimension,
DiagnosticObservation, DiagnosticsObservation, Granularity, HealthStatus,
LatencyPercentiles, ObservabilityPageSize, OpsErrorPage, OpsErrorQuery,
AccountPoolMetrics, AnalyticsState, AttemptMetrics, CostCoverage, CurrencyCost,
DashboardAccountUsage, DashboardObservation, DashboardQuery, DashboardRuntimeSlots,
DiagnosticDimension, DiagnosticObservation, DiagnosticsObservation, Granularity,
HealthStatus, LatencyPercentiles, ObservabilityPageSize, OpsErrorPage, OpsErrorQuery,
PercentileMilliseconds, RequestMetricPoint, RequestMetrics, TimeRange, TrendKind,
UsageBilling, UsageCalculatedBillingFact, UsageDetail, UsageFilter, UsageListRecord,
UsageOverview, UsagePage, UsageQuery,
UsageBilling, UsageDetail, UsageFilter, UsageInsightsObservation, UsageListRecord,
UsageOverview, UsagePage, UsageProjectedCostPoint, UsageQuery,
},
settings::{
AdminApiKey, AdminApiKeyMutation, ReplaceRuntimeSettings, RotationStrategy,
@@ -148,7 +148,7 @@ async fn dashboard_summary_should_derive_first_token_average_from_trend() {
}
#[tokio::test]
async fn health_timeline_should_match_legacy_status_precedence_and_thresholds() {
async fn health_timeline_should_apply_status_precedence_and_thresholds() {
let now = Utc::now();
let current_slot = quarter_hour_start(now);
let store = Arc::new(FixtureObservabilityStore::new(observation_range(now)));
@@ -158,13 +158,28 @@ async fn health_timeline_should_match_legacy_status_precedence_and_thresholds()
RequestMetrics {
failure_count: 1,
cancelled_count: 1,
incomplete_count: 1,
caller_error_count: 1,
..RequestMetrics::default()
},
HealthStatus::NoData,
None,
),
(
RequestMetrics {
incomplete_count: 1,
..RequestMetrics::default()
},
HealthStatus::LowSample,
Some(0.0),
),
(
RequestMetrics {
incomplete_count: 3,
..RequestMetrics::default()
},
HealthStatus::Unavailable,
Some(0.0),
),
(
RequestMetrics {
failure_count: 3,
@@ -184,21 +199,21 @@ async fn health_timeline_should_match_legacy_status_precedence_and_thresholds()
),
(
RequestMetrics {
success_count: 98,
failure_count: 2,
success_count: 94,
failure_count: 6,
..RequestMetrics::default()
},
HealthStatus::Unstable,
Some(98.0),
Some(94.0),
),
(
RequestMetrics {
success_count: 99,
failure_count: 1,
success_count: 95,
incomplete_count: 5,
..RequestMetrics::default()
},
HealthStatus::Stable,
Some(99.0),
Some(95.0),
),
];
@@ -224,6 +239,92 @@ async fn health_timeline_should_match_legacy_status_precedence_and_thresholds()
}
}
#[tokio::test]
async fn health_views_should_include_incomplete_and_keep_final_outcomes_separate() {
let now = Utc::now();
let current_slot = quarter_hour_start(now);
let store = Arc::new(FixtureObservabilityStore::new(observation_range(now)));
let services = observability_services(store.clone()).await;
let fixtures = [
(162, 0, 2, 0, 0, 98.780_487_804_878_05),
(297, 0, 2, 5, 0, 99.331_103_678_929_76),
(467, 1, 5, 0, 0, 98.731_501_057_082_45),
(95, 2, 3, 4, 1, 95.0),
(100, 0, 0, 0, 0, 100.0),
];
for (index, (success, failed, incomplete, cancelled, caller_error, expected_percent)) in
fixtures.into_iter().enumerate()
{
let range = observation_range(
now + Duration::seconds(i64::try_from(index).expect("fixture index fits i64") * 2),
);
let metrics = RequestMetrics {
request_count: success + failed + incomplete + cancelled + caller_error,
success_count: success,
failure_count: failed + caller_error,
incomplete_count: incomplete,
cancelled_count: cancelled,
caller_error_count: caller_error,
..RequestMetrics::default()
};
store.replace_trend(vec![health_metric_point(current_slot, metrics.clone())]);
store.replace_overview(UsageOverview {
analytics: Default::default(),
range,
requests: metrics,
// 中间尝试失败不会增加最终失败请求数
attempts: AttemptMetrics {
failure_count: 2,
..AttemptMetrics::default()
},
providers: Vec::new(),
});
let dashboard = services
.observability()
.dashboard_summary(range, TrendKind::Usage)
.await
.expect("dashboard summary");
let timeline = dashboard.health_timeline;
let point = timeline
.points
.iter()
.find(|point| point.bucket_start == current_slot)
.expect("current health slot");
assert_eq!(point.status, HealthStatus::Stable);
assert_eq!(point.failed_requests, failed);
assert_eq!(point.incomplete_requests, incomplete);
assert_eq!(point.cancelled_requests, cancelled);
assert_eq!(point.caller_error_requests, caller_error);
for reliability in [timeline.reliability_percent, point.reliability_percent] {
assert!((reliability.expect("health reliability") - expected_percent).abs() < 1e-10);
}
for success_rate in [
dashboard.trend.summary.success_rate,
dashboard.trend.points[0].success_rate,
] {
assert!(
(success_rate.expect("trend success rate") * 100.0 - expected_percent).abs()
< 1e-10
);
}
let insights = services
.observability()
.usage_insights(range, UsageFilter::default())
.await
.expect("usage insights");
assert_eq!(insights.health.failed_requests, failed);
assert_eq!(insights.health.incomplete_requests, incomplete);
assert!((insights.health.success_rate * 100.0 - expected_percent).abs() < 1e-10);
assert!(
(insights.health.points[0].error_rate * 100.0 - (100.0 - expected_percent)).abs()
< 1e-10
);
}
}
#[tokio::test]
async fn dashboard_summary_should_project_rebuildable_runtime_slots() {
let now = Utc::now();
@@ -300,7 +401,7 @@ async fn dashboard_summary_should_observe_current_accounts_for_a_historical_rang
}
#[tokio::test]
async fn dashboard_summary_should_coalesce_concurrent_requests_in_one_short_window() {
async fn dashboard_summary_should_keep_distinct_exact_ranges_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");
@@ -330,7 +431,7 @@ async fn dashboard_summary_should_coalesce_concurrent_requests_in_one_short_wind
first.expect("first dashboard summary");
second.expect("second dashboard summary");
assert_eq!(store.dashboard_summary_calls(), 1);
assert_eq!(store.dashboard_summary_calls(), 2);
}
#[tokio::test]
@@ -403,6 +504,7 @@ async fn observability_services_should_calculate_usage_insights_and_diagnostic_s
..RequestMetrics::default()
};
store.replace_overview(UsageOverview {
analytics: Default::default(),
range,
requests: metrics.clone(),
attempts: AttemptMetrics {
@@ -431,23 +533,13 @@ async fn observability_services_should_calculate_usage_insights_and_diagnostic_s
.expect("bucket USD cost"),
}],
}]);
store.replace_calculated_billing_facts(vec![UsageCalculatedBillingFact {
breakdown: None,
store.replace_projected_costs(vec![UsageProjectedCostPoint {
bucket_start: quarter_hour_start(now),
provider_kind: "openai".to_owned(),
upstream_model_id: "gpt-5.5".to_owned(),
service_tier: None,
input_tokens: Some(800),
output_tokens: Some(200),
cached_tokens: Some(0),
cache_write_tokens: Some(0),
total: CurrencyCost {
currency: "USD".to_owned(),
amount: gateway_admin::model::observability::DecimalAmount::from_str("1.25")
.expect("calculated total"),
},
standard_amount: Some("1".parse().expect("standard cost")),
no_cache_amount: Some("1.25".parse().expect("no-cache cost")),
}]);
store.replace_diagnostics(DiagnosticsObservation {
analytics: Default::default(),
total_request_count: 4,
items: vec![diagnostic("openai", 3), diagnostic("xai", 1)],
});
@@ -543,6 +635,7 @@ async fn diagnostics_should_preserve_request_order_when_a_small_sample_fails() {
failure.success_count = 0;
failure.failure_count = 1;
store.replace_diagnostics(DiagnosticsObservation {
analytics: Default::default(),
total_request_count: 101,
items: vec![success, failure],
});
@@ -566,6 +659,7 @@ async fn diagnostics_should_use_the_total_before_truncating_groups() {
let range = observation_range(Utc::now());
let store = Arc::new(FixtureObservabilityStore::new(range));
store.replace_diagnostics(DiagnosticsObservation {
analytics: Default::default(),
total_request_count: 200,
items: (0..100)
.map(|index| diagnostic(&format!("model-{index}"), 1))
@@ -591,6 +685,7 @@ async fn diagnostics_should_count_each_retried_request_once() {
item.retry_count = 6;
item.retried_request_count = 2;
store.replace_diagnostics(DiagnosticsObservation {
analytics: Default::default(),
total_request_count: 8,
items: vec![item],
});
@@ -691,26 +786,16 @@ async fn usage_records_should_enrich_provider_reported_totals_when_pricing_match
}
#[tokio::test]
async fn usage_insights_should_reject_partial_costs_when_billing_stream_fails() {
async fn usage_insights_should_reject_partial_costs_when_insights_fail() {
let now = Utc::now();
let range = observation_range(now);
let store = Arc::new(FixtureObservabilityStore::new(range));
store.replace_calculated_billing_facts(vec![UsageCalculatedBillingFact {
breakdown: None,
store.replace_projected_costs(vec![UsageProjectedCostPoint {
bucket_start: quarter_hour_start(now),
provider_kind: "openai".to_owned(),
upstream_model_id: "gpt-5.5".to_owned(),
service_tier: None,
input_tokens: Some(800),
output_tokens: Some(200),
cached_tokens: Some(0),
cache_write_tokens: Some(0),
total: CurrencyCost {
currency: "USD".to_owned(),
amount: "1.25".parse().expect("cost"),
},
standard_amount: Some("1".parse().expect("standard cost")),
no_cache_amount: Some("1.25".parse().expect("no-cache cost")),
}]);
store.billing_stream_fails.store(true, Ordering::SeqCst);
store.insights_fail.store(true, Ordering::SeqCst);
let services = observability_services_with_calculated_billing(store).await;
assert!(
services
@@ -741,9 +826,9 @@ async fn dashboard_should_preserve_provider_plan_names() {
success_count: 0,
input_tokens: None,
output_tokens: None,
cached_tokens: None,
cache_write_tokens: None,
reasoning_tokens: None,
cached_input_tokens: None,
cache_write_input_tokens: None,
reasoning_output_tokens: None,
image_input_tokens: None,
image_output_tokens: None,
image_request_count: 0,
@@ -783,8 +868,9 @@ async fn dashboard_should_preserve_provider_plan_names() {
struct FixtureObservabilityStore {
trend: Mutex<Vec<RequestMetricPoint>>,
overview: Mutex<UsageOverview>,
calculated_billing_facts: Mutex<Vec<UsageCalculatedBillingFact>>,
billing_stream_fails: AtomicBool,
projected_costs: Mutex<Vec<UsageProjectedCostPoint>>,
insights_fail: AtomicBool,
summary_fail: AtomicBool,
diagnostics: Mutex<DiagnosticsObservation>,
runtime_slots: Mutex<Option<DashboardRuntimeSlots>>,
summary_observed_at: Mutex<Option<DateTime<Utc>>>,
@@ -794,6 +880,7 @@ struct FixtureObservabilityStore {
account_usage: Mutex<Vec<DashboardAccountUsage>>,
dashboard_delay: Mutex<StdDuration>,
dashboard_summary_calls: AtomicUsize,
usage_summary_filters: Mutex<Vec<UsageFilter>>,
}
impl FixtureObservabilityStore {
@@ -801,13 +888,15 @@ impl FixtureObservabilityStore {
Self {
trend: Mutex::new(Vec::new()),
overview: Mutex::new(UsageOverview {
analytics: Default::default(),
range,
requests: RequestMetrics::default(),
attempts: AttemptMetrics::default(),
providers: Vec::new(),
}),
calculated_billing_facts: Mutex::new(Vec::new()),
billing_stream_fails: AtomicBool::new(false),
projected_costs: Mutex::new(Vec::new()),
insights_fail: AtomicBool::new(false),
summary_fail: AtomicBool::new(false),
diagnostics: Mutex::new(DiagnosticsObservation::default()),
runtime_slots: Mutex::new(None),
summary_observed_at: Mutex::new(None),
@@ -817,6 +906,7 @@ impl FixtureObservabilityStore {
account_usage: Mutex::new(Vec::new()),
dashboard_delay: Mutex::new(StdDuration::ZERO),
dashboard_summary_calls: AtomicUsize::new(0),
usage_summary_filters: Mutex::new(Vec::new()),
}
}
@@ -836,11 +926,8 @@ impl FixtureObservabilityStore {
*self.overview.lock().expect("overview") = overview;
}
fn replace_calculated_billing_facts(&self, facts: Vec<UsageCalculatedBillingFact>) {
*self
.calculated_billing_facts
.lock()
.expect("calculated billing facts") = facts;
fn replace_projected_costs(&self, facts: Vec<UsageProjectedCostPoint>) {
*self.projected_costs.lock().expect("projected costs") = facts;
}
fn replace_diagnostics(&self, diagnostics: DiagnosticsObservation) {
@@ -874,6 +961,7 @@ impl ObservabilityStore for FixtureObservabilityStore {
tokio::time::sleep(dashboard_delay).await;
}
Ok(DashboardObservation {
analytics: self.overview.lock().expect("overview").analytics.clone(),
range: query.range,
totals: Default::default(),
provider_accounts: AccountPoolMetrics::default(),
@@ -902,25 +990,29 @@ impl ObservabilityStore for FixtureObservabilityStore {
Ok(self.trend.lock().expect("trend").clone())
}
fn usage_calculated_billing_facts(
async fn analytics_cache_generation(&self) -> AdminStoreResult<u64> {
Ok(self.overview.lock().expect("overview").analytics.generation)
}
async fn usage_insights_observation(
&self,
_: TimeRange,
_: UsageFilter,
granularity: Granularity,
) -> gateway_admin::ports::store::UsageCalculatedBillingStream<'_> {
) -> AdminStoreResult<UsageInsightsObservation> {
self.usage_granularities.lock().unwrap().push(granularity);
let facts = self
.calculated_billing_facts
.lock()
.expect("calculated billing facts")
.clone();
let error = self
.billing_stream_fails
.load(Ordering::SeqCst)
.then(|| Err(super::unavailable("billing stream failed")));
Box::pin(futures::stream::iter(
facts.into_iter().map(Ok).chain(error),
))
if self.insights_fail.load(Ordering::SeqCst) {
return Err(super::unavailable("projection read failed"));
}
Ok(UsageInsightsObservation {
overview: self.overview.lock().expect("overview").clone(),
trend: self.trend.lock().expect("trend").clone(),
costs: self
.projected_costs
.lock()
.expect("projected costs")
.clone(),
})
}
async fn list_usage_records(&self, query: UsageQuery) -> AdminStoreResult<UsagePage> {
@@ -938,8 +1030,26 @@ impl ObservabilityStore for FixtureObservabilityStore {
Err(super::unavailable("usage detail"))
}
async fn usage_summary(&self, _: TimeRange, _: UsageFilter) -> AdminStoreResult<UsageOverview> {
Ok(self.overview.lock().expect("overview").clone())
async fn usage_summary(
&self,
range: TimeRange,
filter: UsageFilter,
) -> AdminStoreResult<UsageOverview> {
self.usage_summary_filters
.lock()
.expect("summary filters")
.push(filter);
let delay = *self.dashboard_delay.lock().expect("query delay");
if !delay.is_zero() {
tokio::time::sleep(delay).await;
}
if self.summary_fail.load(Ordering::SeqCst) {
return Err(super::unavailable("projection summary failed"));
}
Ok(UsageOverview {
range,
..self.overview.lock().expect("overview").clone()
})
}
async fn usage_diagnostics(
@@ -1114,9 +1224,9 @@ fn total_record(
service_tier: None,
input_tokens: Some(800),
output_tokens: Some(200),
cached_tokens: Some(0),
cache_write_tokens: Some(0),
reasoning_tokens: Some(0),
cached_input_tokens: Some(0),
cache_write_input_tokens: Some(0),
reasoning_output_tokens: Some(0),
image_input_tokens: None,
image_output_tokens: None,
total_tokens: Some(1_000),
@@ -1229,9 +1339,204 @@ async fn admin_selects_dashboard_defaults_and_matching_usage_buckets_at_range_bo
.usage_insights(range, UsageFilter::default())
.await
.unwrap();
assert_eq!(
*store.usage_granularities.lock().unwrap(),
vec![expected, expected]
assert_eq!(*store.usage_granularities.lock().unwrap(), vec![expected]);
}
}
#[tokio::test]
async fn usage_summary_cache_should_coalesce_identical_queries_and_invalidate_deleted_generation() {
let range = observation_range(Utc::now());
let store = Arc::new(FixtureObservabilityStore::new(range));
store.set_dashboard_delay(StdDuration::from_millis(20));
let services = observability_services(store.clone()).await;
let service = services.observability();
let filter = UsageFilter {
client_api_key_ref: Some("deleted-key".to_owned()),
..Default::default()
};
let (first, second) = tokio::join!(
service.usage_summary(range, filter.clone()),
service.usage_summary(range, filter.clone()),
);
first.expect("first summary");
second.expect("coalesced summary");
assert_eq!(store.usage_summary_filters.lock().unwrap().len(), 1);
{
let mut overview = store.overview.lock().unwrap();
overview.analytics.generation = 1;
overview.requests.request_count = 7;
}
let refreshed = service
.usage_summary(range, filter)
.await
.expect("generation refresh");
assert_eq!(refreshed.overview.analytics.generation, 1);
assert_eq!(refreshed.overview.requests.request_count, 7);
assert_eq!(store.usage_summary_filters.lock().unwrap().len(), 2);
}
#[tokio::test(start_paused = true)]
async fn usage_summary_cache_should_coalesce_when_generation_advances_during_the_first_load() {
let range = observation_range(Utc::now());
let store = Arc::new(FixtureObservabilityStore::new(range));
store.set_dashboard_delay(StdDuration::from_millis(20));
let services = observability_services(store.clone()).await;
let service = services.observability();
let (first, second, ()) = tokio::join!(
service.usage_summary(range, UsageFilter::default()),
service.usage_summary(range, UsageFilter::default()),
async {
tokio::time::sleep(StdDuration::from_millis(5)).await;
store.overview.lock().unwrap().analytics.generation = 1;
},
);
assert_eq!(
first.expect("first summary").overview.analytics.generation,
1
);
assert_eq!(
second
.expect("coalesced summary")
.overview
.analytics
.generation,
1
);
assert_eq!(store.usage_summary_filters.lock().unwrap().len(), 1);
}
#[tokio::test]
async fn usage_summary_cache_should_isolate_key_identity_and_normalize_group_sets() {
let range = observation_range(Utc::now());
let store = Arc::new(FixtureObservabilityStore::new(range));
let services = observability_services(store.clone()).await;
let service = services.observability();
for (key, groups) in [
("key-a", vec!["deleted-b", "deleted-a", "deleted-b"]),
("key-a", vec!["deleted-a", "deleted-b"]),
("key-b", vec!["deleted-a", "deleted-b"]),
("key-a", vec!["deleted-a"]),
] {
service
.usage_summary(
range,
UsageFilter {
client_api_key_ref: Some(key.to_owned()),
routing_group_refs: groups.into_iter().map(str::to_owned).collect(),
..Default::default()
},
)
.await
.expect("isolated summary");
}
let calls = store.usage_summary_filters.lock().unwrap();
assert_eq!(calls.len(), 3);
assert_eq!(calls[0].routing_group_refs, ["deleted-a", "deleted-b"]);
assert_eq!(calls[1].client_api_key_ref.as_deref(), Some("key-b"));
assert_eq!(calls[2].routing_group_refs, ["deleted-a"]);
}
#[tokio::test(start_paused = true)]
async fn usage_summary_cache_should_bound_waiting_after_failed_refresh_and_allow_recovery() {
let range = observation_range(Utc::now());
let store = Arc::new(FixtureObservabilityStore::new(range));
store.set_dashboard_delay(StdDuration::from_secs(7));
store.summary_fail.store(true, Ordering::SeqCst);
let services = observability_services(store.clone()).await;
let service = services.observability();
let started = tokio::time::Instant::now();
let (first, second, third, fourth) = tokio::join!(
service.usage_summary(range, UsageFilter::default()),
service.usage_summary(range, UsageFilter::default()),
service.usage_summary(range, UsageFilter::default()),
service.usage_summary(range, UsageFilter::default()),
);
assert!(first.is_err() && second.is_err() && third.is_err() && fourth.is_err());
assert!(
started.elapsed() <= StdDuration::from_secs(10),
"同键失败不能使后续请求逐个等待完整数据库超时: {:?}",
started.elapsed()
);
store.summary_fail.store(false, Ordering::SeqCst);
store.set_dashboard_delay(StdDuration::ZERO);
service
.usage_summary(range, UsageFilter::default())
.await
.expect("cancelled cache refresh releases its slot for recovery");
}
#[tokio::test(start_paused = true)]
async fn insights_cache_should_bound_stale_lifetime_without_renewing_failed_refreshes() {
let range = observation_range(Utc::now());
let store = Arc::new(FixtureObservabilityStore::new(range));
{
let mut overview = store.overview.lock().unwrap();
overview.analytics.state = AnalyticsState::Ready;
overview.analytics.generated_at = Utc::now();
overview.analytics.lag_seconds = 7;
overview.requests.request_count = 19;
}
let services = observability_services(store.clone()).await;
let service = services.observability();
service
.usage_insights(range, UsageFilter::default())
.await
.unwrap();
store.insights_fail.store(true, Ordering::SeqCst);
tokio::time::advance(StdDuration::from_secs(16)).await;
let stale = service
.usage_insights(range, UsageFilter::default())
.await
.unwrap();
assert_eq!(stale.analytics.state, AnalyticsState::Stale);
assert!(stale.analytics.lag_seconds >= 7);
assert_eq!(stale.health.total_requests, 19);
tokio::time::advance(StdDuration::from_secs(45)).await;
assert!(
service
.usage_insights(range, UsageFilter::default())
.await
.is_err()
);
}
#[tokio::test(start_paused = true)]
async fn insights_cache_should_reject_deleted_building_and_old_generation_times() {
for reason in ["deleted", "building", "generated-too-long-ago"] {
let range = observation_range(Utc::now());
let store = Arc::new(FixtureObservabilityStore::new(range));
{
let mut overview = store.overview.lock().unwrap();
overview.analytics.state = if reason == "building" {
AnalyticsState::Building
} else {
AnalyticsState::Ready
};
overview.analytics.generated_at = Utc::now()
- Duration::seconds(if reason == "generated-too-long-ago" {
61
} else {
0
});
}
let services = observability_services(store.clone()).await;
let service = services.observability();
service
.usage_insights(range, UsageFilter::default())
.await
.unwrap();
store.insights_fail.store(true, Ordering::SeqCst);
if reason == "deleted" {
store.overview.lock().unwrap().analytics.generation += 1;
}
tokio::time::advance(StdDuration::from_secs(16)).await;
assert!(
service
.usage_insights(range, UsageFilter::default())
.await
.is_err(),
"{reason}"
);
}
}
@@ -100,6 +100,51 @@ async fn client_profile_preview_rejects_unknown_provider_before_loading_settings
assert_eq!(error.kind(), AdminErrorKind::Invalid);
}
#[tokio::test]
async fn pricing_accepts_ultrafast_for_openai_but_rejects_it_for_xai() {
for provider in ["openai", "xai"] {
let store = Arc::new(PricingSettingsStore::default());
let services = super::AdminHarness::new()
.settings(store.clone())
.provider(super::test_provider(provider))
.build()
.await;
let context = MutationContext {
actor: gateway_admin::model::MutationActor::System,
request_id: "request-ultrafast-pricing".to_owned(),
};
for band in ["ultrafast", "long_ultrafast"] {
let pricing = serde_json::from_value(serde_json::json!({
"multiplierBps": 10000,
"bands": {
"standard": {"input":"2", "output":"10", "cacheRead":"0.1", "cacheWrite":"2.5"},
(band): {"input":"12", "output":"60", "cacheRead":"0.6", "cacheWrite":"15"}
}
}))
.expect("Ultrafast override");
let update = UpdatePricing {
provider: provider.to_owned(),
models: vec!["example-model".to_owned()],
change: PricingChange::Replace(pricing),
};
let result = services
.settings()
.update_pricing(&context, update.clone())
.await;
if provider == "openai" {
result.expect("OpenAI accepts Ultrafast overrides");
assert_eq!(store.updated.lock().unwrap().as_ref(), Some(&update));
} else {
assert_eq!(
result.expect_err("xAI rejects Ultrafast").kind(),
AdminErrorKind::Invalid
);
assert!(store.updated.lock().unwrap().is_none());
}
}
}
}
#[tokio::test]
async fn pricing_accepts_native_provider_and_rejects_unknown_sync_provider() {
let store = Arc::new(PricingSettingsStore::default());
@@ -308,8 +308,8 @@ pub(super) fn quota_local_usage(usage: &AccountUsage, time: crate::time::TimePre
"inputTokensDisplay": display_optional_tokens(usage.input_tokens),
"outputTokens": usage.output_tokens.unwrap_or_default(),
"outputTokensDisplay": display_optional_tokens(usage.output_tokens),
"cachedTokens": usage.cached_tokens.unwrap_or_default(),
"cachedTokensDisplay": display_optional_tokens(usage.cached_tokens),
"cachedInputTokens": usage.cached_input_tokens.unwrap_or_default(),
"cachedInputTokensDisplay": display_optional_tokens(usage.cached_input_tokens),
"imageInputTokens": usage.image_input_tokens.unwrap_or_default(),
"imageInputTokensDisplay": display_optional_tokens(usage.image_input_tokens),
"imageOutputTokens": usage.image_output_tokens.unwrap_or_default(),
@@ -355,10 +355,10 @@ pub(super) fn account_usage_view(
input_tokens_display: display_optional_tokens(usage.input_tokens),
output_tokens: usage.output_tokens,
output_tokens_display: display_optional_tokens(usage.output_tokens),
cached_tokens: usage.cached_tokens,
cached_tokens_display: display_optional_tokens(usage.cached_tokens),
reasoning_tokens: usage.reasoning_tokens,
reasoning_tokens_display: display_optional_tokens(usage.reasoning_tokens),
cached_input_tokens: usage.cached_input_tokens,
cached_input_tokens_display: display_optional_tokens(usage.cached_input_tokens),
reasoning_output_tokens: usage.reasoning_output_tokens,
reasoning_output_tokens_display: display_optional_tokens(usage.reasoning_output_tokens),
image_input_tokens: usage.image_input_tokens,
image_input_tokens_display: display_optional_tokens(usage.image_input_tokens),
image_output_tokens: usage.image_output_tokens,
@@ -369,10 +369,8 @@ pub(super) fn account_usage_view(
image_request_failed_count_display: format_number(usage.image_request_failed_count),
total_tokens: usage.total_tokens,
total_tokens_display: display_optional_tokens(usage.total_tokens),
created_tokens: usage.cache_write_tokens,
created_tokens_display: display_optional_tokens(usage.cache_write_tokens),
read_tokens: usage.cached_tokens,
read_tokens_display: display_optional_tokens(usage.cached_tokens),
cache_write_input_tokens: usage.cache_write_input_tokens,
cache_write_input_tokens_display: display_optional_tokens(usage.cache_write_input_tokens),
last_used_at: usage.last_used_at.map(|value| time.rfc3339(&value)),
last_used_at_full_display: usage.last_used_at.map(|value| time.datetime(&value)),
last_used_at_display: usage
@@ -429,8 +427,8 @@ pub(super) fn account_model_usage_view(
input_tokens_display: display_optional_tokens(usage.input_tokens),
output_tokens: usage.output_tokens,
output_tokens_display: display_optional_tokens(usage.output_tokens),
cached_tokens: usage.cached_tokens,
cached_tokens_display: display_optional_tokens(usage.cached_tokens),
cached_input_tokens: usage.cached_input_tokens,
cached_input_tokens_display: display_optional_tokens(usage.cached_input_tokens),
image_input_tokens: usage.image_input_tokens,
image_input_tokens_display: display_optional_tokens(usage.image_input_tokens),
image_output_tokens: usage.image_output_tokens,
@@ -478,10 +476,10 @@ pub(super) fn empty_account_usage() -> AccountUsageView {
input_tokens_display: "—".to_owned(),
output_tokens: None,
output_tokens_display: "—".to_owned(),
cached_tokens: None,
cached_tokens_display: "—".to_owned(),
reasoning_tokens: None,
reasoning_tokens_display: "—".to_owned(),
cached_input_tokens: None,
cached_input_tokens_display: "—".to_owned(),
reasoning_output_tokens: None,
reasoning_output_tokens_display: "—".to_owned(),
image_input_tokens: None,
image_input_tokens_display: "—".to_owned(),
image_output_tokens: None,
@@ -492,10 +490,8 @@ pub(super) fn empty_account_usage() -> AccountUsageView {
image_request_failed_count_display: "—".to_owned(),
total_tokens: None,
total_tokens_display: "—".to_owned(),
created_tokens: None,
created_tokens_display: "—".to_owned(),
read_tokens: None,
read_tokens_display: "—".to_owned(),
cache_write_input_tokens: None,
cache_write_input_tokens_display: "—".to_owned(),
last_used_at: None,
last_used_at_full_display: None,
last_used_at_display: "—".to_owned(),
@@ -435,10 +435,10 @@ pub struct AccountUsageView {
pub input_tokens_display: String,
pub output_tokens: Option<u64>,
pub output_tokens_display: String,
pub cached_tokens: Option<u64>,
pub cached_tokens_display: String,
pub reasoning_tokens: Option<u64>,
pub reasoning_tokens_display: String,
pub cached_input_tokens: Option<u64>,
pub cached_input_tokens_display: String,
pub reasoning_output_tokens: Option<u64>,
pub reasoning_output_tokens_display: String,
pub image_input_tokens: Option<u64>,
pub image_input_tokens_display: String,
pub image_output_tokens: Option<u64>,
@@ -449,10 +449,8 @@ pub struct AccountUsageView {
pub image_request_failed_count_display: String,
pub total_tokens: Option<u64>,
pub total_tokens_display: String,
pub created_tokens: Option<u64>,
pub created_tokens_display: String,
pub read_tokens: Option<u64>,
pub read_tokens_display: String,
pub cache_write_input_tokens: Option<u64>,
pub cache_write_input_tokens_display: String,
pub last_used_at: Option<String>,
pub last_used_at_display: String,
pub last_used_at_full_display: Option<String>,
@@ -477,8 +475,8 @@ pub struct ModelUsageView {
pub input_tokens_display: String,
pub output_tokens: Option<u64>,
pub output_tokens_display: String,
pub cached_tokens: Option<u64>,
pub cached_tokens_display: String,
pub cached_input_tokens: Option<u64>,
pub cached_input_tokens_display: String,
pub image_input_tokens: Option<u64>,
pub image_input_tokens_display: String,
pub image_output_tokens: Option<u64>,
@@ -58,9 +58,9 @@ pub(crate) fn request_metrics_view(metrics: &domain::RequestMetrics) -> RequestM
caller_error_count: metrics.caller_error_count,
input_tokens: metrics.input_tokens,
output_tokens: metrics.output_tokens,
cached_tokens: metrics.cached_tokens,
cache_write_tokens: metrics.cache_write_tokens,
reasoning_tokens: metrics.reasoning_tokens,
cached_input_tokens: metrics.cached_input_tokens,
cache_write_input_tokens: metrics.cache_write_input_tokens,
reasoning_output_tokens: metrics.reasoning_output_tokens,
total_tokens: metrics.total_tokens,
}
}
@@ -103,9 +103,9 @@ pub(crate) fn token_details(record: &domain::UsageRecord) -> 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,
cached_input_tokens: record.cached_input_tokens,
cache_write_input_tokens: record.cache_write_input_tokens,
reasoning_output_tokens: record.reasoning_output_tokens,
image_input_tokens: record.image_input_tokens,
image_output_tokens: record.image_output_tokens,
total_tokens: record.total_tokens,
@@ -115,14 +115,14 @@ pub(crate) fn token_details(record: &domain::UsageRecord) -> TokenDetailsView {
output_tokens_display: record
.output_tokens
.map_or_else(|| "-".to_owned(), format_number),
cached_tokens_display: record
.cached_tokens
cached_input_tokens_display: record
.cached_input_tokens
.map_or_else(|| "-".to_owned(), format_compact_number),
cache_write_tokens_display: record
.cache_write_tokens
cache_write_input_tokens_display: record
.cache_write_input_tokens
.map_or_else(|| "-".to_owned(), format_compact_number),
reasoning_tokens_display: record
.reasoning_tokens
reasoning_output_tokens_display: record
.reasoning_output_tokens
.map_or_else(|| "-".to_owned(), format_number),
image_input_tokens_display: record
.image_input_tokens
@@ -140,9 +140,9 @@ pub(crate) fn usage_list_token_details(record: &domain::UsageListRecord) -> Toke
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,
cached_input_tokens: record.cached_input_tokens,
cache_write_input_tokens: record.cache_write_input_tokens,
reasoning_output_tokens: record.reasoning_output_tokens,
image_input_tokens: record.image_input_tokens,
image_output_tokens: record.image_output_tokens,
total_tokens: record.total_tokens,
@@ -152,14 +152,14 @@ pub(crate) fn usage_list_token_details(record: &domain::UsageListRecord) -> Toke
output_tokens_display: record
.output_tokens
.map_or_else(|| "-".to_owned(), format_number),
cached_tokens_display: record
.cached_tokens
cached_input_tokens_display: record
.cached_input_tokens
.map_or_else(|| "-".to_owned(), format_compact_number),
cache_write_tokens_display: record
.cache_write_tokens
cache_write_input_tokens_display: record
.cache_write_input_tokens
.map_or_else(|| "-".to_owned(), format_compact_number),
reasoning_tokens_display: record
.reasoning_tokens
reasoning_output_tokens_display: record
.reasoning_output_tokens
.map_or_else(|| "-".to_owned(), format_number),
image_input_tokens_display: record
.image_input_tokens
@@ -430,9 +430,9 @@ pub(crate) fn usage_record_view(
latency_ms: record.latency_ms,
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,
cached_input_tokens: record.cached_input_tokens,
cache_write_input_tokens: record.cache_write_input_tokens,
reasoning_output_tokens: record.reasoning_output_tokens,
image_input_tokens: record.image_input_tokens,
image_output_tokens: record.image_output_tokens,
message,
@@ -526,7 +526,7 @@ pub(crate) fn usage_attempt_view(attempt: domain::UsageAttempt) -> UsageAttemptV
estimated_cost_currency: attempt.cost_currency,
input_tokens: attempt.input_tokens,
output_tokens: attempt.output_tokens,
cached_tokens: attempt.cached_tokens,
cached_input_tokens: attempt.cached_input_tokens,
total_tokens: attempt.total_tokens,
first_token_ms: None,
latency_ms: attempt.latency_ms,
@@ -683,8 +683,8 @@ pub(crate) fn trend_point_view(
input_tokens_value: point.metrics.input_tokens,
output_tokens: format_compact_number(point.metrics.output_tokens),
output_tokens_value: point.metrics.output_tokens,
cached_tokens: format_compact_number(point.metrics.cached_tokens),
cached_tokens_value: point.metrics.cached_tokens,
cached_input_tokens: format_compact_number(point.metrics.cached_input_tokens),
cached_input_tokens_value: point.metrics.cached_input_tokens,
cache_hit_rate_value: point.cached_token_rate,
tokens_value: point.metrics.total_tokens,
errors: format_compact_number(point.service_failure_count),
@@ -749,7 +749,7 @@ pub(crate) fn trend_summary_view(
},
TrendSummaryView {
label: "缓存".to_owned(),
value: format_compact_number(summary.cached_tokens),
value: format_compact_number(summary.cached_input_tokens),
ratio: None,
},
],
@@ -800,6 +800,7 @@ pub(crate) fn trend_view(
time: crate::time::TimePresenter,
) -> TrendData {
TrendData {
analytics: analytics_view(trend.analytics, time),
kind,
summary: trend_summary_view(kind, &trend.summary),
points: trend
@@ -913,6 +914,7 @@ pub(crate) fn dashboard_view(
rotation_strategy,
} = result;
let domain::DashboardObservation {
analytics,
runtime_slots: _,
range,
totals,
@@ -982,6 +984,7 @@ pub(crate) fn dashboard_view(
.saturating_sub(provider_accounts.normal);
DashboardDataView {
analytics: analytics_view(analytics, time),
as_of: range.end,
as_of_display: time.datetime(&range.end),
cards: DashboardCardsView {
@@ -1013,7 +1016,7 @@ pub(crate) fn dashboard_view(
today_hit_rate_value: today.observed_cached_token_rate,
yesterday_hit_rate_value: yesterday.observed_cached_token_rate,
total_hit_rate: display_rate(total_cached_token_rate),
total_cached_tokens: format_compact_number(totals.cached_tokens),
total_cached_input_tokens: format_compact_number(totals.cached_input_tokens),
average_first_token_latency_ms: display_duration(average_first_token_latency_ms),
},
},
@@ -1046,14 +1049,18 @@ pub(crate) fn dashboard_view(
}
}
pub(crate) fn usage_summary_view(summary: domain::UsageSummary) -> UsageSummaryView {
pub(crate) fn usage_summary_view(
summary: domain::UsageSummary,
time: crate::time::TimePresenter,
) -> UsageSummaryView {
let overview = summary.overview;
UsageSummaryView {
analytics: analytics_view(overview.analytics, time),
total_requests: format_compact_number(overview.requests.request_count),
input_tokens: format_compact_number(overview.requests.input_tokens),
output_tokens: format_compact_number(overview.requests.output_tokens),
cached_tokens: format_compact_number(overview.requests.cached_tokens),
cache_write_tokens: format_compact_number(overview.requests.cache_write_tokens),
cached_input_tokens: format_compact_number(overview.requests.cached_input_tokens),
cache_write_input_tokens: format_compact_number(overview.requests.cache_write_input_tokens),
total_tokens: format_compact_number(overview.requests.total_tokens),
average_latency_ms: display_duration(summary.average_latency_ms),
logical_requests: request_metrics_view(&overview.requests),
@@ -1139,7 +1146,7 @@ pub(crate) fn usage_insights_view(
label: time.label(point.bucket_start, "%m-%d %H:%M"),
input_tokens: point.input_tokens,
output_tokens: point.output_tokens,
cached_tokens: point.cached_tokens,
cached_input_tokens: point.cached_input_tokens,
total_tokens: point.total_tokens,
estimated_cost: point.estimated_cost.as_ref().map(ToString::to_string),
standard_cost: point.standard_cost.as_ref().map(ToString::to_string),
@@ -1150,6 +1157,7 @@ pub(crate) fn usage_insights_view(
})
.collect();
UsageInsightsOverviewView {
analytics: analytics_view(insights.analytics, time),
granularity: granularity_name(insights.granularity).to_owned(),
health: OverviewHealthView {
total_requests: insights.health.total_requests,
@@ -1270,7 +1278,7 @@ pub(crate) fn usage_insights_view(
cache_hit_request_rate: insights.cost.cache_hit_request_rate,
input_tokens: insights.cost.input_tokens,
output_tokens: insights.cost.output_tokens,
cached_tokens: insights.cost.cached_tokens,
cached_input_tokens: insights.cost.cached_input_tokens,
total_tokens: insights.cost.total_tokens,
points: cost_points,
costs: cost_views(&insights.cost.costs),
@@ -1302,8 +1310,10 @@ pub(crate) fn granularity_name(granularity: domain::Granularity) -> &'static str
pub(crate) fn diagnostics_view(
result: domain::DiagnosticsResult,
dimension: DiagnosticDimension,
time: crate::time::TimePresenter,
) -> DiagnosticsView {
DiagnosticsView {
analytics: analytics_view(result.analytics, time),
dimension: dimension.display_name().to_owned(),
items: result
.items
@@ -1356,3 +1366,22 @@ pub(crate) fn map_wire_error(error: WireValidationError) -> AdminError {
pub(crate) fn map_service_error(error: gateway_admin::model::AdminError) -> AdminError {
map_admin_service_error(error)
}
pub(crate) fn analytics_view(
metadata: domain::AnalyticsMetadata,
time: crate::time::TimePresenter,
) -> AnalyticsView {
AnalyticsView {
coverage_start: metadata.coverage_start,
coverage_start_display: metadata.coverage_start.map(|value| time.datetime(&value)),
coverage_end: metadata.coverage_end,
coverage_end_display: metadata.coverage_end.map(|value| time.datetime(&value)),
generated_at: metadata.generated_at,
generated_at_display: time.datetime(&metadata.generated_at),
lag_seconds: metadata.lag_seconds,
state: metadata.state.as_str().to_owned(),
percentiles_approximate: metadata.percentiles_approximate,
schema_version: metadata.schema_version,
histogram_version: metadata.histogram_version,
}
}
@@ -29,6 +29,7 @@ impl DashboardQuery {
#[derive(Clone, Default, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct UsageQuery {
pub routing_group_refs: Option<String>,
pub current_page: Option<u32>,
pub page_size: Option<u16>,
pub kind: Option<String>,
@@ -76,6 +77,7 @@ impl DetailQuery {
#[derive(Clone, Default, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct DiagnosticsQuery {
pub routing_group_refs: Option<String>,
pub dimension: Option<String>,
pub period: Option<String>,
pub as_of: Option<i64>,
@@ -98,6 +100,7 @@ impl DiagnosticsQuery {
#[derive(Clone, Default, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct OpsQuery {
pub routing_group_refs: Option<String>,
pub current_page: Option<u32>,
pub page_size: Option<u16>,
pub kind: Option<String>,
@@ -298,6 +301,7 @@ pub(crate) fn usage_filter(query: &UsageQuery) -> Result<domain::UsageFilter, Wi
})
});
Ok(domain::UsageFilter {
routing_group_refs: routing_group_refs(query.routing_group_refs.as_deref())?,
client_api_key_ref: non_empty(query.client_api_key_id.clone()),
request_id: non_empty(query.request_id.clone()),
provider_account_ref: non_empty(query.account_id.clone()),
@@ -357,6 +361,7 @@ pub(crate) fn ops_command(
timezone,
)?,
filter: domain::OpsErrorFilter {
routing_group_refs: routing_group_refs(query.routing_group_refs.as_deref())?,
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()),
@@ -395,3 +400,25 @@ fn validate_pagination(
}
Ok((current_page, page_size))
}
/// 逗号分隔的稳定分组身份按集合归一化,不要求分组当前仍存在
pub(crate) fn routing_group_refs(value: Option<&str>) -> Result<Vec<String>, WireValidationError> {
let Some(value) = value else {
return Ok(Vec::new());
};
let mut refs = Vec::new();
for item in value.split(',') {
let item = item.trim();
if item.is_empty()
|| item.len() > 128
|| item.chars().any(char::is_control)
|| refs.len() >= 32
{
return Err(WireValidationError::new("routingGroupRefs"));
}
refs.push(item.to_owned());
}
refs.sort();
refs.dedup();
Ok(refs)
}
@@ -160,7 +160,10 @@ where
.map_err(map_service_error)?;
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(usage_summary_view(result)),
AdminEnvelope::ok(usage_summary_view(
result,
crate::time::TimePresenter::new(state.admin_services().timezone()),
)),
))
}
@@ -212,6 +215,8 @@ where
)
.map_err(map_wire_error)?;
let filter = domain::UsageFilter {
routing_group_refs: routing_group_refs(query.routing_group_refs.as_deref())
.map_err(map_wire_error)?,
provider_kind: non_empty(query.provider),
model: non_empty(query.model),
status_code: parse_status(query.status_code).map_err(map_wire_error)?,
@@ -226,7 +231,11 @@ where
.map_err(map_service_error)?;
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(diagnostics_view(result, dimension)),
AdminEnvelope::ok(diagnostics_view(
result,
dimension,
crate::time::TimePresenter::new(state.admin_services().timezone()),
)),
))
}
@@ -18,17 +18,17 @@ pub struct PageData<T> {
pub struct TokenDetailsView {
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 cached_input_tokens: Option<u64>,
pub cache_write_input_tokens: Option<u64>,
pub reasoning_output_tokens: Option<u64>,
pub image_input_tokens: Option<u64>,
pub image_output_tokens: Option<u64>,
pub total_tokens: Option<u64>,
pub input_tokens_display: String,
pub output_tokens_display: String,
pub cached_tokens_display: String,
pub cache_write_tokens_display: String,
pub reasoning_tokens_display: String,
pub cached_input_tokens_display: String,
pub cache_write_input_tokens_display: String,
pub reasoning_output_tokens_display: String,
pub image_input_tokens_display: String,
pub image_output_tokens_display: String,
pub total_tokens_display: String,
@@ -162,9 +162,9 @@ pub struct UsageRecordView {
pub latency_ms: Option<u64>,
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 cached_input_tokens: Option<u64>,
pub cache_write_input_tokens: Option<u64>,
pub reasoning_output_tokens: Option<u64>,
pub image_input_tokens: Option<u64>,
pub image_output_tokens: Option<u64>,
pub message: String,
@@ -258,7 +258,7 @@ pub struct UsageAttemptView {
pub estimated_cost_currency: Option<String>,
pub input_tokens: Option<u64>,
pub output_tokens: Option<u64>,
pub cached_tokens: Option<u64>,
pub cached_input_tokens: Option<u64>,
pub total_tokens: Option<u64>,
pub first_token_ms: Option<u64>,
pub latency_ms: Option<u64>,
@@ -288,6 +288,7 @@ pub struct UsageRecordDetailView {
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct TrendData {
pub analytics: AnalyticsView,
pub kind: TrendKind,
pub points: Vec<TrendPointView>,
pub summary: Vec<TrendSummaryView>,
@@ -306,8 +307,8 @@ pub struct TrendPointView {
pub input_tokens_value: u64,
pub output_tokens: String,
pub output_tokens_value: u64,
pub cached_tokens: String,
pub cached_tokens_value: u64,
pub cached_input_tokens: String,
pub cached_input_tokens_value: u64,
pub cache_hit_rate_value: f64,
pub tokens_value: u64,
pub errors: String,
@@ -388,7 +389,7 @@ pub struct DashboardCacheCardView {
pub today_hit_rate_value: Option<f64>,
pub yesterday_hit_rate_value: Option<f64>,
pub total_hit_rate: String,
pub total_cached_tokens: String,
pub total_cached_input_tokens: String,
pub average_first_token_latency_ms: String,
}
@@ -446,9 +447,9 @@ pub struct RequestMetricsView {
pub caller_error_count: u64,
pub input_tokens: u64,
pub output_tokens: u64,
pub cached_tokens: u64,
pub cache_write_tokens: u64,
pub reasoning_tokens: u64,
pub cached_input_tokens: u64,
pub cache_write_input_tokens: u64,
pub reasoning_output_tokens: u64,
pub total_tokens: u64,
}
@@ -575,6 +576,7 @@ impl From<domain::DesktopReleaseStatus> for DashboardDesktopReleaseStatusView {
#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
pub struct DashboardDataView {
pub analytics: AnalyticsView,
pub as_of: DateTime<Utc>,
pub as_of_display: String,
pub cards: DashboardCardsView,
@@ -592,11 +594,12 @@ pub struct DashboardDataView {
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct UsageSummaryView {
pub analytics: AnalyticsView,
pub total_requests: String,
pub input_tokens: String,
pub output_tokens: String,
pub cached_tokens: String,
pub cache_write_tokens: String,
pub cached_input_tokens: String,
pub cache_write_input_tokens: String,
pub total_tokens: String,
pub average_latency_ms: String,
pub logical_requests: RequestMetricsView,
@@ -693,7 +696,7 @@ pub struct OverviewCostPointView {
pub label: String,
pub input_tokens: u64,
pub output_tokens: u64,
pub cached_tokens: u64,
pub cached_input_tokens: u64,
pub total_tokens: u64,
pub estimated_cost: Option<String>,
pub standard_cost: Option<String>,
@@ -719,7 +722,7 @@ pub struct OverviewCostView {
pub cache_hit_request_rate: Option<f64>,
pub input_tokens: u64,
pub output_tokens: u64,
pub cached_tokens: u64,
pub cached_input_tokens: u64,
pub total_tokens: u64,
pub points: Vec<OverviewCostPointView>,
pub costs: Vec<CostView>,
@@ -741,6 +744,7 @@ pub struct ProviderOverviewView {
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct UsageInsightsOverviewView {
pub analytics: AnalyticsView,
pub granularity: String,
pub health: OverviewHealthView,
pub performance: OverviewPerformanceView,
@@ -778,6 +782,7 @@ pub struct DiagnosticItemView {
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct DiagnosticsView {
pub analytics: AnalyticsView,
pub dimension: String,
pub items: Vec<DiagnosticItemView>,
}
@@ -852,3 +857,20 @@ pub struct OpsErrorMetadataView {
pub recovery_retry_delay_ms: Option<u64>,
pub recovery_total_latency_ms: Option<u64>,
}
/// 统计覆盖、新鲜度与近似口径,时间文案使用部署时区
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AnalyticsView {
pub coverage_start: Option<DateTime<Utc>>,
pub coverage_start_display: Option<String>,
pub coverage_end: Option<DateTime<Utc>>,
pub coverage_end_display: Option<String>,
pub generated_at: DateTime<Utc>,
pub generated_at_display: String,
pub lag_seconds: u64,
pub state: String,
pub percentiles_approximate: bool,
pub schema_version: u32,
pub histogram_version: u32,
}
@@ -49,6 +49,7 @@ pub(super) fn config(secret: ClientKeySecret) -> ConfigView {
#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
pub(super) struct OverviewView {
analytics: crate::admin::observability::AnalyticsView,
as_of: DateTime<Utc>,
as_of_display: String,
start_time: DateTime<Utc>,
@@ -82,9 +83,9 @@ struct MetricsView {
requests: u64,
input_tokens: u64,
output_tokens: u64,
cached_tokens: u64,
cache_write_tokens: u64,
reasoning_tokens: u64,
cached_input_tokens: u64,
cache_write_input_tokens: u64,
reasoning_output_tokens: u64,
total_tokens: u64,
cost_usd: Option<String>,
cost_incomplete: bool,
@@ -95,9 +96,9 @@ fn metrics(value: &RequestMetrics, costs: &[CurrencyCost], coverage: &CostCovera
requests: value.request_count,
input_tokens: value.input_tokens,
output_tokens: value.output_tokens,
cached_tokens: value.cached_tokens,
cache_write_tokens: value.cache_write_tokens,
reasoning_tokens: value.reasoning_tokens,
cached_input_tokens: value.cached_input_tokens,
cache_write_input_tokens: value.cache_write_input_tokens,
reasoning_output_tokens: value.reasoning_output_tokens,
total_tokens: value.total_tokens,
cost_usd: costs
.iter()
@@ -121,6 +122,10 @@ struct TrendPointView {
pub(super) fn overview(value: KeyUsageOverview, time: crate::time::TimePresenter) -> OverviewView {
let key = value.key;
OverviewView {
analytics: crate::admin::observability::analytics_view(
value.overview.analytics.clone(),
time,
),
as_of: value.overview.range.end,
as_of_display: time.time(&value.overview.range.end),
start_time: value.overview.range.start,
@@ -19,6 +19,9 @@ use gateway_core::{
error::GatewayError,
};
pub(crate) const SSE_KEEPALIVE_INTERVAL: std::time::Duration = std::time::Duration::from_secs(15);
pub(crate) const SSE_KEEPALIVE_FRAME: Bytes = Bytes::from_static(b": keep-alive\n\n");
/// 只读目录/用量与协议错误共用的单帧响应;没有 Provider attempt 或计量副作用
pub(crate) async fn buffered_response(
protocol: &str,
@@ -161,8 +164,8 @@ pub(crate) async fn into_http_response(
let chunk = tokio::select! {
biased;
chunk = stream.next() => chunk?,
() = tokio::time::sleep(std::time::Duration::from_secs(15)), if keep_alive => {
Ok(Bytes::from_static(b": keep-alive\n\n"))
() = tokio::time::sleep(SSE_KEEPALIVE_INTERVAL), if keep_alive => {
Ok(SSE_KEEPALIVE_FRAME)
}
};
Some((chunk, stream))
@@ -4,6 +4,9 @@ mod http;
mod request;
mod session;
pub(crate) use http::{ExpectedBody, buffered_response, error_response, into_http_response};
pub(crate) use http::{
ExpectedBody, SSE_KEEPALIVE_FRAME, SSE_KEEPALIVE_INTERVAL, buffered_response, error_response,
into_http_response,
};
pub(crate) use request::{RequestInput, invoke_request, query_response, request_parts};
pub(crate) use session::{PendingExecution, finalize_session, pending_execution_response};
@@ -86,7 +86,9 @@ async fn models_response(
.iter()
.enumerate()
.map(|(index, entry)| match entry {
PublicModelDescriptor::Adapted(profile) => Ok(codex_model_json(profile, index)),
PublicModelDescriptor::Adapted(profile) => {
Ok(adapted_codex_model_json(profile, index))
}
PublicModelDescriptor::Native { model, payload } => {
if payload.protocol() != "codex" {
return Err(());
@@ -192,7 +194,7 @@ fn openai_model_json(id: &str) -> Value {
})
}
fn codex_model_json(profile: &PublicModelProfile, index: usize) -> Value {
fn adapted_codex_model_json(profile: &PublicModelProfile, index: usize) -> Value {
let id = profile.model().as_str();
let presentation = profile.presentation();
let reasoning_levels = presentation
@@ -265,6 +267,8 @@ fn codex_model_json(profile: &PublicModelProfile, index: usize) -> Value {
"experimental_supported_tools": [],
"input_modalities": input_modalities,
"supports_search_tool": presentation.search_tool(),
// 通用画像没有 Responses Lite 支持依据,避免客户端启用未经声明的协议行为
// 原生目录由 Native 分支保留上游字段,不经过此处
"use_responses_lite": false,
})
}
@@ -141,6 +141,7 @@ pub(crate) async fn execute_prepared_responses(
let cancellation = prepared.cancellation();
let started_request_id = Arc::new(OnceLock::new());
let terminal_request_id = Arc::clone(&started_request_id);
let (admitted, keepalive) = tokio::sync::oneshot::channel();
let mut middleware_headers = headers.clone();
// 中间件正文已经由入口按上限解压;传输编码和旧长度不能与改写后的正文混用
middleware_headers.remove(CONTENT_ENCODING);
@@ -213,37 +214,45 @@ pub(crate) async fn execute_prepared_responses(
.session
.trace()
.capture("client.request.body", &request_body);
let _ = admitted.send(started.session.response_commit());
execution_middleware_response(started, connection_guard, terminal_validation).await
})
}),
);
let middleware = if let Some(authorization) = authorization.as_ref() {
match await_with_authorization(middleware, authorization, &cancellation).await {
Ok(result) => result,
Err(error) => Err(MiddlewareError::Gateway(error)),
let response = async move {
let middleware = if let Some(authorization) = authorization.as_ref() {
match await_with_authorization(middleware, authorization, &cancellation).await {
Ok(result) => result,
Err(error) => Err(MiddlewareError::Gateway(error)),
}
} else {
middleware.await
};
let expected = if stream {
ExpectedBody::Sse
} else {
ExpectedBody::SingleJson
};
match middleware {
Ok(response) => {
let response = if let Some(authorization) = authorization {
authorized_middleware_response(response, authorization, cancellation)
} else {
response
};
into_http_response(
validated_middleware_response(response, validation, stream),
expected,
)
.await
}
Err(error) => error_response(error),
}
} else {
middleware.await
};
let expected = if stream {
ExpectedBody::Sse
let response = if stream {
super::keepalive::response(response, keepalive).await
} else {
ExpectedBody::SingleJson
};
let response = match middleware {
Ok(response) => {
let response = if let Some(authorization) = authorization {
authorized_middleware_response(response, authorization, cancellation)
} else {
response
};
into_http_response(
validated_middleware_response(response, validation, stream),
expected,
)
.await
}
Err(error) => error_response(error),
response.await
};
super::super::with_model_request_id(response, started_request_id.get().unwrap_or(&request_id))
}
@@ -519,16 +528,24 @@ pub async fn stream_execution_response(
session: Box<dyn ExecutionSession>,
connection_guard: Option<Box<dyn ConnectionGuard>>,
) -> Response {
match streaming_execution_middleware_response(
session,
connection_guard,
ResponseValidationFacts::default(),
let (admitted, keepalive) = tokio::sync::oneshot::channel();
let _ = admitted.send(session.response_commit());
super::keepalive::response(
async move {
match streaming_execution_middleware_response(
session,
connection_guard,
ResponseValidationFacts::default(),
)
.await
{
Ok(response) => into_http_response(response, ExpectedBody::Sse).await,
Err(error) => error_response(error),
}
},
keepalive,
)
.await
{
Ok(response) => into_http_response(response, ExpectedBody::Sse).await,
Err(error) => error_response(error),
}
}
fn validated_middleware_response(
@@ -0,0 +1,151 @@
//! 准入成功后覆盖 Responses 首帧等待;已发心跳的 HTTP 错误转为流内失败
use std::future::Future;
use axum::{
body::{Body, BodyDataStream, Bytes, to_bytes},
http::{HeaderValue, StatusCode, header},
response::Response,
};
use futures::{StreamExt as _, future::BoxFuture, stream};
use gateway_core::engine::execution::ResponseCommit;
use gateway_protocol::openai::sse::{
DONE_SSE_FRAME, encode_sse_event, response_failed_sse_data_from_error_event,
response_failed_sse_data_with_id,
};
use serde_json::Value;
use tokio::sync::oneshot;
use crate::openai::middleware::{SSE_KEEPALIVE_FRAME, SSE_KEEPALIVE_INTERVAL};
pub(super) async fn response(
response: impl Future<Output = Response> + Send + 'static,
admitted: oneshot::Receiver<ResponseCommit>,
) -> Response {
let mut response = Box::pin(response);
// 鉴权、模型门禁和准入失败仍保留原 HTTP 状态,不提前建立 SSE
let commit = tokio::select! {
biased;
response = &mut response => return response,
admitted = admitted => match admitted {
Ok(commit) => commit,
Err(_) => return response.await,
},
};
tokio::select! {
biased;
response = &mut response => return response,
() = tokio::time::sleep(SSE_KEEPALIVE_INTERVAL) => {}
}
if !commit.record(Some(StatusCode::OK.as_u16())) {
return response.await;
}
let pending = PendingResponse::Waiting(response);
let stream = stream::once(async { Ok::<_, axum::Error>(SSE_KEEPALIVE_FRAME) })
.chain(stream::try_unfold(pending, PendingResponse::next));
let mut response = Response::new(Body::from_stream(stream));
let headers = response.headers_mut();
headers.insert(
header::CONTENT_TYPE,
HeaderValue::from_static("text/event-stream"),
);
headers.insert(
header::CACHE_CONTROL,
HeaderValue::from_static("no-cache, no-transform"),
);
headers.insert("x-accel-buffering", HeaderValue::from_static("no"));
response
}
enum PendingResponse {
Waiting(BoxFuture<'static, Response>),
Streaming(BodyDataStream),
}
impl PendingResponse {
async fn next(mut self) -> Result<Option<(Bytes, Self)>, axum::Error> {
loop {
match &mut self {
Self::Waiting(response) => {
// 心跳只轮询同一个 future,不取消重建首帧等待或插件调用
let response = tokio::select! {
biased;
response = response => response,
() = tokio::time::sleep(SSE_KEEPALIVE_INTERVAL) => return Ok(Some((SSE_KEEPALIVE_FRAME, self))),
};
let sse = response.status().is_success()
&& response
.headers()
.get(header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.and_then(|value| value.split(';').next())
.is_some_and(|value| {
value.trim().eq_ignore_ascii_case("text/event-stream")
});
let body = if sse {
response.into_body()
} else {
stream_failure(response).await
};
// 正式响应的正文已经带有交付期心跳,后续只转交,不叠加定时器
self = Self::Streaming(body.into_data_stream());
}
Self::Streaming(body) => {
return body
.next()
.await
.transpose()
.map(|chunk| chunk.map(|chunk| (chunk, self)));
}
}
}
}
}
async fn stream_failure(response: Response) -> Body {
let status = response.status();
let retry_after = response
.headers()
.get(header::RETRY_AFTER)
.and_then(|value| value.to_str().ok())
.map(ToOwned::to_owned);
let bytes = to_bytes(response.into_body(), usize::MAX)
.await
.unwrap_or_default();
let json = serde_json::from_slice::<Value>(&bytes).ok();
let raw_message = std::str::from_utf8(&bytes)
.ok()
.map(str::trim)
.filter(|text| !text.is_empty());
// 已完成客户端错误投影的对象原样保留,包括重试建议和策略错误的扩展字段
let mut data = json
.as_ref()
.and_then(|json| response_failed_sse_data_from_error_event(None, None, json))
.unwrap_or_else(|| {
response_failed_sse_data_with_id(
None,
"server_error",
"upstream_error",
raw_message.unwrap_or_else(|| {
status
.canonical_reason()
.unwrap_or("Response stream failed")
}),
)
});
// HTTP 头已提交,按官方 SSE 错误合同保留允许下发的重试建议
if let Some(retry_after) = retry_after
&& let Some(error) = data["response"]["error"].as_object_mut()
&& let Some(headers) = error
.entry("headers")
.or_insert_with(|| serde_json::json!({}))
.as_object_mut()
{
headers
.entry("retry-after")
.or_insert(Value::String(retry_after));
}
let mut event = encode_sse_event("response.failed", &data.to_string());
event.push_str(DONE_SSE_FRAME);
Body::from(event)
}
@@ -2,6 +2,7 @@
mod error;
mod http;
mod keepalive;
mod request;
mod response;
mod validation;
@@ -150,10 +150,12 @@ pub(super) async fn forward_response(
ProviderResponseHeader::new(name, value)
})
.collect::<Vec<_>>();
let mut interrupt = super::interrupt::ResponseInterrupt::default();
let first = match next_body_input(
connection,
body.as_mut(),
response_control,
&mut interrupt,
&request_id,
false,
)
@@ -207,6 +209,7 @@ pub(super) async fn forward_response(
connection,
body.as_mut(),
response_control,
&mut interrupt,
&request_id,
terminal_seen,
)
@@ -274,11 +277,24 @@ pub(super) async fn forward_response(
(true, true) => FramePhase::Terminal,
(true, false) => FramePhase::Data,
};
let created_id = super::interrupt::ResponseInterrupt::created_id(&message);
match connection
.send_text(message, WriteContext::request(&request_id, phase))
.await
{
Ok(super::connection::WriteOutcome::Written) => first_frame_written = true,
Ok(super::connection::WriteOutcome::Written) => {
first_frame_written = true;
if let Some(id) = created_id {
interrupt.observe_created(id, response_control);
tokio::select! {
_ = connection.wait_for_exit() => {
detach_body(body);
return ForwardOutcome::Disconnect;
}
() = interrupt.forward_pending(response_control) => {}
}
}
}
Ok(super::connection::WriteOutcome::Suppressed) => {}
Err(_) => {
detach_body(body);
@@ -315,6 +331,7 @@ async fn next_body_input(
connection: &mut ResponsesWebSocketConnection,
body: &mut dyn MiddlewareBody,
response_control: &ResponseControl,
interrupt: &mut super::interrupt::ResponseInterrupt,
request_id: &Arc<str>,
terminal_seen: bool,
) -> BodyInput {
@@ -331,6 +348,13 @@ async fn next_body_input(
connection.defer(event);
continue;
}
if !terminal_seen && interrupt.retain(payload) {
tokio::select! {
_ = connection.wait_for_exit() => return BodyInput::Disconnect,
() = interrupt.forward_pending(response_control) => {}
}
continue;
}
if send_control(connection, response_control, payload, request_id).await == ForwardOutcome::Disconnect {
return BodyInput::Disconnect;
}
@@ -0,0 +1,91 @@
//! 当前逻辑响应的中断意图,跨重试连接保留并随新响应 ID 转发
use std::{borrow::Cow, collections::BTreeSet};
use gateway_core::engine::response_control::{ResponseControl, ResponseControlBinding};
use serde::Deserialize;
use serde_json::Value;
#[derive(Deserialize)]
struct EventIdentity<'a> {
#[serde(rename = "type", borrow)]
event_type: Cow<'a, str>,
#[serde(borrow)]
response: Option<ResponseIdentity<'a>>,
}
#[derive(Deserialize)]
struct ResponseIdentity<'a> {
#[serde(borrow)]
id: Cow<'a, str>,
}
#[derive(Default)]
pub(super) struct ResponseInterrupt {
delivered_ids: BTreeSet<String>,
current: Option<(String, ResponseControlBinding)>,
pending: Option<Value>,
sent: bool,
}
impl ResponseInterrupt {
pub(super) fn created_id(payload: &str) -> Option<String> {
let event: EventIdentity<'_> = serde_json::from_str(payload).ok()?;
(event.event_type == "response.created").then_some(())?;
Some(event.response?.id.into_owned())
}
/// 只在 created 实际写入下游后更新,不能把未交付的 ID 当成客户端目标
pub(super) fn observe_created(&mut self, id: String, control: &ResponseControl) {
let Some(binding) = control.binding() else {
return;
};
if self
.current
.as_ref()
.is_none_or(|(previous, old)| previous != &id || !old.same_binding(&binding))
{
self.sent = false;
}
self.delivered_ids.insert(id.clone());
self.current = Some((id, binding));
}
/// 未知控制消息、模式和其他逻辑响应的 ID 仍走原始透明通道
pub(super) fn retain(&mut self, payload: &str) -> bool {
let Ok(event) = serde_json::from_str::<Value>(payload) else {
return false;
};
if event["type"] != "response.interrupt"
|| event["mode"] != "discard_partial_items"
|| !event["response_id"]
.as_str()
.is_some_and(|id| self.delivered_ids.contains(id))
{
return false;
}
self.pending = Some(event);
self.sent = false;
true
}
pub(super) async fn forward_pending(&mut self, control: &ResponseControl) {
if self.sent {
return;
}
let (Some(payload), Some((id, binding))) = (&mut self.pending, &self.current) else {
return;
};
// 换号已绑定、created 尚未交付时等待新 ID;连接快照保证 await 期间也不会串号
if !control
.binding()
.is_some_and(|current| binding.same_binding(&current))
{
return;
}
payload["response_id"] = Value::String(id.clone());
if binding.send(&payload.to_string()).await.is_ok() {
self.sent = true;
}
}
}
@@ -2,6 +2,7 @@
pub mod connection;
mod forward;
mod interrupt;
mod protocol;
use std::{
@@ -12,13 +12,8 @@ use super::super::{AdminTestFixture, AdminTestState};
#[tokio::test]
async fn account_cooldown_returns_an_instant_and_server_formatted_recovery() {
use chrono::{DateTime, Duration, Utc};
use gateway_admin::model::{
Revision,
accounts::{AccountCapacity, AccountPageItem, AccountRecord},
};
use gateway_core::account::{
AccountCooldown, AccountCooldownKind, AccountStatusFacts, CredentialState, QuotaState,
resolve_account_status,
};
for zone in ["Asia/Shanghai", "UTC", "Asia/Kathmandu", "America/New_York"] {
@@ -65,41 +60,7 @@ async fn account_cooldown_returns_an_instant_and_server_formatted_recovery() {
last_error_reason: None,
last_error_message: None,
};
*fixture.account.lock().unwrap() = Some(AccountPageItem {
account: AccountRecord {
id: "acct_cooldown".to_owned(),
provider_kind: gateway_core::routing::ProviderKind::new("openai").unwrap(),
groups: Vec::new(),
name: "synthetic cooldown account".to_owned(),
notes: None,
email: None,
upstream_user_id: None,
upstream_account_id: None,
plan_type: None,
authentication_kind: "oauth".to_owned(),
credential_revision: Revision::new(1).unwrap(),
has_refresh_token: true,
access_token_expires_at: None,
next_refresh_at: None,
enabled: true,
concurrency_limit: None,
weight: Default::default(),
model_access: Default::default(),
outbound_proxy: None,
credential_state: facts.credential_state,
credential_observed_at: now,
quota: facts.quota,
last_error_reason: None,
last_error_message: None,
created_at: now,
updated_at: now,
},
capacity: AccountCapacity {
used_slots: None,
total_slots: None,
},
projection: resolve_account_status(&facts, now.into()),
});
*fixture.account.lock().unwrap() = Some(account_page_item(facts, now, "oauth"));
let response = admin::accounts::router::<AdminTestState>()
.with_state(fixture.state())
.oneshot(
@@ -138,6 +99,132 @@ async fn account_cooldown_returns_an_instant_and_server_formatted_recovery() {
}
}
fn account_page_item(
facts: gateway_core::account::AccountStatusFacts,
now: chrono::DateTime<chrono::Utc>,
authentication_kind: &str,
) -> gateway_admin::model::accounts::AccountPageItem {
use gateway_admin::model::{
Revision,
accounts::{AccountCapacity, AccountPageItem, AccountRecord},
};
use gateway_core::account::resolve_account_status;
AccountPageItem {
account: AccountRecord {
id: "acct_cooldown".to_owned(),
provider_kind: gateway_core::routing::ProviderKind::new("openai").unwrap(),
groups: Vec::new(),
name: "synthetic cooldown account".to_owned(),
notes: None,
email: None,
upstream_user_id: None,
upstream_account_id: None,
plan_type: None,
authentication_kind: authentication_kind.to_owned(),
credential_revision: Revision::new(1).unwrap(),
has_refresh_token: authentication_kind == "oauth",
access_token_expires_at: None,
next_refresh_at: None,
enabled: true,
concurrency_limit: None,
weight: Default::default(),
model_access: Default::default(),
outbound_proxy: None,
credential_state: facts.credential_state,
credential_observed_at: now,
quota: facts.quota,
last_error_reason: None,
last_error_message: None,
created_at: now,
updated_at: now,
},
capacity: AccountCapacity {
used_slots: None,
total_slots: None,
},
projection: resolve_account_status(&facts, now.into()),
}
}
#[tokio::test]
async fn account_usage_returns_distinct_cache_read_and_write_tokens_without_aliases() {
use chrono::Utc;
use gateway_admin::model::accounts::AccountUsage;
use gateway_core::account::{AccountStatusFacts, CredentialState, QuotaState};
let fixture = AdminTestFixture::new().await;
fixture.auth.insert_session("valid-session");
let now = Utc::now();
let facts = AccountStatusFacts {
enabled: true,
credential_state: CredentialState::Ready,
access_token_expires_at: None,
quota: QuotaState::unknown(),
cooldown: None,
last_error_reason: None,
last_error_message: None,
};
*fixture.account.lock().unwrap() = Some(account_page_item(facts, now, "api_key"));
for (read, write) in [(Some(13), Some(7)), (None, None)] {
*fixture.account_usage.lock().unwrap() = vec![AccountUsage {
account_id: "acct_cooldown".to_owned(),
request_count: 1,
success_count: 1,
input_tokens: Some(100),
output_tokens: Some(20),
cached_input_tokens: read,
cache_write_input_tokens: write,
reasoning_output_tokens: Some(3),
image_input_tokens: None,
image_output_tokens: None,
image_request_count: 0,
image_request_failed_count: 0,
total_tokens: Some(120),
cost_coverage: Default::default(),
costs: Vec::new(),
last_used_at: Some(now),
request_buckets: Vec::new(),
models: Vec::new(),
}];
let response = admin::accounts::router::<AdminTestState>()
.with_state(fixture.state())
.oneshot(
Request::builder()
.uri("/api/admin/accounts")
.header("x-request-id", "req_account_usage_contract")
.header(header::COOKIE, "cpr_session=valid-session")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
let status = response.status();
let body = to_bytes(response.into_body(), 32768).await.unwrap();
let value: serde_json::Value = serde_json::from_slice(&body).unwrap();
assert_eq!(status, StatusCode::OK, "{value}");
let usage = &value["data"]["items"][0]["usage"];
assert_eq!(usage["cachedInputTokens"], serde_json::json!(read));
assert_eq!(usage["cacheWriteInputTokens"], serde_json::json!(write));
assert_eq!(
usage["cachedInputTokensDisplay"],
if read.is_some() { "13" } else { "—" }
);
assert_eq!(
usage["cacheWriteInputTokensDisplay"],
if write.is_some() { "7" } else { "—" }
);
for alias in [
"createdTokens",
"createdTokensDisplay",
"readTokens",
"readTokensDisplay",
] {
assert!(usage.get(alias).is_none(), "obsolete token alias: {alias}");
}
}
}
#[tokio::test]
async fn standalone_credential_rotation_route_is_not_exposed() {
let fixture = AdminTestFixture::new().await;
@@ -484,10 +484,10 @@ mod response {
input_tokens_display: "-".to_owned(),
output_tokens: None,
output_tokens_display: "-".to_owned(),
cached_tokens: None,
cached_tokens_display: "-".to_owned(),
reasoning_tokens: None,
reasoning_tokens_display: "-".to_owned(),
cached_input_tokens: None,
cached_input_tokens_display: "-".to_owned(),
reasoning_output_tokens: None,
reasoning_output_tokens_display: "-".to_owned(),
image_input_tokens: None,
image_input_tokens_display: "-".to_owned(),
image_output_tokens: None,
@@ -498,10 +498,8 @@ mod response {
image_request_failed_count_display: "-".to_owned(),
total_tokens: None,
total_tokens_display: "-".to_owned(),
created_tokens: None,
created_tokens_display: "-".to_owned(),
read_tokens: None,
read_tokens_display: "-".to_owned(),
cache_write_input_tokens: None,
cache_write_input_tokens_display: "-".to_owned(),
last_used_at: None,
last_used_at_full_display: None,
last_used_at_display: "-".to_owned(),
@@ -516,9 +514,9 @@ mod response {
assert_eq!(value["windowLabelDisplay"], "周/月额度窗口");
assert!(value["inputTokens"].is_null());
assert!(value["totalTokens"].is_null());
assert!(value["reasoningTokens"].is_null());
assert_eq!(value["reasoningTokensDisplay"], "-");
assert_eq!(value["createdTokensDisplay"], "-");
assert!(value["reasoningOutputTokens"].is_null());
assert_eq!(value["reasoningOutputTokensDisplay"], "-");
assert_eq!(value["cacheWriteInputTokensDisplay"], "-");
}
#[test]
+62 -10
View File
@@ -112,6 +112,7 @@ pub(super) struct AdminTestFixture {
pub dashboard_summary_range: Arc<Mutex<Option<TimeRange>>>,
pub provider_error: Arc<Mutex<Option<ProviderAdminError>>>,
pub account: Arc<Mutex<Option<AccountPageItem>>>,
pub account_usage: Arc<Mutex<Vec<AccountUsage>>>,
}
impl AdminTestFixture {
@@ -154,6 +155,7 @@ impl AdminTestFixture {
let dashboard_summary_range = Arc::new(Mutex::new(None));
let provider_error = Arc::new(Mutex::new(None));
let account = Arc::new(Mutex::new(None));
let account_usage = Arc::new(Mutex::new(Vec::new()));
let unused = Arc::new(UnusedStore {
observations: observations.clone(),
usage_records: Arc::clone(&usage_records),
@@ -163,6 +165,7 @@ impl AdminTestFixture {
dashboard_observation: Arc::clone(&dashboard_observation),
dashboard_summary_range: Arc::clone(&dashboard_summary_range),
account: Arc::clone(&account),
account_usage: Arc::clone(&account_usage),
});
let plugin_ports = Arc::new(plugins::TestPluginPorts::default());
let published_snapshot = gateway_core::runtime::RuntimeSnapshotHandle::default();
@@ -232,6 +235,7 @@ impl AdminTestFixture {
dashboard_summary_range,
provider_error,
account,
account_usage,
}
}
@@ -1100,6 +1104,7 @@ struct UnusedStore {
dashboard_observation: Arc<Mutex<Option<DashboardObservation>>>,
dashboard_summary_range: Arc<Mutex<Option<TimeRange>>>,
account: Arc<Mutex<Option<AccountPageItem>>>,
account_usage: Arc<Mutex<Vec<AccountUsage>>>,
}
struct UnusedClientKeyVerifier;
@@ -1161,18 +1166,41 @@ impl AccountStore for UnusedStore {
async fn load_account_usage(
&self,
_: TimeRange,
_: &[String],
account_ids: &[String],
) -> AdminStoreResult<Vec<AccountUsage>> {
if self.account.lock().expect("account").is_some() {
return Ok(Vec::new());
return Ok(self
.account_usage
.lock()
.expect("account usage")
.iter()
.filter(|usage| account_ids.contains(&usage.account_id))
.cloned()
.collect());
}
Err(unavailable("account usage"))
}
async fn load_account_usage_by_windows(
&self,
_: &[AccountUsageWindowQuery],
windows: &[AccountUsageWindowQuery],
) -> AdminStoreResult<Vec<AccountUsageWindowResult>> {
if self.account.lock().expect("account").is_some() {
let usage = self.account_usage.lock().expect("account usage");
return Ok(windows
.iter()
.filter_map(|window| {
usage
.iter()
.find(|usage| usage.account_id == window.account_id)
.map(|usage| AccountUsageWindowResult {
account_id: window.account_id.clone(),
key: window.key.clone(),
usage: usage.clone(),
})
})
.collect());
}
Err(unavailable("account quota window usage"))
}
@@ -1381,15 +1409,39 @@ impl ObservabilityStore for UnusedStore {
data.trend.clone().ok_or_else(|| unavailable("usage trend"))
}
fn usage_calculated_billing_facts(
async fn analytics_cache_generation(&self) -> AdminStoreResult<u64> {
Ok(self
.observations
.lock()
.expect("observations")
.summary
.as_ref()
.map_or(0, |summary| summary.analytics.generation))
}
async fn usage_insights_observation(
&self,
_: TimeRange,
_: UsageFilter,
range: TimeRange,
filter: UsageFilter,
_: Granularity,
) -> gateway_admin::ports::store::UsageCalculatedBillingStream<'_> {
Box::pin(futures::stream::once(async {
Err(unavailable("usage billing facts"))
}))
) -> AdminStoreResult<gateway_admin::model::observability::UsageInsightsObservation> {
let mut data = self.observations.lock().expect("observations");
data.summaries.push((range, filter.clone()));
data.trends.push((range, filter));
Ok(
gateway_admin::model::observability::UsageInsightsObservation {
overview: data
.summary
.clone()
.map(|summary| UsageOverview { range, ..summary })
.ok_or_else(|| unavailable("usage summary"))?,
trend: data
.trend
.clone()
.ok_or_else(|| unavailable("usage trend"))?,
costs: Vec::new(),
},
)
}
async fn list_usage_records(&self, query: UsageQuery) -> AdminStoreResult<UsagePage> {
@@ -25,6 +25,17 @@ async fn dashboard_summary_should_include_quota_exhaustion_in_unavailable_headli
.dashboard_observation
.lock()
.expect("dashboard observation") = Some(DashboardObservation {
analytics: gateway_admin::model::observability::AnalyticsMetadata {
coverage_start: Some(now - Duration::hours(1)),
coverage_end: Some(now),
generated_at: now,
lag_seconds: 12,
state: gateway_admin::model::observability::AnalyticsState::Lagging,
percentiles_approximate: true,
schema_version: 1,
histogram_version: 1,
generation: 42,
},
runtime_slots: None,
range: TimeRange::new(now - Duration::hours(1), now).expect("dashboard range"),
totals: Default::default(),
@@ -59,6 +70,14 @@ async fn dashboard_summary_should_include_quota_exhaustion_in_unavailable_headli
.expect("dashboard response body");
let value: serde_json::Value = serde_json::from_slice(&body).expect("dashboard response JSON");
assert_eq!(status, StatusCode::OK, "{value}");
assert_eq!(value["data"]["analytics"]["state"], "lagging");
assert_eq!(value["data"]["analytics"]["lagSeconds"], 12);
assert_eq!(value["data"]["analytics"]["percentilesApproximate"], true);
assert_eq!(value["data"]["analytics"]["schemaVersion"], 1);
assert_eq!(value["data"]["analytics"]["histogramVersion"], 1);
assert!(value["data"]["analytics"]["coverageStart"].is_string());
assert!(value["data"]["analytics"]["generatedAtDisplay"].is_string());
assert!(value["data"]["analytics"].get("generation").is_none());
assert_eq!(
(
@@ -108,12 +127,13 @@ async fn dashboard_summary_should_use_lifetime_totals_for_card_footers() {
.dashboard_observation
.lock()
.expect("dashboard observation") = Some(DashboardObservation {
analytics: Default::default(),
runtime_slots: None,
range: TimeRange::new(now - Duration::hours(1), now).expect("dashboard range"),
totals: DashboardTotals {
request_count: 42,
input_tokens: 800,
cached_tokens: 200,
cached_input_tokens: 200,
total_tokens: 840,
billing_usd: None,
},
@@ -141,13 +161,21 @@ async fn dashboard_summary_should_use_lifetime_totals_for_card_footers() {
.expect("dashboard response body");
let value: serde_json::Value = serde_json::from_slice(&body).expect("dashboard response JSON");
assert_eq!(status, StatusCode::OK, "{value}");
assert_eq!(
value["data"]["analytics"]["coverageStartDisplay"],
serde_json::Value::Null
);
assert_eq!(
value["data"]["analytics"]["coverageEndDisplay"],
serde_json::Value::Null
);
assert_eq!(
(
&value["data"]["cards"]["traffic"]["totalRequests"],
&value["data"]["cards"]["tokens"]["totalTokens"],
&value["data"]["cards"]["cache"]["totalHitRate"],
&value["data"]["cards"]["cache"]["totalCachedTokens"],
&value["data"]["cards"]["cache"]["totalCachedInputTokens"],
),
(
&serde_json::json!("42"),
@@ -180,6 +208,7 @@ async fn dashboard_summary_should_default_to_current_deployment_day() {
.dashboard_observation
.lock()
.expect("dashboard observation") = Some(DashboardObservation {
analytics: Default::default(),
runtime_slots: None,
range: TimeRange::new(now - Duration::hours(1), now).expect("dashboard range"),
totals: Default::default(),
@@ -282,9 +311,9 @@ async fn usage_detail_should_keep_attempt_snapshot_contract() {
message: Some("limited".to_owned()),
input_tokens: None,
output_tokens: None,
cached_tokens: None,
cache_write_tokens: None,
reasoning_tokens: None,
cached_input_tokens: None,
cache_write_input_tokens: None,
reasoning_output_tokens: None,
total_tokens: None,
cost_source: Some("unavailable".to_owned()),
cost_amount: None,
@@ -401,9 +430,9 @@ async fn zero_attempt_failure_detail_keeps_missing_upstream_facts_and_preparatio
record.error_message = Some("no upstream account is available".to_owned());
record.input_tokens = None;
record.output_tokens = None;
record.cached_tokens = None;
record.cache_write_tokens = None;
record.reasoning_tokens = None;
record.cached_input_tokens = None;
record.cache_write_input_tokens = None;
record.reasoning_output_tokens = None;
record.image_input_tokens = None;
record.image_output_tokens = None;
record.total_tokens = None;
@@ -617,6 +646,7 @@ async fn diagnostics_should_keep_stable_key_and_display_name_contract() {
let fixture = AdminTestFixture::new().await;
fixture.auth.insert_session("valid-session");
*fixture.diagnostics.lock().expect("diagnostics") = DiagnosticsObservation {
analytics: Default::default(),
total_request_count: 4,
items: vec![DiagnosticObservation {
key: "acct_diag".to_owned(),
@@ -726,9 +756,9 @@ fn usage_record_with_account(
retry_after_ms: None,
input_tokens: Some(1),
output_tokens: Some(1),
cached_tokens: Some(0),
cache_write_tokens: Some(0),
reasoning_tokens: Some(0),
cached_input_tokens: Some(0),
cache_write_input_tokens: Some(0),
reasoning_output_tokens: Some(0),
image_input_tokens: Some(0),
image_output_tokens: Some(0),
total_tokens: Some(2),
@@ -847,9 +877,9 @@ async fn usage_route_should_expose_table_facts_without_detail_payload() {
service_tier: Some("default".to_owned()),
input_tokens: Some(1),
output_tokens: Some(1),
cached_tokens: Some(0),
cache_write_tokens: Some(0),
reasoning_tokens: Some(0),
cached_input_tokens: Some(0),
cache_write_input_tokens: Some(0),
reasoning_output_tokens: Some(0),
image_input_tokens: Some(31),
image_output_tokens: Some(9),
total_tokens: Some(2),
@@ -1101,6 +1131,12 @@ async fn calendar_queries_and_display_follow_deployment_timezone_with_one_anchor
let fixture = AdminTestFixture::with_timezone(timezone).await;
fixture.auth.insert_session("valid-session");
*fixture.dashboard_observation.lock().unwrap() = Some(DashboardObservation {
analytics: gateway_admin::model::observability::AnalyticsMetadata {
coverage_start: Some(range.start),
coverage_end: Some(end),
generated_at: end,
..Default::default()
},
runtime_slots: None,
range,
totals: Default::default(),
@@ -1142,6 +1178,15 @@ async fn calendar_queries_and_display_follow_deployment_timezone_with_one_anchor
assert_eq!(queried.end, end);
assert_eq!(value["data"]["asOf"], anchor);
assert_eq!(value["data"]["asOfDisplay"], display);
let analytics = &value["data"]["analytics"];
assert_eq!(analytics["coverageStart"], start);
assert_eq!(analytics["coverageEnd"], anchor);
assert_eq!(
analytics["coverageStartDisplay"],
format!("{} 00:00:00", &display[..10])
);
assert_eq!(analytics["coverageEndDisplay"], display);
assert_eq!(analytics["generatedAtDisplay"], display);
let points = value["data"]["healthTimeline"]["points"]
.as_array()
.unwrap();
@@ -1198,3 +1243,52 @@ async fn calendar_query_rejects_an_anchor_outside_the_local_calendar() {
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
assert!(fixture.dashboard_summary_range.lock().unwrap().is_none());
}
#[tokio::test]
async fn request_and_ops_queries_should_forward_deleted_group_id_sets() {
use crate::admin::{AdminTestFixture, AdminTestState};
use axum::{
body::Body,
http::{Request, StatusCode, header},
};
use gateway_api::admin::observability;
use tower::ServiceExt as _;
let fixture = AdminTestFixture::new().await;
fixture.auth.insert_session("valid-session");
let router = observability::router::<AdminTestState>().with_state(fixture.state());
for path in ["/api/admin/usage/records", "/api/admin/operations/errors"] {
let response = router
.clone()
.oneshot(
Request::builder()
.uri(format!(
"{path}?routingGroupRefs=deleted-b,deleted-a,deleted-b"
))
.header(header::COOKIE, "cpr_session=valid-session")
.header("x-request-id", "req_group_filter")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
let status = response.status();
let body = axum::body::to_bytes(response.into_body(), 1024 * 1024)
.await
.unwrap();
assert_eq!(
status,
StatusCode::OK,
"{path}: {}",
String::from_utf8_lossy(&body)
);
}
let observations = fixture.observations.lock().unwrap();
assert_eq!(
observations.records[0].filter.routing_group_refs,
["deleted-a", "deleted-b"]
);
assert_eq!(
observations.errors[0].filter.routing_group_refs,
["deleted-a", "deleted-b"]
);
}
@@ -48,8 +48,8 @@ pub(super) async fn fixture() -> AdminTestFixture {
failure_count: 1,
input_tokens: 1000,
output_tokens: 100,
cached_tokens: 900,
reasoning_tokens: 40,
cached_input_tokens: 900,
reasoning_output_tokens: 40,
total_tokens: 1100,
..Default::default()
};
@@ -65,6 +65,7 @@ pub(super) async fn fixture() -> AdminTestFixture {
{
let mut observations = fixture.observations.lock().unwrap();
observations.summary = Some(UsageOverview {
analytics: Default::default(),
range: TimeRange {
start: now - Duration::hours(1),
end: now,
@@ -111,9 +112,9 @@ fn usage_record() -> UsageListRecord {
service_tier: Some("private-sentinel".to_owned()),
input_tokens: Some(1000),
output_tokens: Some(100),
cached_tokens: Some(900),
cache_write_tokens: Some(0),
reasoning_tokens: Some(40),
cached_input_tokens: Some(900),
cache_write_input_tokens: Some(0),
reasoning_output_tokens: Some(40),
image_input_tokens: None,
image_output_tokens: None,
total_tokens: Some(1100),
@@ -181,6 +181,7 @@ async fn overview_scopes_every_query_and_projects_only_key_visible_fields() {
"summary",
"trend",
"healthTimeline",
"analytics",
],
);
assert_fields(
@@ -203,8 +204,8 @@ async fn overview_scopes_every_query_and_projects_only_key_visible_fields() {
assert_eq!(data["key"]["dailyUsedUsd"], "0.640001");
assert_eq!(data["key"]["maxConcurrency"], 0);
assert_eq!(data["summary"]["totalTokens"], 1100);
assert_eq!(data["summary"]["reasoningTokens"], 40);
assert_eq!(data["trend"][0]["reasoningTokens"], 40);
assert_eq!(data["summary"]["reasoningOutputTokens"], 40);
assert_eq!(data["trend"][0]["reasoningOutputTokens"], 40);
assert_eq!(data["summary"]["costUsd"], "0.123456");
assert_eq!(data["summary"]["costIncomplete"], true);
assert_eq!(data["trend"][0]["bucketSeconds"], 900);
@@ -300,7 +301,7 @@ async fn records_keep_pagination_and_hide_admin_and_upstream_data() {
assert!(!data.to_string().contains("private-sentinel"));
if kind == "success" {
assert_eq!(record["billing"]["totalAmountDisplay"], "$0.1235");
assert_eq!(record["tokenDetails"]["reasoningTokens"], 40);
assert_eq!(record["tokenDetails"]["reasoningOutputTokens"], 40);
assert_eq!(record["tokenDetails"]["totalTokens"], 1100);
assert_eq!(record["firstTokenLatencyMs"], 210);
assert_fields(
@@ -430,6 +431,7 @@ async fn empty_usage_has_zero_cost_but_unknown_pricing_stays_unknown() {
{
let mut data = fixture.observations.lock().unwrap();
let summary = data.summary.as_mut().unwrap();
summary.analytics.generation += 1;
summary.requests = RequestMetrics {
request_count: requests,
..Default::default()
@@ -13,6 +13,40 @@ use gateway_api::openai::error::{
gateway_error_response, openai_error_response,
};
#[tokio::test]
async fn ordinary_upstream_rejections_preserve_status_body_and_client_retry_advice() {
use bytes::Bytes;
use gateway_core::error::ClientVisibleUpstreamResponse;
use gateway_core::event::ProviderResponseHeader;
let body = Bytes::from_static(br#"{ "error":{"code":"Vendor.Future.v2","type":"server_error","message":"original failure","future":9007199254740993} }"#);
for status in [400, 409, 500, 503, 504] {
let error = EngineError::Provider(
ProviderError::new(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent)
.with_retry_prohibited()
.with_client_visible_upstream_response(
ClientVisibleUpstreamResponse::new(
status,
Some(b"application/problem+json".to_vec()),
body.clone(),
)
.with_headers(vec![ProviderResponseHeader::new(
"retry-after",
Bytes::from_static(b"17"),
)]),
),
);
let response = engine_error_response(&error);
assert_eq!(response.status().as_u16(), status);
assert_eq!(
response.headers()["content-type"],
"application/problem+json"
);
assert_eq!(response.headers()["retry-after"], "17");
assert_eq!(to_bytes(response.into_body(), 4096).await.unwrap(), body);
}
}
#[tokio::test]
async fn capacity_errors_offer_client_retry_without_changing_upstream_facts() {
use bytes::Bytes;
+17 -4
View File
@@ -283,12 +283,25 @@ fn snapshot_with_client_key(
plaintext: &str,
provider_name: &str,
client_key_id: &str,
) -> RuntimeSnapshot {
snapshot_with_accounts(plaintext, provider_name, client_key_id, &["acct_api_test"])
}
fn snapshot_with_accounts(
plaintext: &str,
provider_name: &str,
client_key_id: &str,
accounts: &[&str],
) -> RuntimeSnapshot {
let provider = ProviderKind::new(provider_name).expect("provider");
let account_directory = Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([(
ProviderAccountId::new("acct_api_test").expect("account ID"),
RuntimeAccount::new(provider.clone(), BTreeSet::new()),
)])));
let account_directory = Arc::new(RuntimeAccountDirectory::new(BTreeMap::from_iter(
accounts.iter().map(|id| {
(
ProviderAccountId::new(*id).expect("account ID"),
RuntimeAccount::new(provider.clone(), BTreeSet::new()),
)
}),
)));
let capabilities = ModelCapabilities::new(
BTreeSet::from([gateway_core::operation::OperationKind::Generate]),
Some(16_000),
@@ -23,7 +23,7 @@ use futures::StreamExt;
use futures::future::{BoxFuture, pending};
use gateway_core::engine::execution::{
AuthenticatedClient, ClientAuthenticationError, ExecutionService, ExecutionSession,
StartExecution, StartProviderExecution, StartedExecution,
ResponseCommit, StartExecution, StartProviderExecution, StartedExecution,
};
use gateway_core::engine::middleware::{FrozenMiddlewarePlan, MiddlewareHeader};
use gateway_core::engine::{CommitRequirement, CoordinatedEvent, EngineError, ModelRequestId};
@@ -58,6 +58,7 @@ fn responses_middleware() -> RequestMiddleware {
#[derive(Default)]
struct Trace {
response_commit: ResponseCommit,
events: Mutex<Vec<&'static str>>,
client_statuses: Mutex<Vec<u16>>,
cancelled: AtomicBool,
@@ -77,7 +78,7 @@ impl Trace {
self.client_statuses
.lock()
.expect("client status lock")
.push(status);
.push(self.response_commit.status_code().unwrap_or(status));
}
fn client_statuses(&self) -> Vec<u16> {
@@ -104,6 +105,7 @@ enum NextStep {
Event(CoordinatedEvent),
DelayedEvent(Duration, CoordinatedEvent),
Error(EngineError),
DelayedError(Duration, EngineError),
FinalizeCancelled,
FinalizeSuccess,
End,
@@ -297,6 +299,10 @@ impl FakeSession {
}
impl ExecutionSession for FakeSession {
fn response_commit(&self) -> ResponseCommit {
self.trace.response_commit.clone()
}
fn next_event(&mut self) -> BoxFuture<'_, Result<Option<CoordinatedEvent>, EngineError>> {
Box::pin(async move {
match self.next.pop_front().unwrap_or(NextStep::End) {
@@ -310,6 +316,13 @@ impl ExecutionSession for FakeSession {
self.trace.push("next_event");
Ok(Some(event))
}
NextStep::DelayedError(delay, error) => {
self.trace.push("wait_error");
tokio::time::sleep(delay).await;
self.trace.push("next_error");
self.finalized = true;
Err(error)
}
NextStep::Error(error) => {
self.trace.push("next_error");
self.finalized = true;
@@ -1288,27 +1301,203 @@ async fn dropping_stream_during_keepalive_should_cancel_and_finalize_execution()
}
#[tokio::test(start_paused = true)]
async fn streaming_keepalive_should_wait_for_first_event_before_committing_http_status() {
async fn streaming_keepalive_should_cover_four_minutes_before_first_event() {
let trace = Arc::new(Trace::default());
let session = FakeSession::streaming(
Arc::clone(&trace),
vec![
NextStep::DelayedEvent(
Duration::from_secs(45),
delivery(started(), CommitRequirement::CommitBeforeDelivery),
Duration::from_secs(240),
delivery(completed(), CommitRequirement::CommitBeforeDelivery),
),
NextStep::FinalizeSuccess,
],
);
let start = tokio::time::Instant::now();
let response = tokio::time::timeout(
Duration::from_secs(20),
stream_execution_response(Box::new(session), None),
)
.await
.expect("headers and heartbeat must precede the first upstream event");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(response.headers()[CONTENT_TYPE], "text/event-stream");
assert_eq!(trace.response_commit.status_code(), Some(200));
let mut body = response.into_body().into_data_stream();
for heartbeat in 1..16 {
let chunk = tokio::time::timeout(Duration::from_secs(20), body.next())
.await
.expect("first-event wait must stay alive")
.expect("heartbeat")
.expect("bytes");
assert_eq!(chunk, Bytes::from_static(b": keep-alive\n\n"));
assert_eq!(start.elapsed(), Duration::from_secs(15 * heartbeat));
assert_eq!(trace.snapshot(), vec!["wait_event"]);
}
let terminal = body.next().await.expect("terminal").expect("bytes");
assert!(String::from_utf8_lossy(&terminal).contains("response.completed"));
assert_eq!(start.elapsed(), Duration::from_secs(240));
assert_eq!(
body.next().await.unwrap().unwrap(),
Bytes::from_static(b"data: [DONE]\n\n")
);
assert!(body.next().await.is_none());
assert_eq!(
trace.snapshot(),
vec!["wait_event", "next_event", "commit", "next_end"]
);
assert!(!trace.is_cancelled());
}
#[tokio::test(start_paused = true)]
async fn streaming_error_after_initial_keepalive_should_be_a_terminal_sse_failure() {
let trace = Arc::new(Trace::default());
let session = FakeSession::streaming(
Arc::clone(&trace),
vec![NextStep::DelayedError(
Duration::from_secs(45),
EngineError::Deadline,
)],
)
.with_response_headers(headers(&[("x-codex-turn-state", "private-attempt-state")]));
let response = stream_execution_response(Box::new(session), None).await;
assert_eq!(response.status(), StatusCode::OK);
assert!(!response.headers().contains_key("x-codex-turn-state"));
let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap();
let events = parse_sse_events(std::str::from_utf8(&bytes).unwrap()).unwrap();
assert!(String::from_utf8_lossy(&bytes).starts_with(": keep-alive\n\n"));
let failed = events
.iter()
.find(|event| event.event.as_deref() == Some("response.failed"))
.expect("failure event");
let payload: Value = serde_json::from_str(&failed.data).unwrap();
assert_eq!(payload["response"]["error"]["code"], "request_timeout");
assert!(bytes.ends_with(b"data: [DONE]\n\n"));
assert_eq!(trace.client_statuses(), vec![200]);
assert_eq!(trace.snapshot(), vec!["wait_error", "next_error"]);
}
#[tokio::test(start_paused = true)]
async fn streaming_late_upstream_error_should_preserve_extended_error_fields() {
let mut error = json!({
"type": "invalid_request_error", "code": "misalignment_policy_violation",
"message": "synthetic policy error", "misalignment": {"category": "synthetic"},
"extension": {"opaque": true}
});
let failure = ProviderError::new(ProviderErrorKind::InvalidRequest, UpstreamSendState::Sent)
.with_client_visible_upstream_response(
ClientVisibleUpstreamResponse::new(
400,
Some(b"application/json".to_vec()),
Bytes::from(json!({"error": error}).to_string()),
)
.with_headers(headers(&[("retry-after", "7")])),
);
error["headers"] = json!({"retry-after": "7"});
let session = FakeSession::streaming(
Arc::new(Trace::default()),
vec![NextStep::DelayedError(
Duration::from_secs(30),
EngineError::Provider(failure),
)],
);
let response = stream_execution_response(Box::new(session), None).await;
assert_eq!(response.status(), StatusCode::OK);
let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap();
let events = parse_sse_events(std::str::from_utf8(&bytes).unwrap()).unwrap();
let failed = events
.iter()
.find(|event| event.event.as_deref() == Some("response.failed"))
.unwrap();
let payload: Value = serde_json::from_str(&failed.data).unwrap();
assert_eq!(payload["response"]["error"], error);
assert!(bytes.ends_with(b"data: [DONE]\n\n"));
}
#[tokio::test]
async fn responses_http_connection_should_survive_a_shorter_read_idle_timeout() {
use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
let trace = Arc::new(Trace::default());
let session = FakeSession::streaming(
Arc::clone(&trace),
vec![
NextStep::DelayedEvent(
Duration::from_secs(35),
delivery(completed(), CommitRequirement::CommitBeforeDelivery),
),
NextStep::FinalizeSuccess,
],
);
let execution = Arc::new(SessionExecution {
client: authenticated_client_for_provider("sk_correlation_test", "openai"),
session: Mutex::new(Some(Box::new(session))),
middleware: Some(Arc::new(responses_middleware()).frozen()),
});
let app = api_router(execution).await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
let mut client = tokio::net::TcpStream::connect(address).await.unwrap();
let body = json!({"model": "model-a", "input": "synthetic", "stream": true}).to_string();
client.write_all(format!(
"POST /v1/responses HTTP/1.1\r\nHost: {address}\r\nAuthorization: Bearer sk_correlation_test\r\nUser-Agent: AnotherTerminal/1.0\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len()
).as_bytes()).await.unwrap();
let mut received = Vec::new();
let result = async {
loop {
let mut buffer = [0; 4096];
// 真实 TCP 每次读取都受空闲期限约束,首个业务事件晚于该期限
let read =
tokio::time::timeout(Duration::from_secs(20), client.read(&mut buffer)).await?;
let count = read.unwrap();
if count == 0 {
break;
}
received.extend_from_slice(&buffer[..count]);
}
Ok::<(), tokio::time::error::Elapsed>(())
}
.await;
server.abort();
result.expect("HTTP connection must receive keepalive before its read idle deadline");
let wire = String::from_utf8(received).unwrap();
assert!(wire.starts_with("HTTP/1.1 200 OK\r\n"));
assert!(wire.contains("x-accel-buffering: no\r\n"));
assert_eq!(wire.matches(": keep-alive\n\n").count(), 2);
assert!(wire.contains("response.completed"));
assert!(wire.contains("data: [DONE]\n\n"));
assert!(!trace.is_cancelled());
assert_eq!(
trace.snapshot(),
vec!["wait_event", "next_event", "commit", "next_end"]
);
}
#[tokio::test(start_paused = true)]
async fn dropping_stream_before_first_event_should_cancel_the_original_wait() {
let trace = Arc::new(Trace::default());
let session = FakeSession::streaming(
Arc::clone(&trace),
vec![
NextStep::DelayedEvent(
Duration::from_secs(240),
delivery(completed(), CommitRequirement::CommitBeforeDelivery),
),
NextStep::FinalizeCancelled,
],
);
let response = tokio::spawn(stream_execution_response(Box::new(session), None));
let response = stream_execution_response(Box::new(session), None).await;
let mut body = response.into_body().into_data_stream();
assert_eq!(
body.next().await.unwrap().unwrap(),
Bytes::from_static(b": keep-alive\n\n")
);
drop(body);
tokio::task::yield_now().await;
tokio::time::advance(Duration::from_secs(30)).await;
assert!(!response.is_finished());
assert!(trace.client_statuses().is_empty());
let response = response.await.expect("response task");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(trace.snapshot(), vec!["wait_event", "next_event", "commit"]);
assert!(trace.is_cancelled());
assert_eq!(trace.detached_finalizations(), 1);
assert_eq!(trace.snapshot(), vec!["wait_event", "cancel_finalize"]);
}
#[tokio::test]
@@ -3086,3 +3275,53 @@ async fn streaming_middleware_preserves_heartbeats_before_and_after_rewriting()
assert!(!trace.is_cancelled());
}
}
#[tokio::test(start_paused = true)]
async fn responses_route_should_keep_alive_while_request_middleware_waits_for_first_event() {
let admin = crate::admin::AdminTestFixture::new().await.services;
for invalid_output in [false, true] {
let trace = Arc::new(Trace::default());
let session = FakeSession::streaming(
Arc::clone(&trace),
vec![
NextStep::DelayedEvent(
Duration::from_secs(150),
delivery(started(), CommitRequirement::CommitBeforeDelivery),
),
NextStep::Event(delivery(completed(), CommitRequirement::AlreadyCommitted)),
NextStep::FinalizeSuccess,
],
);
let actions = if invalid_output {
vec![ResponseFrameAction::Replace(Bytes::from_static(
b"data: {invalid-json}\n\n",
))]
} else {
Vec::new()
};
let middleware = Arc::new(responses_middleware().with_response_actions(actions));
let response = tokio::time::timeout(
Duration::from_secs(20),
response_with_middleware(&admin, "openai", session, true, Some(middleware)),
)
.await
.expect("request middleware must not block transport keepalive");
assert_eq!(response.status(), StatusCode::OK);
assert!(response.headers().contains_key("x-gateway-request-id"));
let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap();
let text = std::str::from_utf8(&bytes).unwrap();
assert!(text.starts_with(": keep-alive\n\n"));
assert!(text.ends_with("data: [DONE]\n\n"));
if invalid_output {
assert!(text.contains("response.failed"));
assert!(text.contains("middleware_failed"));
assert!(!text.contains("{invalid-json}"));
assert!(trace.is_cancelled());
} else {
assert!(text.contains("response.completed"));
assert!(!text.contains("response.failed"));
assert!(!trace.is_cancelled());
}
assert_eq!(trace.client_statuses(), vec![200]);
}
}
@@ -62,6 +62,9 @@ impl ResponseControlTransport for ControlTransport {
#[derive(Default)]
struct InterruptProvider {
calls: AtomicUsize,
fail_first: bool,
retry_started: tokio::sync::Notify,
retry_release: tokio::sync::Notify,
transport: Mutex<Option<Arc<ControlTransport>>>,
complete: tokio::sync::Notify,
initial_messages: Mutex<std::collections::VecDeque<ProtocolWireEvent>>,
@@ -87,6 +90,15 @@ impl Provider for InterruptProvider {
context: AttemptContext,
) -> Result<ProviderStream, ProviderError> {
let number = self.calls.fetch_add(1, Ordering::SeqCst);
if number == 1 && self.fail_first {
self.retry_started.notify_one();
self.retry_release.notified().await;
assert!(
context
.excluded_accounts()
.contains(&ProviderAccountId::new("acct_api_test").unwrap())
);
}
let id = format!("resp_interrupt_{number}");
let transport = Arc::new(ControlTransport::new(id));
let owner: Arc<dyn ResponseControlTransport> = transport.clone();
@@ -96,7 +108,12 @@ impl Provider for InterruptProvider {
let metadata = ProviderCallMetadata::new(
candidate.provider().clone(),
candidate.upstream_model().unwrap().clone(),
ProviderAccountId::new("acct_api_test").unwrap(),
ProviderAccountId::new(if self.fail_first && number > 0 {
"acct_api_retry"
} else {
"acct_api_test"
})
.unwrap(),
UpstreamTransport::new("websocket").unwrap(),
);
let body = futures::stream::unfold(
@@ -110,6 +127,18 @@ impl Provider for InterruptProvider {
if let Some(wire) = pending {
return Some((Ok(ProviderEvent::wire(wire)), (phase, transport, provider)));
}
if phase == 1 && provider.fail_first && transport.id == "resp_interrupt_0" {
provider.complete.notified().await;
transport.active.store(false, Ordering::SeqCst);
return Some((
Err(ProviderError::new(
gateway_core::error::ProviderErrorKind::QuotaExhausted,
UpstreamSendState::Sent,
)
.with_replay_safe()),
(2, transport, provider),
));
}
let (wire, events, next_phase) = match phase {
0 => (
json!({"type":"response.created","response":{"id":transport.id,"model":"model-a","status":"in_progress","output":[]}}),
@@ -144,10 +173,13 @@ impl Provider for InterruptProvider {
_ => return None,
};
let event_type = wire["type"].as_str().unwrap().to_owned();
let event = ProviderEvent::canonical_with_wire(
let mut event = ProviderEvent::canonical_with_wire(
events,
ProtocolWireEvent::json("openai", Some(event_type), wire).unwrap(),
);
if phase == 0 && provider.fail_first {
event = event.with_replayable_prelude();
}
Some((Ok(event), (next_phase, transport, provider)))
},
);
@@ -233,7 +265,16 @@ async fn connect(
admissions: Arc<InterruptAdmissions>,
) -> (TestSocket, Server) {
let execution = Arc::new(DefaultExecutionService::new(
RuntimeSnapshotHandle::new(snapshot("sk_ws_interrupt", "openai")),
RuntimeSnapshotHandle::new(if provider.fail_first {
crate::openai::snapshot_with_accounts(
"sk_ws_interrupt",
"openai",
"key_api_test",
&["acct_api_test", "acct_api_retry"],
)
} else {
snapshot("sk_ws_interrupt", "openai")
}),
Arc::new(SettlementPorts::default()),
ProviderRegistry::new([provider as Arc<dyn Provider>]).unwrap(),
admissions,
@@ -448,3 +489,102 @@ async fn active_websocket_delivery_preserves_unparsed_messages_and_business_meta
provider.complete.notify_one();
assert_eq!(next(&mut socket).await["type"], "response.completed");
}
#[tokio::test]
async fn prelude_retry_retains_interrupt_during_account_switch() {
for interrupt_before_new_created in [true, false] {
let provider = Arc::new(InterruptProvider {
fail_first: true,
..Default::default()
});
let admissions = Arc::new(InterruptAdmissions::default());
let (mut socket, _server) = connect(provider.clone(), admissions.clone()).await;
send(
&mut socket,
json!({"type":"response.create","model":"model-a","input":"hello"}),
)
.await;
assert_eq!(
next(&mut socket).await["response"]["id"],
"resp_interrupt_0"
);
let old = provider.transport.lock().unwrap().clone().unwrap();
provider.complete.notify_one();
tokio::time::timeout(Duration::from_secs(5), provider.retry_started.notified())
.await
.unwrap();
if interrupt_before_new_created {
send(
&mut socket,
interrupt("resp_interrupt_0", "discard_partial_items"),
)
.await;
// 等待服务端处理控制帧,随后才允许新 attempt 建立连接
tokio::time::timeout(Duration::from_secs(5), async {
while old.requests.lock().unwrap().is_empty() {
tokio::task::yield_now().await;
}
})
.await
.unwrap();
}
provider.retry_release.notify_one();
assert_eq!(
next(&mut socket).await["response"]["id"],
"resp_interrupt_1"
);
if !interrupt_before_new_created {
// 迟到的旧 ID 也应作用于当前逻辑响应
send(
&mut socket,
interrupt("resp_interrupt_0", "discard_partial_items"),
)
.await;
}
let terminal = next(&mut socket).await;
assert_eq!(terminal["type"], "response.incomplete");
assert_eq!(terminal["response"]["id"], "resp_interrupt_1");
let current = provider.transport.lock().unwrap().clone().unwrap();
let requests = current.requests.lock().unwrap().clone();
assert_eq!(requests.len(), 1);
assert_eq!(
serde_json::from_str::<Value>(&requests[0]).unwrap(),
interrupt("resp_interrupt_1", "discard_partial_items")
);
assert_eq!(provider.calls.load(Ordering::SeqCst), 2);
socket.close(None).await.unwrap();
}
}
#[tokio::test]
async fn non_replayable_delivery_blocks_retry_even_after_a_replayable_prelude() {
for event in [
json!({"type":"response.output_text.delta","delta":"visible","output_index":0,"content_index":0}),
json!({"type":"response.future_progress","marker":"preserved"}),
] {
let provider = Arc::new(InterruptProvider {
fail_first: true,
..Default::default()
});
provider.initial_messages.lock().unwrap().push_back(
ProtocolWireEvent::json(
"openai",
Some(event["type"].as_str().unwrap().to_owned()),
event.clone(),
)
.unwrap(),
);
let (mut socket, _server) = connect(provider.clone(), Arc::default()).await;
send(
&mut socket,
json!({"type":"response.create","model":"model-a","input":"hello"}),
)
.await;
assert_eq!(next(&mut socket).await["type"], "response.created");
assert_eq!(next(&mut socket).await, event);
provider.complete.notify_one();
assert_eq!(next(&mut socket).await["type"], "error");
assert_eq!(provider.calls.load(Ordering::SeqCst), 1);
socket.close(None).await.unwrap();
}
}
@@ -750,12 +750,40 @@ const fn status_projection(status: AccountStatus) -> AccountStatusProjection {
}
}
/// 选中账号时冻结的历史展示事实,不包含凭据与运行状态
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AccountHistorySnapshot {
pub name: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub email: Option<String>,
pub authentication_kind: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub plan_type_snapshot: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub notes_snapshot: Option<String>,
}
impl AccountHistorySnapshot {
#[must_use]
pub fn from_account(account: &crate::account::ProviderAccount) -> Self {
Self {
name: account.name().to_owned(),
email: account.email().map(str::to_owned),
authentication_kind: account.authentication_kind().to_owned(),
plan_type_snapshot: account.plan_type().map(str::to_owned),
notes_snapshot: account.notes().map(str::to_owned),
}
}
}
/// 账号持久事实;代理认证信息只通过显式 secret accessor 读取
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderAccount {
id: ProviderAccountId,
provider: ProviderKind,
name: String,
notes: Option<String>,
email: Option<String>,
upstream_user_id: Option<String>,
upstream_account_id: Option<String>,
@@ -793,6 +821,7 @@ impl ProviderAccount {
id,
provider,
name,
notes: None,
email: None,
upstream_user_id,
upstream_account_id: None,
@@ -815,6 +844,17 @@ impl ProviderAccount {
}
}
#[must_use]
pub fn with_notes(mut self, notes: Option<String>) -> Self {
self.notes = notes;
self
}
#[must_use]
pub fn notes(&self) -> Option<&str> {
self.notes.as_deref()
}
#[must_use]
pub fn with_profile(
mut self,
@@ -18,6 +18,7 @@ pub struct OperationalFailure {
pub correlation_id: Option<String>,
pub provider_kind: Option<ProviderKind>,
pub account_id: Option<ProviderAccountId>,
pub account_snapshot: Option<crate::account::AccountHistorySnapshot>,
pub upstream_status: Option<u16>,
pub upstream_code: Option<OpaqueUpstreamValue>,
pub details: Option<ErrorDetails>,
@@ -40,6 +41,7 @@ impl OperationalFailure {
correlation_id: None,
provider_kind: None,
account_id: None,
account_snapshot: None,
upstream_status: None,
upstream_code: None,
details: None,
@@ -307,6 +307,7 @@ where
request_observation,
budget_prior_attempts_usd: Decimal::ZERO,
budget_attempt_already_counted: false,
response_commit: super::execution::ResponseCommit::new(trace.clone()),
trace,
deadline,
deadline_timer: deadline.wait().fuse(),
@@ -345,7 +346,8 @@ where
request_profiles: BTreeMap::new(),
current: None,
send_state_watermark: UpstreamSendState::NotSent,
downstream_committed_at: None,
delivery_committed: false,
replay_barrier_reached: false,
client_status_code: None,
delivery_pending: false,
upstream_complete: false,
@@ -482,7 +484,11 @@ pub struct ResponseExecutionSession<S: ?Sized> {
current: Option<CurrentAttempt>,
/// 请求级发送状态水位;跨 attempt 单调不降,终态写回不得低于此档
send_state_watermark: UpstreamSendState,
downstream_committed_at: Option<SystemTime>,
response_commit: super::execution::ResponseCommit,
// 传输心跳可先于首个业务批次提交,两者不能共用重放屏障
delivery_committed: bool,
/// 已交付的实际输出禁止重放,可替换前导事件不跨越此屏障
replay_barrier_reached: bool,
client_status_code: Option<u16>,
delivery_pending: bool,
upstream_complete: bool,
@@ -518,6 +524,11 @@ where
self
}
/// 传输首字节事实可在 next_event 等待期间由协议层提交
pub fn response_commit(&self) -> super::execution::ResponseCommit {
self.response_commit.clone()
}
/// 当前请求共享的诊断上下文
pub fn trace(&self) -> TraceContext {
self.trace.clone()
@@ -548,7 +559,11 @@ where
loop {
match self.pull().await? {
PullOutcome::Events(events) => {
let requirement = if self.downstream_committed_at.is_some() {
self.replay_barrier_reached |= self.middleware.is_some()
|| self.client_transport != super::execution::ClientTransport::WebSocket
|| self.response_control.is_none()
|| events.iter().any(|event| !event.is_replayable_prelude());
let requirement = if self.delivery_committed {
CommitRequirement::AlreadyCommitted
} else {
self.delivery_pending = true;
@@ -562,18 +577,19 @@ where
error,
send_state,
} => {
let requirement = if self.downstream_committed_at.is_some() {
let requirement = if self.delivery_committed {
CommitRequirement::AlreadyCommitted
} else {
self.delivery_pending = true;
CommitRequirement::CommitBeforeDelivery
};
self.replay_barrier_reached = true;
self.pending_terminal_failure =
Some(PendingTerminalFailure { error, send_state });
return CoordinatedEvent::try_batch(events, requirement).map(Some);
}
PullOutcome::End => {
if self.downstream_committed_at.is_some() {
if self.delivery_committed {
self.finish_success().await?;
}
return Ok(None);
@@ -589,7 +605,7 @@ where
/// 会话已提交、已有待提交结果或执行失败时返回错误
pub async fn collect_uncommitted(&mut self) -> Result<Vec<ProviderEvent>, EngineError> {
self.resume_finalization().await;
if self.downstream_committed_at.is_some() || self.delivery_pending {
if self.delivery_committed || self.delivery_pending {
return Err(EngineError::InvalidDeliveryState);
}
if self.is_finalized() {
@@ -628,14 +644,15 @@ where
&mut self,
client_status_code: Option<u16>,
) -> Result<(), EngineError> {
if !self.delivery_pending || self.downstream_committed_at.is_some() || self.is_finalized() {
if !self.delivery_pending || self.delivery_committed || self.is_finalized() {
return Err(EngineError::InvalidDeliveryState);
}
let committed_at = SystemTime::now();
self.trace.record(
"downstream.committed",
json!({"status": client_status_code}),
);
self.response_commit.record(client_status_code);
let committed_at = self
.response_commit
.committed_at()
.ok_or(EngineError::InvalidDeliveryState)?;
let client_status_code = self.response_commit.status_code();
if self.request_persisted {
best_effort_store_write(
"mark_downstream_committed",
@@ -648,7 +665,7 @@ where
)
.await;
}
self.downstream_committed_at = Some(committed_at);
self.delivery_committed = true;
self.client_status_code = client_status_code;
self.delivery_pending = false;
if self.upstream_complete {
@@ -661,13 +678,14 @@ where
/// 原始 Provider facts 已经观察,不回滚计量,也不把丢弃误记为客户端 commit
pub fn discard_pending_delivery(&mut self) -> Result<(), EngineError> {
if !self.delivery_pending
|| self.downstream_committed_at.is_some()
|| self.delivery_committed
|| self.pending_terminal_failure.is_some()
|| self.is_finalized()
{
return Err(EngineError::InvalidDeliveryState);
}
self.delivery_pending = false;
self.replay_barrier_reached = false;
Ok(())
}
@@ -683,7 +701,10 @@ where
if self.client_status_code.is_some() {
return Err(EngineError::InvalidDeliveryState);
}
if self.request_persisted {
let committed_status = self.response_commit.status_code();
let client_status_code = committed_status.unwrap_or(client_status_code);
// 保活后的失败终结已写入实际 HTTP 状态,不能再触发仅补空值的 Store 更新
if self.request_persisted && !(self.is_finalized() && committed_status.is_some()) {
best_effort_store_write(
"record_client_status",
&self.request_id,
@@ -1316,11 +1337,11 @@ where
));
}
let attempt_record = AttemptRecord {
account_snapshot: metadata.account_snapshot().cloned(),
request_id: self.request_id.clone(),
attempt_count: next_attempt,
trigger,
provider_kind: metadata.provider().clone(),
provider_account_id: Some(metadata.provider_account_id().clone()),
provider_account_ref: Some(metadata.provider_account_id().clone()),
upstream_model_id: metadata.upstream_model().cloned(),
upstream_transport: metadata.transport().as_str().to_owned(),
@@ -1497,7 +1518,7 @@ where
let account_rotation_retry = !execution_effect_observed
&& self.account_selection.required_account().is_none()
&& self.continuation_attempt == ContinuationAttempt::None
&& self.downstream_committed_at.is_none()
&& !self.replay_barrier_reached
&& !self.delivery_pending
&& attempt_send_state != UpstreamSendState::Ambiguous
&& matches!(
@@ -1511,11 +1532,13 @@ where
error.pre_delivery_retry(),
Some(crate::error::PreDeliveryRetry::SameAccountConnectionRetry { .. })
);
// Provider 的有界同账号传输恢复是尽力重试策略,不要求把 Sent/Ambiguous
// 改写成安全证明;外部副作用、取消、预算与下游提交仍由 Core 拦截
let transport_recovery = match error.pre_delivery_retry() {
Some(crate::error::PreDeliveryRetry::SameAccountConnectionRetry { transport })
if !execution_effect_observed
&& !self.connection_budget.exhausted()
&& self.downstream_committed_at.is_none()
&& !self.replay_barrier_reached
&& !self.delivery_pending
&& attempt_send_state == UpstreamSendState::NotSent
&& self.current_send_state() == UpstreamSendState::NotSent
@@ -1533,18 +1556,16 @@ where
delay,
}) if !execution_effect_observed
&& !self.connection_budget.exhausted()
&& self.downstream_committed_at.is_none()
&& !self.delivery_pending
&& attempt_send_state != UpstreamSendState::Ambiguous =>
&& !self.replay_barrier_reached
&& !self.delivery_pending =>
{
Some((AttemptTransport::Retry(retry_index), delay))
}
Some(crate::error::PreDeliveryRetry::SameAccountTransportFallback)
if !execution_effect_observed
&& !self.connection_budget.exhausted()
&& self.downstream_committed_at.is_none()
&& !self.delivery_pending
&& attempt_send_state != UpstreamSendState::Ambiguous =>
&& !self.replay_barrier_reached
&& !self.delivery_pending =>
{
Some((
AttemptTransport::Fallback,
@@ -1558,7 +1579,7 @@ where
&& !execution_effect_observed
&& self.account_selection.required_account().is_none()
&& self.continuation_attempt == ContinuationAttempt::None
&& self.downstream_committed_at.is_none()
&& !self.replay_barrier_reached
&& !self.delivery_pending
&& attempt_send_state != UpstreamSendState::Ambiguous
&& provider_proved_replay_safe
@@ -1590,7 +1611,7 @@ where
let same_account_retry = !execution_effect_observed
&& error.retries_same_account()
&& provider_proved_replay_safe
&& self.downstream_committed_at.is_none()
&& !self.replay_barrier_reached
&& !self.delivery_pending
&& attempt_send_state != UpstreamSendState::Ambiguous
&& self.routing_attempts < self.plan.max_attempts().get()
@@ -1655,7 +1676,8 @@ where
"connectionBudgetRemainingMs": self.connection_budget.remaining().map(duration_ms),
"executionEffectObserved": execution_effect_observed,
"delayMs": transient_retry.or(transport_recovery.map(|(_, delay)| delay)).map(duration_ms),
"downstreamCommitted": self.downstream_committed_at.is_some(),
"downstreamCommitted": self.response_commit.committed_at().is_some(),
"replayBarrierReached": self.replay_barrier_reached,
"sendState": format!("{attempt_send_state:?}"),
}));
if retryable {
@@ -1701,6 +1723,7 @@ where
self.engine
.store()
.record_intermediate_failure(IntermediateFailure {
account_snapshot: current.metadata.account_snapshot().cloned(),
request_id: self.request_id.clone(),
attempt_index: current.index,
trigger: current.trigger,
@@ -1813,7 +1836,7 @@ where
) -> bool {
if self.account_selection.required_account().is_some()
|| self.continuation_attempt == ContinuationAttempt::None
|| self.downstream_committed_at.is_some()
|| self.replay_barrier_reached
|| self.delivery_pending
|| send_state == UpstreamSendState::Ambiguous
|| !provider_proved_replay_safe
@@ -1952,8 +1975,11 @@ where
outcome: ExecutionOutcome::Succeeded,
send_state: UpstreamSendState::Sent,
attempt_count: self.attempts,
downstream_committed_at: self.downstream_committed_at,
client_status_code: self.client_status_code,
downstream_committed_at: self.response_commit.committed_at(),
client_status_code: self
.response_commit
.status_code()
.or(self.client_status_code),
upstream_status_code,
client_response_id: self.observation.client_response_id.clone(),
upstream_request_id,
@@ -2001,7 +2027,7 @@ where
let send_state = self.raise_send_watermark(send_state);
let outcome = if error.kind() == ProviderErrorKind::Cancelled {
ExecutionOutcome::Cancelled
} else if self.downstream_committed_at.is_some() {
} else if self.response_commit.committed_at().is_some() {
ExecutionOutcome::Incomplete
} else {
ExecutionOutcome::Failed
@@ -2026,13 +2052,15 @@ where
}
async fn finish_interruption(&mut self, error: &EngineError) -> Result<(), EngineError> {
// 结算 I/O 可能继续等待,先冻结交付事实以保持终态与实际 HTTP 状态一致
self.response_commit.finish();
let (outcome, gateway_error) = match error {
EngineError::Cancelled => (
ExecutionOutcome::Cancelled,
GatewayError::new(GatewayErrorKind::Cancelled, "request was cancelled"),
),
EngineError::Deadline => (
if self.downstream_committed_at.is_some() {
if self.response_commit.committed_at().is_some() {
ExecutionOutcome::Incomplete
} else {
ExecutionOutcome::Failed
@@ -2046,7 +2074,7 @@ where
};
let send_state = if self.attempts == 0 {
UpstreamSendState::NotSent
} else if self.downstream_committed_at.is_some()
} else if self.delivery_committed
|| self
.current
.as_ref()
@@ -2117,8 +2145,11 @@ where
outcome: finalization.outcome,
send_state: finalization.send_state,
attempt_count: self.attempts,
downstream_committed_at: self.downstream_committed_at,
client_status_code: self.client_status_code,
downstream_committed_at: self.response_commit.committed_at(),
client_status_code: self
.response_commit
.status_code()
.or(self.client_status_code),
upstream_status_code: finalization.upstream_status_code.or(observed_status_code),
client_response_id: self.observation.client_response_id.clone(),
upstream_request_id: finalization.upstream_request_id.or(upstream_request_id),
@@ -2151,6 +2182,7 @@ where
}
async fn persist_finalization(&mut self, finalization: ModelRequestFinalization) {
self.response_commit.finish();
self.finalized_at = Some(finalization.completed_at);
let observed_provider = self
.current
@@ -297,6 +297,11 @@ pub struct StartedExecution {
}
pub trait ExecutionSession: Send {
/// 允许协议层在等待业务事件期间记录传输首字节,保活不提交业务批次
fn response_commit(&self) -> super::ResponseCommit {
super::ResponseCommit::default()
}
fn trace(&self) -> crate::diagnostics::TraceContext {
crate::diagnostics::TraceContext::default()
}
@@ -0,0 +1,65 @@
//! 下游首字节提交事实;传输保活不代表业务事件已交付
use std::sync::{Arc, OnceLock};
use std::time::SystemTime;
use crate::diagnostics::TraceContext;
#[derive(Clone, Default)]
pub struct ResponseCommit {
first: Arc<OnceLock<TransportState>>,
trace: TraceContext,
}
enum TransportState {
Committed { at: SystemTime, status: Option<u16> },
Finished,
}
impl ResponseCommit {
pub(crate) fn new(trace: TraceContext) -> Self {
Self {
trace,
..Self::default()
}
}
/// 只记录首次实际交付;已结束且未提交的执行返回 false,不改变业务重放屏障
pub fn record(&self, status: Option<u16>) -> bool {
if self
.first
.set(TransportState::Committed {
at: SystemTime::now(),
status,
})
.is_ok()
{
self.trace.record(
"downstream.committed",
serde_json::json!({"status": status}),
);
}
self.committed_at().is_some()
}
pub(crate) fn finish(&self) {
// 执行已结束且未提交时保留原 HTTP 错误,不能因外层插件迟迟返回再改成 200
let _ = self.first.set(TransportState::Finished);
}
#[must_use]
pub fn committed_at(&self) -> Option<SystemTime> {
match self.first.get() {
Some(TransportState::Committed { at, .. }) => Some(*at),
_ => None,
}
}
#[must_use]
pub fn status_code(&self) -> Option<u16> {
match self.first.get() {
Some(TransportState::Committed { status, .. }) => *status,
_ => None,
}
}
}
@@ -2,6 +2,8 @@
mod admission;
mod contract;
mod delivery;
pub use delivery::ResponseCommit;
mod failure;
pub use contract::{
AuthenticatedClient, ClientApiKeyUsageSink, ClientAuthenticationError, ClientKeyVerifier,
@@ -807,7 +807,7 @@ impl DefaultExecutionService {
}
let new_request = NewModelRequest {
id: request_id.clone(),
client_api_key_id: Some(client.policy.key_id().clone()),
client_api_key_name_snapshot: client.policy.name_snapshot().map(str::to_owned),
client_api_key_ref: client.policy.key_id().clone(),
config_revision: plan.config_revision(),
routing: authorization.account_scope.routing_snapshot(),
@@ -1347,7 +1347,7 @@ impl DefaultExecutionService {
.map_err(|_| GatewayError::new(GatewayErrorKind::Internal, "invalid admin actor"))?;
let new_request = NewModelRequest {
id: request_id,
client_api_key_id: None,
client_api_key_name_snapshot: None,
client_api_key_ref: actor,
config_revision: plan.config_revision(),
routing: crate::routing::AccountRoutingSnapshot::all(),
@@ -106,6 +106,10 @@ impl Drop for DefaultExecutionSession {
}
impl ExecutionSession for DefaultExecutionSession {
fn response_commit(&self) -> super::ResponseCommit {
self.core.response_commit()
}
fn trace(&self) -> crate::diagnostics::TraceContext {
self.core.trace()
}
+11 -2
View File
@@ -559,6 +559,14 @@ impl AttemptContext {
self.request.client_transport
}
/// 原生双向通道可提前交付可替换前导事件;中间件可能改变其输出语义
#[must_use]
pub fn allows_replayable_prelude(&self) -> bool {
self.client_transport() == execution::ClientTransport::WebSocket
&& self.request.middleware.is_none()
&& self.response_control().is_some()
}
#[must_use]
pub fn response_control(&self) -> Option<&response_control::ResponseControl> {
self.request.response_control.as_ref()
@@ -815,8 +823,8 @@ impl AttemptContext {
/// 创建唯一 `model_requests` 行所需的入口事实
#[derive(Debug, Clone)]
pub struct NewModelRequest {
pub client_api_key_name_snapshot: Option<String>,
pub id: ModelRequestId,
pub client_api_key_id: Option<ClientApiKeyId>,
pub client_api_key_ref: ClientApiKeyId,
pub config_revision: ConfigRevision,
pub routing: crate::routing::AccountRoutingSnapshot,
@@ -844,11 +852,11 @@ pub struct NewModelRequest {
/// 每次真实上游发送前对同一 `model_requests` 行的更新
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AttemptRecord {
pub account_snapshot: Option<crate::account::AccountHistorySnapshot>,
pub request_id: ModelRequestId,
pub attempt_count: NonZeroU32,
pub trigger: AttemptTrigger,
pub provider_kind: ProviderKind,
pub provider_account_id: Option<ProviderAccountId>,
pub provider_account_ref: Option<ProviderAccountId>,
pub upstream_model_id: Option<UpstreamModelId>,
pub upstream_transport: String,
@@ -863,6 +871,7 @@ pub struct AttemptRecord {
/// 需要解释换号的中间失败
#[derive(Debug)]
pub struct IntermediateFailure {
pub account_snapshot: Option<crate::account::AccountHistorySnapshot>,
pub request_id: ModelRequestId,
pub attempt_index: NonZeroU32,
pub trigger: AttemptTrigger,
@@ -13,7 +13,8 @@ use futures::{Stream, StreamExt, future::BoxFuture};
use thiserror::Error;
use crate::account::{
AccountAttemptFeedback, AccountCapacitySnapshot, AccountFeedbackStats, ProviderAccountId,
AccountAttemptFeedback, AccountCapacitySnapshot, AccountFeedbackStats, AccountHistorySnapshot,
ProviderAccountId,
};
use crate::engine::AttemptContext;
use crate::engine::execution::ClientTransport;
@@ -39,7 +40,7 @@ use crate::upstream::{UpstreamSendState, UpstreamTransport};
pub type EventStream =
Pin<Box<dyn Stream<Item = Result<ProviderEvent, ProviderError>> + Send + 'static>>;
/// Provider 选定单个 credential 后返回的事实
/// Provider 选定单个 credential 后返回的调用事实
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderCallMetadata {
provider: ProviderKind,
@@ -48,6 +49,7 @@ pub struct ProviderCallMetadata {
upstream_request_id: Option<OpaqueUpstreamValue>,
transport: UpstreamTransport,
selection_observation: Option<ProviderSelectionObservation>,
account_snapshot: Option<AccountHistorySnapshot>,
}
/// Provider 账号选择阶段输出的中立运行压力事实
@@ -96,6 +98,7 @@ impl ProviderCallMetadata {
upstream_request_id: None,
transport,
selection_observation: None,
account_snapshot: None,
}
}
@@ -113,9 +116,21 @@ impl ProviderCallMetadata {
upstream_request_id: None,
transport,
selection_observation: None,
account_snapshot: None,
}
}
#[must_use]
pub fn with_account_snapshot(mut self, snapshot: Option<AccountHistorySnapshot>) -> Self {
self.account_snapshot = snapshot;
self
}
#[must_use]
pub const fn account_snapshot(&self) -> Option<&AccountHistorySnapshot> {
self.account_snapshot.as_ref()
}
/// 设置 adapter 已分类为非 bearer 的 request ID
#[must_use]
pub fn with_upstream_request_id(mut self, request_id: OpaqueUpstreamValue) -> Self {
@@ -24,6 +24,25 @@ pub struct ResponseControl {
target: Arc<Mutex<Option<Weak<dyn ResponseControlTransport>>>>,
}
/// 某次已交付响应所属的弱连接快照,重绑后不能把旧消息发送到新连接
#[derive(Clone)]
pub struct ResponseControlBinding(Weak<dyn ResponseControlTransport>);
impl ResponseControlBinding {
#[must_use]
pub fn same_binding(&self, other: &Self) -> bool {
self.0.ptr_eq(&other.0)
}
pub async fn send(&self, payload: &str) -> Result<(), ResponseControlUnavailable> {
self.0
.upgrade()
.ok_or(ResponseControlUnavailable)?
.send(payload)
.await
}
}
impl fmt::Debug for ResponseControl {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str("ResponseControl([REDACTED])")
@@ -31,6 +50,17 @@ impl fmt::Debug for ResponseControl {
}
impl ResponseControl {
/// 仅保存弱引用,不延长旧 attempt 的连接生命周期
#[must_use]
pub fn binding(&self) -> Option<ResponseControlBinding> {
self.target
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.as_ref()
.cloned()
.map(ResponseControlBinding)
}
/// Provider 在请求成功写入上游后绑定;连接 owner 持有强引用并负责撤销
pub fn bind(&self, transport: &Arc<dyn ResponseControlTransport>) {
*self
+4 -2
View File
@@ -135,6 +135,7 @@ pub enum PreDeliveryRetry {
max_delay: Duration,
},
/// 固定本次账号,并按 Provider 给出的序号重试当前传输
/// 可用于发送结果不明确时的有界尽力恢复,不等于上游未执行的证明
SameAccountTransportRetry {
/// Provider-owned 传输重试序号,从 1 开始
retry_index: NonZeroU32,
@@ -630,8 +631,8 @@ impl ProviderError {
/// 允许 Core 在首个客户端事件前固定原账号并改用 Provider 备用传输
///
/// 与普通换号恢复相同,Core 仍会拒绝 continuation、指定账号、预算耗尽或
/// 下游已提交后的隐藏重放
/// Provider 必须确认备用传输满足当前 continuation scope 并限制恢复次数
/// Core 保持原账号,拒绝预算耗尽、外部副作用或下游已提交后的隐藏重放
#[must_use]
pub fn with_pre_delivery_transport_fallback(mut self) -> Self {
self.set_pre_delivery_transport_fallback();
@@ -647,6 +648,7 @@ impl ProviderError {
///
/// 重试预算和退避策略由 Provider 所属协议定义;Core 只负责同账号钉选、
/// deadline/取消边界以及下游 commit barrier
/// 该意图允许发送结果不明确时尽力恢复,但不修改发送事实或安全重放证明
#[must_use]
pub fn with_pre_delivery_transport_retry(
mut self,
+19
View File
@@ -834,6 +834,7 @@ pub struct ProviderEvent {
observation: Option<Box<ProviderResponseObservation>>,
session_update: Option<Box<ProviderSessionState>>,
middleware_transformed: bool,
replayable_prelude: bool,
middleware_origin_wire: Option<Arc<ProtocolWireEvent>>,
}
@@ -847,6 +848,7 @@ impl ProviderEvent {
observation: None,
session_update: None,
middleware_transformed: false,
replayable_prelude: false,
middleware_origin_wire: None,
}
}
@@ -867,6 +869,7 @@ impl ProviderEvent {
observation: None,
session_update: None,
middleware_transformed: false,
replayable_prelude: false,
middleware_origin_wire: None,
}
}
@@ -880,6 +883,7 @@ impl ProviderEvent {
observation: None,
session_update: None,
middleware_transformed: false,
replayable_prelude: false,
middleware_origin_wire: None,
}
}
@@ -893,10 +897,24 @@ impl ProviderEvent {
observation: Some(Box::new(observation)),
session_update: None,
middleware_transformed: false,
replayable_prelude: false,
middleware_origin_wire: None,
}
}
/// Provider 确认该事件仅为可替换的前导信息,客户端支持后续 attempt 的新前导事件
/// 此证明不代替错误分类、发送状态、预算或续接恢复检查
#[must_use]
pub fn with_replayable_prelude(mut self) -> Self {
self.replayable_prelude = true;
self
}
#[must_use]
pub const fn is_replayable_prelude(&self) -> bool {
self.replayable_prelude && !self.middleware_transformed
}
/// 返回协议表达是否由中间件替换、展开或丢弃映射产生
///
/// 宿主最终协议边界据此启用严格复核,插件可读取其快照
@@ -970,6 +988,7 @@ impl ProviderEvent {
/// 仅替换客户端协议表达;canonical facts、响应观察和会话状态保持原始事实
pub fn replace_wire(&mut self, wire: Option<ProtocolWireEvent>) {
self.wire = wire.map(Box::new);
self.replayable_prelude = false;
}
/// 取出仅供 Core 持久化的响应观察
+13 -13
View File
@@ -194,9 +194,9 @@ impl Money {
pub struct Usage {
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 cached_input_tokens: Option<u64>,
pub cache_write_input_tokens: Option<u64>,
pub reasoning_output_tokens: Option<u64>,
pub image_input_tokens: Option<u64>,
pub image_output_tokens: Option<u64>,
/// Provider/协议报告的独立事实,不从其他列相加推导
@@ -209,9 +209,9 @@ impl Usage {
Self {
input_tokens: None,
output_tokens: None,
cached_tokens: None,
cache_write_tokens: None,
reasoning_tokens: None,
cached_input_tokens: None,
cache_write_input_tokens: None,
reasoning_output_tokens: None,
image_input_tokens: None,
image_output_tokens: None,
total_tokens: None,
@@ -226,14 +226,14 @@ impl Usage {
if newer.output_tokens.is_some() {
self.output_tokens = newer.output_tokens;
}
if newer.cached_tokens.is_some() {
self.cached_tokens = newer.cached_tokens;
if newer.cached_input_tokens.is_some() {
self.cached_input_tokens = newer.cached_input_tokens;
}
if newer.cache_write_tokens.is_some() {
self.cache_write_tokens = newer.cache_write_tokens;
if newer.cache_write_input_tokens.is_some() {
self.cache_write_input_tokens = newer.cache_write_input_tokens;
}
if newer.reasoning_tokens.is_some() {
self.reasoning_tokens = newer.reasoning_tokens;
if newer.reasoning_output_tokens.is_some() {
self.reasoning_output_tokens = newer.reasoning_output_tokens;
}
if newer.image_input_tokens.is_some() {
self.image_input_tokens = newer.image_input_tokens;
@@ -332,7 +332,7 @@ pub struct CalculatedCostBreakdown {
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ImageCostBreakdown {
pub input_tokens: u64,
pub cached_tokens: u64,
pub cached_input_tokens: u64,
pub input_amount: Money,
pub cache_read_amount: Money,
pub input_price_per_million: Money,
@@ -69,9 +69,11 @@ impl ModelPriceOverride {
band.as_str(),
"standard"
| "fast"
| "ultrafast"
| "flex"
| "long_standard"
| "long_fast"
| "long_ultrafast"
| "long_flex"
| "image"
)
@@ -119,6 +119,7 @@ pub struct ClientSettings {
/// 从 `client_api_keys` 冻结的公开准入事实
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ClientPolicy {
name_snapshot: Option<String>,
defaults: Arc<ClientSettings>,
key_id: ClientApiKeyId,
plaintext_key: PlaintextClientApiKey,
@@ -163,6 +164,7 @@ impl ClientPolicy {
limits: RateLimits,
) -> Self {
Self {
name_snapshot: None,
defaults: Arc::new(ClientSettings {
request_profiles: account_scope.request_profiles().clone(),
fast_mode: account_scope.fast_mode(),
@@ -176,6 +178,17 @@ impl ClientPolicy {
}
}
#[must_use]
pub fn with_name_snapshot(mut self, name: Option<String>) -> Self {
self.name_snapshot = name;
self
}
#[must_use]
pub fn name_snapshot(&self) -> Option<&str> {
self.name_snapshot.as_deref()
}
#[must_use]
pub const fn key_id(&self) -> &ClientApiKeyId {
&self.key_id
@@ -29,6 +29,7 @@ type ModelCatalogAccounts = BTreeMap<ProviderKind, BTreeMap<String, BTreeSet<Pro
/// Store 读取到的一个启用 Client API Key 策略事实
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SnapshotClientPolicyFacts {
name_snapshot: Option<String>,
request_profiles: BTreeMap<ProviderKind, crate::account::OpaqueProviderData>,
key_id: ClientApiKeyId,
plaintext_key: PlaintextClientApiKey,
@@ -37,6 +38,12 @@ pub struct SnapshotClientPolicyFacts {
}
impl SnapshotClientPolicyFacts {
#[must_use]
pub fn with_name_snapshot(mut self, name: Option<String>) -> Self {
self.name_snapshot = name;
self
}
#[must_use]
pub fn with_request_profiles(
mut self,
@@ -55,6 +62,7 @@ impl SnapshotClientPolicyFacts {
) -> Self {
Self {
key_id,
name_snapshot: None,
request_profiles: BTreeMap::new(),
plaintext_key,
group_ids,
@@ -268,7 +276,7 @@ impl RuntimeSnapshotCompiler {
/// 读取一个 revision,并为已注册 Provider 查询实时模型目录
pub async fn compile(&self) -> Result<RuntimeSnapshot, RuntimeSnapshotCompileError> {
self.compile_inner(None).await
self.compile_inner(None, true).await
}
/// 配置提交复用同一扩展集合的已发布目录,目录留待对账刷新
@@ -276,12 +284,21 @@ impl RuntimeSnapshotCompiler {
&self,
previous: &RuntimeSnapshot,
) -> Result<RuntimeSnapshot, RuntimeSnapshotCompileError> {
self.compile_inner(Some(previous)).await
self.compile_inner(Some(previous), false).await
}
/// 刷新目录,单个 Provider 读取失败时保留同一扩展集合的已知模型归属
pub(crate) async fn compile_with_catalog_fallback(
&self,
previous: &RuntimeSnapshot,
) -> Result<RuntimeSnapshot, RuntimeSnapshotCompileError> {
self.compile_inner(Some(previous), true).await
}
async fn compile_inner(
&self,
previous: Option<&RuntimeSnapshot>,
refresh_catalog: bool,
) -> Result<RuntimeSnapshot, RuntimeSnapshotCompileError> {
// 同一配置 revision 的重试复用扩展候选,保留插件策略的发布一致性
let mut prepared_extensions: Option<(
@@ -330,6 +347,7 @@ impl RuntimeSnapshotCompiler {
catalogs.as_ref(),
provider_kinds,
cached,
refresh_catalog,
self.privacy_compiler.clone(),
)
.await?;
@@ -346,7 +364,7 @@ impl RuntimeSnapshotCompiler {
return Err(RuntimeSnapshotCompileError::RevisionChanged);
}
let snapshot = snapshot
.with_provider_catalog_generations(if cached.is_some() {
.with_provider_catalog_generations(if cached.is_some() && !refresh_catalog {
BTreeMap::new()
} else {
observed_generations
@@ -365,14 +383,20 @@ async fn compile_runtime_snapshot(
catalogs: &dyn ProviderCatalogPort,
provider_kinds: Vec<ProviderKind>,
previous: Option<&RuntimeSnapshot>,
refresh_catalog: bool,
privacy_compiler: Option<Arc<dyn crate::settings::privacy::PrivacyPolicyCompiler>>,
) -> Result<RuntimeSnapshot, RuntimeSnapshotCompileError> {
// 只有成功取得的完整目录能证明模型缺项;发现型目录和查询失败均交由上游验证
// 目录读取失败不能抹掉已知归属;成功返回空目录则仍是权威的模型撤回
let mut provider_models = Vec::new();
let mut catalog_accounts = BTreeMap::new();
let mut exhaustive_provider_catalogs = BTreeSet::new();
for provider in &provider_kinds {
if let Some(previous) = previous {
let models = if refresh_catalog || previous.is_none() {
catalogs.query_model_capabilities(provider).await.ok()
} else {
None
};
if let Some(previous) = previous.filter(|_| models.is_none()) {
if let Some(accounts) = previous.model_catalog_accounts.get(provider) {
catalog_accounts.insert(provider.clone(), accounts.clone());
}
@@ -395,7 +419,7 @@ async fn compile_runtime_snapshot(
}
continue;
}
let Ok(models) = catalogs.query_model_capabilities(provider).await else {
let Some(models) = models else {
continue;
};
if catalogs.model_catalog_is_exhaustive(provider) {
@@ -509,17 +533,20 @@ async fn compile_runtime_snapshot(
.map_err(|_| RuntimeSnapshotCompileError::InvalidData)?,
)
};
client_policies.push(ClientPolicy::new(
policy.key_id,
policy.plaintext_key,
Arc::new(
account_scope
.with_fast_mode(fast_mode)
.with_request_profiles(policy.request_profiles),
),
true,
policy.limits,
));
client_policies.push(
ClientPolicy::new(
policy.key_id,
policy.plaintext_key,
Arc::new(
account_scope
.with_fast_mode(fast_mode)
.with_request_profiles(policy.request_profiles),
),
true,
policy.limits,
)
.with_name_snapshot(policy.name_snapshot),
);
}
RuntimeSnapshot::new_with_privacy_compiler(
@@ -935,7 +962,7 @@ impl RuntimeSnapshot {
.collect()
}
/// 完整目录按映射后的模型判定;发现型或不可用目录不作为能力白名单
/// 已知模型只属于明确声明它的 Provider;未知模型仍可由非完整目录透传
#[must_use]
pub fn contains_public_model_for_provider(
&self,
@@ -951,13 +978,22 @@ impl RuntimeSnapshot {
{
return false;
}
if !self.exhaustive_provider_catalogs.contains(provider) {
return true;
}
let upstream_model = self.mapped_model(public_model.as_str());
let Ok(upstream_model) =
UpstreamModelId::from_client_wire(self.mapped_model(public_model.as_str()))
else {
return false;
};
self.provider_models
.get(provider)
.is_some_and(|models| models.keys().any(|model| model.as_str() == upstream_model))
.is_some_and(|models| models.contains_key(&upstream_model))
|| (!self.exhaustive_provider_catalogs.contains(provider)
&& !self.is_known_upstream_model(&upstream_model))
}
fn is_known_upstream_model(&self, upstream_model: &UpstreamModelId) -> bool {
self.provider_models
.values()
.any(|models| models.contains_key(upstream_model))
}
#[must_use]
@@ -1095,6 +1131,23 @@ impl RuntimeSnapshot {
return Err(RoutingError::EmptyAccountScope);
}
let requested_model = public_model.as_str();
let mapped_model = self.mapped_model(requested_model);
let upstream_model = if self
.settings
.values
.model_mappings
.contains_key(requested_model)
|| self.model_alias(requested_model).is_some()
{
UpstreamModelId::new(mapped_model)
} else {
UpstreamModelId::from_client_wire(mapped_model)
}
.map_err(|_| RoutingError::InvalidIdentifier)?;
// 归属先于账号范围、能力与路由约束判定,原 Provider 被排除不能变成跨平台透传
let known_model = self.is_known_upstream_model(&upstream_model);
let providers = context.required_provider.as_ref().map_or_else(
|| account_scope.provider_kinds().clone(),
|provider| BTreeSet::from([provider.clone()]),
@@ -1117,38 +1170,26 @@ impl RuntimeSnapshot {
{
continue;
}
let requested_model = public_model.as_str();
let mapped_model = self.mapped_model(requested_model);
let upstream_model = if self
.settings
.values
.model_mappings
.contains_key(requested_model)
|| self.model_alias(requested_model).is_some()
{
UpstreamModelId::new(mapped_model)
} else {
UpstreamModelId::from_client_wire(mapped_model)
}
.map_err(|_| RoutingError::InvalidIdentifier)?;
// 强制 Provider 路由也不能扩大 Key 授权;管理诊断的目标账号由探测入口固定
if enforce_account_model_policy
&& !account_scope.allows_provider_model(provider, upstream_model.as_str())
{
continue;
}
let emulated_features = match self
let capabilities = self
.provider_models
.get(provider)
.and_then(|models| models.get(&upstream_model))
{
.and_then(|models| models.get(&upstream_model));
let emulated_features = match capabilities {
Some(capabilities) => {
let Some(emulated) = capabilities.match_requirements(&requirements) else {
continue;
};
emulated
}
None if self.exhaustive_provider_catalogs.contains(provider) => continue,
None if known_model || self.exhaustive_provider_catalogs.contains(provider) => {
continue;
}
None => BTreeSet::new(),
};
candidates.push(ProviderCandidate {
@@ -1158,12 +1199,19 @@ impl RuntimeSnapshot {
.get(provider)
.and_then(|models| models.get(&upstream_model))
.cloned(),
upstream_model: Some(upstream_model),
upstream_model: Some(upstream_model.clone()),
emulated_features,
account_scope: Arc::clone(&account_scope),
});
}
// 冷启动或目录不可读时不能按 Provider 名称猜测模型归属
if !known_model && candidates.len() > 1 {
return Err(RoutingError::NoCapableProvider {
model: public_model.as_str().to_owned(),
});
}
if candidates.is_empty() {
// 模型存在性必须包含被请求路由限制排除的 Provider;目录未知时不能断言模型不存在
let mut scoped_providers = providers.intersection(&self.providers).peekable();
@@ -234,7 +234,14 @@ impl RuntimeSnapshotPublisher {
None => self.compiler.compile().await,
}
} else {
preempt_for_committed_refresh(cancellation, self.compiler.compile()).await?
let previous = self.snapshots.snapshot_for_diagnostics();
preempt_for_committed_refresh(cancellation, async {
match previous {
Some(previous) => self.compiler.compile_with_catalog_fallback(&previous).await,
None => self.compiler.compile().await,
}
})
.await?
};
let snapshot = match compiled {
Ok(snapshot) => snapshot,
+4 -1
View File
@@ -20,6 +20,7 @@ pub enum WorkerKind {
RuntimeSnapshotReconciliation,
RuntimeChangeSubscription,
StaleModelRequestRecovery,
RequestAnalytics,
Retention,
OpsFlush,
Backup,
@@ -28,12 +29,13 @@ pub enum WorkerKind {
}
impl WorkerKind {
pub const ALL: [Self; 10] = [
pub const ALL: [Self; 11] = [
Self::OAuthRefresh,
Self::QuotaCatalogHealth,
Self::RuntimeSnapshotReconciliation,
Self::RuntimeChangeSubscription,
Self::StaleModelRequestRecovery,
Self::RequestAnalytics,
Self::Retention,
Self::OpsFlush,
Self::Backup,
@@ -48,6 +50,7 @@ impl WorkerKind {
Self::QuotaCatalogHealth => "quota_catalog_health",
Self::RuntimeSnapshotReconciliation => "runtime_snapshot_reconciliation",
Self::RuntimeChangeSubscription => "runtime_change_subscription",
Self::RequestAnalytics => "request_analytics",
Self::StaleModelRequestRecovery => "stale_model_request_recovery",
Self::Retention => "retention",
Self::OpsFlush => "ops_flush",
@@ -860,7 +860,7 @@ fn model_request(operation: &Operation, deadline: SystemTime) -> NewModelRequest
NewModelRequest {
admission_decision_ms: None,
id: ModelRequestId::new("req_core_1").expect("request id"),
client_api_key_id: Some(client_key.clone()),
client_api_key_name_snapshot: None,
client_api_key_ref: client_key,
config_revision: ConfigRevision::new(1).expect("config revision"),
routing: AccountRoutingSnapshot::all(),
@@ -3493,7 +3493,7 @@ fn ambiguous_pre_delivery_retry_marker_does_not_rotate_account() {
}
#[test]
fn ambiguous_transport_retry_is_rejected_even_with_provider_session_state() {
fn ambiguous_transport_retry_keeps_the_account_without_claiming_replay_safety() {
let operation = generate_operation();
let route_plan = plan(&operation);
let checkpoint = ProviderSessionState::new(
@@ -3513,7 +3513,6 @@ fn ambiguous_transport_retry_is_rejected_even_with_provider_session_state() {
Ok(checkpoint_event),
Err(
ProviderError::new(ProviderErrorKind::Transport, UpstreamSendState::Ambiguous)
.with_replay_safe()
.with_pre_delivery_transport_retry(retry_index, Duration::ZERO),
),
],
@@ -3533,14 +3532,25 @@ fn ambiguous_transport_retry_is_rejected_even_with_provider_session_state() {
CancellationToken::new(),
))
.expect("start execution");
let error =
block_on(session.collect_uncommitted()).expect_err("ambiguous send cannot be replayed");
assert!(matches!(error, EngineError::Provider(_)));
assert_eq!(provider.contexts.lock().expect("contexts lock").len(), 1);
block_on(session.collect_uncommitted()).expect("bounded transport recovery succeeds");
block_on(session.commit_downstream(Some(200))).expect("commit winning response");
let contexts = provider.contexts.lock().expect("contexts lock");
assert_eq!(contexts.len(), 2);
assert_eq!(
contexts[1].transport(),
AttemptTransport::Retry(retry_index)
);
assert_eq!(
contexts[1].required_account(),
Some(&ProviderAccountId::new("acct_first").unwrap())
);
assert!(contexts[1].excluded_accounts().is_empty());
let state = store.state.lock().expect("store lock");
assert_eq!(state.intermediate_failures, 0);
assert_eq!(state.finalizations[0].attempt_count, 1);
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Failed);
assert_eq!(state.intermediate_failures, 1);
assert_eq!(state.finalizations[0].attempt_count, 2);
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Succeeded);
assert!(state.send_states.contains(&UpstreamSendState::Ambiguous));
assert_eq!(state.finalizations[0].send_state, UpstreamSendState::Sent);
}
#[test]
@@ -5844,3 +5854,144 @@ fn zero_rotation_budget_still_allows_retrying_the_same_account() {
Some(&ProviderAccountId::new("acct_first").unwrap())
);
}
#[test]
fn transport_keepalive_preserves_pre_output_recovery_and_first_business_commit() {
let operation = generate_operation();
let route_plan = plan(&operation);
let (coordinator, store, provider) = coordinator(vec![
Script::Stream {
account_id: "acct_first",
items: vec![Err(transient_rejection())],
},
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("session");
let transport = session.response_commit();
transport.record(Some(200));
let first_commit = transport.committed_at();
let event = block_on(session.next_event())
.expect("recovery")
.expect("event");
assert_eq!(
event.commit_requirement(),
CommitRequirement::CommitBeforeDelivery
);
block_on(session.commit_downstream(Some(200))).expect("business commit");
while block_on(session.next_event())
.expect("remaining stream")
.is_some()
{}
assert_eq!(transport.committed_at(), first_commit);
assert_eq!(provider.contexts.lock().unwrap().len(), 2);
let state = store.state.lock().unwrap();
assert_eq!(state.finalizations.len(), 1);
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Succeeded);
assert_eq!(state.finalizations[0].client_status_code, Some(200));
assert!(state.finalizations[0].committed);
}
#[test]
fn transport_keepalive_records_late_failure_as_incomplete_http_200() {
let operation = generate_operation();
let route_plan = plan(&operation);
let (coordinator, store, _) = coordinator(vec![Script::Stream {
account_id: "acct_first",
items: vec![Err(ProviderError::new(
ProviderErrorKind::Timeout,
UpstreamSendState::Sent,
)
.with_status(504))],
}]);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
None,
CancellationToken::new(),
))
.expect("session");
session.response_commit().record(Some(200));
block_on(session.next_event()).expect_err("upstream failure");
block_on(session.record_client_status(504))
.expect("late logical status cannot replace HTTP 200");
let state = store.state.lock().unwrap();
assert_eq!(state.finalizations.len(), 1);
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Incomplete);
assert_eq!(state.finalizations[0].client_status_code, Some(200));
assert_eq!(state.finalizations[0].upstream_status_code, Some(504));
assert!(state.finalizations[0].committed);
assert!(state.recorded_statuses.is_empty());
}
#[test]
fn transport_keepalive_cancellation_retains_commit_without_inventing_a_send() {
let operation = generate_operation();
let route_plan = plan(&operation);
let (coordinator, store, provider) = coordinator(Vec::new());
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
None,
CancellationToken::new(),
))
.expect("session");
session.response_commit().record(Some(200));
block_on(session.cancel_and_finalize()).expect("cancel");
assert!(provider.contexts.lock().unwrap().is_empty());
let state = store.state.lock().unwrap();
assert_eq!(state.finalizations.len(), 1);
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Cancelled);
assert_eq!(state.finalizations[0].client_status_code, Some(200));
assert_eq!(
state.finalizations[0].send_state,
UpstreamSendState::NotSent
);
assert!(state.finalizations[0].committed);
}
#[test]
fn transport_keepalive_cannot_change_an_already_finalized_failure_to_http_200() {
let operation = generate_operation();
let route_plan = plan(&operation);
let (coordinator, store, _) = coordinator(vec![Script::Stream {
account_id: "acct_first",
items: vec![Err(ProviderError::new(
ProviderErrorKind::Timeout,
UpstreamSendState::Sent,
)
.with_status(504))],
}]);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
None,
CancellationToken::new(),
))
.unwrap();
let transport = session.response_commit();
block_on(session.next_event()).expect_err("original failure");
assert!(!transport.record(Some(200)));
assert!(transport.committed_at().is_none());
block_on(session.record_client_status(504)).unwrap();
let state = store.state.lock().unwrap();
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Failed);
assert!(!state.finalizations[0].committed);
assert_eq!(state.recorded_statuses, vec![504]);
}
@@ -579,10 +579,7 @@ fn assert_early_failure_recorded(
let request = &requests[0];
assert_eq!(request.id, *request_id);
assert_eq!(request.client_api_key_ref.as_str(), "key_start_test");
assert_eq!(
request.client_api_key_id.as_ref(),
Some(&request.client_api_key_ref)
);
assert_eq!(request.client_api_key_name_snapshot, None);
assert_eq!(request.protocol, "openai");
assert_eq!(request.endpoint, "/v1/responses");
assert_eq!(request.client_transport, transport.as_str());
@@ -66,3 +66,28 @@ fn replacing_or_clearing_a_binding_does_not_share_control_between_executions() {
);
});
}
#[test]
fn captured_binding_never_redirects_old_control_to_a_replacement() {
block_on(async {
let control = ResponseControl::default();
let old = Arc::new(Transport::default());
let old_owner: Arc<dyn ResponseControlTransport> = old.clone();
control.bind(&old_owner);
let binding = control.binding().unwrap();
let new = Arc::new(Transport::default());
let new_owner: Arc<dyn ResponseControlTransport> = new.clone();
control.bind(&new_owner);
assert!(!binding.same_binding(&control.binding().unwrap()));
binding.send("old ID").await.unwrap();
assert_eq!(*old.0.lock().unwrap(), ["old ID"]);
assert!(new.0.lock().unwrap().is_empty());
drop(old_owner);
drop(old);
assert_eq!(
binding.send("expired ID").await,
Err(ResponseControlUnavailable)
);
assert!(new.0.lock().unwrap().is_empty());
});
}
@@ -37,6 +37,19 @@ fn override_validation_rejects_unsupported_bands_and_excessive_multiplier() {
assert!(price.validate().is_err());
}
#[test]
fn override_validation_accepts_ultrafast_bands() {
let pricing: ModelPriceOverride = serde_json::from_value(json!({
"multiplierBps": 10000,
"bands": {
"ultrafast": {"input":"12", "output":"60", "cacheRead":"0.6", "cacheWrite":"15"},
"long_ultrafast": {"input":"24", "output":"90", "cacheRead":"1.2", "cacheWrite":"30"}
}
}))
.unwrap();
assert!(pricing.validate().is_ok());
}
#[test]
fn sync_layer_cannot_erase_manual_bands_or_multiplier() {
let synced: PricingOverrides = serde_json::from_value(json!({"openai":{"model":{
@@ -282,6 +282,83 @@ impl ProviderCatalogPort for TestCatalog {
}
}
struct MixedCatalog {
xai_available: bool,
}
impl ProviderCatalogPort for MixedCatalog {
fn catalog_generations(&self) -> BTreeMap<ProviderKind, ProviderCatalogGeneration> {
["openai", "xai"]
.into_iter()
.map(|provider| {
(
ProviderKind::new(provider).unwrap(),
ProviderCatalogGeneration::new(1),
)
})
.collect()
}
fn model_catalog_is_exhaustive(&self, provider: &ProviderKind) -> bool {
provider.as_str() != "openai"
}
fn query_model_capabilities(
&self,
provider: &ProviderKind,
) -> BoxFuture<'_, Result<Vec<ProviderModelCapabilities>, ProviderCatalogUnavailable>> {
let model = match provider.as_str() {
"openai" => "gpt-5.5",
"xai" if self.xai_available => "grok-4.7",
_ => return Box::pin(async { Err(ProviderCatalogUnavailable) }),
};
Box::pin(async move {
Ok([model, "shared-model"]
.into_iter()
.map(|model| {
ProviderModelCapabilities::new(
UpstreamModelId::new(model).unwrap(),
super::capabilities(),
)
})
.collect())
})
}
}
fn mixed_catalog_snapshot() -> gateway_core::routing::RuntimeSnapshot {
mixed_catalog_snapshot_with_xai(true)
}
fn mixed_catalog_snapshot_with_xai(xai_available: bool) -> gateway_core::routing::RuntimeSnapshot {
let facts = SnapshotFacts::new(
revision(1),
revision(1),
super::settings(),
vec![],
vec![],
["openai", "xai"]
.into_iter()
.map(|provider| {
SnapshotProviderAccountFacts::new(
gateway_core::account::ProviderAccountId::new(format!("acct_{provider}"))
.unwrap(),
provider,
)
})
.collect(),
vec![],
);
block_on(
RuntimeSnapshotCompiler::new(
Arc::new(TestSnapshotStore::new(Ok(facts))),
Arc::new(MixedCatalog { xai_available }),
)
.compile(),
)
.unwrap()
}
fn catalog_generations(generation: u64) -> BTreeMap<ProviderKind, ProviderCatalogGeneration> {
BTreeMap::from([(
ProviderKind::new("alpha").expect("provider"),
@@ -606,6 +683,205 @@ fn discovery_catalog_should_not_reject_an_unlisted_model() {
.expect("upstream decides model availability");
}
#[test]
fn mixed_catalog_should_bind_known_model_to_its_declared_provider() {
let snapshot = mixed_catalog_snapshot();
let model = PublicModelId::new("grok-4.7").unwrap();
assert!(
snapshot
.public_models_for_scope(&snapshot.all_account_scope())
.contains(&model)
);
let plan = snapshot
.plan(
&model,
&super::operation(),
snapshot.all_account_scope(),
&Default::default(),
)
.unwrap();
assert_eq!(
plan.candidates()
.iter()
.map(|candidate| candidate.provider().as_str())
.collect::<Vec<_>>(),
["xai"]
);
}
#[test]
fn mixed_catalog_should_bind_the_mapped_upstream_model() {
let snapshot = mixed_catalog_snapshot();
let settings = snapshot
.settings()
.clone()
.with_model_mappings(BTreeMap::from([(
"public-model".to_owned(),
"grok-4.7".to_owned(),
)]));
let snapshot = snapshot.with_settings(&settings).unwrap();
let plan = snapshot
.plan(
&PublicModelId::new("public-model").unwrap(),
&super::operation(),
snapshot.all_account_scope(),
&Default::default(),
)
.unwrap();
assert_eq!(plan.candidates().len(), 1);
assert_eq!(plan.candidates()[0].provider().as_str(), "xai");
assert_eq!(
plan.candidates()[0].upstream_model().unwrap().as_str(),
"grok-4.7"
);
}
#[test]
fn mixed_catalog_should_keep_discovery_routes_for_unlisted_models() {
let snapshot = mixed_catalog_snapshot();
let plan = snapshot
.plan(
&PublicModelId::new("unlisted-model").unwrap(),
&super::operation(),
snapshot.all_account_scope(),
&Default::default(),
)
.unwrap();
assert_eq!(plan.candidates().len(), 1);
assert_eq!(plan.candidates()[0].provider().as_str(), "openai");
}
#[test]
fn mixed_catalog_should_preserve_provider_order_for_shared_models() {
let snapshot = mixed_catalog_snapshot();
let plan = snapshot
.plan(
&PublicModelId::new("shared-model").unwrap(),
&super::operation(),
snapshot.all_account_scope(),
&Default::default(),
)
.unwrap();
assert_eq!(
plan.candidates()
.iter()
.map(|candidate| candidate.provider().as_str())
.collect::<Vec<_>>(),
["openai", "xai"]
);
}
#[test]
fn mixed_catalog_should_reject_forcing_a_model_to_another_provider() {
let snapshot = mixed_catalog_snapshot();
let plan = snapshot.plan(
&PublicModelId::new("grok-4.7").unwrap(),
&super::operation(),
snapshot.all_account_scope(),
&gateway_core::routing::RoutingContext {
required_provider: Some(ProviderKind::new("openai").unwrap()),
..Default::default()
},
);
assert!(matches!(
plan,
Err(gateway_core::validation::RoutingError::NoCapableProvider { .. })
));
}
#[test]
fn mixed_catalog_should_not_escape_a_blocked_model_provider() {
let snapshot = mixed_catalog_snapshot();
let plan = snapshot.plan(
&PublicModelId::new("grok-4.7").unwrap(),
&super::operation(),
snapshot.all_account_scope(),
&gateway_core::routing::RoutingContext {
blocked_providers: std::collections::BTreeSet::from(
[ProviderKind::new("xai").unwrap()],
),
..Default::default()
},
);
assert!(matches!(
plan,
Err(gateway_core::validation::RoutingError::NoCapableProvider { .. })
));
}
#[test]
fn mixed_catalog_should_not_escape_model_capability_requirements() {
let snapshot = mixed_catalog_snapshot();
let plan = snapshot.plan(
&PublicModelId::new("grok-4.7").unwrap(),
&super::image_operation(),
snapshot.all_account_scope(),
&Default::default(),
);
assert!(matches!(
plan,
Err(gateway_core::validation::RoutingError::NoCapableProvider { .. })
));
}
#[test]
fn mixed_catalog_should_not_escape_the_authorized_account_scope() {
let snapshot = mixed_catalog_snapshot();
let scope = Arc::new(gateway_core::routing::FrozenAccountScope::new(
Arc::new(gateway_core::routing::RuntimeAccountDirectory::new(
BTreeMap::from([(
gateway_core::account::ProviderAccountId::new("acct_openai").unwrap(),
gateway_core::routing::RuntimeAccount::new(
ProviderKind::new("openai").unwrap(),
Default::default(),
),
)]),
)),
gateway_core::routing::ClientRoutingScope::all_accounts(),
));
let model = PublicModelId::new("grok-4.7").unwrap();
assert!(!snapshot.contains_public_model_for_scope(&model, &scope));
let plan = snapshot.plan(&model, &super::operation(), scope, &Default::default());
assert!(matches!(
plan,
Err(gateway_core::validation::RoutingError::NoCapableProvider { .. })
));
}
#[test]
fn mixed_catalog_should_reject_ambiguous_models_when_xai_catalog_is_unavailable() {
let snapshot = mixed_catalog_snapshot_with_xai(false);
let plan = snapshot.plan(
&PublicModelId::new("grok-4.7").unwrap(),
&super::operation(),
snapshot.all_account_scope(),
&Default::default(),
);
assert!(matches!(
plan,
Err(gateway_core::validation::RoutingError::NoCapableProvider { .. })
));
let plan = snapshot
.plan(
&PublicModelId::new("gpt-5.5").unwrap(),
&super::operation(),
snapshot.all_account_scope(),
&Default::default(),
)
.unwrap();
assert_eq!(plan.candidates().len(), 1);
assert_eq!(plan.candidates()[0].provider().as_str(), "openai");
}
#[test]
fn known_empty_catalog_should_report_model_not_found() {
let compiler = RuntimeSnapshotCompiler::new(
@@ -2,7 +2,7 @@
use std::collections::BTreeMap;
mod publications;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use futures::channel::{mpsc, oneshot};
@@ -612,6 +612,85 @@ fn catalog_only_reconciliation_should_keep_snapshot_on_failure_and_retry_generat
});
}
#[test]
fn catalog_failure_should_preserve_model_facts_without_restoring_revoked_accounts() {
block_on(async {
let catalog = Arc::new(TestCatalog::default());
let provider = ProviderKind::new("audit").unwrap();
let model = gateway_core::routing::PublicModelId::new("audit-model").unwrap();
catalog.models.lock().unwrap().push(
ProviderModelCapabilities::new(
UpstreamModelId::new("audit-model").unwrap(),
ModelCapabilities::new(
std::collections::BTreeSet::from([
gateway_core::operation::OperationKind::Generate,
]),
Some(16_000),
),
)
.with_presentation(gateway_core::routing::ModelPresentation::new(
Some("Audit model".to_owned()),
None,
))
.with_catalog_accounts(std::collections::BTreeSet::from([ProviderAccountId::new(
"acct_audit_removed",
)
.unwrap()])),
);
let store = Arc::new(TestSnapshotStore::new(Ok(scoped_facts(10, true))));
let handle = RuntimeSnapshotHandle::default();
let publisher = RuntimeSnapshotPublisher::new(
Arc::new(catalog_compiler(store.clone(), catalog.clone())),
handle.clone(),
Arc::new(TestSnapshotSubscriptions::default()),
);
publisher.refresh().await.unwrap();
let initial = handle.acquire().unwrap();
let initial_profiles = initial.public_model_profiles_for_provider(&provider);
catalog.unavailable.store(true, Ordering::SeqCst);
catalog.generation.store(1, Ordering::SeqCst);
*store.facts.lock().unwrap() = Ok(scoped_facts(11, false));
*store.current_revision.lock().unwrap() = Ok(revision(11));
publisher.refresh().await.unwrap();
let retained = handle.acquire().unwrap();
assert_eq!(retained.revision(), revision(11));
assert!(!removed_account_allowed(&retained));
assert_eq!(
retained.public_model_profiles_for_provider(&provider),
initial_profiles
);
assert!(retained.contains_public_model_for_provider(&model, &provider));
assert!(!retained.contains_public_model_for_provider(
&gateway_core::routing::PublicModelId::new("unlisted-model").unwrap(),
&provider,
));
assert!(retained.public_models_for_scope(
retained.client_policies().next().unwrap().account_scope(),
).is_empty());
assert!(removed_account_allowed(&initial));
assert_eq!(
retained.provider_catalog_generations(),
&catalog.catalog_generations()
);
// 上游恢复并明确撤回模型后,旧目录不能继续作为已知归属
catalog.unavailable.store(false, Ordering::SeqCst);
catalog.models.lock().unwrap().clear();
catalog.generation.store(2, Ordering::SeqCst);
let (task, context) = reconciliation_task(&publisher);
task.run_cycle(context).await.unwrap();
let refreshed = handle.acquire().unwrap();
assert!(refreshed.public_models_for_provider(&provider).is_empty());
assert!(!refreshed.contains_public_model_for_provider(&model, &provider));
assert_eq!(
refreshed.provider_catalog_generations(),
&catalog.catalog_generations()
);
});
}
#[test]
fn dropping_inflight_refresh_should_release_publication_for_next_refresh() {
block_on(async {
@@ -688,6 +767,7 @@ fn reconciliation_task(
#[derive(Default)]
struct TestCatalog {
generation: AtomicU64,
unavailable: AtomicBool,
queries: AtomicUsize,
next_query: Mutex<Option<oneshot::Receiver<()>>>,
models: Mutex<Vec<ProviderModelCapabilities>>,
@@ -711,6 +791,9 @@ impl ProviderCatalogPort for TestCatalog {
if let Some(gate) = gate {
gate.await.expect("release catalog query");
}
if self.unavailable.load(Ordering::SeqCst) {
return Err(ProviderCatalogUnavailable);
}
Ok(self.models.lock().expect("models lock").clone())
})
}
+2 -1
View File
@@ -19,12 +19,13 @@ use gateway_core::task::{
use gateway_host::workers::{WorkerStartError, WorkerSupervisor};
use tokio::sync::Notify;
const ACTIVE_KINDS: [WorkerKind; 9] = [
const ACTIVE_KINDS: [WorkerKind; 10] = [
WorkerKind::OAuthRefresh,
WorkerKind::QuotaCatalogHealth,
WorkerKind::RuntimeSnapshotReconciliation,
WorkerKind::RuntimeChangeSubscription,
WorkerKind::StaleModelRequestRecovery,
WorkerKind::RequestAnalytics,
WorkerKind::Retention,
WorkerKind::Backup,
WorkerKind::AccountImport,
@@ -48,8 +48,10 @@ pub(super) async fn invoke(
Err(MiddlewareError::Fault)
};
};
let remaining = context.deadline().bounded(maximum_timeout);
if remaining.is_zero() || context.cancellation().is_cancelled() {
if maximum_timeout.is_zero()
|| context.deadline().is_elapsed()
|| context.cancellation().is_cancelled()
{
return Err(MiddlewareError::Fault);
}
let invocation = MiddlewareInvocation::new(&context, request, next, entry.instance_id.clone());
@@ -65,7 +67,7 @@ pub(super) async fn invoke(
MiddlewareMount::Request => Stage::Request,
MiddlewareMount::Attempt => Stage::Attempt,
},
remaining,
maximum_timeout,
);
call_context.request_id = Some(context.request_id().as_str().to_owned());
call_context.attempt_id = context
@@ -129,6 +131,10 @@ pub(super) async fn invoke(
let stream = tokio::select! {
biased;
() = context.cancellation().cancelled() => return Err(MiddlewareError::Fault),
() = context.deadline().wait() => return Err(gateway_core::error::GatewayError::new(
gateway_core::error::GatewayErrorKind::Timeout,
"request deadline elapsed",
).into()),
result = invocation_ports.session.call_stream(
HANDLE_METHOD,
call_context,
@@ -495,9 +495,9 @@ fn wire_observation(observation: &RequestObservation) -> RequestCompleted {
RequestUsage {
input_tokens: usage.input_tokens,
output_tokens: usage.output_tokens,
cached_tokens: usage.cached_tokens,
cache_write_tokens: usage.cache_write_tokens,
reasoning_tokens: usage.reasoning_tokens,
cached_tokens: usage.cached_input_tokens,
cache_write_tokens: usage.cache_write_input_tokens,
reasoning_tokens: usage.reasoning_output_tokens,
image_input_tokens: usage.image_input_tokens,
image_output_tokens: usage.image_output_tokens,
total_tokens: usage.total_tokens,
@@ -165,9 +165,9 @@ impl EventDecoder {
let usage = Usage {
input_tokens: usage.input_tokens,
output_tokens: usage.output_tokens,
cached_tokens: usage.cached_tokens,
cache_write_tokens: usage.cache_write_tokens,
reasoning_tokens: usage.reasoning_tokens,
cached_input_tokens: usage.cached_tokens,
cache_write_input_tokens: usage.cache_write_tokens,
reasoning_output_tokens: usage.reasoning_tokens,
image_input_tokens: usage.image_input_tokens,
image_output_tokens: usage.image_output_tokens,
total_tokens: usage.total_tokens,
@@ -3,16 +3,15 @@
use super::{CallbackScope, MiddlewareCallback, denied, http_middleware, model};
use gateway_core::lifecycle::CancellationToken;
use gateway_host::outbound::HttpBody;
use gateway_plugin_sdk::{CallContext, ErrorCode, PluginFault};
use gateway_plugin_sdk::{CallContext, ErrorCode, PluginFault, client::CallDeadline};
use std::{
collections::BTreeMap,
sync::{Arc, Mutex},
time::Duration,
};
use tokio::time::Instant;
pub(super) struct CallResources {
pub(super) deadline: Instant,
pub(super) deadline: CallDeadline,
operation_timeout: Duration,
pub(super) cancellation: CancellationToken,
pub(super) scope: Arc<CallbackScope>,
@@ -26,11 +25,12 @@ pub(super) struct CallResources {
impl CallResources {
pub(super) fn new(
context: &CallContext,
deadline: &CallDeadline,
scope: Arc<CallbackScope>,
middleware: Option<Arc<dyn MiddlewareCallback>>,
) -> Self {
Self {
deadline: Instant::now() + Duration::from_millis(context.timeout_ms),
deadline: deadline.clone(),
operation_timeout: Duration::from_millis(context.timeout_ms),
cancellation: CancellationToken::new(),
scope,
@@ -46,13 +46,6 @@ impl CallResources {
}
}
pub(super) fn set_streaming(&self, streaming: bool) {
self.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.resource_stream = streaming;
}
pub(super) fn close(&self) {
self.cancellation.cancel();
let mut state = self
@@ -94,20 +87,20 @@ impl CallResources {
if state.closed {
return Err(denied());
}
if state.resource_stream {
return Ok(self.operation_timeout);
let remaining = self.deadline.remaining().unwrap_or(self.operation_timeout);
if remaining.is_zero() {
return Err(PluginFault::new(
ErrorCode::Timeout,
"callback deadline elapsed",
));
}
self.deadline
.checked_duration_since(Instant::now())
.filter(|timeout| !timeout.is_zero())
.ok_or_else(|| PluginFault::new(ErrorCode::Timeout, "callback deadline elapsed"))
Ok(remaining)
}
}
#[derive(Default)]
pub(super) struct CallState {
pub(super) closed: bool,
pub(super) resource_stream: bool,
pub(super) middleware: Option<Arc<dyn MiddlewareCallback>>,
pub(super) streams: BTreeMap<String, Arc<HttpStream>>,
pub(super) model_streams: BTreeMap<String, Arc<model::ModelStream>>,
@@ -86,7 +86,7 @@ fn breakdown(value: &CalculatedCostBreakdown) -> facts::CostBreakdown {
long_context_billing_applied: value.long_context_billing_applied(),
image: value.image().map(|value| facts::ImageCostBreakdown {
input_tokens: value.input_tokens,
cached_tokens: value.cached_tokens,
cached_tokens: value.cached_input_tokens,
input_amount: money(value.input_amount),
cache_read_amount: money(value.cache_read_amount),
input_price_per_million: money(value.input_price_per_million),
@@ -234,9 +234,9 @@ fn project_fact(event: GatewayEvent) -> Option<Result<CanonicalEvent, PluginFaul
usage: WireUsage {
input_tokens: usage.input_tokens,
output_tokens: usage.output_tokens,
cached_tokens: usage.cached_tokens,
cache_write_tokens: usage.cache_write_tokens,
reasoning_tokens: usage.reasoning_tokens,
cached_tokens: usage.cached_input_tokens,
cache_write_tokens: usage.cache_write_input_tokens,
reasoning_tokens: usage.reasoning_output_tokens,
image_input_tokens: usage.image_input_tokens,
image_output_tokens: usage.image_output_tokens,
total_tokens: usage.total_tokens,
@@ -30,8 +30,7 @@ use futures::future::BoxFuture;
use gateway_admin::model::AdminError;
use gateway_core::{account::OutboundProxy, lifecycle::CancellationToken};
use gateway_host::outbound::{HttpClient, NetworkPolicy};
use gateway_plugin_sdk::{CallContext, ErrorCode, PluginFault};
use tokio::time::Instant;
use gateway_plugin_sdk::{CallContext, ErrorCode, PluginFault, client::CallDeadline};
use crate::{CallbackHandler, RpcReply};
pub(crate) use accounts::PluginAccountPortSlot;
@@ -329,7 +328,7 @@ impl PluginCallbacks {
}
impl CallbackHandler for PluginCallbacks {
fn begin(&self, context: &CallContext) {
fn begin(&self, context: &CallContext, deadline: &CallDeadline) {
let scope = self
.scopes
.lock()
@@ -347,21 +346,10 @@ impl CallbackHandler for PluginCallbacks {
.unwrap_or_else(std::sync::PoisonError::into_inner)
.insert(
context.call_id,
Arc::new(CallResources::new(context, scope, middleware)),
Arc::new(CallResources::new(context, deadline, scope, middleware)),
);
}
fn streaming(&self, context: &CallContext) {
if let Some(call) = self
.calls
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.get(&context.call_id)
{
call.set_streaming(context.resource_stream);
}
}
fn finish(&self, context: &CallContext) {
if let Some(call) = self
.calls
@@ -406,7 +394,7 @@ impl CallbackHandler for PluginCallbacks {
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if state.closed || (!state.resource_stream && Instant::now() >= call.deadline) {
if state.closed || call.deadline.is_elapsed() {
return Err(denied());
}
if method == "host.log" {
@@ -4,6 +4,7 @@ use std::sync::Arc;
use gateway_plugin_sdk::{
ErrorCode, Frame, Message, PluginFault,
call::middleware::NEXT_METHOD,
client::{read_frame, validate_frame, write_frame},
};
use tokio::{
@@ -128,12 +129,16 @@ async fn read_loop<R: AsyncRead + Unpin>(
let handler = Arc::clone(&callbacks);
let response = Arc::clone(&shared);
let (ready, started) = oneshot::channel();
let pause = (method == NEXT_METHOD)
.then(|| shared.pause_next(parent_id))
.flatten();
let task = tokio::spawn(async move {
// 先把任务归属登记到父调用,再允许执行宿主操作
if started.await.is_err() {
return;
}
let result = handler.call(context, method, params, frame.payload).await;
drop(pause);
let mut reply = match result {
Ok(reply) => Frame {
message: Message::Result {
@@ -13,7 +13,7 @@ use futures::future::BoxFuture;
use gateway_host::process::{ProcessControl, ProcessSpec, ProcessStartError, ProcessSupervisor};
use gateway_plugin_sdk::{
CallContext, Frame, Handshake, Message, PROTOCOL_VERSION, PluginFault, Stage,
client::{read_frame, validate_frame, write_frame},
client::{CallDeadline, CallDeadlinePause, read_frame, validate_frame, write_frame},
};
use tokio::sync::{Mutex as AsyncMutex, Semaphore, mpsc, oneshot, watch};
@@ -155,7 +155,7 @@ pub struct RpcReply {
/// 具体回调在领域适配器校验账号、出站与执行身份;上下文只取自宿主在途调用表
pub trait CallbackHandler: Send + Sync {
/// 在 Call 可见之前建立资源范围;不得执行 I/O 或调用 RPC
fn begin(&self, _context: &CallContext) {}
fn begin(&self, _context: &CallContext, _deadline: &CallDeadline) {}
/// 调用结束或取消时立即撤销其句柄;实现必须幂等且不能阻塞
fn finish(&self, _context: &CallContext) {}
@@ -224,6 +224,7 @@ struct State {
}
struct PendingCall {
deadline: CallDeadline,
handler: Arc<dyn CallbackHandler>,
transmitted: bool,
cancellation: Option<RpcError>,
@@ -293,7 +294,14 @@ impl Shared {
if call.cancellation.is_some() {
return Ok(());
}
if matches!(&result, Err(RpcError::Remote(fault)) if fault.code == gateway_plugin_sdk::ErrorCode::Timeout)
{
log_timeout(&call.context, "plugin", "initial_reply", "remote_timeout");
}
if call.stream.is_some() && result.is_ok() {
if call.context.resource_stream {
call.deadline.clear();
}
let initial = call.result.take().ok_or(RpcError::Protocol(None))?;
call.handler.streaming(&call.context);
let _ = initial.send(result);
@@ -375,6 +383,14 @@ impl Shared {
.map(|call| call.context.clone())
}
pub fn pause_next(&self, id: u64) -> Option<CallDeadlinePause> {
self.state()
.calls
.get(&id)
.filter(|call| call.cancellation.is_none())
.and_then(|call| call.deadline.pause())
}
pub fn is_cancelling(&self, id: u64) -> bool {
self.state()
.calls
@@ -403,6 +419,18 @@ impl Shared {
if call.cancellation.is_some() {
return;
}
if matches!(error, RpcError::Timeout) {
log_timeout(
&call.context,
"runtime",
if call.result.is_some() {
"initial_reply"
} else {
"response_stream"
},
"processing_budget_exhausted",
);
}
if !call.transmitted {
if let Some(call) = state.calls.remove(&id) {
call.finish(Err(error));
@@ -759,10 +787,12 @@ impl RpcSession {
if state.calls.len() >= self.limits.maximum_calls {
return Err(RpcError::Capacity);
}
self.callbacks.begin(&context);
let deadline = CallDeadline::new(deadline);
self.callbacks.begin(&context, &deadline);
state.calls.insert(
id,
PendingCall {
deadline: deadline.clone(),
handler: self.callbacks.clone(),
transmitted: false,
cancellation: None,
@@ -817,7 +847,7 @@ impl RpcSession {
id: started.id,
armed: true,
};
let sent = tokio::time::timeout_at(started.deadline, async {
let send = async {
self.data
.send(frame)
.await
@@ -834,8 +864,12 @@ impl RpcSession {
.map_err(|_| RpcError::Closed(None))?;
}
Ok::<(), RpcError>(())
})
.await;
};
let sent = tokio::select! {
biased;
() = started.deadline.wait() => Err(()),
result = send => Ok(result),
};
drop(next_call);
match sent {
Ok(Ok(())) => {}
@@ -850,7 +884,11 @@ impl RpcSession {
return Err(RpcError::Timeout);
}
}
let result = tokio::time::timeout_at(started.deadline, started.received).await;
let result = tokio::select! {
biased;
() = started.deadline.wait() => Err(()),
result = started.received => Ok(result),
};
let reply = match result {
Ok(Ok(result)) => {
guard.armed = false;
@@ -900,7 +938,6 @@ impl RpcSession {
let bytes = self.limits.maximum_buffered_body_bytes.min(256 * 1024) as u32;
let frames = 32;
let (stream, chunks, terminal) = StreamIngress::new(bytes, frames);
let resource_stream = context.resource_stream;
let completed = self
.send_call(
method,
@@ -915,7 +952,7 @@ impl RpcSession {
initial: completed.reply,
chunks,
terminal: Some(terminal),
deadline: (!resource_stream).then_some(completed.deadline),
deadline: completed.deadline,
id: completed.id,
shared: Arc::clone(&self.shared),
_slot: slot,
@@ -980,12 +1017,33 @@ impl Drop for CallGuard {
struct StartedCall {
id: u64,
deadline: tokio::time::Instant,
deadline: CallDeadline,
received: oneshot::Receiver<Result<RpcReply, RpcError>>,
}
struct CompletedCall {
id: u64,
deadline: tokio::time::Instant,
deadline: CallDeadline,
reply: RpcReply,
}
fn log_timeout(
context: &CallContext,
owner: &'static str,
phase: &'static str,
reason: &'static str,
) {
tracing::warn!(
event = "plugin.rpc.timeout",
owner,
phase,
instance_id = context.instance_id,
call_id = context.call_id,
request_id = context.request_id,
attempt_id = context.attempt_id,
stage = ?context.stage,
budget_ms = context.timeout_ms,
reason,
"插件 RPC 调用超时"
);
}
@@ -2,10 +2,9 @@
use std::sync::Arc;
use tokio::{
sync::{OwnedSemaphorePermit, mpsc, oneshot},
time::Instant,
};
use tokio::sync::{OwnedSemaphorePermit, mpsc, oneshot};
use gateway_plugin_sdk::client::CallDeadline;
use crate::{RpcError, RpcReply, RpcSession, rpc::Shared};
@@ -60,7 +59,7 @@ pub struct RpcStream {
pub initial: RpcReply,
pub(crate) chunks: mpsc::Receiver<Vec<u8>>,
pub(crate) terminal: Option<oneshot::Receiver<Result<(), RpcError>>>,
pub(crate) deadline: Option<Instant>,
pub(crate) deadline: CallDeadline,
pub(crate) id: u64,
pub(crate) shared: Arc<Shared>,
pub(crate) _slot: OwnedSemaphorePermit,
@@ -72,9 +71,10 @@ impl RpcStream {
if self.terminal.is_none() {
return Ok(None);
}
let result = match self.deadline {
Some(deadline) => tokio::time::timeout_at(deadline, self.chunks.recv()).await,
None => Ok(self.chunks.recv().await),
let result = tokio::select! {
biased;
() = self.deadline.wait() => Err(()),
chunk = self.chunks.recv() => Ok(chunk),
};
match result {
Ok(Some(chunk)) => {
@@ -1,4 +1,4 @@
//! 验证跨集合复用实例时旧请求和宿主回调仍持有各自的冻结计划
//! 验证中间件续体的调用预算、取消回收及跨集合冻结计划
use super::*;
use gateway_admin::ports::plugins::PluginPreparation;
@@ -124,3 +124,206 @@ async fn partial_update_preserves_inflight_middleware_and_uses_new_configuration
drop(new);
super::super::wait_until_empty(cache.path()).await;
}
async fn sdk_budget_runtime(stage: Stage, count: usize) -> (tempfile::TempDir, PluginRuntime) {
let worker = std::fs::read(env!("CARGO_BIN_EXE_gateway-plugin-test-middleware")).unwrap();
let package = crate::support::package_with_contributions(
&worker,
Contributions::from([crate::support::contribution(
Capability::Middleware,
vec![stage],
vec!["openai".into()],
vec!["openai".into()],
)]),
);
let mount = if stage == Stage::Attempt {
"attempt"
} else {
"request"
};
let instances = ["budget-outer", "budget-inner"]
.into_iter()
.take(count)
.enumerate()
.map(|(order, id)| InstanceFixture {
id,
configuration: serde_json::json!({"mode":"map", "attempt":stage == Stage::Attempt}),
bindings: vec![binding(
MIDDLEWARE_CONTRIBUTION,
mount,
order as i32,
PluginFailurePolicy::Reject,
)],
})
.collect();
setup_package_with_limits(
instances,
package,
RpcLimits {
maximum_call_timeout: Duration::from_millis(300),
..RpcLimits::default()
},
)
.await
}
fn budget_request() -> MiddlewareRequest {
MiddlewareRequest::new(
"openai",
middleware_headers(),
Bytes::from_static(br#"{"input":"hello"}"#),
)
}
#[tokio::test]
async fn nested_sdk_middleware_waits_beyond_budget_and_can_read_response_callbacks() {
for stage in [Stage::Request, Stage::Attempt] {
let (cache, runtime) = sdk_budget_runtime(stage, 2).await;
let generation = prepare(&runtime).await;
let plan = runtime
.execution_registry()
.middleware(&generation)
.unwrap();
let context = if stage == Stage::Attempt {
attempt_middleware_context(Arc::default())
} else {
middleware_context_with_deadline(ClientTransport::HttpJson, Default::default())
};
let calls = Arc::new(AtomicUsize::new(0));
let downstream = Downstream {
calls: calls.clone(),
closes: Arc::default(),
reads: Arc::default(),
error: false,
};
let response = tokio::time::timeout(
Duration::from_secs(3),
plan.handle(
context,
budget_request(),
gateway_core::middleware::compose(Vec::new(), move |request| {
Box::pin(async move {
tokio::time::sleep(Duration::from_millis(900)).await;
downstream.into_next().run(request).await
})
}),
),
)
.await
.unwrap()
.expect("host next must not consume either plugin's processing budget");
let (_, status, _, mut body, _) = response.into_parts();
assert_eq!(status, 200);
// 正文改写还会发起新的 body_read,证明 callback 有效期也随暂停恢复
let frame = body.next_frame().await.unwrap().unwrap();
assert_eq!(
frame.into_bytes(),
Bytes::from_static(b"{\"downstream\":true} ")
);
assert!(body.next_frame().await.unwrap().is_none());
body.close().await;
assert_eq!(calls.load(Ordering::Relaxed), 1);
assert!(generation.is_ready());
drop(plan);
drop(generation);
super::super::wait_until_empty(cache.path()).await;
}
}
struct NotifyOnDrop(Option<tokio::sync::oneshot::Sender<()>>);
impl Drop for NotifyOnDrop {
fn drop(&mut self) {
if let Some(sender) = self.0.take() {
let _ = sender.send(());
}
}
}
#[tokio::test]
async fn sdk_waiting_next_obeys_request_deadline_and_client_cancellation() {
for cancel_client in [false, true] {
let (cache, runtime) = sdk_budget_runtime(Stage::Request, 1).await;
let generation = prepare(&runtime).await;
let plan = runtime
.execution_registry()
.middleware(&generation)
.unwrap();
let deadline = if cancel_client {
Default::default()
} else {
(SystemTime::now() + Duration::from_millis(900)).into()
};
let context = middleware_context_with_deadline(ClientTransport::HttpJson, deadline);
let cancellation = context.cancellation().clone();
let (started, entered) = tokio::sync::oneshot::channel();
let (closed, dropped) = tokio::sync::oneshot::channel();
let pending = tokio::spawn(async move {
plan.handle(
context,
budget_request(),
gateway_core::middleware::compose(Vec::new(), move |_| {
Box::pin(async move {
let _closed = NotifyOnDrop(Some(closed));
let _ = started.send(());
futures::future::pending().await
})
}),
)
.await
});
tokio::time::timeout(Duration::from_secs(2), entered)
.await
.unwrap()
.unwrap();
tokio::time::sleep(Duration::from_millis(600)).await;
assert!(
!pending.is_finished(),
"plugin budget must stay paused while waiting for next"
);
if cancel_client {
cancellation.cancel();
}
let result = tokio::time::timeout(Duration::from_secs(2), pending)
.await
.unwrap()
.unwrap();
match result {
Err(MiddlewareError::Fault) if cancel_client => {}
Err(MiddlewareError::Gateway(error)) if !cancel_client => {
assert_eq!(error.kind(), gateway_core::error::GatewayErrorKind::Timeout);
}
_ => panic!("request termination must end the suspended middleware"),
}
tokio::time::timeout(Duration::from_secs(2), dropped)
.await
.unwrap()
.unwrap();
// 取消单次调用后,同一进程仍能完成新的中间件及正文回调
let plan = runtime
.execution_registry()
.middleware(&generation)
.unwrap();
let response = plan
.handle(
middleware_context(ClientTransport::HttpJson),
budget_request(),
Downstream {
calls: Arc::default(),
closes: Arc::default(),
reads: Arc::default(),
error: false,
}
.into_next(),
)
.await
.unwrap();
let (_, _, _, mut body, _) = response.into_parts();
assert!(body.next_frame().await.unwrap().is_some());
body.close().await;
assert!(generation.is_ready());
drop(plan);
drop(generation);
super::super::wait_until_empty(cache.path()).await;
}
}
@@ -681,6 +681,16 @@ impl Downstream {
}
fn middleware_context(transport: ClientTransport) -> MiddlewareContext {
middleware_context_with_deadline(
transport,
(SystemTime::now() + Duration::from_secs(5)).into(),
)
}
fn middleware_context_with_deadline(
transport: ClientTransport,
deadline: gateway_core::lifecycle::Deadline,
) -> MiddlewareContext {
MiddlewareContext::new(
MiddlewareTarget {
request_id: ModelRequestId::new("req_middleware").unwrap(),
@@ -697,7 +707,7 @@ fn middleware_context(transport: ClientTransport) -> MiddlewareContext {
client_key_id: ClientApiKeyId::new("key-middleware").unwrap(),
account_group_ids: Arc::<[AccountGroupId]>::from([]),
cancellation: CancellationToken::new(),
deadline: (SystemTime::now() + Duration::from_secs(5)).into(),
deadline,
extension_scope: ExtensionCallScope::default(),
execution_effects: None,
},
@@ -97,7 +97,7 @@ impl CallbackHandler for BlockingCallbacks {
}
impl CallbackHandler for LifecycleCallbacks {
fn begin(&self, _context: &CallContext) {
fn begin(&self, _context: &CallContext, _deadline: &gateway_plugin_sdk::client::CallDeadline) {
self.begun.fetch_add(1, Ordering::Relaxed);
}
@@ -122,7 +122,7 @@ impl CallbackHandler for LifecycleCallbacks {
}
impl CallbackHandler for GatedFirstBegin {
fn begin(&self, context: &CallContext) {
fn begin(&self, context: &CallContext, _deadline: &gateway_plugin_sdk::client::CallDeadline) {
if context.call_id == 1 {
self.entered.notify_one();
let released = self.released.lock().unwrap();
@@ -865,3 +865,53 @@ async fn acknowledged_timeout_does_not_interrupt_another_call() {
assert!(session.is_ready());
session.shutdown(Duration::from_secs(1)).await;
}
struct SlowNextCallbacks;
impl CallbackHandler for SlowNextCallbacks {
fn call(
&self,
_context: CallContext,
method: String,
params: serde_json::Value,
payload: Vec<u8>,
) -> BoxFuture<'static, Result<RpcReply, PluginFault>> {
Box::pin(async move {
assert_eq!(method, gateway_plugin_sdk::call::middleware::NEXT_METHOD);
tokio::time::sleep(Duration::from_millis(600)).await;
Ok(RpcReply {
result: params,
payload,
})
})
}
}
#[tokio::test]
async fn managed_next_suspends_only_its_parent_and_not_a_concurrent_call() {
let (_cache, session) = session(Arc::new(SlowNextCallbacks)).await;
let delegated = session.call(
"callback_method",
session.context(Stage::Request, Duration::from_millis(200)),
json!({"method": gateway_plugin_sdk::call::middleware::NEXT_METHOD}),
Vec::new(),
);
let hanging = session.call(
"hang",
session.context(Stage::Request, Duration::from_millis(100)),
json!({}),
Vec::new(),
);
let (delegated, hanging) = tokio::time::timeout(Duration::from_secs(2), async {
tokio::join!(delegated, hanging)
})
.await
.unwrap();
assert!(
delegated.is_ok(),
"host waiting for next must outlive its processing budget"
);
assert!(matches!(hanging, Err(RpcError::Timeout)));
assert!(session.is_ready());
session.shutdown(Duration::from_secs(1)).await;
}
+1 -1
View File
@@ -19,7 +19,7 @@ thiserror = "2.0.18"
tokio = { version = "1.52.3", optional = true, default-features = false, features = ["io-util", "macros", "rt", "sync", "time"] }
[dev-dependencies]
tokio = { version = "1.52.3", default-features = false, features = ["io-util", "macros", "rt", "sync", "time"] }
tokio = { version = "1.52.3", default-features = false, features = ["io-util", "macros", "rt", "sync", "time", "test-util"] }
[[test]]
name = "main"
@@ -242,6 +242,11 @@ SDK 不自动重试。超时或断连不能证明写入未提交;重试上限
`PluginBuilder::middleware` 与 `MiddlewarePlugin::new` 共用 `MiddlewareInput` / `MiddlewareOutput`。
`next.run(input)` 消费自身,最多执行一次;结果按相反顺序经过外层。重试或独立子请求通过宿主执行服务发起
调用上下文的 `timeout_ms` 限制插件自身处理时间,默认上限为 120 秒。等待宿主 `next` 时暂停这份预算,
收到结果或错误后恢复剩余时间;入站处理与返回处理合计使用同一预算,普通宿主回调不会暂停计时。
请求总期限、下游操作超时和取消信号独立生效,客户端取消或插件退出仍会结束等待并回收资源。
宿主与插件 SDK 必须同时支持此语义;仍使用固定绝对期限的 SDK 构建的插件需重新构建,单独升级宿主不能解除插件侧期限
HTTP 总入口统一包裹原生路由,`dispatch_http` 复用同一路由与组合能力;无需为每个 HTTP 接口登记新的挂载点。
`service` 用于不经过 HTTP 的类型化主动调用,具体公开范围见下节
@@ -28,9 +28,9 @@ pub use plugin::{
AuthorError, ComposedPlugin, Empty, Method, PluginBuilder, TypedCall, TypedReply, methods,
};
pub use session::{
CallCancellation, CallFuture, CallReply, HostClient, HostReply, PluginCall, PluginHandler,
PluginSession, PullResponseFuture, PullResponseStream, ResponseStream, SessionConfig,
SessionError, StreamSender,
CallCancellation, CallDeadline, CallDeadlinePause, CallFuture, CallReply, HostClient,
HostReply, PluginCall, PluginHandler, PluginSession, PullResponseFuture, PullResponseStream,
ResponseStream, SessionConfig, SessionError, StreamSender,
};
pub use upstream_adapter::{UpstreamWebSocket, UpstreamWebSocketUpgrade};
@@ -2,9 +2,7 @@
use std::sync::Mutex;
use tokio::time::Instant;
use super::{CallCancellation, SessionError, lifecycle::wait_deadline};
use super::{CallCancellation, CallDeadline, SessionError};
#[derive(Default)]
pub(super) struct CreditWindow {
@@ -53,7 +51,7 @@ impl CreditWindow {
pub(super) async fn initial_window(
&self,
deadline: Option<Instant>,
deadline: &CallDeadline,
cancellation: &CallCancellation,
) -> Result<u64, SessionError> {
loop {
@@ -70,7 +68,7 @@ impl CreditWindow {
tokio::select! {
biased;
() = cancellation.cancelled() => return Err(SessionError::Cancelled),
() = wait_deadline(deadline) => return Err(SessionError::Timeout),
() = deadline.wait() => return Err(SessionError::Timeout),
() = changed => {}
}
}
@@ -79,7 +77,7 @@ impl CreditWindow {
pub(super) async fn take(
&self,
bytes: u64,
deadline: Option<Instant>,
deadline: &CallDeadline,
cancellation: &CallCancellation,
) -> Result<(), SessionError> {
loop {
@@ -101,7 +99,7 @@ impl CreditWindow {
tokio::select! {
biased;
() = cancellation.cancelled() => return Err(SessionError::Cancelled),
() = wait_deadline(deadline) => return Err(SessionError::Timeout),
() = deadline.wait() => return Err(SessionError::Timeout),
() = changed => {}
}
}
@@ -0,0 +1,106 @@
//! 宿主与插件共用的调用处理预算,受管续体等待期间暂停并保留剩余额度
use std::time::Duration;
use tokio::{sync::watch, time::Instant};
/// 同一次调用的所有等待点共享此期限,避免回调恢复后仍使用旧截止时间
#[derive(Clone)]
pub struct CallDeadline {
state: watch::Sender<State>,
}
#[derive(Clone, Copy)]
enum State {
Running(Instant),
Paused(Duration),
Unlimited,
}
/// 随受管续体释放而恢复计时;取消和错误出口同样归还暂停状态
pub struct CallDeadlinePause {
deadline: CallDeadline,
}
impl CallDeadline {
#[must_use]
pub fn new(deadline: Instant) -> Self {
let (state, _) = watch::channel(State::Running(deadline));
Self { state }
}
#[must_use]
pub fn remaining(&self) -> Option<Duration> {
match *self.state.borrow() {
State::Running(at) => Some(at.saturating_duration_since(Instant::now())),
State::Paused(remaining) => Some(remaining),
State::Unlimited => None,
}
}
#[must_use]
pub fn is_elapsed(&self) -> bool {
self.remaining()
.is_some_and(|remaining| remaining.is_zero())
}
/// 只有尚未耗尽的运行预算可以暂停;重复 next 不能重置或复活预算
#[must_use]
pub fn pause(&self) -> Option<CallDeadlinePause> {
self.state
.send_if_modified(|state| {
let State::Running(at) = *state else {
return false;
};
let remaining = at.saturating_duration_since(Instant::now());
if remaining.is_zero() {
return false;
}
*state = State::Paused(remaining);
true
})
.then(|| CallDeadlinePause {
deadline: self.clone(),
})
}
/// 初始资源流响应已交付,后续期限由资源 owner 持有
pub fn clear(&self) {
self.state.send_replace(State::Unlimited);
}
pub async fn wait(&self) {
let mut state = self.state.subscribe();
loop {
let current = *state.borrow_and_update();
let elapsed = async {
match current {
State::Running(at) => tokio::time::sleep_until(at).await,
State::Paused(_) | State::Unlimited => std::future::pending().await,
}
};
tokio::select! {
biased;
_ = state.changed() => {}
() = elapsed => {
// 暂停可与旧计时器唤醒并发,返回前重新核对当前预算
if self.is_elapsed() {
return;
}
}
}
}
}
}
impl Drop for CallDeadlinePause {
fn drop(&mut self) {
self.deadline.state.send_if_modified(|state| {
let State::Paused(remaining) = *state else {
return false;
};
*state = State::Running(Instant::now() + remaining);
true
});
}
}
@@ -16,9 +16,10 @@ use crate::{
use super::super::frame::{read_frame, validate_frame, write_frame};
use super::{
CallCancellation, HostClient, PluginCall, PluginHandler, SessionConfig, SessionError,
CallCancellation, CallDeadline, HostClient, PluginCall, PluginHandler, SessionConfig,
SessionError,
credit::CreditWindow,
lifecycle::{CancellationSource, LifecycleHooks, wait_deadline},
lifecycle::{CancellationSource, LifecycleHooks},
output::{Output, writer_loop},
registry::CallbackRegistry,
stream::validate_stream_chunk,
@@ -340,13 +341,12 @@ impl SessionDriver {
let deadline = Instant::now()
.checked_add(Duration::from_millis(context.timeout_ms))
.ok_or(SessionError::Protocol)?;
let deadline = Some(deadline);
let (deadline_source, callback_deadline) = tokio::sync::watch::channel(deadline);
let deadline = CallDeadline::new(deadline);
let credits = Arc::new(CreditWindow::default());
let (cancellation, cancellation_signal) = CancellationSource::new();
let host = HostClient {
parent_id: id,
deadline: callback_deadline,
deadline: deadline.clone(),
cancellation: cancellation_signal.clone(),
callbacks: Arc::clone(&self.callbacks),
output: self.output.clone(),
@@ -379,7 +379,6 @@ impl SessionDriver {
output,
lifecycle,
deadline,
deadline_source,
maximum_stream_chunk_bytes,
maximum_chunks,
)
@@ -453,8 +452,7 @@ async fn run_call(
credits: Arc<CreditWindow>,
output: Output,
lifecycle: LifecycleHooks,
deadline: Option<Instant>,
deadline_source: tokio::sync::watch::Sender<Option<Instant>>,
deadline: CallDeadline,
maximum_stream_chunk_bytes: usize,
maximum_chunks: usize,
) -> Result<(), SessionError> {
@@ -463,7 +461,7 @@ async fn run_call(
let response = tokio::select! {
biased;
() = cancellation.cancelled() => return Ok(()),
() = wait_deadline(deadline) => {
() = deadline.wait() => {
lifecycle.cancel(context.clone());
output.send_control(fault_frame(
id,
@@ -481,7 +479,7 @@ async fn run_call(
bounded_fault_frame(id, error),
&context,
&lifecycle,
deadline,
&deadline,
&cancellation,
)
.await?;
@@ -511,13 +509,13 @@ async fn run_call(
initial,
&context,
&lifecycle,
deadline,
&deadline,
&cancellation,
)
.await?;
return Ok(());
};
let window_bytes = match credits.initial_window(deadline, &cancellation).await {
let window_bytes = match credits.initial_window(&deadline, &cancellation).await {
Ok(window) => window,
Err(SessionError::Timeout) => {
lifecycle.cancel(context.clone());
@@ -541,25 +539,22 @@ async fn run_call(
initial,
&context,
&lifecycle,
deadline,
&deadline,
&cancellation,
)
.await?
{
return Ok(());
}
let deadline = if context.resource_stream {
None
} else {
deadline
};
deadline_source.send_replace(deadline);
if context.resource_stream {
deadline.clear();
}
let mut sequence = 0_u64;
loop {
let item = tokio::select! {
biased;
() = cancellation.cancelled() => return Ok(()),
() = wait_deadline(deadline) => {
() = deadline.wait() => {
lifecycle.cancel(context.clone());
output.send_control(end_frame(
id,
@@ -587,7 +582,7 @@ async fn run_call(
return Ok(());
}
match credits
.take(payload.len() as u64, deadline, &cancellation)
.take(payload.len() as u64, &deadline, &cancellation)
.await
{
Ok(()) => {}
@@ -615,7 +610,7 @@ async fn run_call(
id,
&context,
&lifecycle,
deadline,
&deadline,
&cancellation,
)
.await?
@@ -631,7 +626,7 @@ async fn send_initial(
frame: Frame,
context: &CallContext,
lifecycle: &LifecycleHooks,
deadline: Option<Instant>,
deadline: &CallDeadline,
cancellation: &CallCancellation,
) -> Result<bool, SessionError> {
match output.send_data_until(frame, deadline, cancellation).await {
@@ -655,7 +650,7 @@ async fn send_stream_data(
id: u64,
context: &CallContext,
lifecycle: &LifecycleHooks,
deadline: Option<Instant>,
deadline: &CallDeadline,
cancellation: &CallCancellation,
) -> Result<bool, SessionError> {
match output.send_data_until(frame, deadline, cancellation).await {
@@ -2,7 +2,7 @@
use std::sync::Arc;
use tokio::{sync::watch, time::Instant};
use tokio::sync::watch;
use crate::CallContext;
@@ -95,22 +95,3 @@ impl LifecycleHooks {
let _ = self.events.try_send(LifecycleEvent::Shutdown);
}
}
pub(super) async fn wait_deadline(deadline: Option<Instant>) {
match deadline {
Some(deadline) => tokio::time::sleep_until(deadline).await,
None => std::future::pending().await,
}
}
pub(super) async fn callback_deadline(mut deadline: tokio::sync::watch::Receiver<Option<Instant>>) {
loop {
let current = *deadline.borrow_and_update();
tokio::select! {
() = wait_deadline(current) => return,
changed = deadline.changed() => {
if changed.is_err() { return wait_deadline(current).await; }
}
}
}
}
@@ -3,6 +3,7 @@
mod call;
mod config;
mod credit;
mod deadline;
mod driver;
mod error;
mod lifecycle;
@@ -12,6 +13,7 @@ mod stream;
pub use call::{CallFuture, CallReply, PluginCall, PluginHandler};
pub use config::SessionConfig;
pub use deadline::{CallDeadline, CallDeadlinePause};
pub use driver::PluginSession;
pub use error::SessionError;
pub use lifecycle::CallCancellation;
@@ -5,13 +5,12 @@ use std::sync::Arc;
use tokio::{
io::AsyncWrite,
sync::{OwnedSemaphorePermit, Semaphore, mpsc},
time::Instant,
};
use crate::Frame;
use super::super::frame::{validate_frame, write_frame};
use super::{CallCancellation, SessionError, lifecycle::wait_deadline};
use super::{CallCancellation, CallDeadline, SessionError};
pub(super) struct Output {
pub(super) sender: mpsc::Sender<Outbound>,
@@ -46,13 +45,13 @@ impl Output {
pub(super) async fn send_data_until(
&self,
frame: Frame,
deadline: Option<Instant>,
deadline: &CallDeadline,
cancellation: &CallCancellation,
) -> Result<(), SessionError> {
tokio::select! {
biased;
() = cancellation.cancelled() => Err(SessionError::Cancelled),
() = wait_deadline(deadline) => Err(SessionError::Timeout),
() = deadline.wait() => Err(SessionError::Timeout),
result = self.send_data(frame) => result,
}
}
@@ -6,13 +6,13 @@ use std::{
};
use serde_json::Value;
use tokio::{sync::oneshot, time::Instant};
use tokio::sync::oneshot;
use crate::{Frame, Message};
use crate::{Frame, Message, call::middleware::NEXT_METHOD};
use super::super::frame::validate_frame;
use super::{
CallCancellation, SessionError, credit::CreditWindow, lifecycle::callback_deadline,
CallCancellation, CallDeadline, CallDeadlinePause, SessionError, credit::CreditWindow,
output::Output,
};
@@ -26,7 +26,7 @@ pub struct HostReply {
#[derive(Clone)]
pub struct HostClient {
pub(super) parent_id: u64,
pub(super) deadline: tokio::sync::watch::Receiver<Option<Instant>>,
pub(super) deadline: CallDeadline,
pub(super) cancellation: CallCancellation,
pub(super) callbacks: Arc<CallbackRegistry>,
pub(super) output: Output,
@@ -60,7 +60,13 @@ impl HostClient {
if method.is_empty() || method.len() > 128 {
return Err(SessionError::Protocol);
}
let ordered = self.callbacks.order.lock().await;
let ordered = tokio::select! {
biased;
() = self.cancellation.cancelled() => return Err(SessionError::Cancelled),
() = self.deadline.wait() => return Err(SessionError::Timeout),
ordered = self.callbacks.order.lock() => ordered,
};
let managed_next = method == NEXT_METHOD;
let (id, received) = self.callbacks.begin(self.parent_id)?;
let frame = Frame {
message: Message::Callback {
@@ -78,13 +84,16 @@ impl HostClient {
let sent = tokio::select! {
biased;
() = self.cancellation.cancelled() => Err(SessionError::Cancelled),
() = callback_deadline(self.deadline.clone()) => Err(SessionError::Timeout),
() = self.deadline.wait() => Err(SessionError::Timeout),
result = self.output.send_data(frame) => result,
};
if let Err(error) = sent {
self.callbacks.retire(id);
return Err(error);
}
if managed_next {
self.callbacks.pause_next(id, &self.deadline);
}
drop(ordered);
tokio::select! {
biased;
@@ -92,7 +101,7 @@ impl HostClient {
self.callbacks.retire(id);
Err(SessionError::Cancelled)
}
() = callback_deadline(self.deadline.clone()) => {
() = self.deadline.wait() => {
self.callbacks.retire(id);
Err(SessionError::Timeout)
}
@@ -117,6 +126,7 @@ struct CallbackState {
struct PendingCallback {
parent_id: u64,
pause: Option<CallDeadlinePause>,
result: oneshot::Sender<Result<HostReply, SessionError>>,
}
@@ -149,12 +159,27 @@ impl CallbackRegistry {
let id = state.next_id;
state.next_id = state.next_id.checked_add(2).ok_or(SessionError::Capacity)?;
let (result, received) = oneshot::channel();
state
.pending
.insert(id, PendingCallback { parent_id, result });
state.pending.insert(
id,
PendingCallback {
parent_id,
result,
pause: None,
},
);
Ok((id, received))
}
fn pause_next(&self, id: u64, deadline: &CallDeadline) {
let mut state = self
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if let Some(pending) = state.pending.get_mut(&id) {
pending.pause = deadline.pause();
}
}
pub(super) fn resolve(&self, frame: Frame) -> Result<(), SessionError> {
let id = match &frame.message {
Message::Result { id, .. } | Message::Error { id, .. } => *id,
@@ -174,7 +199,7 @@ impl CallbackRegistry {
return Err(SessionError::Protocol);
}
};
let Some(pending) = pending else {
let Some(mut pending) = pending else {
return Ok(());
};
let result = match frame.message {
@@ -187,6 +212,8 @@ impl CallbackRegistry {
}
_ => return Err(SessionError::Protocol),
};
// 收到宿主结果就恢复预算,不依赖插件再次轮询 next future
pending.pause.take();
let _ = pending.result.send(result);
Ok(())
}
@@ -41,6 +41,7 @@ pub struct CallContext {
pub generation: u64,
pub incarnation: String,
pub stage: Stage,
/// 插件处理预算;等待宿主 middleware next 时暂停,请求总期限由宿主单独执行
pub timeout_ms: u64,
/// 首个流式结果返回后由资源 owner 控制取消,不再沿用初始调用期限
pub resource_stream: bool,
@@ -0,0 +1,52 @@
//! 验证续体等待不计入预算,恢复、到期和资源流切换不会重置已消耗时间
use std::time::Duration;
use gateway_plugin_sdk::client::CallDeadline;
use tokio::time::{Instant, advance};
#[tokio::test(start_paused = true)]
async fn suspension_preserves_remaining_budget_and_wakes_existing_waiters() {
let deadline = CallDeadline::new(Instant::now() + Duration::from_secs(10));
let waiter = tokio::spawn({
let deadline = deadline.clone();
async move { deadline.wait().await }
});
tokio::task::yield_now().await;
advance(Duration::from_secs(3)).await;
let pause = deadline.pause().unwrap();
advance(Duration::from_secs(30)).await;
assert!(!waiter.is_finished());
assert_eq!(deadline.remaining(), Some(Duration::from_secs(7)));
drop(pause);
advance(Duration::from_secs(6)).await;
assert!(!waiter.is_finished());
advance(Duration::from_secs(1)).await;
waiter.await.unwrap();
assert!(deadline.is_elapsed());
}
#[tokio::test(start_paused = true)]
async fn exhausted_budget_cannot_be_suspended_or_revived() {
let deadline = CallDeadline::new(Instant::now() + Duration::from_secs(1));
advance(Duration::from_secs(1)).await;
assert!(deadline.pause().is_none());
deadline.wait().await;
assert!(deadline.is_elapsed());
}
#[tokio::test(start_paused = true)]
async fn resource_stream_clear_survives_an_outstanding_pause_guard() {
let deadline = CallDeadline::new(Instant::now() + Duration::from_secs(1));
let pause = deadline.pause().unwrap();
assert!(deadline.pause().is_none());
deadline.clear();
drop(pause);
advance(Duration::from_secs(10)).await;
assert_eq!(deadline.remaining(), None);
assert!(
tokio::time::timeout(Duration::from_secs(1), deadline.wait())
.await
.is_err()
);
}
@@ -110,3 +110,203 @@ async fn quiesce_rejects_new_calls_and_shutdown_closes_the_session() {
shutdown(&mut host, task).await;
}
async fn send_middleware_call(host: &mut HostPeer, id: u64, resource_stream: bool) {
let mut context = middleware_context(id);
context.resource_stream = resource_stream;
write_frame(
&mut host.writer,
&Frame {
message: Message::Call {
id,
method: HANDLE_METHOD.into(),
context,
params: serde_json::to_value(middleware_request()).unwrap(),
},
payload: Vec::new(),
},
)
.await
.unwrap();
send_control(
host,
Message::Credit {
id,
bytes: 1024,
frames: 2,
},
)
.await;
}
async fn receive_next(host: &mut HostPeer, parent: u64) -> u64 {
let Message::Callback {
id,
parent_id,
method,
..
} = receive(host).await.message
else {
panic!("expected middleware next")
};
assert_eq!(parent_id, parent);
assert_eq!(method, NEXT_METHOD);
id
}
async fn reply_next(host: &mut HostPeer, id: u64) {
send_control(
host,
Message::Result {
id,
result: serde_json::to_value(MiddlewareNextResponse {
metadata: None,
response: format!("response-{id}"),
protocol: "openai".into(),
status: 200,
headers: Vec::new(),
body: None,
})
.unwrap(),
},
)
.await;
}
#[tokio::test(start_paused = true)]
async fn managed_next_can_outlive_the_initial_budget_with_or_without_resource_stream() {
for resource_stream in [false, true] {
let (mut host, task) = start_session(middleware_plugin(false)).await;
send_middleware_call(&mut host, 1, resource_stream).await;
let next = receive_next(&mut host, 1).await;
tokio::time::advance(Duration::from_secs(10)).await;
reply_next(&mut host, next).await;
assert!(matches!(
receive(&mut host).await.message,
Message::Result { id: 1, .. }
));
assert!(matches!(
receive(&mut host).await.message,
Message::End { id: 1, error: None }
));
shutdown(&mut host, task).await;
}
}
#[tokio::test(start_paused = true)]
async fn response_processing_uses_only_the_budget_left_before_next() {
let plugin = MiddlewarePlugin::new(
&middleware_contributions(),
|call: RequestCall| async move {
tokio::time::sleep(Duration::from_millis(200)).await;
let response = call.next.run(call.request).await?;
tokio::time::sleep(Duration::from_millis(850)).await;
Ok(response)
},
)
.unwrap();
let (mut host, task) = start_session(plugin).await;
send_middleware_call(&mut host, 1, false).await;
let next = receive_next(&mut host, 1).await;
tokio::time::advance(Duration::from_secs(10)).await;
reply_next(&mut host, next).await;
assert!(matches!(
receive(&mut host).await.message,
Message::Error {
id: 1,
error: PluginFault {
code: ErrorCode::Timeout,
..
}
}
));
shutdown(&mut host, task).await;
}
#[tokio::test(start_paused = true)]
async fn waiting_next_does_not_suspend_an_independent_handlers_budget() {
let plugin = MiddlewarePlugin::new(
&middleware_contributions(),
|call: RequestCall| async move {
if call.context.call_id == 3 {
return pending().await;
}
call.next.run(call.request).await
},
)
.unwrap();
let (mut host, task) = start_session(plugin).await;
send_middleware_call(&mut host, 1, false).await;
let next = receive_next(&mut host, 1).await;
send_middleware_call(&mut host, 3, false).await;
// 先让第二个 Call 进入 SDK,再推进到其预算之外
tokio::task::yield_now().await;
tokio::time::advance(Duration::from_secs(2)).await;
assert!(matches!(
receive(&mut host).await.message,
Message::Error {
id: 3,
error: PluginFault {
code: ErrorCode::Timeout,
..
}
}
));
reply_next(&mut host, next).await;
assert!(matches!(
receive(&mut host).await.message,
Message::Result { id: 1, .. }
));
shutdown(&mut host, task).await;
}
#[tokio::test(start_paused = true)]
async fn cancelling_managed_next_retires_its_late_reply_and_keeps_session_available() {
let (mut host, task) = start_session(middleware_plugin(false)).await;
send_middleware_call(&mut host, 1, false).await;
let next = receive_next(&mut host, 1).await;
tokio::time::advance(Duration::from_secs(10)).await;
send_control(&mut host, Message::Cancel { id: 1 }).await;
assert!(matches!(
receive(&mut host).await.message,
Message::Cancelled { id: 1 }
));
reply_next(&mut host, next).await;
send_middleware_call(&mut host, 3, false).await;
let next = receive_next(&mut host, 3).await;
reply_next(&mut host, next).await;
assert!(matches!(
receive(&mut host).await.message,
Message::Result { id: 3, .. }
));
shutdown(&mut host, task).await;
}
#[tokio::test(start_paused = true)]
async fn ordinary_callbacks_still_consume_the_parent_budget() {
let (mut host, task) = start_session(TestHandler::default()).await;
send_call_with_timeout(
&mut host,
1,
"callback",
json!({}),
Vec::new(),
Duration::from_millis(100),
)
.await;
assert!(matches!(
receive(&mut host).await.message,
Message::Callback { parent_id: 1, .. }
));
assert!(matches!(
receive(&mut host).await.message,
Message::Error {
id: 1,
error: PluginFault {
code: ErrorCode::Timeout,
..
}
}
));
shutdown(&mut host, task).await;
}
@@ -1,6 +1,7 @@
//! 插件会话测试入口与双工传输、处理器和宿主回调夹具
mod credit;
mod deadline;
mod driver;
mod lifecycle;
mod registry;
@@ -66,11 +66,11 @@ pub struct TokenUsage {
/// 输出 token 数
pub output_tokens: u64,
/// 命中缓存的输入 token 数
pub cached_tokens: u64,
pub cached_input_tokens: u64,
/// 写入 prompt cache 的输入 token 数
pub cache_write_tokens: u64,
pub cache_write_input_tokens: u64,
/// 输出 token 中的 reasoning token 数
pub reasoning_tokens: u64,
pub reasoning_output_tokens: u64,
/// 图片工具输入 token 数
pub image_input_tokens: u64,
/// 图片工具输出 token 数
@@ -94,18 +94,19 @@ pub fn extract_usage(body: &Value) -> Option<TokenUsage> {
let output_tokens = number_field(usage, "output_tokens")
.or_else(|| number_field(usage, "completion_tokens"))
.unwrap_or_default();
let cached_tokens = nested_number_field(usage, &["input_tokens_details", "cached_tokens"])
.or_else(|| nested_number_field(usage, &["prompt_tokens_details", "cached_tokens"]))
.or_else(|| number_field(usage, "cached_tokens"))
.unwrap_or_default();
let cache_write_tokens =
let cached_input_tokens =
nested_number_field(usage, &["input_tokens_details", "cached_tokens"])
.or_else(|| nested_number_field(usage, &["prompt_tokens_details", "cached_tokens"]))
.or_else(|| number_field(usage, "cached_tokens"))
.unwrap_or_default();
let cache_write_input_tokens =
nested_number_field(usage, &["input_tokens_details", "cache_write_tokens"])
.or_else(|| {
nested_number_field(usage, &["prompt_tokens_details", "cache_write_tokens"])
})
.or_else(|| number_field(usage, "cache_write_tokens"))
.unwrap_or_default();
let reasoning_tokens =
let reasoning_output_tokens =
nested_number_field(usage, &["output_tokens_details", "reasoning_tokens"])
.or_else(|| {
nested_number_field(usage, &["completion_tokens_details", "reasoning_tokens"])
@@ -142,9 +143,9 @@ pub fn extract_usage(body: &Value) -> Option<TokenUsage> {
has_usage.then_some(TokenUsage {
input_tokens,
output_tokens,
cached_tokens,
cache_write_tokens,
reasoning_tokens,
cached_input_tokens,
cache_write_input_tokens,
reasoning_output_tokens,
image_input_tokens,
image_output_tokens,
total_tokens,
@@ -187,11 +188,11 @@ pub fn billable_usage_is_complete(response: &Value, usage: TokenUsage) -> bool {
input == Some(usage.input_tokens)
&& output == Some(usage.output_tokens)
&& cached == Some(usage.cached_tokens)
&& written == Some(usage.cache_write_tokens)
&& cached == Some(usage.cached_input_tokens)
&& written == Some(usage.cache_write_input_tokens)
&& usage
.cached_tokens
.checked_add(usage.cache_write_tokens)
.cached_input_tokens
.checked_add(usage.cache_write_input_tokens)
.is_some_and(|cached| cached <= usage.input_tokens)
&& total == Some(usage.total_tokens)
&& raw
@@ -165,9 +165,9 @@ fn extract_usage_should_read_codex_usage_shape() {
Some(TokenUsage {
input_tokens: 12,
output_tokens: 5,
cached_tokens: 3,
cache_write_tokens: 4,
reasoning_tokens: 0,
cached_input_tokens: 3,
cache_write_input_tokens: 4,
reasoning_output_tokens: 0,
image_input_tokens: 0,
image_output_tokens: 0,
total_tokens: 17,
@@ -191,9 +191,9 @@ fn extract_usage_should_read_openai_usage_shape() {
Some(TokenUsage {
input_tokens: 8,
output_tokens: 4,
cached_tokens: 2,
cache_write_tokens: 0,
reasoning_tokens: 3,
cached_input_tokens: 2,
cache_write_input_tokens: 0,
reasoning_output_tokens: 3,
image_input_tokens: 0,
image_output_tokens: 0,
total_tokens: 12,
@@ -217,9 +217,9 @@ fn extract_usage_should_preserve_explicit_total_and_reasoning_tokens() {
Some(TokenUsage {
input_tokens: 12,
output_tokens: 8,
cached_tokens: 0,
cache_write_tokens: 0,
reasoning_tokens: 6,
cached_input_tokens: 0,
cache_write_input_tokens: 0,
reasoning_output_tokens: 6,
image_input_tokens: 0,
image_output_tokens: 0,
total_tokens: 20,
@@ -248,9 +248,9 @@ fn extract_usage_should_keep_image_tool_usage_separate_from_model_total() {
Some(TokenUsage {
input_tokens: 12,
output_tokens: 5,
cached_tokens: 0,
cache_write_tokens: 0,
reasoning_tokens: 0,
cached_input_tokens: 0,
cache_write_input_tokens: 0,
reasoning_output_tokens: 0,
image_input_tokens: 31,
image_output_tokens: 9,
total_tokens: 17,
@@ -278,9 +278,9 @@ fn extract_sse_usage_should_prefer_completed_response_usage() {
Some(TokenUsage {
input_tokens: 3,
output_tokens: 5,
cached_tokens: 1,
cache_write_tokens: 0,
reasoning_tokens: 0,
cached_input_tokens: 1,
cache_write_input_tokens: 0,
reasoning_output_tokens: 0,
image_input_tokens: 0,
image_output_tokens: 0,
total_tokens: 8,
@@ -302,9 +302,9 @@ fn extract_sse_usage_should_fall_back_to_last_visible_usage() {
Some(TokenUsage {
input_tokens: 4,
output_tokens: 6,
cached_tokens: 0,
cache_write_tokens: 0,
reasoning_tokens: 0,
cached_input_tokens: 0,
cache_write_input_tokens: 0,
reasoning_output_tokens: 0,
image_input_tokens: 0,
image_output_tokens: 0,
total_tokens: 10,
@@ -324,9 +324,9 @@ fn extract_sse_usage_should_preserve_completed_image_tool_usage() {
Some(TokenUsage {
input_tokens: 12,
output_tokens: 5,
cached_tokens: 3,
cache_write_tokens: 0,
reasoning_tokens: 0,
cached_input_tokens: 3,
cache_write_input_tokens: 0,
reasoning_output_tokens: 0,
image_input_tokens: 31,
image_output_tokens: 9,
total_tokens: 17,
@@ -348,9 +348,9 @@ fn extract_sse_usage_should_ignore_non_json_control_events() {
Some(TokenUsage {
input_tokens: 2,
output_tokens: 3,
cached_tokens: 0,
cache_write_tokens: 0,
reasoning_tokens: 0,
cached_input_tokens: 0,
cache_write_input_tokens: 0,
reasoning_output_tokens: 0,
image_input_tokens: 0,
image_output_tokens: 0,
total_tokens: 5,

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