mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-07 16:05:27 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
ae5eaa46b4 | ||
|
|
65d5b32eba | ||
|
|
6c158447be | ||
|
|
81d939bd03 | ||
|
|
4a22479dff | ||
|
|
3ee454b750 | ||
|
|
b5bd63b585 | ||
|
|
286031a441 | ||
|
|
5358de53dd | ||
|
|
d814582c09 | ||
|
|
a9ce01f91f | ||
|
|
0fc25bfce9 | ||
|
|
eed8c93ae2 | ||
|
|
d865a00768 | ||
|
|
a0671fb70b | ||
|
|
fce543d4ce | ||
|
|
219888d5c2 | ||
|
|
5b80340135 | ||
|
|
b2dab66a80 | ||
|
|
f0cf8e7198 | ||
|
|
a9a86071bf | ||
|
|
27cef052c9 | ||
|
|
ef40088303 | ||
|
|
05ba116eef | ||
|
|
3ef4743e1a | ||
|
|
9e02e52fec | ||
|
|
42ff5b74d0 | ||
|
|
eb2eefd51e | ||
|
|
abd60ce60a | ||
|
|
2f96c02ea0 | ||
|
|
0f575e279a | ||
|
|
634b8fc674 | ||
|
|
00bfb1699d | ||
|
|
956ca70a14 | ||
|
|
24dd7b94f4 | ||
|
|
94c769fc9b |
Generated
+20
-19
@@ -1681,6 +1681,7 @@ dependencies = [
|
||||
"gateway-core",
|
||||
"heck",
|
||||
"hex",
|
||||
"hkdf",
|
||||
"prettyplease",
|
||||
"quote",
|
||||
"rand_core 0.6.4",
|
||||
@@ -1980,9 +1981,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 +2447,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 +2475,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",
|
||||
@@ -3401,9 +3402,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "quinn-proto"
|
||||
version = "0.11.16"
|
||||
version = "0.11.18"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2f4bfc015262b9df63c8845072ce59068853ff5872180c2ce2f13038b970e560"
|
||||
checksum = "a9746dbde176634f4f2f1faf2404e30a31b2bc1e9cafb5329c95d8177a18c9fc"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"getrandom 0.4.3",
|
||||
@@ -3587,9 +3588,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 +3605,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 +4505,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 +5007,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 +5547,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",
|
||||
]
|
||||
|
||||
@@ -21,6 +21,7 @@ futures = "0.3.32"
|
||||
gateway-core = { path = "../gateway-core" }
|
||||
heck = "0.5.0"
|
||||
hex = "0.4.3"
|
||||
hkdf = "0.13.0"
|
||||
rand_core = { version = "0.6.4", features = ["getrandom"] }
|
||||
serde = { version = "1.0.228", features = ["derive"] }
|
||||
serde_json = { version = "1.0.150", features = ["arbitrary_precision", "preserve_order"] }
|
||||
|
||||
@@ -491,11 +491,13 @@ impl BackupTask {
|
||||
)
|
||||
.await
|
||||
{
|
||||
let _ = store.delete_object(&config, &record.object_key).await;
|
||||
self.fail_task(
|
||||
self.fail_uploaded_task(
|
||||
record,
|
||||
error.code(),
|
||||
format!("上传到对象存储失败:{}", error.message()),
|
||||
&config,
|
||||
BackupError::new(
|
||||
error.code(),
|
||||
format!("上传到对象存储失败:{}", error.message()),
|
||||
),
|
||||
)
|
||||
.await?;
|
||||
return Ok(());
|
||||
@@ -533,16 +535,19 @@ impl BackupTask {
|
||||
sha256_matches,
|
||||
"备份远端对象校验不一致"
|
||||
);
|
||||
let _ = store.delete_object(&config, &record.object_key).await;
|
||||
self.fail_task(
|
||||
self.fail_uploaded_task(
|
||||
record,
|
||||
code::REMOTE_VERIFICATION_FAILED,
|
||||
match (size_matches, sha256_matches) {
|
||||
(false, true) => "远端对象大小不一致",
|
||||
(true, false) => "远端对象 sha256 校验值不一致",
|
||||
(false, false) => "远端对象大小和 sha256 校验值均不一致",
|
||||
(true, true) => "远端对象校验状态异常",
|
||||
},
|
||||
&config,
|
||||
BackupError::new(
|
||||
code::REMOTE_VERIFICATION_FAILED,
|
||||
match (size_matches, sha256_matches) {
|
||||
(false, true) => "远端对象大小不一致",
|
||||
(true, false) => "远端对象 sha256 校验值不一致",
|
||||
(false, false) => "远端对象大小和 sha256 校验值均不一致",
|
||||
(true, true) => "远端对象校验状态异常",
|
||||
}
|
||||
.to_owned(),
|
||||
),
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
@@ -556,11 +561,13 @@ impl BackupTask {
|
||||
.await?;
|
||||
}
|
||||
Err(error) => {
|
||||
let _ = store.delete_object(&config, &record.object_key).await;
|
||||
self.fail_task(
|
||||
self.fail_uploaded_task(
|
||||
record,
|
||||
error.code(),
|
||||
format!("远端对象校验失败:{}", error.message()),
|
||||
&config,
|
||||
BackupError::new(
|
||||
error.code(),
|
||||
format!("远端对象校验失败:{}", error.message()),
|
||||
),
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
@@ -568,6 +575,24 @@ impl BackupTask {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 清理失败并入原终态,后续清理沿用备份记录的到期策略
|
||||
async fn fail_uploaded_task(
|
||||
&self,
|
||||
record: &BackupRecord,
|
||||
config: &BackupStorageConfig,
|
||||
mut error: BackupError,
|
||||
) -> Result<(), WorkerTaskError> {
|
||||
if let Err(cleanup) = self
|
||||
.object_store
|
||||
.delete_object(config, &record.object_key)
|
||||
.await
|
||||
{
|
||||
error =
|
||||
error.with_cleanup_failure("DeleteObject", &cleanup, record.expires_at.is_some());
|
||||
}
|
||||
self.fail_task(record, error.code(), error.message()).await
|
||||
}
|
||||
|
||||
/// 把活跃任务写入失败终态并清理暂存
|
||||
async fn fail_task(
|
||||
&self,
|
||||
|
||||
@@ -405,9 +405,7 @@ async fn initialize_inner(
|
||||
let registry = providers;
|
||||
|
||||
let auth = Arc::new(DefaultAuthService::new(
|
||||
config.default_username,
|
||||
config.session_ttl_minutes,
|
||||
config.session_absolute_ttl_minutes,
|
||||
&config,
|
||||
client_config.session_ttl_minutes,
|
||||
store.auth(),
|
||||
client_key_verifier.clone(),
|
||||
|
||||
@@ -66,7 +66,7 @@ pub struct AccountGroupCapacity {
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AccountGroupUsage {
|
||||
pub today_usd: DecimalAmount,
|
||||
/// 当前 usage retention 窗口内、按请求发生时 routing group 快照归属的累计成本
|
||||
/// 当前 usage retention 窗口内、当前组成员实际完成请求的累计成本
|
||||
pub retained_total_usd: DecimalAmount,
|
||||
}
|
||||
|
||||
|
||||
@@ -76,18 +76,23 @@ impl std::fmt::Debug for LoginCommand {
|
||||
}
|
||||
}
|
||||
|
||||
/// 登录成功后返回的会话事实
|
||||
/// 登录或刷新后仅向当前浏览器交付的会话凭据
|
||||
#[derive(Clone, PartialEq, Eq)]
|
||||
pub struct LoginResult {
|
||||
pub struct SessionGrant {
|
||||
pub session_id: String,
|
||||
pub refresh_token: Option<String>,
|
||||
pub session: AuthSession,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for LoginResult {
|
||||
impl std::fmt::Debug for SessionGrant {
|
||||
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
formatter
|
||||
.debug_struct("LoginResult")
|
||||
.debug_struct("SessionGrant")
|
||||
.field("session_id", &"[REDACTED]")
|
||||
.field(
|
||||
"refresh_token",
|
||||
&self.refresh_token.as_ref().map(|_| "[REDACTED]"),
|
||||
)
|
||||
.field("session", &self.session)
|
||||
.finish()
|
||||
}
|
||||
@@ -138,12 +143,41 @@ impl std::fmt::Debug for ChangePassword {
|
||||
}
|
||||
}
|
||||
|
||||
/// 会话有效期与管理员续期上限;没有上限字段的旧会话保持固定有效期
|
||||
/// 访问有效期与逻辑会话;没有续期凭据的会话保持固定期限
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AuthSession {
|
||||
pub subject: SessionSubject,
|
||||
pub expires_at: DateTime<Utc>,
|
||||
pub absolute_expires_at: Option<DateTime<Utc>>,
|
||||
pub renewal: Option<SessionRenewal>,
|
||||
}
|
||||
|
||||
/// 长效登录的凭据摘要与续期窗口;访问到期不删除整段登录
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct SessionRenewal {
|
||||
pub expires_at: DateTime<Utc>,
|
||||
pub access_digest: [u8; 32],
|
||||
pub refresh_digest: [u8; 32],
|
||||
pub previous: Option<SessionRefreshReplay>,
|
||||
}
|
||||
|
||||
/// 仅保存上一代摘要及派生盐;响应丢失时须持有上一代凭据才能重建同一响应
|
||||
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct SessionRefreshReplay {
|
||||
pub digest: [u8; 32],
|
||||
pub salt: [u8; 32],
|
||||
pub expires_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
impl AuthSession {
|
||||
#[must_use]
|
||||
pub fn retained_until(&self) -> DateTime<Utc> {
|
||||
self.renewal
|
||||
.as_ref()
|
||||
.map_or(self.expires_at, |renewal| renewal.expires_at)
|
||||
}
|
||||
}
|
||||
|
||||
/// 安全审计事件类型
|
||||
|
||||
@@ -474,6 +474,26 @@ impl BackupError {
|
||||
pub fn message(&self) -> &str {
|
||||
&self.message
|
||||
}
|
||||
|
||||
/// 清理失败附加到原始终态,不覆盖上传失败或主动取消的分类
|
||||
#[must_use]
|
||||
pub fn with_cleanup_failure(
|
||||
mut self,
|
||||
stage: &'static str,
|
||||
error: &Self,
|
||||
retry_on_expiration: bool,
|
||||
) -> Self {
|
||||
let recovery = if retry_on_expiration {
|
||||
"等待备份到期清理"
|
||||
} else {
|
||||
"未安排自动重试"
|
||||
};
|
||||
self.message.push_str(&format!(
|
||||
";清理阶段 {stage} 失败({}),{recovery}",
|
||||
error.code()
|
||||
));
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
/// 构造稳定错误码常量
|
||||
|
||||
@@ -157,7 +157,6 @@ impl OtherRequestOutcome {
|
||||
/// 用量记录过滤条件
|
||||
#[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>,
|
||||
@@ -211,7 +210,6 @@ 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>,
|
||||
|
||||
@@ -142,7 +142,7 @@ pub fn select_forecast_sample(
|
||||
};
|
||||
history.push(current);
|
||||
// 重置边界不在查询起点时,从首个新观测重新计数,避免把旧周期累计带入总量
|
||||
let mut cycle_baseline = interrupted.then(|| &history[0]);
|
||||
let cycle_baseline = interrupted.then(|| &history[0]);
|
||||
let mut anchors: Vec<&QuotaForecastPoint> = Vec::new();
|
||||
let mut high_water: Option<&QuotaForecastPoint> = None;
|
||||
for point in &history {
|
||||
@@ -155,9 +155,9 @@ pub fn select_forecast_sample(
|
||||
if previous.used_percent - point.used_percent <= 1.0 {
|
||||
continue;
|
||||
}
|
||||
// 同窗口回落不能证明周期重置,只重建采样并保留已记录累计
|
||||
sample.discontinuous = true;
|
||||
anchors.clear();
|
||||
cycle_baseline = Some(point);
|
||||
}
|
||||
}
|
||||
high_water = Some(point);
|
||||
|
||||
@@ -51,7 +51,6 @@ pub struct RuntimeSettingsValues {
|
||||
pub refresh_margin_seconds: u64,
|
||||
pub refresh_concurrency: u32,
|
||||
pub max_concurrent_per_account: u32,
|
||||
pub request_interval_ms: u64,
|
||||
pub max_waiting_per_key: u32,
|
||||
pub max_waiting_per_account: u32,
|
||||
pub concurrency_wait_timeout_seconds: u32,
|
||||
@@ -212,10 +211,6 @@ impl RuntimeSettingsValues {
|
||||
"refresh_margin_seconds",
|
||||
),
|
||||
(self.refresh_concurrency > 0, "refresh_concurrency"),
|
||||
(
|
||||
i64::try_from(self.request_interval_ms).is_ok(),
|
||||
"request_interval_ms",
|
||||
),
|
||||
] {
|
||||
if !valid {
|
||||
return Err(field);
|
||||
|
||||
@@ -294,12 +294,12 @@ pub trait AuthStore: Send + Sync {
|
||||
|
||||
async fn store_session(&self, session_id: &str, session: &AuthSession) -> AdminStoreResult<()>;
|
||||
|
||||
/// 只延长仍存在且匹配的会话;并发续期返回当前值,已退出或过期时不重建
|
||||
async fn renew_session(
|
||||
/// 原子替换或撤销仍匹配的会话并返回当前值;过期或退出后不重新创建
|
||||
async fn compare_exchange_session(
|
||||
&self,
|
||||
session_id: &str,
|
||||
expected: &AuthSession,
|
||||
expires_at: chrono::DateTime<chrono::Utc>,
|
||||
replacement: Option<&AuthSession>,
|
||||
) -> AdminStoreResult<Option<AuthSession>>;
|
||||
|
||||
async fn delete_session(&self, session_id: &str) -> AdminStoreResult<Option<AuthSession>>;
|
||||
@@ -432,6 +432,9 @@ pub trait UsageBillingProjector: Send + Sync {
|
||||
/// 用量、趋势、诊断与运维错误的只读能力
|
||||
#[async_trait]
|
||||
pub trait ObservabilityStore: Send + Sync {
|
||||
/// 读取实现声明的总预算,缓存等待沿用同一配置,避免外层提前截断
|
||||
fn query_timeout(&self) -> std::time::Duration;
|
||||
|
||||
/// 缓存必须与删除和后台投影提交的版本一致
|
||||
async fn analytics_cache_generation(&self) -> AdminStoreResult<u64>;
|
||||
|
||||
|
||||
@@ -0,0 +1,139 @@
|
||||
//! 控制面认证端口、会话策略与访问和续期共用的身份校验
|
||||
|
||||
use std::{net::IpAddr, sync::Arc};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
|
||||
use chrono::{Duration, Utc};
|
||||
use gateway_core::engine::execution::ClientKeyVerifier;
|
||||
use sha2::{Digest as _, Sha256};
|
||||
|
||||
use super::map_store_error;
|
||||
use crate::{
|
||||
model::{
|
||||
AdminError,
|
||||
auth::{
|
||||
AuthSession, ChangePassword, LoginCommand, LoginError, SessionGrant, SessionSubject,
|
||||
},
|
||||
},
|
||||
ports::store::AuthStore,
|
||||
};
|
||||
|
||||
mod service;
|
||||
mod session;
|
||||
|
||||
/// 所有控制面接口消费同一个会话服务,权限由服务端身份决定
|
||||
#[async_trait]
|
||||
pub trait AuthService: Send + Sync {
|
||||
async fn change_password(
|
||||
&self,
|
||||
session_id: Option<&str>,
|
||||
command: ChangePassword,
|
||||
source_ip: IpAddr,
|
||||
) -> Result<(), AdminError>;
|
||||
async fn ensure_default_admin(&self, password: &str) -> Result<bool, AdminError>;
|
||||
async fn session(&self, session_id: Option<&str>) -> Result<Option<AuthSession>, AdminError>;
|
||||
async fn renew_session(
|
||||
&self,
|
||||
refresh_token: Option<&str>,
|
||||
access_token: Option<&str>,
|
||||
) -> Result<Option<SessionGrant>, AdminError>;
|
||||
async fn resolve_admin_user_id(
|
||||
&self,
|
||||
session_id: Option<&str>,
|
||||
) -> Result<Option<String>, AdminError>;
|
||||
async fn verify_admin_api_key(&self, key: &str) -> Result<bool, AdminError>;
|
||||
async fn login(
|
||||
&self,
|
||||
command: LoginCommand,
|
||||
source_ip: IpAddr,
|
||||
previous_session_id: Option<&str>,
|
||||
) -> Result<SessionGrant, LoginError>;
|
||||
async fn logout(&self, session_id: &str) -> Result<(), AdminError>;
|
||||
}
|
||||
|
||||
const MAX_SESSION_TTL_MINUTES: i64 = 366 * 24 * 60;
|
||||
|
||||
pub(crate) struct DefaultAuthService {
|
||||
default_admin_user_id: String,
|
||||
admin_access_ttl: Duration,
|
||||
admin_session_ttl: Duration,
|
||||
admin_session_absolute_ttl: Duration,
|
||||
key_session_ttl: Duration,
|
||||
store: Arc<dyn AuthStore>,
|
||||
verifier: Arc<dyn ClientKeyVerifier>,
|
||||
}
|
||||
|
||||
impl DefaultAuthService {
|
||||
#[must_use]
|
||||
pub(crate) fn new(
|
||||
config: &crate::AdminConfig,
|
||||
key_session_ttl_minutes: u64,
|
||||
store: Arc<dyn AuthStore>,
|
||||
verifier: Arc<dyn ClientKeyVerifier>,
|
||||
) -> Self {
|
||||
let admin_session_ttl = session_ttl(config.session_ttl_minutes);
|
||||
Self {
|
||||
default_admin_user_id: config.default_username.clone(),
|
||||
// 访问凭据最多有效 24 小时,较短的不活动期限也要留出续期窗口
|
||||
admin_access_ttl: Duration::hours(24).min(admin_session_ttl / 4),
|
||||
admin_session_ttl,
|
||||
admin_session_absolute_ttl: session_ttl(config.session_absolute_ttl_minutes),
|
||||
key_session_ttl: session_ttl(key_session_ttl_minutes),
|
||||
store,
|
||||
verifier,
|
||||
}
|
||||
}
|
||||
|
||||
async fn load_live_session(&self, id: &str) -> Result<Option<AuthSession>, AdminError> {
|
||||
let Some(session) = self
|
||||
.store
|
||||
.load_session(id)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "authentication session"))?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
if session.retained_until() <= Utc::now()
|
||||
|| session
|
||||
.absolute_expires_at
|
||||
.is_some_and(|expiry| expiry <= Utc::now())
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
let valid = match &session.subject {
|
||||
SessionSubject::Admin {
|
||||
admin_user_id,
|
||||
credential_fingerprint,
|
||||
} => self
|
||||
.store
|
||||
.load_password_hash(admin_user_id)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "administrator session"))?
|
||||
.is_some_and(|hash| password_fingerprint(&hash) == *credential_fingerprint),
|
||||
SessionSubject::Key { client_key_id } => self
|
||||
.store
|
||||
.client_key_enabled(client_key_id)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "session key"))?,
|
||||
};
|
||||
if !valid {
|
||||
let _ = self.store.delete_session(id).await;
|
||||
return Ok(None);
|
||||
}
|
||||
Ok(Some(session))
|
||||
}
|
||||
}
|
||||
|
||||
fn session_ttl(minutes: u64) -> Duration {
|
||||
Duration::minutes(
|
||||
i64::try_from(minutes)
|
||||
.unwrap_or(MAX_SESSION_TTL_MINUTES)
|
||||
.clamp(1, MAX_SESSION_TTL_MINUTES),
|
||||
)
|
||||
}
|
||||
|
||||
// 指纹只绑定已加盐的密码哈希,不把密码或原始哈希复制到 Redis 会话
|
||||
fn password_fingerprint(password_hash: &str) -> String {
|
||||
URL_SAFE_NO_PAD.encode(Sha256::digest(password_hash.as_bytes()))
|
||||
}
|
||||
+80
-156
@@ -1,93 +1,35 @@
|
||||
//! 控制面登录、会话恢复与身份校验的唯一 owner
|
||||
//! 管理员和 Key 的登录、退出、密码策略与会话访问
|
||||
|
||||
use std::{net::IpAddr, sync::Arc, time::Duration as StdDuration};
|
||||
use std::{net::IpAddr, time::Duration as StdDuration};
|
||||
|
||||
use argon2::{Argon2, PasswordHash, PasswordHasher, PasswordVerifier};
|
||||
use async_trait::async_trait;
|
||||
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
|
||||
use chrono::{Duration, Utc};
|
||||
use gateway_core::engine::execution::ClientKeyVerifier;
|
||||
use chrono::Utc;
|
||||
use rand_core::{OsRng, RngCore as _};
|
||||
use sha2::{Digest, Sha256};
|
||||
use subtle::ConstantTimeEq as _;
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::{
|
||||
AuthService, DefaultAuthService, password_fingerprint,
|
||||
session::{self, Credential, storage_id},
|
||||
};
|
||||
use crate::{
|
||||
model::{
|
||||
AdminError, AdminErrorKind,
|
||||
auth::{
|
||||
AdminAuditEvent, AuditActorKind, AuthSession, ChangePassword, LoginCommand, LoginError,
|
||||
LoginResult, SessionSubject,
|
||||
SessionGrant, SessionSubject,
|
||||
},
|
||||
},
|
||||
ports::store::AuthStore,
|
||||
use_case::map_store_error,
|
||||
};
|
||||
|
||||
use super::map_store_error;
|
||||
|
||||
/// 所有控制面接口消费同一个会话服务,权限由服务端身份决定
|
||||
#[async_trait]
|
||||
pub trait AuthService: Send + Sync {
|
||||
async fn change_password(
|
||||
&self,
|
||||
session_id: Option<&str>,
|
||||
command: ChangePassword,
|
||||
source_ip: IpAddr,
|
||||
) -> Result<(), AdminError>;
|
||||
async fn ensure_default_admin(&self, password: &str) -> Result<bool, AdminError>;
|
||||
async fn session(&self, session_id: Option<&str>) -> Result<Option<AuthSession>, AdminError>;
|
||||
async fn renew_session(
|
||||
&self,
|
||||
session_id: Option<&str>,
|
||||
) -> Result<Option<AuthSession>, AdminError>;
|
||||
async fn resolve_admin_user_id(
|
||||
&self,
|
||||
session_id: Option<&str>,
|
||||
) -> Result<Option<String>, AdminError>;
|
||||
async fn verify_admin_api_key(&self, key: &str) -> Result<bool, AdminError>;
|
||||
async fn login(
|
||||
&self,
|
||||
command: LoginCommand,
|
||||
source_ip: IpAddr,
|
||||
previous_session_id: Option<&str>,
|
||||
) -> Result<LoginResult, LoginError>;
|
||||
async fn logout(&self, session_id: &str) -> Result<(), AdminError>;
|
||||
}
|
||||
|
||||
const MAX_SESSION_TTL_MINUTES: i64 = 366 * 24 * 60;
|
||||
const LOGIN_WINDOW: StdDuration = StdDuration::from_secs(60);
|
||||
const LOGIN_ATTEMPTS_PER_SOURCE: u32 = 10;
|
||||
const LOGIN_ATTEMPTS_GLOBAL: u32 = 200;
|
||||
|
||||
pub(crate) struct DefaultAuthService {
|
||||
default_admin_user_id: String,
|
||||
admin_session_ttl: Duration,
|
||||
admin_session_absolute_ttl: Duration,
|
||||
key_session_ttl: Duration,
|
||||
store: Arc<dyn AuthStore>,
|
||||
verifier: Arc<dyn ClientKeyVerifier>,
|
||||
}
|
||||
|
||||
impl DefaultAuthService {
|
||||
#[must_use]
|
||||
pub(crate) fn new(
|
||||
default_admin_user_id: impl Into<String>,
|
||||
admin_session_ttl_minutes: u64,
|
||||
admin_session_absolute_ttl_minutes: u64,
|
||||
key_session_ttl_minutes: u64,
|
||||
store: Arc<dyn AuthStore>,
|
||||
verifier: Arc<dyn ClientKeyVerifier>,
|
||||
) -> Self {
|
||||
Self {
|
||||
default_admin_user_id: default_admin_user_id.into(),
|
||||
admin_session_ttl: session_ttl(admin_session_ttl_minutes),
|
||||
admin_session_absolute_ttl: session_ttl(admin_session_absolute_ttl_minutes),
|
||||
key_session_ttl: session_ttl(key_session_ttl_minutes),
|
||||
store,
|
||||
verifier,
|
||||
}
|
||||
}
|
||||
|
||||
fn auth_audit(&self, action: &str, admin_user_id: &str) -> AdminAuditEvent {
|
||||
AdminAuditEvent {
|
||||
id: format!("audit_{}", Uuid::now_v7().simple()),
|
||||
@@ -228,50 +170,23 @@ impl AuthService for DefaultAuthService {
|
||||
}
|
||||
|
||||
async fn session(&self, session_id: Option<&str>) -> Result<Option<AuthSession>, AdminError> {
|
||||
let Some(session_id) = session_id.filter(|value| !value.is_empty()) else {
|
||||
let Some(raw) = session_id.filter(|value| !value.is_empty()) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(session) = self
|
||||
.store
|
||||
.load_session(session_id)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "authentication session"))?
|
||||
else {
|
||||
let Some(session) = self.load_live_session(storage_id(raw)).await? else {
|
||||
return Ok(None);
|
||||
};
|
||||
if session.expires_at <= Utc::now()
|
||||
|| session
|
||||
.absolute_expires_at
|
||||
.is_some_and(|expiry| expiry <= Utc::now())
|
||||
{
|
||||
// Redis 的 TTL 负责过期清理,避免旧读取删除另一个请求刚续期的会话
|
||||
if session.expires_at <= Utc::now() {
|
||||
return Ok(None);
|
||||
}
|
||||
if let SessionSubject::Admin {
|
||||
admin_user_id,
|
||||
credential_fingerprint,
|
||||
} = &session.subject
|
||||
{
|
||||
let password_hash = self
|
||||
.store
|
||||
.load_password_hash(admin_user_id)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "administrator session"))?;
|
||||
if password_hash
|
||||
.is_none_or(|hash| password_fingerprint(&hash) != *credential_fingerprint)
|
||||
{
|
||||
let _ = self.store.delete_session(session_id).await;
|
||||
if let Some(renewal) = &session.renewal {
|
||||
let Some(credential) = Credential::parse(raw).filter(|value| !value.refresh) else {
|
||||
return Ok(None);
|
||||
};
|
||||
if !credential.owns(renewal) {
|
||||
return Ok(None);
|
||||
}
|
||||
}
|
||||
if let SessionSubject::Key { client_key_id } = &session.subject
|
||||
&& !self
|
||||
.store
|
||||
.client_key_enabled(client_key_id)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "session key"))?
|
||||
{
|
||||
let _ = self.store.delete_session(session_id).await;
|
||||
} else if Credential::parse(raw).is_some() {
|
||||
return Ok(None);
|
||||
}
|
||||
Ok(Some(session))
|
||||
@@ -279,33 +194,25 @@ impl AuthService for DefaultAuthService {
|
||||
|
||||
async fn renew_session(
|
||||
&self,
|
||||
session_id: Option<&str>,
|
||||
) -> Result<Option<AuthSession>, AdminError> {
|
||||
let Some(session_id) = session_id else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(session) = self.session(Some(session_id)).await? else {
|
||||
return Ok(None);
|
||||
};
|
||||
// Key 与升级前的会话保留原有期限,不能凭续期获得新的最长登录窗口
|
||||
let Some(absolute) = session
|
||||
.absolute_expires_at
|
||||
.filter(|_| matches!(session.subject, SessionSubject::Admin { .. }))
|
||||
else {
|
||||
return Ok(Some(session));
|
||||
};
|
||||
let now = Utc::now();
|
||||
let expires_at = (now + self.admin_session_ttl).min(absolute);
|
||||
let interval = (self.admin_session_ttl / 4).min(Duration::minutes(1));
|
||||
if expires_at <= session.expires_at
|
||||
|| session.expires_at - now > self.admin_session_ttl - interval
|
||||
{
|
||||
return Ok(Some(session));
|
||||
refresh_token: Option<&str>,
|
||||
access_token: Option<&str>,
|
||||
) -> Result<Option<SessionGrant>, AdminError> {
|
||||
if let Some(refresh) = refresh_token {
|
||||
return self.refresh(refresh).await;
|
||||
}
|
||||
self.store
|
||||
.renew_session(session_id, &session, expires_at)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "authentication session renewal"))
|
||||
// Key 与升级前的单凭据仅返回仍有效的原会话,不授予新的续期能力
|
||||
let Some(session) = self
|
||||
.session(access_token)
|
||||
.await?
|
||||
.filter(|session| session.renewal.is_none())
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
Ok(Some(SessionGrant {
|
||||
session_id: access_token.unwrap_or_default().to_owned(),
|
||||
refresh_token: None,
|
||||
session,
|
||||
}))
|
||||
}
|
||||
|
||||
async fn resolve_admin_user_id(
|
||||
@@ -346,7 +253,7 @@ impl AuthService for DefaultAuthService {
|
||||
command: LoginCommand,
|
||||
source_ip: IpAddr,
|
||||
previous_session_id: Option<&str>,
|
||||
) -> Result<LoginResult, LoginError> {
|
||||
) -> Result<SessionGrant, LoginError> {
|
||||
if let Some(retry_after) = self
|
||||
.store
|
||||
.consume_login_attempt(
|
||||
@@ -370,14 +277,31 @@ impl AuthService for DefaultAuthService {
|
||||
let now = Utc::now();
|
||||
let absolute_expires_at = matches!(subject, SessionSubject::Admin { .. })
|
||||
.then_some(now + self.admin_session_absolute_ttl);
|
||||
let session = AuthSession {
|
||||
subject,
|
||||
expires_at: (now + ttl).min(absolute_expires_at.unwrap_or(now + ttl)),
|
||||
absolute_expires_at,
|
||||
let refresh_expiry = (now + ttl).min(absolute_expires_at.unwrap_or(now + ttl));
|
||||
let grant = if let Some(absolute) = absolute_expires_at {
|
||||
session::issue(
|
||||
subject,
|
||||
(now + self.admin_access_ttl).min(refresh_expiry),
|
||||
refresh_expiry,
|
||||
absolute,
|
||||
)
|
||||
.map_err(|_| LoginError::Unavailable)?
|
||||
} else {
|
||||
SessionGrant {
|
||||
session_id: random_session_token(),
|
||||
refresh_token: None,
|
||||
session: AuthSession {
|
||||
subject,
|
||||
expires_at: refresh_expiry,
|
||||
absolute_expires_at: None,
|
||||
renewal: None,
|
||||
},
|
||||
}
|
||||
};
|
||||
let session_id = random_session_token();
|
||||
let session_id = storage_id(&grant.session_id);
|
||||
let session = &grant.session;
|
||||
self.store
|
||||
.store_session(&session_id, &session)
|
||||
.store_session(session_id, session)
|
||||
.await
|
||||
.map_err(|_| LoginError::Unavailable)?;
|
||||
if let SessionSubject::Admin { admin_user_id, .. } = &session.subject
|
||||
@@ -387,26 +311,39 @@ impl AuthService for DefaultAuthService {
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
let _ = self.store.delete_session(&session_id).await;
|
||||
let _ = self.store.delete_session(session_id).await;
|
||||
return Err(LoginError::Unavailable);
|
||||
}
|
||||
// 新身份验证成功后才撤销旧会话;撤销失败时不向浏览器提交新会话
|
||||
if let Some(previous) = previous_session_id.filter(|value| !value.is_empty())
|
||||
&& self.logout(previous).await.is_err()
|
||||
{
|
||||
let _ = self.store.delete_session(&session_id).await;
|
||||
let _ = self.store.delete_session(session_id).await;
|
||||
return Err(LoginError::Unavailable);
|
||||
}
|
||||
Ok(LoginResult {
|
||||
session_id,
|
||||
session,
|
||||
})
|
||||
Ok(grant)
|
||||
}
|
||||
|
||||
async fn logout(&self, session_id: &str) -> Result<(), AdminError> {
|
||||
let id = storage_id(session_id);
|
||||
let Some(stored) = self
|
||||
.store
|
||||
.load_session(id)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "authentication session"))?
|
||||
else {
|
||||
return Ok(());
|
||||
};
|
||||
if let Some(renewal) = &stored.renewal {
|
||||
if !Credential::parse(session_id).is_some_and(|credential| credential.owns(renewal)) {
|
||||
return Ok(());
|
||||
}
|
||||
} else if Credential::parse(session_id).is_some() {
|
||||
return Ok(());
|
||||
}
|
||||
let session = self
|
||||
.store
|
||||
.delete_session(session_id)
|
||||
.delete_session(id)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "authentication session"))?;
|
||||
if let Some(AuthSession {
|
||||
@@ -423,14 +360,6 @@ impl AuthService for DefaultAuthService {
|
||||
}
|
||||
}
|
||||
|
||||
fn session_ttl(minutes: u64) -> Duration {
|
||||
Duration::minutes(
|
||||
i64::try_from(minutes)
|
||||
.unwrap_or(MAX_SESSION_TTL_MINUTES)
|
||||
.clamp(1, MAX_SESSION_TTL_MINUTES),
|
||||
)
|
||||
}
|
||||
|
||||
fn validate_new_password(password: &str) -> Result<(), AdminError> {
|
||||
if password.trim().chars().count() < 12
|
||||
|| password.len() > 1024
|
||||
@@ -470,8 +399,3 @@ fn valid_admin_api_key_shape(value: &str) -> bool {
|
||||
&& value.starts_with("admin-")
|
||||
&& value[6..].bytes().all(|byte| byte.is_ascii_hexdigit())
|
||||
}
|
||||
|
||||
// 指纹只绑定已加盐的密码哈希,不把密码或原始哈希复制到 Redis 会话
|
||||
fn password_fingerprint(password_hash: &str) -> String {
|
||||
URL_SAFE_NO_PAD.encode(Sha256::digest(password_hash.as_bytes()))
|
||||
}
|
||||
@@ -0,0 +1,214 @@
|
||||
//! 双凭据会话的签发、摘要校验与短窗口重试,不持久化可直接使用的凭据
|
||||
|
||||
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
|
||||
use chrono::{DateTime, Duration, Utc};
|
||||
use rand_core::{OsRng, RngCore as _};
|
||||
use sha2::{Digest as _, Sha256};
|
||||
use subtle::ConstantTimeEq as _;
|
||||
|
||||
use super::DefaultAuthService;
|
||||
use crate::{
|
||||
model::{
|
||||
AdminError, AdminErrorKind,
|
||||
auth::{AuthSession, SessionGrant, SessionRefreshReplay, SessionRenewal, SessionSubject},
|
||||
},
|
||||
use_case::map_store_error,
|
||||
};
|
||||
|
||||
const REPLAY_WINDOW: Duration = Duration::seconds(30);
|
||||
|
||||
pub(super) struct Credential<'a> {
|
||||
pub id: &'a str,
|
||||
pub secret: &'a str,
|
||||
pub refresh: bool,
|
||||
}
|
||||
|
||||
impl<'a> Credential<'a> {
|
||||
pub fn parse(value: &'a str) -> Option<Self> {
|
||||
let mut parts = value.split('.');
|
||||
let refresh = match parts.next()? {
|
||||
"a2" => false,
|
||||
"r2" => true,
|
||||
_ => return None,
|
||||
};
|
||||
let id = parts.next()?;
|
||||
let secret = parts.next()?;
|
||||
if parts.next().is_some() || !valid_secret(id) || !valid_secret(secret) {
|
||||
return None;
|
||||
}
|
||||
Some(Self {
|
||||
id,
|
||||
secret,
|
||||
refresh,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn digest(&self) -> [u8; 32] {
|
||||
digest(self.secret)
|
||||
}
|
||||
|
||||
pub fn owns(&self, renewal: &SessionRenewal) -> bool {
|
||||
if self.refresh {
|
||||
matches_digest(&renewal.refresh_digest, &self.digest())
|
||||
|| renewal
|
||||
.previous
|
||||
.as_ref()
|
||||
.is_some_and(|previous| matches_digest(&previous.digest, &self.digest()))
|
||||
} else {
|
||||
matches_digest(&renewal.access_digest, &self.digest())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn valid_secret(value: &str) -> bool {
|
||||
value.len() == 43
|
||||
&& URL_SAFE_NO_PAD
|
||||
.decode(value)
|
||||
.is_ok_and(|bytes| bytes.len() == 32)
|
||||
}
|
||||
|
||||
pub(super) fn storage_id(value: &str) -> &str {
|
||||
Credential::parse(value).map_or(value, |credential| credential.id)
|
||||
}
|
||||
|
||||
fn matches_digest(left: &[u8; 32], right: &[u8; 32]) -> bool {
|
||||
bool::from(left.ct_eq(right))
|
||||
}
|
||||
|
||||
fn digest(secret: &str) -> [u8; 32] {
|
||||
Sha256::digest(secret.as_bytes()).into()
|
||||
}
|
||||
|
||||
fn random_bytes() -> [u8; 32] {
|
||||
let mut bytes = [0; 32];
|
||||
OsRng.fill_bytes(&mut bytes);
|
||||
bytes
|
||||
}
|
||||
|
||||
// 使用标准 HKDF 的独立用途标签;只有盐和摘要留在 Store,重试仍须持有原始凭据
|
||||
fn derive(secret: &str, salt: Option<&[u8]>, purpose: &[u8]) -> Result<String, AdminError> {
|
||||
let mut output = [0; 32];
|
||||
hkdf::Hkdf::<Sha256>::new(salt, secret.as_bytes())
|
||||
.expand(purpose, &mut output)
|
||||
.map_err(|_| AdminError::internal("会话凭据派生失败"))?;
|
||||
Ok(URL_SAFE_NO_PAD.encode(output))
|
||||
}
|
||||
|
||||
fn access_secret(refresh: &str) -> Result<String, AdminError> {
|
||||
derive(refresh, None, b"cpr-session-v2/access")
|
||||
}
|
||||
|
||||
pub(super) fn issue(
|
||||
subject: SessionSubject,
|
||||
access_expiry: DateTime<Utc>,
|
||||
refresh_expiry: DateTime<Utc>,
|
||||
absolute: DateTime<Utc>,
|
||||
) -> Result<SessionGrant, AdminError> {
|
||||
let id = URL_SAFE_NO_PAD.encode(random_bytes());
|
||||
let secret = URL_SAFE_NO_PAD.encode(random_bytes());
|
||||
let session = AuthSession {
|
||||
subject,
|
||||
expires_at: access_expiry,
|
||||
absolute_expires_at: Some(absolute),
|
||||
renewal: Some(SessionRenewal {
|
||||
expires_at: refresh_expiry,
|
||||
access_digest: digest(&access_secret(&secret)?),
|
||||
refresh_digest: digest(&secret),
|
||||
previous: None,
|
||||
}),
|
||||
};
|
||||
grant(&id, &secret, session)
|
||||
}
|
||||
|
||||
fn grant(id: &str, refresh: &str, session: AuthSession) -> Result<SessionGrant, AdminError> {
|
||||
Ok(SessionGrant {
|
||||
session_id: format!("a2.{id}.{}", access_secret(refresh)?),
|
||||
refresh_token: Some(format!("r2.{id}.{refresh}")),
|
||||
session,
|
||||
})
|
||||
}
|
||||
|
||||
impl DefaultAuthService {
|
||||
pub(super) async fn refresh(&self, raw: &str) -> Result<Option<SessionGrant>, AdminError> {
|
||||
let Some(credential) = Credential::parse(raw).filter(|value| value.refresh) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let mut current = self.load_live_session(credential.id).await?;
|
||||
// CAS 冲突只重新解释同一请求的凭据,不对并发退出进行补建
|
||||
for _ in 0..3 {
|
||||
let Some(session) = current else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(renewal) = &session.renewal else {
|
||||
return Ok(None);
|
||||
};
|
||||
let now = Utc::now();
|
||||
if session.retained_until() <= now
|
||||
|| session.absolute_expires_at.is_none_or(|limit| limit <= now)
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
if matches_digest(&renewal.refresh_digest, &credential.digest()) {
|
||||
// 其他标签页已经刷新或客户端重试时复用当前代,避免 Cookie 响应相互覆盖
|
||||
if session.expires_at > now {
|
||||
return grant(credential.id, credential.secret, session).map(Some);
|
||||
}
|
||||
let absolute = session
|
||||
.absolute_expires_at
|
||||
.ok_or_else(|| AdminError::internal("缺少会话最长有效期"))?;
|
||||
let salt = random_bytes();
|
||||
let next_secret =
|
||||
derive(credential.secret, Some(&salt), b"cpr-session-v2/refresh")?;
|
||||
let refresh_expiry = (now + self.admin_session_ttl).min(absolute);
|
||||
let access_expiry = (now + self.admin_access_ttl).min(refresh_expiry);
|
||||
let replacement = AuthSession {
|
||||
expires_at: access_expiry,
|
||||
renewal: Some(SessionRenewal {
|
||||
expires_at: refresh_expiry,
|
||||
access_digest: digest(&access_secret(&next_secret)?),
|
||||
refresh_digest: digest(&next_secret),
|
||||
previous: Some(SessionRefreshReplay {
|
||||
digest: credential.digest(),
|
||||
salt,
|
||||
expires_at: (now + REPLAY_WINDOW).min(access_expiry),
|
||||
}),
|
||||
}),
|
||||
..session.clone()
|
||||
};
|
||||
current = self
|
||||
.store
|
||||
.compare_exchange_session(credential.id, &session, Some(&replacement))
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "authentication session rotation"))?;
|
||||
if current.as_ref() == Some(&replacement) {
|
||||
return grant(credential.id, &next_secret, replacement).map(Some);
|
||||
}
|
||||
} else if let Some(previous) = renewal
|
||||
.previous
|
||||
.as_ref()
|
||||
.filter(|previous| matches_digest(&previous.digest, &credential.digest()))
|
||||
{
|
||||
if previous.expires_at > now {
|
||||
let secret = derive(
|
||||
credential.secret,
|
||||
Some(&previous.salt),
|
||||
b"cpr-session-v2/refresh",
|
||||
)?;
|
||||
return grant(credential.id, &secret, session).map(Some);
|
||||
}
|
||||
// 已知上一代在重试窗口外再次使用,原子撤销整个会话
|
||||
current = self
|
||||
.store
|
||||
.compare_exchange_session(credential.id, &session, None)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "authentication session revocation"))?;
|
||||
} else {
|
||||
return Ok(None);
|
||||
}
|
||||
}
|
||||
Err(AdminError::new(
|
||||
AdminErrorKind::Unavailable,
|
||||
"会话正在变更,请重试",
|
||||
))
|
||||
}
|
||||
}
|
||||
@@ -37,7 +37,7 @@ 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 HEALTH_TIMELINE_STABLE_RELIABILITY: f64 = 95.0;
|
||||
mod store_cache;
|
||||
pub(crate) use store_cache::CachedObservabilityStore;
|
||||
|
||||
@@ -379,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,
|
||||
@@ -398,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(),
|
||||
@@ -956,7 +959,7 @@ fn trend(
|
||||
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,
|
||||
@@ -1023,7 +1026,10 @@ fn trend_point(point: RequestMetricPoint) -> TrendPoint {
|
||||
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,
|
||||
@@ -1123,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) {
|
||||
@@ -1155,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 {
|
||||
@@ -1178,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)
|
||||
}
|
||||
|
||||
@@ -12,7 +12,6 @@ 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> {
|
||||
@@ -72,7 +71,7 @@ impl<K: Eq + Hash, T: AnalyticsValue> Cache<K, T> {
|
||||
{
|
||||
// 刷新失败后的同键等待也计入预算,避免串行重试把延迟放大为队列长度倍数
|
||||
tokio::time::timeout(
|
||||
REQUEST_TIMEOUT,
|
||||
store.query_timeout(),
|
||||
self.get_within_budget(key, store, ttl, load),
|
||||
)
|
||||
.await
|
||||
@@ -174,12 +173,6 @@ fn with_age<T: AnalyticsValue>(mut value: T, stale: bool) -> T {
|
||||
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>,
|
||||
@@ -202,6 +195,10 @@ impl CachedObservabilityStore {
|
||||
|
||||
#[async_trait]
|
||||
impl ObservabilityStore for CachedObservabilityStore {
|
||||
fn query_timeout(&self) -> Duration {
|
||||
self.inner.query_timeout()
|
||||
}
|
||||
|
||||
async fn analytics_cache_generation(&self) -> AdminStoreResult<u64> {
|
||||
self.inner.analytics_cache_generation().await
|
||||
}
|
||||
@@ -241,7 +238,6 @@ impl ObservabilityStore for CachedObservabilityStore {
|
||||
filter: UsageFilter,
|
||||
granularity: Granularity,
|
||||
) -> AdminStoreResult<UsageInsightsObservation> {
|
||||
let filter = normalized(filter);
|
||||
self.insights
|
||||
.get(
|
||||
((range, filter.clone()), granularity),
|
||||
@@ -259,7 +255,6 @@ impl ObservabilityStore for CachedObservabilityStore {
|
||||
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)
|
||||
@@ -273,7 +268,6 @@ impl ObservabilityStore for CachedObservabilityStore {
|
||||
dimension: DiagnosticDimension,
|
||||
limit: u16,
|
||||
) -> AdminStoreResult<DiagnosticsObservation> {
|
||||
let filter = normalized(filter);
|
||||
self.diagnostics
|
||||
.get(
|
||||
((range, filter.clone()), dimension, limit),
|
||||
|
||||
@@ -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 不支持此价格档位"));
|
||||
}
|
||||
|
||||
@@ -536,6 +536,9 @@ impl DatabaseDumpPort for FakeDumpPort {
|
||||
pub(crate) struct FakeObjectStore {
|
||||
pub(crate) objects: Mutex<std::collections::HashMap<String, Vec<u8>>>,
|
||||
upload_error: Option<BackupError>,
|
||||
head_error: Option<BackupError>,
|
||||
delete_error: Option<BackupError>,
|
||||
connection_error: Option<BackupError>,
|
||||
}
|
||||
|
||||
impl FakeObjectStore {
|
||||
@@ -543,6 +546,9 @@ impl FakeObjectStore {
|
||||
Self {
|
||||
objects: Mutex::new(std::collections::HashMap::new()),
|
||||
upload_error: None,
|
||||
head_error: None,
|
||||
delete_error: None,
|
||||
connection_error: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -552,6 +558,21 @@ impl FakeObjectStore {
|
||||
self
|
||||
}
|
||||
|
||||
pub(crate) fn fail_head(mut self, code: &'static str) -> Self {
|
||||
self.head_error = Some(BackupError::new(code, "HEAD failed".into()));
|
||||
self
|
||||
}
|
||||
|
||||
pub(crate) fn fail_delete(mut self, code: &'static str) -> Self {
|
||||
self.delete_error = Some(BackupError::new(code, "DELETE failed".into()));
|
||||
self
|
||||
}
|
||||
|
||||
pub(crate) fn fail_connection(mut self, code: &'static str) -> Self {
|
||||
self.connection_error = Some(BackupError::new(code, "DELETE failed".into()));
|
||||
self
|
||||
}
|
||||
|
||||
pub(crate) fn object(&self, key: &str) -> Option<Vec<u8>> {
|
||||
self.objects.lock().expect("objects").get(key).cloned()
|
||||
}
|
||||
@@ -563,6 +584,14 @@ impl BackupObjectStorePort for FakeObjectStore {
|
||||
&self,
|
||||
_config: &BackupStorageConfig,
|
||||
) -> Result<ConnectionTestResult, gateway_admin::model::backup::BackupError> {
|
||||
if let Some(error) = &self.connection_error {
|
||||
return Ok(ConnectionTestResult {
|
||||
ok: false,
|
||||
stage: "deleteObject",
|
||||
code: Some(error.code()),
|
||||
message: error.message().into(),
|
||||
});
|
||||
}
|
||||
Ok(ConnectionTestResult {
|
||||
ok: true,
|
||||
stage: "deleteObject",
|
||||
@@ -597,6 +626,9 @@ impl BackupObjectStorePort for FakeObjectStore {
|
||||
_config: &BackupStorageConfig,
|
||||
object_key: &str,
|
||||
) -> Result<Option<BackupObjectMetadata>, gateway_admin::model::backup::BackupError> {
|
||||
if let Some(error) = &self.head_error {
|
||||
return Err(error.clone());
|
||||
}
|
||||
let Some(content) = self.object(object_key) else {
|
||||
return Ok(None);
|
||||
};
|
||||
@@ -613,6 +645,9 @@ impl BackupObjectStorePort for FakeObjectStore {
|
||||
_config: &BackupStorageConfig,
|
||||
object_key: &str,
|
||||
) -> Result<(), gateway_admin::model::backup::BackupError> {
|
||||
if let Some(error) = &self.delete_error {
|
||||
return Err(error.clone());
|
||||
}
|
||||
self.objects.lock().expect("objects").remove(object_key);
|
||||
Ok(())
|
||||
}
|
||||
@@ -690,11 +725,11 @@ impl AuthStore for FakeAuthStore {
|
||||
) -> AdminStoreResult<()> {
|
||||
Ok(())
|
||||
}
|
||||
async fn renew_session(
|
||||
async fn compare_exchange_session(
|
||||
&self,
|
||||
_: &str,
|
||||
_: &gateway_admin::model::auth::AuthSession,
|
||||
_: chrono::DateTime<chrono::Utc>,
|
||||
_: Option<&gateway_admin::model::auth::AuthSession>,
|
||||
) -> AdminStoreResult<Option<gateway_admin::model::auth::AuthSession>> {
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
@@ -324,3 +324,58 @@ async fn completed_recovery_removes_staging_for_dumping_and_uploading_tasks() {
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn manual_backup_preserves_primary_and_final_delete_failures() {
|
||||
for (upload_fails, expiring) in [(false, false), (true, false), (false, true)] {
|
||||
let repository = Arc::new(FakeBackupRepository::new(configured_settings()));
|
||||
let dump = Arc::new(FakeDumpPort::new());
|
||||
let store = FakeObjectStore::new().fail_delete(code::S3_PERMISSION_DENIED);
|
||||
let store = Arc::new(if upload_fails {
|
||||
store.fail_upload(code::S3_UPLOAD_FAILED, "上传失败")
|
||||
} else {
|
||||
store.fail_head(code::REMOTE_VERIFICATION_FAILED)
|
||||
});
|
||||
let seed = BackupRecordSeed {
|
||||
id: backup_id("cleanup-failure"),
|
||||
trigger_kind: BackupTriggerKind::Manual,
|
||||
scheduled_at: None,
|
||||
object_key: "codex/cleanup-failure.dump".into(),
|
||||
expires_at: expiring.then(|| Utc::now() + Duration::hours(1)),
|
||||
};
|
||||
repository.insert_backup_record(seed.clone()).await.unwrap();
|
||||
let task = BackupTask::new(repository.clone(), dump.clone(), store.clone());
|
||||
task.run_cycle(&CancellationToken::new()).await.unwrap();
|
||||
let records = repository.all_records();
|
||||
let record = &records[0];
|
||||
assert_eq!(record.status, BackupStatus::Failed);
|
||||
assert_eq!(
|
||||
record.error_code.as_deref(),
|
||||
Some(if upload_fails {
|
||||
code::S3_UPLOAD_FAILED
|
||||
} else {
|
||||
code::REMOTE_VERIFICATION_FAILED
|
||||
})
|
||||
);
|
||||
let message = record.error_message.as_deref().unwrap();
|
||||
assert!(message.contains(if upload_fails {
|
||||
"上传失败"
|
||||
} else {
|
||||
"HEAD failed"
|
||||
}));
|
||||
assert!(message.contains("DeleteObject"));
|
||||
assert!(message.contains(code::S3_PERMISSION_DENIED));
|
||||
assert!(message.contains(if expiring {
|
||||
"等待备份到期清理"
|
||||
} else {
|
||||
"未安排自动重试"
|
||||
}));
|
||||
assert_eq!(store.object(&seed.object_key).is_some(), !upload_fails);
|
||||
assert!(dump.inspect_staging(&seed.id).await.unwrap().is_none());
|
||||
task.run_cycle(&CancellationToken::new()).await.unwrap();
|
||||
assert_eq!(
|
||||
repository.all_records()[0].error_message,
|
||||
record.error_message
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -73,14 +73,14 @@ fn snapshot_at_a_block_boundary_keeps_only_three_recent_blocks() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn large_percentage_reversal_restarts_sampling_without_crossing_old_consumption() {
|
||||
fn large_percentage_reversal_restarts_only_sampling_without_crossing_old_consumption() {
|
||||
let sample = select(
|
||||
vec![point(1, 40.0, 100), point(2, 50.0, 200), point(3, 1.0, 210)],
|
||||
point(4, 8.0, 280),
|
||||
);
|
||||
assert_eq!(sample.baseline_percent, 1.0);
|
||||
assert_eq!(sample.usage.tokens, 70);
|
||||
assert_eq!(sample.cycle_usage.tokens, 70);
|
||||
assert_eq!(sample.cycle_usage.tokens, 280);
|
||||
assert_eq!(sample.sampled_percent, 7.0);
|
||||
assert!(!sample.discontinuous);
|
||||
let waiting = select(
|
||||
@@ -90,6 +90,30 @@ fn large_percentage_reversal_restarts_sampling_without_crossing_old_consumption(
|
||||
assert!(waiting.discontinuous);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn same_window_reversal_preserves_all_recorded_cycle_usage() {
|
||||
let current = point(3, 35.0, 140);
|
||||
let recorded = current.usage.clone();
|
||||
let sample = select(vec![point(1, 32.0, 100), point(2, 30.0, 120)], current);
|
||||
assert_eq!(sample.cycle_usage, recorded);
|
||||
assert_eq!(sample.usage.tokens, 20);
|
||||
assert!((sample.usage.usd - 0.02).abs() < 1e-10);
|
||||
assert!(!sample.discontinuous);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn small_reversals_after_full_usage_preserve_cycle_totals_and_cumulative_forecasts() {
|
||||
let current = point(3, 99.5, 10_020);
|
||||
let recorded = current.usage.clone();
|
||||
let sample = select(
|
||||
vec![point(1, 100.0, 10_000), point(2, 99.0, 10_010)],
|
||||
current,
|
||||
);
|
||||
assert_eq!(sample.cycle_usage, recorded);
|
||||
assert_eq!(sample.method, QuotaForecastMethod::Cumulative);
|
||||
assert!(!sample.discontinuous);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn one_point_reversal_is_skipped_without_counting_an_extra_block() {
|
||||
let sample = select(
|
||||
@@ -128,17 +152,41 @@ fn counter_regressions_fail_closed_and_future_points_are_not_used() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn exhausted_history_is_retained_to_detect_a_reset_with_an_unchanged_deadline() {
|
||||
fn same_window_recovery_restarts_sampling_but_preserves_recorded_cycle_usage() {
|
||||
let sample = select(
|
||||
vec![point(1, 100.0, 10_000), point(2, 0.0, 10_000)],
|
||||
point(3, 10.0, 10_100),
|
||||
);
|
||||
assert_eq!(sample.cycle_usage.tokens, 100);
|
||||
assert_eq!(sample.cycle_usage.tokens, 10_100);
|
||||
assert!((sample.cycle_usage.usd - 10.1).abs() < 1e-10);
|
||||
assert_eq!(sample.usage.tokens, 100);
|
||||
assert_eq!(sample.sampled_percent, 10.0);
|
||||
assert!(!sample.discontinuous);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_single_non_exhausted_observation_does_not_restart_cycle_usage() {
|
||||
let sample = select(vec![point(1, 100.0, 10_000)], point(2, 0.0, 10_000));
|
||||
assert_eq!(sample.cycle_usage.tokens, 10_000);
|
||||
assert!(sample.discontinuous);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn repeated_same_window_recoveries_preserve_recorded_cycle_usage() {
|
||||
let sample = select(
|
||||
vec![
|
||||
point(1, 100.0, 10_000),
|
||||
point(2, 0.0, 10_000),
|
||||
point(3, 100.0, 10_100),
|
||||
point(4, 0.0, 10_100),
|
||||
],
|
||||
point(5, 10.0, 10_200),
|
||||
);
|
||||
assert_eq!(sample.cycle_usage.tokens, 10_200);
|
||||
assert_eq!(sample.usage.tokens, 100);
|
||||
assert!(!sample.discontinuous);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn changed_window_restarts_cycle_totals_at_first_matching_observation() {
|
||||
let sample = select_forecast_sample(
|
||||
|
||||
@@ -1121,7 +1121,6 @@ impl SettingsStore for StaticSettingsStore {
|
||||
refresh_margin_seconds: 300,
|
||||
refresh_concurrency: 2,
|
||||
max_concurrent_per_account: 1,
|
||||
request_interval_ms: 0,
|
||||
max_waiting_per_key: 0,
|
||||
max_waiting_per_account: 0,
|
||||
concurrency_wait_timeout_seconds: 30,
|
||||
@@ -2399,6 +2398,93 @@ async fn quota_forecast_only_queries_the_two_selected_source_windows() {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn quota_forecast_same_window_reversal_keeps_recorded_usage_in_the_report() {
|
||||
use gateway_admin::model::quota_forecast_sampling::QuotaForecastHistoryPoint;
|
||||
|
||||
for (percents, tokens, expected_tokens, expected_usd) in [
|
||||
([32.0, 30.0, 35.0], [100, 120, 140], 400, 0.4),
|
||||
(
|
||||
[100.0, 99.0, 99.5],
|
||||
[10_000, 10_010, 10_020],
|
||||
10_070,
|
||||
10.02 / 0.995,
|
||||
),
|
||||
] {
|
||||
let now = Utc::now();
|
||||
let observed = now - TimeDelta::minutes(1);
|
||||
let reset = now + TimeDelta::days(1);
|
||||
let provider = FakeProviderAdmin::new("openai", events());
|
||||
provider.set_quota(ProviderQuota {
|
||||
plan_type: Some("pro".to_owned()),
|
||||
observed_at: Some(observed),
|
||||
windows: vec![ProviderQuotaWindow {
|
||||
key: "week".to_owned(),
|
||||
group: "shortTerm".to_owned(),
|
||||
label: "周额度".to_owned(),
|
||||
limit_id: None,
|
||||
limit_name: None,
|
||||
role: None,
|
||||
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
|
||||
window_seconds: Some(604_800),
|
||||
used_percent: Some(percents[2]),
|
||||
reset_at: Some(reset),
|
||||
limit_reached: false,
|
||||
local_usage: None,
|
||||
provider_data: None,
|
||||
}],
|
||||
..empty_quota()
|
||||
});
|
||||
let mut account = account_record("openai");
|
||||
account.created_at = now - TimeDelta::days(7);
|
||||
let store = FakeAccountStore::with_account(account, events());
|
||||
let usage = |requests, tokens| QuotaForecastUsage {
|
||||
request_count: requests,
|
||||
tokens,
|
||||
input_tokens: tokens,
|
||||
known_cost_count: requests,
|
||||
usd: tokens as f64 / 1_000.0,
|
||||
..Default::default()
|
||||
};
|
||||
*store.quota_forecast_history.lock().unwrap() = QuotaForecastHistory {
|
||||
points: [(3, percents[0], tokens[0]), (2, percents[1], tokens[1])]
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.map(
|
||||
|(index, (minutes, percent, tokens))| QuotaForecastHistoryPoint {
|
||||
started_at: now - TimeDelta::minutes(5),
|
||||
completed_at: now - TimeDelta::minutes(minutes),
|
||||
usage: usage(index as u64 + 1, tokens),
|
||||
provider_observation: ProviderDocument::new(OpaqueProviderData::new(
|
||||
serde_json::json!({
|
||||
"percent": percent, "reset": reset.to_rfc3339(), "plan": "pro",
|
||||
})
|
||||
.as_object()
|
||||
.unwrap()
|
||||
.clone(),
|
||||
)),
|
||||
},
|
||||
)
|
||||
.collect(),
|
||||
usage: usage(3, tokens[2]),
|
||||
pending_request_count: 0,
|
||||
};
|
||||
let services = accounts_service(provider, store).await;
|
||||
let result = services
|
||||
.accounts()
|
||||
.quota_forecast(&ProviderAccountId::new("acct_test").unwrap())
|
||||
.await
|
||||
.unwrap();
|
||||
let forecast = &result.forecasts[0];
|
||||
let source = forecast.source.as_ref().unwrap();
|
||||
assert_eq!(source.tokens, Some(tokens[2]));
|
||||
assert!((source.usd.unwrap() - tokens[2] as f64 / 1_000.0).abs() < 1e-10);
|
||||
assert_eq!(forecast.estimated_tokens, Some(expected_tokens));
|
||||
assert!((forecast.estimated_usd.unwrap() - expected_usd).abs() < 1e-10);
|
||||
assert!(forecast.unavailable_reason.is_none());
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn quota_forecast_mid_cycle_sampling_accepts_small_reset_jitter_but_not_a_changed_bucket() {
|
||||
use gateway_admin::model::quota_forecast_sampling::QuotaForecastHistoryPoint;
|
||||
|
||||
+47
-110
@@ -1,4 +1,4 @@
|
||||
//! 验证管理员认证、密码变更、会话绑定与限流行为
|
||||
//! 验证管理员认证、密码变更与限流,提供会话轮换的内存存储辅助
|
||||
|
||||
use std::{collections::BTreeMap, sync::Mutex};
|
||||
|
||||
@@ -69,6 +69,7 @@ async fn password_change_obeys_rate_limit_and_rejects_unbound_legacy_sessions()
|
||||
admin_user_id: "admin".into(),
|
||||
credential_fingerprint: String::new(),
|
||||
},
|
||||
renewal: None,
|
||||
absolute_expires_at: None,
|
||||
expires_at: Utc::now() + TimeDelta::hours(1),
|
||||
},
|
||||
@@ -138,19 +139,24 @@ impl AuthStore for MemoryAuthStore {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn renew_session(
|
||||
async fn compare_exchange_session(
|
||||
&self,
|
||||
session_id: &str,
|
||||
expected: &AuthSession,
|
||||
expires_at: chrono::DateTime<chrono::Utc>,
|
||||
replacement: Option<&AuthSession>,
|
||||
) -> AdminStoreResult<Option<AuthSession>> {
|
||||
let mut sessions = self.sessions.lock().unwrap();
|
||||
if let Some(session) = sessions.get_mut(session_id) {
|
||||
if session.expires_at <= chrono::Utc::now() {
|
||||
if session.retained_until() <= chrono::Utc::now() {
|
||||
return Ok(None);
|
||||
}
|
||||
if session == expected {
|
||||
session.expires_at = expires_at;
|
||||
if let Some(replacement) = replacement {
|
||||
*session = replacement.clone();
|
||||
} else {
|
||||
sessions.remove(session_id);
|
||||
return Ok(None);
|
||||
}
|
||||
}
|
||||
return Ok(Some(session.clone()));
|
||||
}
|
||||
@@ -190,8 +196,13 @@ impl AuthStore for MemoryAuthStore {
|
||||
#[tokio::test]
|
||||
async fn successful_login_should_create_expiring_session_and_audit() {
|
||||
let store = std::sync::Arc::new(MemoryAuthStore::default());
|
||||
let services = super::AdminHarness::new().auth(store.clone()).build().await;
|
||||
let services = super::AdminHarness::new()
|
||||
.session_ttl_minutes(7 * 24 * 60)
|
||||
.auth(store.clone())
|
||||
.build()
|
||||
.await;
|
||||
|
||||
let before = Utc::now();
|
||||
let result = services
|
||||
.auth()
|
||||
.login(
|
||||
@@ -205,6 +216,17 @@ async fn successful_login_should_create_expiring_session_and_audit() {
|
||||
.await
|
||||
.expect("login");
|
||||
|
||||
let after = Utc::now();
|
||||
for (expiry, ttl) in [
|
||||
(result.session.expires_at, TimeDelta::hours(24)),
|
||||
(result.session.retained_until(), TimeDelta::days(7)),
|
||||
(
|
||||
result.session.absolute_expires_at.unwrap(),
|
||||
TimeDelta::days(30),
|
||||
),
|
||||
] {
|
||||
assert!((before + ttl..=after + ttl).contains(&expiry));
|
||||
}
|
||||
assert!(
|
||||
services
|
||||
.auth()
|
||||
@@ -293,6 +315,7 @@ async fn expired_sessions_are_rejected_and_store_outages_are_not_treated_as_logo
|
||||
credential_fingerprint: String::new(),
|
||||
admin_user_id: "admin".to_owned(),
|
||||
},
|
||||
renewal: None,
|
||||
absolute_expires_at: None,
|
||||
expires_at: Utc::now() - TimeDelta::seconds(1),
|
||||
},
|
||||
@@ -362,7 +385,7 @@ async fn failed_rotation_discards_the_new_session_and_leaves_the_old_session_ret
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(!format!("{original:?}").contains(&original.session_id));
|
||||
*store.reject_delete.lock().unwrap() = Some(original.session_id.clone());
|
||||
*store.reject_delete.lock().unwrap() = Some(stored_id(&original.session_id).to_owned());
|
||||
assert_eq!(
|
||||
services
|
||||
.auth()
|
||||
@@ -391,109 +414,23 @@ async fn failed_rotation_discards_the_new_session_and_leaves_the_old_session_ret
|
||||
assert!(store.sessions.lock().unwrap().is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn renewal_is_throttled_and_capped_by_the_original_absolute_expiry() {
|
||||
let store = std::sync::Arc::new(MemoryAuthStore::default());
|
||||
let services = super::AdminHarness::new()
|
||||
.auth(store.clone())
|
||||
.session_ttl_minutes(60)
|
||||
.build()
|
||||
.await;
|
||||
let auth = services.auth();
|
||||
let login = auth
|
||||
.login(
|
||||
LoginCommand::Admin {
|
||||
username: None,
|
||||
password: "strong-test-password".into(),
|
||||
},
|
||||
std::net::Ipv4Addr::LOCALHOST.into(),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
auth.renew_session(Some(&login.session_id)).await.unwrap(),
|
||||
Some(login.session.clone())
|
||||
);
|
||||
let limit = Utc::now() + TimeDelta::minutes(30);
|
||||
{
|
||||
let mut sessions = store.sessions.lock().unwrap();
|
||||
let session = sessions.get_mut(&login.session_id).unwrap();
|
||||
session.expires_at = Utc::now() + TimeDelta::minutes(1);
|
||||
session.absolute_expires_at = Some(limit);
|
||||
}
|
||||
let renewed = auth
|
||||
.renew_session(Some(&login.session_id))
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(renewed.expires_at, limit);
|
||||
assert_eq!(renewed.absolute_expires_at, Some(limit));
|
||||
assert_eq!(
|
||||
auth.renew_session(Some(&login.session_id)).await.unwrap(),
|
||||
Some(renewed)
|
||||
);
|
||||
fn stored_id(token: &str) -> &str {
|
||||
token.split('.').nth(1).unwrap_or(token)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn renewal_never_revives_expired_revoked_or_legacy_sessions() {
|
||||
let store = std::sync::Arc::new(MemoryAuthStore::default());
|
||||
let services = super::AdminHarness::new().auth(store.clone()).build().await;
|
||||
let auth = services.auth();
|
||||
let login = auth
|
||||
.login(
|
||||
LoginCommand::Admin {
|
||||
username: None,
|
||||
password: "strong-test-password".into(),
|
||||
},
|
||||
std::net::Ipv4Addr::LOCALHOST.into(),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let mut session = login.session.clone();
|
||||
session.absolute_expires_at = None;
|
||||
session.expires_at = Utc::now() + TimeDelta::seconds(30);
|
||||
store
|
||||
.sessions
|
||||
.lock()
|
||||
.unwrap()
|
||||
.insert(login.session_id.clone(), session.clone());
|
||||
assert_eq!(
|
||||
auth.renew_session(Some(&login.session_id)).await.unwrap(),
|
||||
Some(session.clone())
|
||||
);
|
||||
session.expires_at = Utc::now() - TimeDelta::seconds(1);
|
||||
store
|
||||
.sessions
|
||||
.lock()
|
||||
.unwrap()
|
||||
.insert(login.session_id.clone(), session);
|
||||
assert!(
|
||||
auth.renew_session(Some(&login.session_id))
|
||||
.await
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
let mut session = login.session;
|
||||
session.absolute_expires_at = Some(Utc::now() - TimeDelta::seconds(1));
|
||||
store
|
||||
.sessions
|
||||
.lock()
|
||||
.unwrap()
|
||||
.insert(login.session_id.clone(), session);
|
||||
assert!(
|
||||
auth.renew_session(Some(&login.session_id))
|
||||
.await
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
auth.logout(&login.session_id).await.unwrap();
|
||||
assert!(
|
||||
auth.renew_session(Some(&login.session_id))
|
||||
.await
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
assert!(store.sessions.lock().unwrap().is_empty());
|
||||
async fn admin_login(
|
||||
auth: &dyn gateway_admin::AuthService,
|
||||
) -> gateway_admin::model::auth::SessionGrant {
|
||||
auth.login(
|
||||
LoginCommand::Admin {
|
||||
username: None,
|
||||
password: "strong-test-password".into(),
|
||||
},
|
||||
std::net::Ipv4Addr::LOCALHOST.into(),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
mod session;
|
||||
@@ -0,0 +1,268 @@
|
||||
//! 验证双凭据续期、响应重试、并发撤销与固定期限边界
|
||||
|
||||
use super::super::AdminHarness;
|
||||
use super::{MemoryAuthStore, admin_login, stored_id};
|
||||
use chrono::{TimeDelta, Utc};
|
||||
use gateway_admin::model::auth::AuthSession;
|
||||
|
||||
#[tokio::test]
|
||||
async fn expired_access_is_restored_by_refresh_and_duplicate_refresh_replays_the_same_grant() {
|
||||
let store = std::sync::Arc::new(MemoryAuthStore::default());
|
||||
let services = AdminHarness::new().auth(store.clone()).build().await;
|
||||
let auth = services.auth();
|
||||
let login = admin_login(auth).await;
|
||||
let refresh = login.refresh_token.as_deref().unwrap();
|
||||
assert!(auth.session(Some(refresh)).await.unwrap().is_none());
|
||||
assert!(
|
||||
auth.renew_session(Some(&login.session_id), None)
|
||||
.await
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
assert!(
|
||||
auth.renew_session(None, Some(&login.session_id))
|
||||
.await
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
assert_eq!(
|
||||
auth.renew_session(Some(refresh), None).await.unwrap(),
|
||||
Some(login.clone())
|
||||
);
|
||||
store
|
||||
.sessions
|
||||
.lock()
|
||||
.unwrap()
|
||||
.get_mut(stored_id(&login.session_id))
|
||||
.unwrap()
|
||||
.expires_at = Utc::now() - TimeDelta::seconds(1);
|
||||
assert!(
|
||||
auth.session(Some(&login.session_id))
|
||||
.await
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
let (first, concurrent) = tokio::join!(
|
||||
auth.renew_session(Some(refresh), None),
|
||||
auth.renew_session(Some(refresh), None)
|
||||
);
|
||||
let first = first.unwrap().unwrap();
|
||||
assert_eq!(concurrent.unwrap().as_ref(), Some(&first));
|
||||
assert_ne!(first.session_id, login.session_id);
|
||||
assert_ne!(first.refresh_token, login.refresh_token);
|
||||
assert_eq!(
|
||||
first.session.absolute_expires_at,
|
||||
login.session.absolute_expires_at
|
||||
);
|
||||
assert!(
|
||||
auth.session(Some(&first.session_id))
|
||||
.await
|
||||
.unwrap()
|
||||
.is_some()
|
||||
);
|
||||
assert_eq!(
|
||||
auth.renew_session(first.refresh_token.as_deref(), None)
|
||||
.await
|
||||
.unwrap(),
|
||||
Some(first.clone())
|
||||
);
|
||||
assert_eq!(
|
||||
auth.renew_session(Some(refresh), None).await.unwrap(),
|
||||
Some(first.clone())
|
||||
);
|
||||
let persisted = format!("{:?}", store.sessions.lock().unwrap());
|
||||
for token in [
|
||||
&first.session_id,
|
||||
first.refresh_token.as_ref().unwrap(),
|
||||
&login.session_id,
|
||||
login.refresh_token.as_ref().unwrap(),
|
||||
] {
|
||||
assert!(!persisted.contains(token));
|
||||
assert!(!persisted.contains(token.rsplit('.').next().unwrap()));
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn late_refresh_reuse_revokes_the_family_but_unknown_secrets_do_not() {
|
||||
let store = std::sync::Arc::new(MemoryAuthStore::default());
|
||||
let services = AdminHarness::new().auth(store.clone()).build().await;
|
||||
let auth = services.auth();
|
||||
let login = admin_login(auth).await;
|
||||
store
|
||||
.sessions
|
||||
.lock()
|
||||
.unwrap()
|
||||
.get_mut(stored_id(&login.session_id))
|
||||
.unwrap()
|
||||
.expires_at = Utc::now() - TimeDelta::seconds(1);
|
||||
let next = auth
|
||||
.renew_session(login.refresh_token.as_deref(), None)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
let unknown = format!("r2.{}.{}", stored_id(&login.session_id), "A".repeat(43));
|
||||
assert!(
|
||||
auth.renew_session(Some(&unknown), None)
|
||||
.await
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
auth.logout(&unknown).await.unwrap();
|
||||
auth.logout(stored_id(&login.session_id)).await.unwrap();
|
||||
assert!(
|
||||
auth.session(Some(&next.session_id))
|
||||
.await
|
||||
.unwrap()
|
||||
.is_some()
|
||||
);
|
||||
store
|
||||
.sessions
|
||||
.lock()
|
||||
.unwrap()
|
||||
.get_mut(stored_id(&login.session_id))
|
||||
.unwrap()
|
||||
.renewal
|
||||
.as_mut()
|
||||
.unwrap()
|
||||
.previous
|
||||
.as_mut()
|
||||
.unwrap()
|
||||
.expires_at = Utc::now() - TimeDelta::seconds(1);
|
||||
assert!(
|
||||
auth.renew_session(login.refresh_token.as_deref(), None)
|
||||
.await
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
assert!(
|
||||
auth.session(Some(&next.session_id))
|
||||
.await
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
assert!(
|
||||
auth.renew_session(next.refresh_token.as_deref(), None)
|
||||
.await
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn refresh_respects_idle_and_absolute_expiry_and_logout_after_access_expiry() {
|
||||
let store = std::sync::Arc::new(MemoryAuthStore::default());
|
||||
let services = AdminHarness::new().auth(store.clone()).build().await;
|
||||
let auth = services.auth();
|
||||
for absolute in [false, true] {
|
||||
let login = admin_login(auth).await;
|
||||
{
|
||||
let mut sessions = store.sessions.lock().unwrap();
|
||||
let session = sessions.get_mut(stored_id(&login.session_id)).unwrap();
|
||||
if absolute {
|
||||
session.absolute_expires_at = Some(Utc::now() - TimeDelta::seconds(1));
|
||||
} else {
|
||||
session.renewal.as_mut().unwrap().expires_at = Utc::now() - TimeDelta::seconds(1);
|
||||
}
|
||||
}
|
||||
assert!(
|
||||
auth.renew_session(login.refresh_token.as_deref(), None)
|
||||
.await
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
assert!(
|
||||
auth.session(Some(&login.session_id))
|
||||
.await
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
}
|
||||
let login = admin_login(auth).await;
|
||||
store
|
||||
.sessions
|
||||
.lock()
|
||||
.unwrap()
|
||||
.get_mut(stored_id(&login.session_id))
|
||||
.unwrap()
|
||||
.expires_at = Utc::now() - TimeDelta::seconds(1);
|
||||
auth.logout(login.refresh_token.as_deref().unwrap())
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
auth.renew_session(login.refresh_token.as_deref(), None)
|
||||
.await
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn refreshed_expiry_is_capped_and_store_outages_preserve_the_session() {
|
||||
let store = std::sync::Arc::new(MemoryAuthStore::default());
|
||||
let services = AdminHarness::new().auth(store.clone()).build().await;
|
||||
let auth = services.auth();
|
||||
let login = admin_login(auth).await;
|
||||
let limit = Utc::now() + TimeDelta::minutes(1);
|
||||
{
|
||||
let mut sessions = store.sessions.lock().unwrap();
|
||||
let session = sessions.get_mut(stored_id(&login.session_id)).unwrap();
|
||||
session.absolute_expires_at = Some(limit);
|
||||
session.expires_at = Utc::now() - TimeDelta::seconds(1);
|
||||
}
|
||||
store
|
||||
.unavailable
|
||||
.store(true, std::sync::atomic::Ordering::SeqCst);
|
||||
assert!(
|
||||
auth.renew_session(login.refresh_token.as_deref(), None)
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
store
|
||||
.unavailable
|
||||
.store(false, std::sync::atomic::Ordering::SeqCst);
|
||||
let refreshed = auth
|
||||
.renew_session(login.refresh_token.as_deref(), None)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(refreshed.session.expires_at, limit);
|
||||
assert_eq!(refreshed.session.retained_until(), limit);
|
||||
assert_eq!(refreshed.session.absolute_expires_at, Some(limit));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fixed_sessions_cannot_acquire_refresh_credentials() {
|
||||
let store = std::sync::Arc::new(MemoryAuthStore::default());
|
||||
let services = AdminHarness::new().auth(store.clone()).build().await;
|
||||
let auth = services.auth();
|
||||
let login = admin_login(auth).await;
|
||||
let fixed = AuthSession {
|
||||
renewal: None,
|
||||
..login.session
|
||||
};
|
||||
store
|
||||
.sessions
|
||||
.lock()
|
||||
.unwrap()
|
||||
.insert("fixed".into(), fixed.clone());
|
||||
let grant = auth
|
||||
.renew_session(None, Some("fixed"))
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(grant.session, fixed);
|
||||
assert!(grant.refresh_token.is_none());
|
||||
store
|
||||
.sessions
|
||||
.lock()
|
||||
.unwrap()
|
||||
.get_mut("fixed")
|
||||
.unwrap()
|
||||
.expires_at = Utc::now() - TimeDelta::seconds(1);
|
||||
assert!(
|
||||
auth.renew_session(None, Some("fixed"))
|
||||
.await
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
}
|
||||
@@ -285,19 +285,24 @@ impl AuthStore for MemoryKeyAuthStore {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn renew_session(
|
||||
async fn compare_exchange_session(
|
||||
&self,
|
||||
session_id: &str,
|
||||
expected: &AuthSession,
|
||||
expires_at: chrono::DateTime<chrono::Utc>,
|
||||
replacement: Option<&AuthSession>,
|
||||
) -> AdminStoreResult<Option<AuthSession>> {
|
||||
let mut sessions = self.sessions.lock().unwrap();
|
||||
if let Some(session) = sessions.get_mut(session_id) {
|
||||
if session.expires_at <= chrono::Utc::now() {
|
||||
if session.retained_until() <= chrono::Utc::now() {
|
||||
return Ok(None);
|
||||
}
|
||||
if session == expected {
|
||||
session.expires_at = expires_at;
|
||||
if let Some(replacement) = replacement {
|
||||
*session = replacement.clone();
|
||||
} else {
|
||||
sessions.remove(session_id);
|
||||
return Ok(None);
|
||||
}
|
||||
}
|
||||
return Ok(Some(session.clone()));
|
||||
}
|
||||
|
||||
@@ -102,3 +102,51 @@ async fn backup_service_rejects_download_of_non_completed_record() {
|
||||
.expect_err("non-completed record must be rejected");
|
||||
assert_eq!(error.kind(), gateway_admin::model::AdminErrorKind::Conflict);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn failed_delete_probe_does_not_refresh_verification() {
|
||||
use gateway_admin::{model::backup::code, ports::backup::BackupRepository as _};
|
||||
for succeeds in [false, true] {
|
||||
let mut settings = configured_settings();
|
||||
settings.last_verified_at = None;
|
||||
let repository = Arc::new(FakeBackupRepository::new(settings));
|
||||
let auth = Arc::new(FakeAuthStore::new());
|
||||
let store = FakeObjectStore::new();
|
||||
let store = if succeeds {
|
||||
store
|
||||
} else {
|
||||
store.fail_connection(code::S3_PERMISSION_DENIED)
|
||||
};
|
||||
let services = AdminHarness::new()
|
||||
.auth(auth.clone())
|
||||
.backup(BackupStorePorts::new(
|
||||
repository.clone(),
|
||||
Arc::new(FakeDumpPort::new()),
|
||||
Arc::new(store),
|
||||
))
|
||||
.build()
|
||||
.await;
|
||||
let result = services
|
||||
.backups()
|
||||
.test_connection(&system_context())
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(result.ok, succeeds);
|
||||
assert_eq!(
|
||||
repository
|
||||
.load_settings()
|
||||
.await
|
||||
.unwrap()
|
||||
.last_verified_at
|
||||
.is_some(),
|
||||
succeeds
|
||||
);
|
||||
assert_eq!(
|
||||
auth.audit_actions()
|
||||
.iter()
|
||||
.filter(|action| action.as_str() == "backup.s3_connection_tested")
|
||||
.count(),
|
||||
1
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -38,7 +38,6 @@ fn runtime_settings(enabled: bool, probe_enabled: bool, adaptive: bool) -> Runti
|
||||
refresh_margin_seconds: 300,
|
||||
refresh_concurrency: 2,
|
||||
max_concurrent_per_account: 5,
|
||||
request_interval_ms: 0,
|
||||
max_waiting_per_key: 0,
|
||||
max_waiting_per_account: 0,
|
||||
concurrency_wait_timeout_seconds: 30,
|
||||
|
||||
@@ -369,11 +369,11 @@ impl AuthStore for BootstrapAuthStore {
|
||||
Err(unavailable("admin session"))
|
||||
}
|
||||
|
||||
async fn renew_session(
|
||||
async fn compare_exchange_session(
|
||||
&self,
|
||||
_: &str,
|
||||
_: &gateway_admin::model::auth::AuthSession,
|
||||
_: chrono::DateTime<chrono::Utc>,
|
||||
_: Option<&gateway_admin::model::auth::AuthSession>,
|
||||
) -> AdminStoreResult<Option<gateway_admin::model::auth::AuthSession>> {
|
||||
Ok(None)
|
||||
}
|
||||
@@ -747,6 +747,10 @@ impl ClientKeyStore for UnavailableStore {
|
||||
|
||||
#[async_trait]
|
||||
impl ObservabilityStore for UnavailableStore {
|
||||
fn query_timeout(&self) -> std::time::Duration {
|
||||
std::time::Duration::from_secs(10)
|
||||
}
|
||||
|
||||
async fn dashboard_summary(
|
||||
&self,
|
||||
_: DashboardQuery,
|
||||
|
||||
@@ -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();
|
||||
@@ -844,6 +945,10 @@ impl FixtureObservabilityStore {
|
||||
|
||||
#[async_trait]
|
||||
impl ObservabilityStore for FixtureObservabilityStore {
|
||||
fn query_timeout(&self) -> std::time::Duration {
|
||||
std::time::Duration::from_secs(10)
|
||||
}
|
||||
|
||||
async fn dashboard_summary(
|
||||
&self,
|
||||
query: DashboardQuery,
|
||||
@@ -1003,7 +1108,6 @@ impl SettingsStore for FixtureSettingsStore {
|
||||
refresh_margin_seconds: 300,
|
||||
refresh_concurrency: 2,
|
||||
max_concurrent_per_account: self.max_concurrent_per_account,
|
||||
request_interval_ms: 0,
|
||||
max_waiting_per_key: 0,
|
||||
max_waiting_per_account: 0,
|
||||
concurrency_wait_timeout_seconds: 30,
|
||||
@@ -1305,23 +1409,23 @@ async fn usage_summary_cache_should_coalesce_when_generation_advances_during_the
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn usage_summary_cache_should_isolate_key_identity_and_normalize_group_sets() {
|
||||
async fn usage_summary_cache_should_isolate_key_identity_and_filters() {
|
||||
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"]),
|
||||
for (key, model) in [
|
||||
("key-a", None),
|
||||
("key-a", None),
|
||||
("key-b", None),
|
||||
("key-a", Some("model-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(),
|
||||
model: model.map(str::to_owned),
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
@@ -1330,16 +1434,16 @@ async fn usage_summary_cache_should_isolate_key_identity_and_normalize_group_set
|
||||
}
|
||||
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[0].client_api_key_ref.as_deref(), Some("key-a"));
|
||||
assert_eq!(calls[1].client_api_key_ref.as_deref(), Some("key-b"));
|
||||
assert_eq!(calls[2].routing_group_refs, ["deleted-a"]);
|
||||
assert_eq!(calls[2].model.as_deref(), Some("model-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.set_dashboard_delay(store.query_timeout().mul_f64(0.7));
|
||||
store.summary_fail.store(true, Ordering::SeqCst);
|
||||
let services = observability_services(store.clone()).await;
|
||||
let service = services.observability();
|
||||
@@ -1352,7 +1456,7 @@ async fn usage_summary_cache_should_bound_waiting_after_failed_refresh_and_allow
|
||||
);
|
||||
assert!(first.is_err() && second.is_err() && third.is_err() && fourth.is_err());
|
||||
assert!(
|
||||
started.elapsed() <= StdDuration::from_secs(10),
|
||||
started.elapsed() <= store.query_timeout(),
|
||||
"同键失败不能使后续请求逐个等待完整数据库超时: {:?}",
|
||||
started.elapsed()
|
||||
);
|
||||
|
||||
@@ -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());
|
||||
@@ -286,7 +331,6 @@ async fn settings_should_reject_zero_refresh_margin_before_store_call() {
|
||||
refresh_margin_seconds: 0,
|
||||
refresh_concurrency: 1,
|
||||
max_concurrent_per_account: 1,
|
||||
request_interval_ms: 0,
|
||||
max_waiting_per_key: 0,
|
||||
max_waiting_per_account: 0,
|
||||
concurrency_wait_timeout_seconds: 30,
|
||||
|
||||
@@ -29,7 +29,6 @@ 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>,
|
||||
@@ -77,7 +76,6 @@ 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>,
|
||||
@@ -86,7 +84,6 @@ pub struct DiagnosticsQuery {
|
||||
pub provider: Option<String>,
|
||||
pub model: Option<String>,
|
||||
pub status_code: Option<i64>,
|
||||
pub search: Option<String>,
|
||||
}
|
||||
|
||||
impl DiagnosticsQuery {
|
||||
@@ -100,7 +97,6 @@ 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>,
|
||||
@@ -301,7 +297,6 @@ 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()),
|
||||
@@ -318,6 +313,15 @@ pub(crate) fn usage_filter(query: &UsageQuery) -> Result<domain::UsageFilter, Wi
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn analytics_filter(
|
||||
query: &UsageQuery,
|
||||
) -> Result<domain::UsageFilter, WireValidationError> {
|
||||
if query.search.is_some() {
|
||||
return Err(WireValidationError::new("search"));
|
||||
}
|
||||
usage_filter(query)
|
||||
}
|
||||
|
||||
pub(crate) fn usage_command(
|
||||
query: &UsageQuery,
|
||||
timezone: gateway_core::time::DeploymentTimeZone,
|
||||
@@ -361,7 +365,6 @@ 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()),
|
||||
@@ -400,25 +403,3 @@ 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)
|
||||
}
|
||||
|
||||
@@ -151,7 +151,7 @@ where
|
||||
state.admin_services().timezone(),
|
||||
)
|
||||
.map_err(map_wire_error)?;
|
||||
let filter = usage_filter(&query).map_err(map_wire_error)?;
|
||||
let filter = analytics_filter(&query).map_err(map_wire_error)?;
|
||||
let result = state
|
||||
.admin_services()
|
||||
.observability()
|
||||
@@ -184,7 +184,7 @@ where
|
||||
state.admin_services().timezone(),
|
||||
)
|
||||
.map_err(map_wire_error)?;
|
||||
let filter = usage_filter(&query).map_err(map_wire_error)?;
|
||||
let filter = analytics_filter(&query).map_err(map_wire_error)?;
|
||||
let result = state
|
||||
.admin_services()
|
||||
.observability()
|
||||
@@ -215,12 +215,9 @@ 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)?,
|
||||
search: non_empty(query.search),
|
||||
..domain::UsageFilter::default()
|
||||
};
|
||||
let result = state
|
||||
|
||||
@@ -42,7 +42,6 @@ pub struct RuntimeSettingsFields {
|
||||
pub refresh_margin_seconds: u64,
|
||||
pub refresh_concurrency: u64,
|
||||
pub max_concurrent_per_account: u64,
|
||||
pub request_interval_ms: u64,
|
||||
pub max_waiting_per_key: u32,
|
||||
pub max_waiting_per_account: u32,
|
||||
pub concurrency_wait_timeout_seconds: u32,
|
||||
@@ -735,7 +734,6 @@ impl RuntimeSettingsFields {
|
||||
.map_err(|_| WireValidationError::new("settingsRefreshConcurrencyOverflow"))?,
|
||||
max_concurrent_per_account: u32::try_from(self.max_concurrent_per_account)
|
||||
.map_err(|_| WireValidationError::new("settingsMaxConcurrencyOverflow"))?,
|
||||
request_interval_ms: self.request_interval_ms,
|
||||
max_waiting_per_key: self.max_waiting_per_key,
|
||||
max_waiting_per_account: self.max_waiting_per_account,
|
||||
concurrency_wait_timeout_seconds: self.concurrency_wait_timeout_seconds,
|
||||
@@ -779,7 +777,6 @@ impl From<gateway_admin::model::settings::RuntimeSettingsValues> for RuntimeSett
|
||||
refresh_margin_seconds: settings.refresh_margin_seconds,
|
||||
refresh_concurrency: u64::from(settings.refresh_concurrency),
|
||||
max_concurrent_per_account: u64::from(settings.max_concurrent_per_account),
|
||||
request_interval_ms: settings.request_interval_ms,
|
||||
max_waiting_per_key: settings.max_waiting_per_key,
|
||||
max_waiting_per_account: settings.max_waiting_per_account,
|
||||
concurrency_wait_timeout_seconds: settings.concurrency_wait_timeout_seconds,
|
||||
@@ -835,7 +832,6 @@ fn settings_validation_error(field: &'static str) -> WireValidationError {
|
||||
"ops_event_retention_days" => "opsEventRetentionDays",
|
||||
"refresh_concurrency" => "refreshConcurrency",
|
||||
"refresh_margin_seconds" => "refreshMarginSeconds",
|
||||
"request_interval_ms" => "requestIntervalMs",
|
||||
"request_location" => "requestLocation",
|
||||
"request_location_enabled" => "requestLocationEnabled",
|
||||
"responses_max_decompressed_body_bytes" => "responsesMaxDecompressedBodyBytes",
|
||||
|
||||
@@ -138,7 +138,9 @@ where
|
||||
.login(
|
||||
payload.into(),
|
||||
peer.ip(),
|
||||
session_cookie::value(&headers).as_deref(),
|
||||
session_cookie::refresh_value(&headers)
|
||||
.or_else(|| session_cookie::value(&headers))
|
||||
.as_deref(),
|
||||
)
|
||||
.await
|
||||
.map_err(map_login_error)?;
|
||||
@@ -149,8 +151,7 @@ where
|
||||
)
|
||||
.into_response(),
|
||||
&headers,
|
||||
&result.session_id,
|
||||
&result.session,
|
||||
&result,
|
||||
)
|
||||
}
|
||||
|
||||
@@ -177,16 +178,18 @@ async fn refresh_session<S>(
|
||||
where
|
||||
S: SessionState + Send + Sync,
|
||||
{
|
||||
let session_id = session_cookie::value(&headers);
|
||||
let session = state
|
||||
let grant = state
|
||||
.admin_services()
|
||||
.auth()
|
||||
.renew_session(session_id.as_deref())
|
||||
.renew_session(
|
||||
session_cookie::refresh_value(&headers).as_deref(),
|
||||
session_cookie::value(&headers).as_deref(),
|
||||
)
|
||||
.await
|
||||
.map_err(map_admin_service_error)?;
|
||||
let response = session_response(session.as_ref());
|
||||
if let (Some(session_id), Some(session)) = (session_id, session) {
|
||||
return set_session_cookie(response, &headers, &session_id, &session);
|
||||
let response = session_response(grant.as_ref().map(|grant| &grant.session));
|
||||
if let Some(grant) = grant {
|
||||
return set_session_cookie(response, &headers, &grant);
|
||||
}
|
||||
// 失效响应不清 Cookie,避免晚到的请求覆盖另一个标签页刚建立的会话
|
||||
Ok(response)
|
||||
@@ -206,32 +209,72 @@ fn session_response(session: Option<&AuthSession>) -> Response {
|
||||
fn set_session_cookie(
|
||||
mut response: Response,
|
||||
headers: &HeaderMap,
|
||||
session_id: &str,
|
||||
session: &AuthSession,
|
||||
grant: &gateway_admin::model::auth::SessionGrant,
|
||||
) -> Result<Response, AdminError> {
|
||||
let max_age = session
|
||||
.expires_at
|
||||
.signed_duration_since(chrono::Utc::now())
|
||||
.num_seconds()
|
||||
.max(1);
|
||||
let expires = session.expires_at.format("%a, %d %b %Y %H:%M:%S GMT");
|
||||
let cookie = format!(
|
||||
"{}={session_id}; {}; Max-Age={max_age}; Expires={expires}",
|
||||
append_cookie(
|
||||
&mut response,
|
||||
headers,
|
||||
session_cookie::NAME,
|
||||
session_cookie::attributes(headers)
|
||||
);
|
||||
response.headers_mut().insert(
|
||||
header::SET_COOKIE,
|
||||
HeaderValue::from_str(&cookie).map_err(|_| AdminError::internal())?,
|
||||
);
|
||||
&grant.session_id,
|
||||
Some(grant.session.expires_at),
|
||||
)?;
|
||||
append_cookie(
|
||||
&mut response,
|
||||
headers,
|
||||
session_cookie::REFRESH_NAME,
|
||||
grant.refresh_token.as_deref().unwrap_or_default(),
|
||||
grant
|
||||
.session
|
||||
.renewal
|
||||
.as_ref()
|
||||
.map(|renewal| renewal.expires_at),
|
||||
)?;
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
fn append_cookie(
|
||||
response: &mut Response,
|
||||
headers: &HeaderMap,
|
||||
name: &str,
|
||||
value: &str,
|
||||
expires_at: Option<chrono::DateTime<chrono::Utc>>,
|
||||
) -> Result<(), AdminError> {
|
||||
let attributes = session_cookie::attributes(headers);
|
||||
let path = if name == session_cookie::REFRESH_NAME {
|
||||
"/api/auth"
|
||||
} else {
|
||||
"/"
|
||||
};
|
||||
let lifetime = expires_at.map_or_else(
|
||||
|| "Max-Age=0".to_owned(),
|
||||
|expiry| {
|
||||
format!(
|
||||
"Max-Age={}; Expires={}",
|
||||
expiry
|
||||
.signed_duration_since(chrono::Utc::now())
|
||||
.num_seconds()
|
||||
.max(1),
|
||||
expiry.format("%a, %d %b %Y %H:%M:%S GMT")
|
||||
)
|
||||
},
|
||||
);
|
||||
response.headers_mut().append(
|
||||
header::SET_COOKIE,
|
||||
HeaderValue::from_str(&format!(
|
||||
"{name}={value}; Path={path}; {attributes}; {lifetime}"
|
||||
))
|
||||
.map_err(|_| AdminError::internal())?,
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn logout<S>(State(state): State<S>, headers: HeaderMap) -> Result<Response, AdminError>
|
||||
where
|
||||
S: SessionState + Send + Sync,
|
||||
{
|
||||
if let Some(session_id) = session_cookie::value(&headers) {
|
||||
if let Some(session_id) =
|
||||
session_cookie::refresh_value(&headers).or_else(|| session_cookie::value(&headers))
|
||||
{
|
||||
state
|
||||
.admin_services()
|
||||
.auth()
|
||||
@@ -285,15 +328,9 @@ fn clear_session_cookie(
|
||||
mut response: Response,
|
||||
headers: &HeaderMap,
|
||||
) -> Result<Response, AdminError> {
|
||||
let cookie = format!(
|
||||
"{}=; {}; Max-Age=0",
|
||||
session_cookie::NAME,
|
||||
session_cookie::attributes(headers)
|
||||
);
|
||||
response.headers_mut().insert(
|
||||
header::SET_COOKIE,
|
||||
HeaderValue::from_str(&cookie).map_err(|_| AdminError::internal())?,
|
||||
);
|
||||
for name in [session_cookie::NAME, session_cookie::REFRESH_NAME] {
|
||||
append_cookie(&mut response, headers, name, "", None)?;
|
||||
}
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
|
||||
@@ -228,6 +228,8 @@ async fn handle(
|
||||
Ok(response.map(|body| body.map_err(|error| Box::new(error) as _).boxed_unsync()))
|
||||
})
|
||||
});
|
||||
let request_id = context.request_id.clone();
|
||||
let call_id = context.call_id.clone();
|
||||
let result = match plan {
|
||||
Some(plan) => plan.handle_http(context, request, next).await,
|
||||
None => next.run(request).await,
|
||||
@@ -251,7 +253,18 @@ async fn handle(
|
||||
tokio::select! {
|
||||
biased;
|
||||
() = cancellation.cancelled() => {},
|
||||
_ = task => {},
|
||||
result = task => {
|
||||
if let Err(contract::upgrade::SessionError::Failed { kind, plugin_instance_id }) = result {
|
||||
tracing::warn!(
|
||||
event = "middleware.session.failed",
|
||||
request_id,
|
||||
call_id,
|
||||
plugin_instance_id,
|
||||
kind,
|
||||
"升级后的中间件会话失败"
|
||||
);
|
||||
}
|
||||
},
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
@@ -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};
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -4,6 +4,7 @@ use axum::http::{HeaderMap, header::ORIGIN};
|
||||
use url::Url;
|
||||
|
||||
pub(crate) const NAME: &str = "cpr_session";
|
||||
pub(crate) const REFRESH_NAME: &str = "cpr_refresh";
|
||||
|
||||
pub(crate) fn attributes(headers: &HeaderMap) -> &'static str {
|
||||
// HTTPS 反代回源 HTTP 不改变浏览器 Origin
|
||||
@@ -19,16 +20,24 @@ pub(crate) fn attributes(headers: &HeaderMap) -> &'static str {
|
||||
})
|
||||
});
|
||||
if http_origin {
|
||||
"Path=/; HttpOnly; SameSite=Lax"
|
||||
"HttpOnly; SameSite=Lax"
|
||||
} else {
|
||||
"Path=/; Secure; HttpOnly; SameSite=Lax"
|
||||
"Secure; HttpOnly; SameSite=Lax"
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn value(headers: &HeaderMap) -> Option<String> {
|
||||
named_value(headers, NAME)
|
||||
}
|
||||
|
||||
pub(crate) fn refresh_value(headers: &HeaderMap) -> Option<String> {
|
||||
named_value(headers, REFRESH_NAME)
|
||||
}
|
||||
|
||||
fn named_value(headers: &HeaderMap, name: &str) -> Option<String> {
|
||||
let cookie = headers.get("cookie")?.to_str().ok()?;
|
||||
cookie.split(';').find_map(|part| {
|
||||
let (candidate, value) = part.trim().split_once('=')?;
|
||||
(candidate == NAME).then(|| value.to_owned())
|
||||
(candidate == name && !value.is_empty()).then(|| value.to_owned())
|
||||
})
|
||||
}
|
||||
|
||||
@@ -321,6 +321,7 @@ impl MemoryAuthStore {
|
||||
},
|
||||
admin_user_id: "admin_1".to_owned(),
|
||||
},
|
||||
renewal: None,
|
||||
absolute_expires_at: None,
|
||||
expires_at: Utc::now() + Duration::hours(1),
|
||||
},
|
||||
@@ -339,7 +340,7 @@ impl MemoryAuthStore {
|
||||
self.sessions
|
||||
.lock()
|
||||
.unwrap()
|
||||
.get_mut(session_id)
|
||||
.get_mut(session_id.split('.').nth(1).unwrap_or(session_id))
|
||||
.unwrap()
|
||||
.expires_at = expires_at;
|
||||
}
|
||||
@@ -415,19 +416,24 @@ impl AuthStore for MemoryAuthStore {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn renew_session(
|
||||
async fn compare_exchange_session(
|
||||
&self,
|
||||
session_id: &str,
|
||||
expected: &AuthSession,
|
||||
expires_at: chrono::DateTime<chrono::Utc>,
|
||||
replacement: Option<&AuthSession>,
|
||||
) -> AdminStoreResult<Option<AuthSession>> {
|
||||
let mut sessions = self.sessions.lock().unwrap();
|
||||
if let Some(session) = sessions.get_mut(session_id) {
|
||||
if session.expires_at <= chrono::Utc::now() {
|
||||
if session.retained_until() <= chrono::Utc::now() {
|
||||
return Ok(None);
|
||||
}
|
||||
if session == expected {
|
||||
session.expires_at = expires_at;
|
||||
if let Some(replacement) = replacement {
|
||||
*session = replacement.clone();
|
||||
} else {
|
||||
sessions.remove(session_id);
|
||||
return Ok(None);
|
||||
}
|
||||
}
|
||||
return Ok(Some(session.clone()));
|
||||
}
|
||||
@@ -610,7 +616,6 @@ impl SettingsStore for MemorySettingsStore {
|
||||
refresh_margin_seconds: command.values.refresh_margin_seconds,
|
||||
refresh_concurrency: command.values.refresh_concurrency,
|
||||
max_concurrent_per_account: command.values.max_concurrent_per_account,
|
||||
request_interval_ms: command.values.request_interval_ms,
|
||||
max_waiting_per_key: command.values.max_waiting_per_key,
|
||||
max_waiting_per_account: command.values.max_waiting_per_account,
|
||||
concurrency_wait_timeout_seconds: command.values.concurrency_wait_timeout_seconds,
|
||||
@@ -1374,6 +1379,10 @@ impl AccountRuntimeStore for UnusedStore {
|
||||
|
||||
#[async_trait]
|
||||
impl ObservabilityStore for UnusedStore {
|
||||
fn query_timeout(&self) -> std::time::Duration {
|
||||
std::time::Duration::from_secs(10)
|
||||
}
|
||||
|
||||
async fn dashboard_summary(
|
||||
&self,
|
||||
query: DashboardQuery,
|
||||
@@ -1726,7 +1735,6 @@ fn test_runtime_settings() -> RuntimeSettings {
|
||||
refresh_margin_seconds: 3_600,
|
||||
refresh_concurrency: 2,
|
||||
max_concurrent_per_account: 3,
|
||||
request_interval_ms: 50,
|
||||
max_waiting_per_key: 0,
|
||||
max_waiting_per_account: 0,
|
||||
concurrency_wait_timeout_seconds: 30,
|
||||
|
||||
@@ -1245,7 +1245,7 @@ async fn calendar_query_rejects_an_anchor_outside_the_local_calendar() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn request_and_ops_queries_should_forward_deleted_group_id_sets() {
|
||||
async fn observability_queries_reject_group_filters_and_analytics_search() {
|
||||
use crate::admin::{AdminTestFixture, AdminTestState};
|
||||
use axum::{
|
||||
body::Body,
|
||||
@@ -1256,39 +1256,50 @@ async fn request_and_ops_queries_should_forward_deleted_group_id_sets() {
|
||||
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"] {
|
||||
for (path, parameter) in [
|
||||
("/api/admin/usage/records", "routingGroupRefs=group-a"),
|
||||
("/api/admin/operations/errors", "routingGroupRefs=group-a"),
|
||||
(
|
||||
"/api/admin/usage/records/summary",
|
||||
"routingGroupRefs=group-a",
|
||||
),
|
||||
(
|
||||
"/api/admin/usage/insights/overview",
|
||||
"routingGroupRefs=group-a",
|
||||
),
|
||||
(
|
||||
"/api/admin/usage/insights/diagnostics",
|
||||
"routingGroupRefs=group-a",
|
||||
),
|
||||
("/api/admin/usage/records/summary", "search=request-prefix"),
|
||||
(
|
||||
"/api/admin/usage/insights/overview",
|
||||
"search=request-prefix",
|
||||
),
|
||||
(
|
||||
"/api/admin/usage/insights/diagnostics",
|
||||
"search=request-prefix",
|
||||
),
|
||||
] {
|
||||
let response = router
|
||||
.clone()
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.uri(format!(
|
||||
"{path}?routingGroupRefs=deleted-b,deleted-a,deleted-b"
|
||||
))
|
||||
.uri(format!("{path}?{parameter}"))
|
||||
.header(header::COOKIE, "cpr_session=valid-session")
|
||||
.header("x-request-id", "req_group_filter")
|
||||
.header("x-request-id", "req_rejected_observability_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)
|
||||
response.status(),
|
||||
StatusCode::BAD_REQUEST,
|
||||
"{path}?{parameter}"
|
||||
);
|
||||
}
|
||||
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"]
|
||||
);
|
||||
assert!(observations.records.is_empty());
|
||||
assert!(observations.errors.is_empty());
|
||||
}
|
||||
|
||||
@@ -395,14 +395,7 @@ async fn published_fixture() -> AdminTestFixture {
|
||||
fixture.published_snapshot.publish(
|
||||
RuntimeSnapshot::new(
|
||||
ConfigRevision::new(7).unwrap(),
|
||||
gateway_core::settings::SettingsValues::new(
|
||||
1,
|
||||
0,
|
||||
"smart",
|
||||
Default::default(),
|
||||
None,
|
||||
None,
|
||||
),
|
||||
gateway_core::settings::SettingsValues::new(1, "smart", Default::default(), None, None),
|
||||
vec![],
|
||||
vec![],
|
||||
vec![],
|
||||
|
||||
@@ -50,7 +50,6 @@ fn update_body() -> Value {
|
||||
"refreshMarginSeconds": 1800,
|
||||
"refreshConcurrency": 4,
|
||||
"maxConcurrentPerAccount": 5,
|
||||
"requestIntervalMs": 25,
|
||||
"maxWaitingPerKey": 0,
|
||||
"maxWaitingPerAccount": 0,
|
||||
"openaiGuardianReservedConcurrency": 0,
|
||||
@@ -287,7 +286,6 @@ fn settings_response_should_cover_the_full_runtime_settings_contract() {
|
||||
refresh_margin_seconds: 1800,
|
||||
refresh_concurrency: 4,
|
||||
max_concurrent_per_account: 5,
|
||||
request_interval_ms: 25,
|
||||
max_waiting_per_key: 0,
|
||||
max_waiting_per_account: 0,
|
||||
concurrency_wait_timeout_seconds: 30,
|
||||
@@ -337,7 +335,6 @@ fn settings_response_should_cover_the_full_runtime_settings_contract() {
|
||||
"refreshMarginSeconds": 1800,
|
||||
"refreshConcurrency": 4,
|
||||
"maxConcurrentPerAccount": 5,
|
||||
"requestIntervalMs": 25,
|
||||
"maxWaitingPerKey": 0,
|
||||
"maxWaitingPerAccount": 0,
|
||||
"openaiGuardianReservedConcurrency": 0,
|
||||
@@ -415,7 +412,6 @@ fn settings_request_and_response_fields_should_stay_in_lockstep() {
|
||||
refresh_concurrency: u32::try_from(request.values.refresh_concurrency).expect("u32"),
|
||||
max_concurrent_per_account: u32::try_from(request.values.max_concurrent_per_account)
|
||||
.expect("u32"),
|
||||
request_interval_ms: request.values.request_interval_ms,
|
||||
max_waiting_per_key: 0,
|
||||
max_waiting_per_account: 0,
|
||||
concurrency_wait_timeout_seconds: 30,
|
||||
|
||||
@@ -23,6 +23,27 @@ fn session_cookie(response: &axum::response::Response) -> String {
|
||||
.to_owned()
|
||||
}
|
||||
|
||||
fn refresh_cookie(response: &axum::response::Response) -> String {
|
||||
response
|
||||
.headers()
|
||||
.get_all(header::SET_COOKIE)
|
||||
.iter()
|
||||
.filter_map(|value| value.to_str().ok())
|
||||
.find(|value| value.starts_with("cpr_refresh="))
|
||||
.expect("refresh cookie")
|
||||
.split(';')
|
||||
.next()
|
||||
.unwrap()
|
||||
.to_owned()
|
||||
}
|
||||
|
||||
async fn refresh(app: &axum::Router, cookie: &str) -> axum::response::Response {
|
||||
app.clone()
|
||||
.oneshot(cookie_request(Method::POST, "/api/auth/refresh", cookie))
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
async fn login(
|
||||
app: &axum::Router,
|
||||
body: Value,
|
||||
@@ -537,64 +558,129 @@ async fn key_auth_should_issue_restore_and_clear_a_unified_cookie() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn refresh_renews_admin_cookie_and_preserves_fixed_key_expiry() {
|
||||
async fn refresh_restores_expired_access_with_only_the_refresh_cookie() {
|
||||
let (app, fixture) = auth_app().await;
|
||||
for body in [
|
||||
let response = login(
|
||||
&app,
|
||||
json!({"mode": "admin", "password": "strong-admin-password"}),
|
||||
json!({"mode": "key", "apiKey": RAW_KEY}),
|
||||
] {
|
||||
let response = login(&app, body.clone(), None).await;
|
||||
let cookie = session_cookie(&response);
|
||||
let original = response_json(response).await["data"]["expiresAt"]
|
||||
.as_str()
|
||||
.unwrap()
|
||||
.to_owned();
|
||||
if body["mode"] == "admin" {
|
||||
fixture.set_session_expiry(
|
||||
cookie.split_once('=').unwrap().1,
|
||||
chrono::Utc::now() + chrono::Duration::minutes(1),
|
||||
);
|
||||
}
|
||||
let response = app
|
||||
.clone()
|
||||
.oneshot(cookie_request(Method::POST, "/api/auth/refresh", &cookie))
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
let access = session_cookie(&response);
|
||||
let renewal = refresh_cookie(&response);
|
||||
let cookies: Vec<_> = response
|
||||
.headers()
|
||||
.get_all(header::SET_COOKIE)
|
||||
.iter()
|
||||
.map(|value| value.to_str().unwrap())
|
||||
.collect();
|
||||
assert_eq!(cookies.len(), 2);
|
||||
assert!(cookies[0].contains("Path=/;"));
|
||||
assert!(cookies[1].contains("Path=/api/auth;"));
|
||||
assert!(
|
||||
cookies
|
||||
.iter()
|
||||
.all(|value| value.contains("HttpOnly") && value.contains("SameSite=Lax"))
|
||||
);
|
||||
assert_eq!(
|
||||
response_json(refresh(&app, &access).await).await["data"]["authenticated"],
|
||||
false
|
||||
);
|
||||
fixture.set_session_expiry(
|
||||
access.split_once('=').unwrap().1,
|
||||
chrono::Utc::now() - chrono::Duration::seconds(1),
|
||||
);
|
||||
assert_eq!(
|
||||
get(&app, "/api/admin/system/version", &access)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(response.headers()[header::CACHE_CONTROL], "no-store");
|
||||
assert_eq!(session_cookie(&response), cookie);
|
||||
assert!(
|
||||
response.headers()[header::SET_COOKIE]
|
||||
.to_str()
|
||||
.unwrap()
|
||||
.contains("HttpOnly")
|
||||
);
|
||||
let data = response_json(response).await;
|
||||
assert_eq!(data["data"]["authenticated"], true);
|
||||
let expiry = data["data"]["session"]["expiresAt"].as_str().unwrap();
|
||||
if body["mode"] == "key" {
|
||||
assert_eq!(expiry, original);
|
||||
} else {
|
||||
assert!(
|
||||
chrono::DateTime::parse_from_rfc3339(expiry).unwrap()
|
||||
> chrono::Utc::now() + chrono::Duration::minutes(1)
|
||||
);
|
||||
}
|
||||
}
|
||||
.status(),
|
||||
StatusCode::UNAUTHORIZED
|
||||
);
|
||||
assert_eq!(
|
||||
response_json(get(&app, "/api/auth/status", &renewal).await).await["data"]["authenticated"],
|
||||
false
|
||||
);
|
||||
|
||||
let response = refresh(&app, &renewal).await;
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(response.headers()[header::CACHE_CONTROL], "no-store");
|
||||
let next_access = session_cookie(&response);
|
||||
let next_renewal = refresh_cookie(&response);
|
||||
assert_ne!(next_access, access);
|
||||
assert_ne!(next_renewal, renewal);
|
||||
assert_eq!(response_json(response).await["data"]["authenticated"], true);
|
||||
assert_eq!(
|
||||
get(&app, "/api/admin/system/version", &next_access)
|
||||
.await
|
||||
.status(),
|
||||
StatusCode::OK
|
||||
);
|
||||
assert_eq!(
|
||||
get(&app, "/api/admin/system/version", &access)
|
||||
.await
|
||||
.status(),
|
||||
StatusCode::UNAUTHORIZED
|
||||
);
|
||||
// 模拟响应丢失后的重试,返回同一代 Cookie
|
||||
let retry = refresh(&app, &renewal).await;
|
||||
assert_eq!(session_cookie(&retry), next_access);
|
||||
assert_eq!(refresh_cookie(&retry), next_renewal);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn key_refresh_preserves_fixed_expiry_and_clears_admin_refresh_cookie() {
|
||||
let (app, _) = auth_app().await;
|
||||
let admin = login(
|
||||
&app,
|
||||
json!({"mode": "admin", "password": "strong-admin-password"}),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
let renewal = refresh_cookie(&admin);
|
||||
let response = login(
|
||||
&app,
|
||||
json!({"mode": "key", "apiKey": RAW_KEY}),
|
||||
Some(&renewal),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(refresh_cookie(&response), "cpr_refresh=");
|
||||
assert!(
|
||||
response
|
||||
.headers()
|
||||
.get_all(header::SET_COOKIE)
|
||||
.iter()
|
||||
.any(|value| {
|
||||
let value = value.to_str().unwrap();
|
||||
value.starts_with("cpr_refresh=")
|
||||
&& value.contains("Max-Age=0")
|
||||
&& value.contains("Path=/api/auth;")
|
||||
})
|
||||
);
|
||||
let cookie = session_cookie(&response);
|
||||
let original = response_json(response).await["data"]["expiresAt"].clone();
|
||||
let response = refresh(&app, &cookie).await;
|
||||
assert_eq!(session_cookie(&response), cookie);
|
||||
let data = response_json(response).await;
|
||||
assert_eq!(data["data"]["authenticated"], true);
|
||||
assert_eq!(data["data"]["session"]["expiresAt"], original);
|
||||
assert_eq!(
|
||||
response_json(refresh(&app, &renewal).await).await["data"]["authenticated"],
|
||||
false
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn refresh_after_logout_or_password_change_cannot_recreate_a_session() {
|
||||
let (app, _) = auth_app().await;
|
||||
let (app, fixture) = auth_app().await;
|
||||
for password_change in [false, true] {
|
||||
let cookie = session_cookie(
|
||||
&login(
|
||||
&app,
|
||||
json!({"mode": "admin", "password": "strong-admin-password"}),
|
||||
None,
|
||||
)
|
||||
.await,
|
||||
);
|
||||
let login_response = login(
|
||||
&app,
|
||||
json!({"mode": "admin", "password": "strong-admin-password"}),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
let cookie = session_cookie(&login_response);
|
||||
let renewal = refresh_cookie(&login_response);
|
||||
let response = if password_change {
|
||||
change_password(
|
||||
&app,
|
||||
@@ -604,15 +690,34 @@ async fn refresh_after_logout_or_password_change_cannot_recreate_a_session() {
|
||||
)
|
||||
.await
|
||||
} else {
|
||||
fixture.set_session_expiry(
|
||||
cookie.split_once('=').unwrap().1,
|
||||
chrono::Utc::now() - chrono::Duration::seconds(1),
|
||||
);
|
||||
app.clone()
|
||||
.oneshot(cookie_request(Method::POST, "/api/auth/logout", &cookie))
|
||||
.oneshot(cookie_request(Method::POST, "/api/auth/logout", &renewal))
|
||||
.await
|
||||
.unwrap()
|
||||
};
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response
|
||||
.headers()
|
||||
.get_all(header::SET_COOKIE)
|
||||
.iter()
|
||||
.count(),
|
||||
2
|
||||
);
|
||||
assert!(
|
||||
response
|
||||
.headers()
|
||||
.get_all(header::SET_COOKIE)
|
||||
.iter()
|
||||
.all(|value| value.to_str().unwrap().contains("Max-Age=0"))
|
||||
);
|
||||
let response = app
|
||||
.clone()
|
||||
.oneshot(cookie_request(Method::POST, "/api/auth/refresh", &cookie))
|
||||
.oneshot(cookie_request(Method::POST, "/api/auth/refresh", &renewal))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
//! HTTP 与 WebSocket 中间件测试入口及请求生命周期替身
|
||||
|
||||
mod session;
|
||||
mod websocket;
|
||||
|
||||
use std::{
|
||||
@@ -40,6 +41,9 @@ enum Mode {
|
||||
idle: bool,
|
||||
block_messages: bool,
|
||||
},
|
||||
SessionEnd {
|
||||
kind: Option<&'static str>,
|
||||
},
|
||||
Service,
|
||||
Echo,
|
||||
Slow {
|
||||
@@ -120,7 +124,7 @@ impl MiddlewarePlan for Plan {
|
||||
assert!(input.is_null());
|
||||
Box::pin(async move {
|
||||
let mut result = next.run(input).await?;
|
||||
result["request_interval_ms"] = serde_json::json!(4321);
|
||||
result["refresh_margin_seconds"] = serde_json::json!(4321);
|
||||
Ok(result)
|
||||
})
|
||||
}
|
||||
@@ -161,13 +165,43 @@ impl MiddlewarePlan for Plan {
|
||||
session.send(Message { kind: Kind::Text, payload: Bytes::from_static(b"ready") }).await?;
|
||||
let message = receive.await?.unwrap();
|
||||
session.send(message).await?;
|
||||
session.send(Message { kind: Kind::Close { code: Some(1000) }, payload: Bytes::from_static(b"done") }).await
|
||||
session.send(Message { kind: Kind::Close { code: Some(1000) }, payload: Bytes::from_static(b"done") }).await.map_err(Into::into)
|
||||
})));
|
||||
if reject {
|
||||
*response.status_mut() = StatusCode::FORBIDDEN;
|
||||
}
|
||||
Ok(response)
|
||||
}
|
||||
Mode::SessionEnd { kind } => {
|
||||
let upgrade = request
|
||||
.extensions_mut()
|
||||
.remove::<Arc<dyn contract::upgrade::WebSocketUpgrade>>()
|
||||
.unwrap();
|
||||
let (parts, _) = request.into_parts();
|
||||
let (mut response, session) = upgrade.accept(parts, Vec::new()).await?;
|
||||
response
|
||||
.extensions_mut()
|
||||
.insert(contract::upgrade::Upgraded::new(
|
||||
StatusCode::SWITCHING_PROTOCOLS,
|
||||
Box::pin(async move {
|
||||
let _session = session.await.map_err(|_| MiddlewareError::Fault)?;
|
||||
match kind {
|
||||
None => Ok(()),
|
||||
Some("cancelled") => {
|
||||
Err(contract::upgrade::SessionError::Cancelled)
|
||||
}
|
||||
Some("reported") => {
|
||||
Err(contract::upgrade::SessionError::Reported)
|
||||
}
|
||||
Some(kind) => Err(contract::upgrade::SessionError::Failed {
|
||||
kind,
|
||||
plugin_instance_id: Some("plugin-fixture".into()),
|
||||
}),
|
||||
}
|
||||
}),
|
||||
));
|
||||
Ok(response)
|
||||
}
|
||||
Mode::Rewrite => {
|
||||
assert_eq!(request.uri(), "/unrecognized/client/path?raw=%2f%FF");
|
||||
assert_eq!(request.headers()["authorization"], "Bearer fixture-secret");
|
||||
@@ -389,7 +423,7 @@ async fn admin_http_does_not_invoke_service_middleware_for_native_business_metho
|
||||
.await
|
||||
.unwrap();
|
||||
let result: serde_json::Value = serde_json::from_slice(&bytes).unwrap();
|
||||
assert_ne!(result["data"]["requestIntervalMs"], 4321);
|
||||
assert_ne!(result["data"]["refreshMarginSeconds"], 4321);
|
||||
assert_eq!(plan.calls.load(Ordering::SeqCst), 0);
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,79 @@
|
||||
//! 验证升级后终态只记录一次安全分类,正常结束及已报告失败不重复输出
|
||||
|
||||
use super::*;
|
||||
use tracing::{
|
||||
Event, Metadata, Subscriber,
|
||||
span::{Attributes, Id, Record},
|
||||
};
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
struct Recorder(Arc<Mutex<Vec<String>>>);
|
||||
|
||||
impl Subscriber for Recorder {
|
||||
fn enabled(&self, _: &Metadata<'_>) -> bool {
|
||||
true
|
||||
}
|
||||
fn new_span(&self, _: &Attributes<'_>) -> Id {
|
||||
Id::from_u64(1)
|
||||
}
|
||||
fn record(&self, _: &Id, _: &Record<'_>) {}
|
||||
fn record_follows_from(&self, _: &Id, _: &Id) {}
|
||||
fn enter(&self, _: &Id) {}
|
||||
fn exit(&self, _: &Id) {}
|
||||
fn event(&self, event: &Event<'_>) {
|
||||
struct Fields(String);
|
||||
impl tracing::field::Visit for Fields {
|
||||
fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
|
||||
use std::fmt::Write as _;
|
||||
write!(&mut self.0, "{}={value:?} ", field.name()).unwrap();
|
||||
}
|
||||
}
|
||||
let mut fields = Fields(String::new());
|
||||
event.record(&mut fields);
|
||||
if fields.0.contains("middleware.session.failed") {
|
||||
self.0.lock().unwrap().push(fields.0);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn upgraded_session_records_one_correlated_failure_and_ignores_expected_endings() {
|
||||
use tokio_tungstenite::{connect_async, tungstenite::client::IntoClientRequest as _};
|
||||
for kind in [
|
||||
None,
|
||||
Some("cancelled"),
|
||||
Some("reported"),
|
||||
Some("permission_denied"),
|
||||
] {
|
||||
let recorder = Recorder::default();
|
||||
let _subscriber = tracing::subscriber::set_default(recorder.clone());
|
||||
let (router, plan) = router(Mode::SessionEnd { kind }, None).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, router).await.unwrap();
|
||||
});
|
||||
let mut request = format!("ws://{address}/plugin-owned/session")
|
||||
.into_client_request()
|
||||
.unwrap();
|
||||
request
|
||||
.headers_mut()
|
||||
.insert("x-request-id", "session-regression".parse().unwrap());
|
||||
let (_socket, response) = connect_async(request).await.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::SWITCHING_PROTOCOLS);
|
||||
let cancellation = plan.cancellation.lock().unwrap().clone().unwrap();
|
||||
tokio::time::timeout(Duration::from_secs(3), cancellation.cancelled())
|
||||
.await
|
||||
.unwrap();
|
||||
let events = recorder.0.lock().unwrap();
|
||||
if kind == Some("permission_denied") {
|
||||
assert_eq!(events.len(), 1, "{events:?}");
|
||||
assert!(events[0].contains("session-regression"));
|
||||
assert!(events[0].contains("plugin-fixture"));
|
||||
assert!(events[0].contains("permission_denied"));
|
||||
} else {
|
||||
assert!(events.is_empty(), "{events:?}");
|
||||
}
|
||||
server.abort();
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
@@ -308,7 +308,7 @@ fn snapshot_with_accounts(
|
||||
);
|
||||
RuntimeSnapshot::new(
|
||||
ConfigRevision::new(1).expect("revision"),
|
||||
gateway_core::settings::SettingsValues::new(2, 1, "smart", Default::default(), None, None),
|
||||
gateway_core::settings::SettingsValues::new(2, "smart", Default::default(), None, None),
|
||||
vec![provider.clone()],
|
||||
["model-a", "model-b"]
|
||||
.into_iter()
|
||||
|
||||
@@ -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]);
|
||||
}
|
||||
}
|
||||
|
||||
+4
-4
@@ -221,7 +221,7 @@ async fn turns_refresh_host_settings_and_key_identity_while_preserving_handshake
|
||||
.unwrap();
|
||||
let baseline = settings.runtime.as_ref().unwrap();
|
||||
let mut values = serde_json::to_value(baseline.values()).unwrap();
|
||||
values["request_interval_ms"] = json!(0);
|
||||
values["max_account_rotations"] = json!(0);
|
||||
settings.runtime = Some(
|
||||
baseline
|
||||
.replace(serde_json::from_value(values).unwrap(), "handshake")
|
||||
@@ -292,7 +292,7 @@ async fn turns_refresh_host_settings_and_key_identity_while_preserving_handshake
|
||||
for round in 0..3 {
|
||||
if round == 1 {
|
||||
let mut settings = serde_json::to_value(initial.settings()).unwrap();
|
||||
settings["request_interval_ms"] = json!(50);
|
||||
settings["max_account_rotations"] = json!(5);
|
||||
settings["responses_max_decompressed_body_bytes"] = json!(2048);
|
||||
snapshots.publish(
|
||||
initial
|
||||
@@ -337,8 +337,8 @@ async fn turns_refresh_host_settings_and_key_identity_while_preserving_handshake
|
||||
observed[1]["runtime"]["responses_max_decompressed_body_bytes"],
|
||||
2048
|
||||
);
|
||||
assert_eq!(observed[0]["runtime"]["request_interval_ms"], 0);
|
||||
assert_eq!(observed[1]["runtime"]["request_interval_ms"], 0);
|
||||
assert_eq!(observed[0]["runtime"]["max_account_rotations"], 0);
|
||||
assert_eq!(observed[1]["runtime"]["max_account_rotations"], 0);
|
||||
}
|
||||
socket.close(None).await.unwrap();
|
||||
server.abort();
|
||||
|
||||
@@ -89,7 +89,6 @@ pub struct AccountSelectionPolicy {
|
||||
strategy: RotationStrategy,
|
||||
smart_scheduling: SmartSchedulingConfig,
|
||||
max_concurrent_per_account: AccountConcurrency,
|
||||
request_interval: Duration,
|
||||
queue_policy: ConcurrencyQueuePolicy,
|
||||
openai_guardian_reserved_concurrency: u32,
|
||||
openai_account_affinity: AccountAffinity,
|
||||
@@ -102,13 +101,11 @@ impl AccountSelectionPolicy {
|
||||
pub fn new(
|
||||
strategy: RotationStrategy,
|
||||
max_concurrent_per_account: impl Into<AccountConcurrency>,
|
||||
request_interval: Duration,
|
||||
) -> Self {
|
||||
Self {
|
||||
strategy,
|
||||
smart_scheduling: SmartSchedulingConfig::default(),
|
||||
max_concurrent_per_account: max_concurrent_per_account.into(),
|
||||
request_interval,
|
||||
openai_guardian_reserved_concurrency: 0,
|
||||
openai_account_affinity: AccountAffinity::default(),
|
||||
max_account_rotations: 3,
|
||||
@@ -196,11 +193,6 @@ impl AccountSelectionPolicy {
|
||||
pub const fn max_concurrent_per_account(self) -> AccountConcurrency {
|
||||
self.max_concurrent_per_account
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn request_interval(self) -> Duration {
|
||||
self.request_interval
|
||||
}
|
||||
}
|
||||
|
||||
/// Store 提供并发事实,Provider 叠加自己解释的额度事实;全部信号均可重建
|
||||
@@ -543,8 +535,8 @@ impl AccountSelectionContext {
|
||||
/// 选择账号时是否执行本地调度资格投影
|
||||
///
|
||||
/// 管理端对指定账号执行诊断时,会直接向上游确认实际状态;该模式跳过
|
||||
/// `enabled`、可用性、冷却和 token 到期等本地投影,仍保留租约、并发和
|
||||
/// 请求间隔约束,且只能与固定账号约束组合使用
|
||||
/// `enabled`、可用性、冷却和 token 到期等本地投影,仍保留租约和并发
|
||||
/// 约束,且只能与固定账号约束组合使用
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
||||
pub enum AccountEligibilityPolicy {
|
||||
#[default]
|
||||
@@ -566,7 +558,6 @@ pub enum AccountSchedulingBlocker {
|
||||
LocalAvailability,
|
||||
Excluded,
|
||||
ConcurrencyLimit,
|
||||
RequestInterval,
|
||||
LowerWeight,
|
||||
}
|
||||
|
||||
@@ -791,7 +782,7 @@ impl AccountSelector {
|
||||
}
|
||||
}
|
||||
|
||||
/// 只有本地并发/调度间隔可等待;账号权限、失效、额度与上游冷却仍立即排除
|
||||
/// 只有本地并发可等待;账号权限、失效、额度与上游冷却仍立即排除
|
||||
#[must_use]
|
||||
pub fn wait_candidates(
|
||||
&self,
|
||||
@@ -803,10 +794,7 @@ impl AccountSelector {
|
||||
.filter(|candidate| {
|
||||
matches!(
|
||||
self.scheduling_blocker(candidate, context),
|
||||
None | Some(
|
||||
AccountSchedulingBlocker::ConcurrencyLimit
|
||||
| AccountSchedulingBlocker::RequestInterval
|
||||
)
|
||||
None | Some(AccountSchedulingBlocker::ConcurrencyLimit)
|
||||
)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
@@ -879,18 +867,6 @@ impl AccountSelector {
|
||||
{
|
||||
return Some(AccountSchedulingBlocker::ConcurrencyLimit);
|
||||
}
|
||||
if candidate
|
||||
.signals
|
||||
.last_started_at
|
||||
.is_some_and(|last_started| {
|
||||
!context
|
||||
.now
|
||||
.duration_since(last_started)
|
||||
.is_ok_and(|elapsed| elapsed >= context.policy.request_interval())
|
||||
})
|
||||
{
|
||||
return Some(AccountSchedulingBlocker::RequestInterval);
|
||||
}
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
@@ -219,7 +219,7 @@ where
|
||||
|
||||
/// 对固定账号执行诊断请求
|
||||
///
|
||||
/// 仅跳过该账号的本地可用性投影;账号租约、并发和请求间隔仍必须满足
|
||||
/// 仅跳过该账号的本地可用性投影;账号租约和并发仍必须满足
|
||||
pub async fn start_diagnostic(
|
||||
&self,
|
||||
request: NewModelRequest,
|
||||
@@ -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,7 @@ 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,
|
||||
@@ -483,7 +484,9 @@ 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>,
|
||||
@@ -521,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()
|
||||
@@ -552,10 +560,8 @@ where
|
||||
match self.pull().await? {
|
||||
PullOutcome::Events(events) => {
|
||||
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.downstream_committed_at.is_some() {
|
||||
let requirement = if self.delivery_committed {
|
||||
CommitRequirement::AlreadyCommitted
|
||||
} else {
|
||||
self.delivery_pending = true;
|
||||
@@ -569,7 +575,7 @@ 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;
|
||||
@@ -581,7 +587,7 @@ where
|
||||
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);
|
||||
@@ -597,7 +603,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() {
|
||||
@@ -636,14 +642,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",
|
||||
@@ -656,7 +663,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 {
|
||||
@@ -669,7 +676,7 @@ 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()
|
||||
{
|
||||
@@ -692,7 +699,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,
|
||||
@@ -1664,7 +1674,7 @@ 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:?}"),
|
||||
}));
|
||||
@@ -1963,8 +1973,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,
|
||||
@@ -2012,7 +2025,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
|
||||
@@ -2037,13 +2050,21 @@ 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"),
|
||||
GatewayError::new(
|
||||
GatewayErrorKind::Cancelled,
|
||||
self.cancellation
|
||||
.reason()
|
||||
.unwrap_or(crate::lifecycle::CancellationReason::Requested)
|
||||
.message(),
|
||||
),
|
||||
),
|
||||
EngineError::Deadline => (
|
||||
if self.downstream_committed_at.is_some() {
|
||||
if self.response_commit.committed_at().is_some() {
|
||||
ExecutionOutcome::Incomplete
|
||||
} else {
|
||||
ExecutionOutcome::Failed
|
||||
@@ -2057,7 +2078,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()
|
||||
@@ -2128,8 +2149,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),
|
||||
@@ -2162,6 +2186,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,
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
|
||||
@@ -8,7 +8,42 @@ use crate::{engine::middleware::MiddlewareError, engine::middleware::websocket::
|
||||
|
||||
pub type PendingSession =
|
||||
Shared<BoxFuture<'static, Result<Arc<dyn Session>, Arc<MiddlewareError>>>>;
|
||||
pub type SessionTask = BoxFuture<'static, Result<(), MiddlewareError>>;
|
||||
pub type SessionTask = BoxFuture<'static, Result<(), SessionError>>;
|
||||
|
||||
/// 升级后的终态由传输 owner 消费,协议适配器只提供安全分类与诊断归属
|
||||
#[derive(Debug)]
|
||||
pub enum SessionError {
|
||||
/// 预期取消,不作为异常终态输出
|
||||
Cancelled,
|
||||
/// 下层已记录同一失败,传输 owner 不重复输出
|
||||
Reported,
|
||||
Failed {
|
||||
kind: &'static str,
|
||||
plugin_instance_id: Option<String>,
|
||||
},
|
||||
}
|
||||
|
||||
impl From<MiddlewareError> for SessionError {
|
||||
fn from(error: MiddlewareError) -> Self {
|
||||
use crate::error::GatewayErrorKind;
|
||||
if matches!(&error, MiddlewareError::Gateway(error) if error.kind() == GatewayErrorKind::Cancelled)
|
||||
|| matches!(
|
||||
error,
|
||||
MiddlewareError::Engine(crate::engine::EngineError::Cancelled)
|
||||
)
|
||||
{
|
||||
return Self::Cancelled;
|
||||
}
|
||||
Self::Failed {
|
||||
kind: if error.is_rejected() {
|
||||
"rejected"
|
||||
} else {
|
||||
"middleware"
|
||||
},
|
||||
plugin_instance_id: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 具体服务器负责握手与 socket;Core 不依赖 HTTP 服务器实现
|
||||
pub trait WebSocketUpgrade: Send + Sync {
|
||||
|
||||
@@ -559,12 +559,10 @@ 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()
|
||||
self.request.middleware.is_none()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
|
||||
@@ -1,16 +1,69 @@
|
||||
//! 请求与连接的截止、租约、取消和 drain 契约
|
||||
|
||||
use std::fmt;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::{Arc, OnceLock};
|
||||
use std::time::{Duration, SystemTime};
|
||||
|
||||
use event_listener::Event;
|
||||
use futures::future::{BoxFuture, pending};
|
||||
use futures_timer::Delay;
|
||||
|
||||
/// 参与业务并发约束的请求租约类别,观测续期不属于此合同
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum RequestLeaseKind {
|
||||
ClientAdmission,
|
||||
Credential,
|
||||
}
|
||||
|
||||
/// 续期最终结束的原因,不包含可恢复的中间失败
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum LeaseFailure {
|
||||
Lost,
|
||||
RenewalTimedOut,
|
||||
Expired,
|
||||
}
|
||||
|
||||
/// 只携带稳定原因,不保存凭据、Redis key 或任意基础设施错误文本
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum CancellationReason {
|
||||
Requested,
|
||||
Lease {
|
||||
kind: RequestLeaseKind,
|
||||
failure: LeaseFailure,
|
||||
},
|
||||
}
|
||||
|
||||
impl CancellationReason {
|
||||
#[must_use]
|
||||
pub const fn message(self) -> &'static str {
|
||||
match self {
|
||||
Self::Requested => "request was cancelled",
|
||||
Self::Lease {
|
||||
kind: RequestLeaseKind::ClientAdmission,
|
||||
failure,
|
||||
} => match failure {
|
||||
LeaseFailure::Lost => "request was cancelled: client admission lease was lost",
|
||||
LeaseFailure::RenewalTimedOut => {
|
||||
"request was cancelled: client admission lease renewal timed out"
|
||||
}
|
||||
LeaseFailure::Expired => "request was cancelled: client admission lease expired",
|
||||
},
|
||||
Self::Lease {
|
||||
kind: RequestLeaseKind::Credential,
|
||||
failure,
|
||||
} => match failure {
|
||||
LeaseFailure::Lost => "request was cancelled: credential lease was lost",
|
||||
LeaseFailure::RenewalTimedOut => {
|
||||
"request was cancelled: credential lease renewal timed out"
|
||||
}
|
||||
LeaseFailure::Expired => "request was cancelled: credential lease expired",
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct CancellationState {
|
||||
cancelled: AtomicBool,
|
||||
reason: OnceLock<CancellationReason>,
|
||||
event: Event,
|
||||
}
|
||||
|
||||
@@ -41,7 +94,7 @@ impl CancellationToken {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
state: Arc::new(CancellationState {
|
||||
cancelled: AtomicBool::new(false),
|
||||
reason: OnceLock::new(),
|
||||
event: Event::new(),
|
||||
}),
|
||||
ancestors: Arc::from([]),
|
||||
@@ -56,7 +109,7 @@ impl CancellationToken {
|
||||
.collect::<Vec<_>>();
|
||||
Self {
|
||||
state: Arc::new(CancellationState {
|
||||
cancelled: AtomicBool::new(false),
|
||||
reason: OnceLock::new(),
|
||||
event: Event::new(),
|
||||
}),
|
||||
ancestors: ancestors.into(),
|
||||
@@ -64,19 +117,29 @@ impl CancellationToken {
|
||||
}
|
||||
|
||||
pub fn cancel(&self) {
|
||||
if self.state.cancelled.swap(true, Ordering::AcqRel) {
|
||||
self.cancel_with_reason(CancellationReason::Requested);
|
||||
}
|
||||
|
||||
/// 首次取消确定原因,后续资源回收不覆盖已观察到的取消事实
|
||||
pub fn cancel_with_reason(&self, reason: CancellationReason) {
|
||||
if self.is_cancelled() || self.state.reason.set(reason).is_err() {
|
||||
return;
|
||||
}
|
||||
self.state.event.notify(usize::MAX);
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn is_cancelled(&self) -> bool {
|
||||
self.state.cancelled.load(Ordering::Acquire)
|
||||
|| self
|
||||
.ancestors
|
||||
pub fn reason(&self) -> Option<CancellationReason> {
|
||||
self.state.reason.get().copied().or_else(|| {
|
||||
self.ancestors
|
||||
.iter()
|
||||
.any(|state| state.cancelled.load(Ordering::Acquire))
|
||||
.find_map(|state| state.reason.get().copied())
|
||||
})
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn is_cancelled(&self) -> bool {
|
||||
self.reason().is_some()
|
||||
}
|
||||
|
||||
pub async fn cancelled(&self) {
|
||||
@@ -92,13 +155,13 @@ impl CancellationToken {
|
||||
|
||||
fn wait_for_state(state: Arc<CancellationState>) -> BoxFuture<'static, ()> {
|
||||
Box::pin(async move {
|
||||
if state.cancelled.load(Ordering::Acquire) {
|
||||
if state.reason.get().is_some() {
|
||||
return;
|
||||
}
|
||||
// listener 随等待 future 丢弃而注销,避免 select 败选分支在自身及祖先状态中积累
|
||||
let listener = state.event.listen();
|
||||
// 注册后复查,覆盖取消发生在首次检查与注册之间的竞态
|
||||
if state.cancelled.load(Ordering::Acquire) {
|
||||
if state.reason.get().is_some() {
|
||||
return;
|
||||
}
|
||||
listener.await;
|
||||
|
||||
@@ -69,9 +69,11 @@ impl ModelPriceOverride {
|
||||
band.as_str(),
|
||||
"standard"
|
||||
| "fast"
|
||||
| "ultrafast"
|
||||
| "flex"
|
||||
| "long_standard"
|
||||
| "long_fast"
|
||||
| "long_ultrafast"
|
||||
| "long_flex"
|
||||
| "image"
|
||||
)
|
||||
|
||||
@@ -55,7 +55,6 @@ pub struct ProviderSchedulingLeaseRequest {
|
||||
credential_revision: CredentialRevision,
|
||||
max_concurrent: AccountConcurrency,
|
||||
concurrency_pool: ProviderConcurrencyPool,
|
||||
request_interval: Duration,
|
||||
deadline: crate::lifecycle::Deadline,
|
||||
cancellation: crate::lifecycle::CancellationToken,
|
||||
}
|
||||
@@ -67,7 +66,6 @@ impl ProviderSchedulingLeaseRequest {
|
||||
account_id: ProviderAccountId,
|
||||
credential_revision: CredentialRevision,
|
||||
max_concurrent: impl Into<AccountConcurrency>,
|
||||
request_interval: Duration,
|
||||
deadline: impl Into<crate::lifecycle::Deadline>,
|
||||
) -> Self {
|
||||
Self {
|
||||
@@ -76,7 +74,6 @@ impl ProviderSchedulingLeaseRequest {
|
||||
credential_revision,
|
||||
max_concurrent: max_concurrent.into(),
|
||||
concurrency_pool: ProviderConcurrencyPool::Shared,
|
||||
request_interval,
|
||||
deadline: deadline.into(),
|
||||
cancellation: crate::lifecycle::CancellationToken::new(),
|
||||
}
|
||||
@@ -124,11 +121,6 @@ impl ProviderSchedulingLeaseRequest {
|
||||
self.max_concurrent
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn request_interval(&self) -> Duration {
|
||||
self.request_interval
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn deadline(&self) -> crate::lifecycle::Deadline {
|
||||
self.deadline
|
||||
@@ -166,7 +158,7 @@ pub enum ProviderLeaseRequest {
|
||||
}
|
||||
|
||||
pub trait ProviderLeasePort: Send + Sync {
|
||||
/// 按目标容量池读取在途数,账号最小请求间隔仍跨池共享
|
||||
/// 按目标容量池读取在途数,最近调度时间仍跨池共享
|
||||
fn load_state<'a>(
|
||||
&'a self,
|
||||
client_api_key_id: &'a ClientApiKeyId,
|
||||
|
||||
@@ -276,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
|
||||
}
|
||||
|
||||
/// 配置提交复用同一扩展集合的已发布目录,目录留待对账刷新
|
||||
@@ -284,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<(
|
||||
@@ -338,6 +347,7 @@ impl RuntimeSnapshotCompiler {
|
||||
catalogs.as_ref(),
|
||||
provider_kinds,
|
||||
cached,
|
||||
refresh_catalog,
|
||||
self.privacy_compiler.clone(),
|
||||
)
|
||||
.await?;
|
||||
@@ -354,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
|
||||
@@ -373,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());
|
||||
}
|
||||
@@ -403,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) {
|
||||
@@ -946,7 +962,7 @@ impl RuntimeSnapshot {
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// 完整目录按映射后的模型判定;发现型或不可用目录不作为能力白名单
|
||||
/// 已知模型只属于明确声明它的 Provider;未知模型仍可由非完整目录透传
|
||||
#[must_use]
|
||||
pub fn contains_public_model_for_provider(
|
||||
&self,
|
||||
@@ -962,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]
|
||||
@@ -1101,12 +1126,28 @@ impl RuntimeSnapshot {
|
||||
) -> Result<RoutingPlan, RoutingError> {
|
||||
let requirements = operation.capability_requirements();
|
||||
let mut candidates = Vec::new();
|
||||
let mut unlisted_candidates = Vec::new();
|
||||
|
||||
if context.required_provider.is_none() && account_scope.provider_kinds().is_empty() {
|
||||
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()]),
|
||||
@@ -1129,20 +1170,6 @@ 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())
|
||||
@@ -1160,28 +1187,30 @@ impl RuntimeSnapshot {
|
||||
};
|
||||
emulated
|
||||
}
|
||||
None if self.exhaustive_provider_catalogs.contains(provider) => continue,
|
||||
None if known_model || self.exhaustive_provider_catalogs.contains(provider) => {
|
||||
continue;
|
||||
}
|
||||
None => BTreeSet::new(),
|
||||
};
|
||||
let candidate = ProviderCandidate {
|
||||
candidates.push(ProviderCandidate {
|
||||
provider: provider.clone(),
|
||||
model_presentation: self
|
||||
.provider_model_presentations
|
||||
.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),
|
||||
};
|
||||
if capabilities.is_some() {
|
||||
candidates.push(candidate);
|
||||
} else {
|
||||
unlisted_candidates.push(candidate);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
// 冷启动或目录不可读时不能按 Provider 名称猜测模型归属
|
||||
if !known_model && candidates.len() > 1 {
|
||||
return Err(RoutingError::NoCapableProvider {
|
||||
model: public_model.as_str().to_owned(),
|
||||
});
|
||||
}
|
||||
// 已知模型优先;发现型目录的未知模型仍可透传,但不能抢在明确匹配的 Provider 前面
|
||||
candidates.extend(unlisted_candidates);
|
||||
|
||||
if candidates.is_empty() {
|
||||
// 模型存在性必须包含被请求路由限制排除的 Provider;目录未知时不能断言模型不存在
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -73,7 +73,6 @@ impl CompiledSettings {
|
||||
let account_selection_policy = AccountSelectionPolicy::new(
|
||||
strategy,
|
||||
AccountConcurrency::new(settings.max_concurrent_per_account),
|
||||
Duration::from_millis(settings.request_interval_ms),
|
||||
)
|
||||
.with_openai_guardian_reserved_concurrency(settings.openai_guardian_reserved_concurrency)
|
||||
.with_openai_account_affinity(settings.openai_account_affinity)
|
||||
|
||||
@@ -22,7 +22,6 @@ pub struct SettingsValues {
|
||||
pub(crate) max_account_rotations: u32,
|
||||
pub(crate) openai_session_affinity_ttl_hours: u32,
|
||||
pub(crate) responses_max_decompressed_body_bytes: u64,
|
||||
pub(crate) request_interval_ms: u64,
|
||||
pub(crate) smart_scheduling: crate::account::SmartSchedulingConfig,
|
||||
pub(crate) rotation_strategy: String,
|
||||
pub(crate) model_mappings: Arc<BTreeMap<String, String>>,
|
||||
@@ -129,7 +128,6 @@ impl SettingsValues {
|
||||
#[must_use]
|
||||
pub fn new(
|
||||
max_concurrent_per_account: u32,
|
||||
request_interval_ms: u64,
|
||||
rotation_strategy: impl Into<String>,
|
||||
model_mappings: BTreeMap<String, String>,
|
||||
min_codex_desktop_version: Option<String>,
|
||||
@@ -150,7 +148,6 @@ impl SettingsValues {
|
||||
max_account_rotations: 3,
|
||||
openai_session_affinity_ttl_hours: 24,
|
||||
responses_max_decompressed_body_bytes: 64 * 1024 * 1024,
|
||||
request_interval_ms,
|
||||
smart_scheduling: crate::account::SmartSchedulingConfig::default(),
|
||||
rotation_strategy: rotation_strategy.into(),
|
||||
model_mappings: Arc::new(model_mappings),
|
||||
|
||||
@@ -99,11 +99,7 @@ fn account_scheduling_should_default_and_resolve_concurrency_override() {
|
||||
|
||||
fn context(strategy: RotationStrategy) -> AccountSelectionContext {
|
||||
AccountSelectionContext {
|
||||
policy: AccountSelectionPolicy::new(
|
||||
strategy,
|
||||
NonZeroU32::new(3).expect("positive"),
|
||||
Duration::ZERO,
|
||||
),
|
||||
policy: AccountSelectionPolicy::new(strategy, NonZeroU32::new(3).expect("positive")),
|
||||
now: SystemTime::now(),
|
||||
excluded_accounts: BTreeSet::new(),
|
||||
preferred_account: None,
|
||||
@@ -471,7 +467,6 @@ fn highest_priority_affinity_should_still_obey_existing_scheduling_constraints()
|
||||
for blocker in [
|
||||
AccountSchedulingBlocker::Excluded,
|
||||
AccountSchedulingBlocker::ConcurrencyLimit,
|
||||
AccountSchedulingBlocker::RequestInterval,
|
||||
AccountSchedulingBlocker::LocalAvailability,
|
||||
] {
|
||||
let mut candidates = [
|
||||
@@ -488,14 +483,6 @@ fn highest_priority_affinity_should_still_obey_existing_scheduling_constraints()
|
||||
.insert(candidates[0].account.id().clone());
|
||||
}
|
||||
AccountSchedulingBlocker::ConcurrencyLimit => candidates[0].signals.in_flight = 3,
|
||||
AccountSchedulingBlocker::RequestInterval => {
|
||||
candidates[0].signals.last_started_at = Some(selection.now);
|
||||
selection.policy = AccountSelectionPolicy::new(
|
||||
RotationStrategy::Smart,
|
||||
NonZeroU32::new(3).expect("limit"),
|
||||
Duration::from_secs(1),
|
||||
);
|
||||
}
|
||||
AccountSchedulingBlocker::LocalAvailability => {
|
||||
candidates[0].signals.cooldown =
|
||||
Some((selection.now + Duration::from_secs(60)).into())
|
||||
@@ -788,25 +775,25 @@ fn selector_should_reject_account_at_concurrency_limit() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selector_should_enforce_request_interval_until_boundary() {
|
||||
fn selector_should_allow_immediate_reuse_regardless_of_last_started_time() {
|
||||
let now = SystemTime::now();
|
||||
let mut cooling = candidate("acct_cooling", 0, Some(100));
|
||||
cooling.signals.last_started_at = Some(now - Duration::from_millis(9));
|
||||
let mut ready = candidate("acct_ready", 0, Some(1));
|
||||
ready.signals.last_started_at = Some(now - Duration::from_millis(10));
|
||||
let candidates = vec![cooling, ready];
|
||||
let mut context = context(RotationStrategy::Smart);
|
||||
context.now = now;
|
||||
context.policy = AccountSelectionPolicy::new(
|
||||
RotationStrategy::Smart,
|
||||
NonZeroU32::new(3).expect("positive"),
|
||||
Duration::from_millis(10),
|
||||
);
|
||||
let selected = AccountSelector
|
||||
.select(&candidates, &context)
|
||||
.expect("candidate available");
|
||||
|
||||
assert_eq!(selected.candidate().account.id().as_str(), "acct_ready");
|
||||
for last_started in [
|
||||
now - Duration::from_millis(1),
|
||||
now,
|
||||
now + Duration::from_secs(30),
|
||||
] {
|
||||
let mut candidate = candidate("acct_recent", 1, None);
|
||||
candidate.signals.last_started_at = Some(last_started);
|
||||
let candidates = [candidate];
|
||||
let mut context = context(RotationStrategy::Sticky);
|
||||
context.now = now;
|
||||
context.preferred_account = Some(candidates[0].account.id().clone());
|
||||
let selected = AccountSelector
|
||||
.select(&candidates, &context)
|
||||
.expect("available slot");
|
||||
assert_eq!(selected.candidate().account.id().as_str(), "acct_recent");
|
||||
assert_eq!(selected.preferred(), PreferredAccountSelection::Hit);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -1136,7 +1123,6 @@ fn reserved_capacity_stays_bounded_when_normal_concurrency_is_unlimited() {
|
||||
normal.policy = AccountSelectionPolicy::new(
|
||||
RotationStrategy::Smart,
|
||||
gateway_core::account::AccountConcurrency::Unlimited,
|
||||
Duration::ZERO,
|
||||
);
|
||||
assert!(AccountSelector.select(&candidates, &normal).is_some());
|
||||
assert!(
|
||||
|
||||
@@ -13,7 +13,7 @@ use super::{candidate, candidate_with_concurrency, context};
|
||||
const FAILURE_RATE_HALF_LIFE: Duration = Duration::from_secs(15 * 60);
|
||||
|
||||
#[test]
|
||||
fn unlimited_account_concurrency_preserves_overrides_intervals_and_finite_scores() {
|
||||
fn unlimited_account_concurrency_preserves_overrides_and_finite_scores() {
|
||||
use gateway_core::account::{AccountConcurrency, AccountSelectionPolicy};
|
||||
|
||||
for strategy in [
|
||||
@@ -23,8 +23,7 @@ fn unlimited_account_concurrency_preserves_overrides_intervals_and_finite_scores
|
||||
RotationStrategy::QuotaResetPriority,
|
||||
] {
|
||||
let mut context = context(strategy);
|
||||
context.policy =
|
||||
AccountSelectionPolicy::new(strategy, AccountConcurrency::Unlimited, Duration::ZERO);
|
||||
context.policy = AccountSelectionPolicy::new(strategy, AccountConcurrency::Unlimited);
|
||||
let candidates = [
|
||||
candidate("acct_unlimited", u32::MAX, None),
|
||||
candidate_with_concurrency("acct_limited", 2, 2),
|
||||
@@ -60,14 +59,9 @@ fn unlimited_account_concurrency_preserves_overrides_intervals_and_finite_scores
|
||||
.is_finite()
|
||||
);
|
||||
}
|
||||
context.policy = AccountSelectionPolicy::new(
|
||||
strategy,
|
||||
AccountConcurrency::Unlimited,
|
||||
Duration::from_secs(1),
|
||||
);
|
||||
let mut recent = candidate("acct_unlimited", u32::MAX, None);
|
||||
recent.signals.last_started_at = Some(context.now);
|
||||
assert!(AccountSelector.select(&[recent], &context).is_none());
|
||||
assert!(AccountSelector.select(&[recent], &context).is_some());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -77,11 +71,8 @@ fn unavailable_unlimited_account_does_not_hide_the_finite_pool_capacity() {
|
||||
AccountConcurrency, AccountSelectionPolicy, CredentialState, QuotaState,
|
||||
};
|
||||
let mut context = context(RotationStrategy::Smart);
|
||||
context.policy = AccountSelectionPolicy::new(
|
||||
RotationStrategy::Smart,
|
||||
AccountConcurrency::Unlimited,
|
||||
Duration::ZERO,
|
||||
);
|
||||
context.policy =
|
||||
AccountSelectionPolicy::new(RotationStrategy::Smart, AccountConcurrency::Unlimited);
|
||||
let mut unavailable = candidate("acct_disabled", 0, None);
|
||||
unavailable.account = unavailable.account.with_account_facts(
|
||||
false,
|
||||
|
||||
@@ -53,6 +53,7 @@ use serde_json::{Map, Value, json};
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
struct FinalState {
|
||||
error_message: Option<String>,
|
||||
error_kind: Option<GatewayErrorKind>,
|
||||
diagnostic_trace_json: Option<String>,
|
||||
outcome: ExecutionOutcome,
|
||||
@@ -228,6 +229,10 @@ impl ExecutionStore for FakeStore {
|
||||
.expect("store lock")
|
||||
.finalizations
|
||||
.push(FinalState {
|
||||
error_message: finalization
|
||||
.error
|
||||
.as_ref()
|
||||
.map(|error| error.safe_message().to_owned()),
|
||||
error_kind: finalization.error.as_ref().map(|error| error.kind()),
|
||||
diagnostic_trace_json: finalization.diagnostic_trace_json,
|
||||
outcome: finalization.outcome,
|
||||
@@ -671,7 +676,6 @@ fn plan(operation: &Operation) -> RoutingPlan {
|
||||
AccountSelectionPolicy::new(
|
||||
RotationStrategy::Smart,
|
||||
NonZeroU32::new(2).expect("account concurrency"),
|
||||
Duration::from_millis(50),
|
||||
),
|
||||
)
|
||||
}
|
||||
@@ -752,7 +756,6 @@ fn plan_with_profiles(
|
||||
ConfigRevision::new(1).expect("config revision"),
|
||||
gateway_core::settings::SettingsValues::new(
|
||||
account_selection_policy.max_concurrent_per_account().get(),
|
||||
u64::try_from(account_selection_policy.request_interval().as_millis()).unwrap(),
|
||||
account_selection_policy.strategy().as_str(),
|
||||
Default::default(),
|
||||
None,
|
||||
@@ -824,7 +827,7 @@ fn dual_provider_plan(operation: &Operation) -> RoutingPlan {
|
||||
let directory = Arc::new(RuntimeAccountDirectory::new(entries));
|
||||
let snapshot = RuntimeSnapshot::new(
|
||||
ConfigRevision::new(1).expect("config revision"),
|
||||
gateway_core::settings::SettingsValues::new(2, 50, "smart", Default::default(), None, None),
|
||||
gateway_core::settings::SettingsValues::new(2, "smart", Default::default(), None, None),
|
||||
vec![openai.clone(), xai.clone()],
|
||||
vec![
|
||||
ProviderModel::new(
|
||||
@@ -2156,7 +2159,6 @@ fn unavailable_native_continuation_uses_provider_recovery_with_the_configured_ac
|
||||
let policy = AccountSelectionPolicy::new(
|
||||
RotationStrategy::RoundRobin,
|
||||
NonZeroU32::new(2).expect("account concurrency"),
|
||||
Duration::from_millis(50),
|
||||
);
|
||||
let route_plan = plan_with_policy(&operation, policy);
|
||||
let (coordinator, store, provider) = coordinator(vec![
|
||||
@@ -4284,6 +4286,85 @@ fn structural_event_before_replay_safe_failure_should_switch_account_before_comm
|
||||
assert_eq!(state.finalizations[0].client_status_code, Some(200));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn delivered_prelude_allows_retry_without_control_until_non_replayable_delivery() {
|
||||
for close_window in [false, true] {
|
||||
let operation = generate_operation();
|
||||
let route_plan = plan(&operation);
|
||||
let prelude = || {
|
||||
ProviderEvent::wire(
|
||||
ProtocolWireEvent::json(
|
||||
"openai",
|
||||
Some("response.created".to_owned()),
|
||||
json!({"type":"response.created","response":{"id":"response-1","output":[]}}),
|
||||
)
|
||||
.unwrap(),
|
||||
)
|
||||
.with_replayable_prelude()
|
||||
};
|
||||
let mut items = vec![Ok(prelude())];
|
||||
if close_window {
|
||||
items.push(Ok(ProviderEvent::wire(
|
||||
ProtocolWireEvent::json(
|
||||
"openai",
|
||||
Some("response.future".to_owned()),
|
||||
json!({"type":"response.future"}),
|
||||
)
|
||||
.unwrap(),
|
||||
)));
|
||||
// 后续前导不能撤销已经交付未知事件的事实
|
||||
items.push(Ok(prelude()));
|
||||
}
|
||||
items.push(Err(ProviderError::new(
|
||||
ProviderErrorKind::Transport,
|
||||
UpstreamSendState::Sent,
|
||||
)
|
||||
.with_replay_safe()));
|
||||
let (coordinator, store, provider) = coordinator(vec![
|
||||
Script::ObservedStream {
|
||||
account_id: "acct_first",
|
||||
items,
|
||||
},
|
||||
Script::Stream {
|
||||
account_id: "acct_second",
|
||||
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(),
|
||||
))
|
||||
.unwrap();
|
||||
let first = block_on(session.next_event()).unwrap().unwrap();
|
||||
assert!(first.into_provider_events()[0].is_replayable_prelude());
|
||||
block_on(session.commit_downstream(Some(200))).unwrap();
|
||||
let result = loop {
|
||||
match block_on(session.next_event()) {
|
||||
Ok(Some(_)) => {}
|
||||
terminal => break terminal,
|
||||
}
|
||||
};
|
||||
assert_eq!(result.is_err(), close_window);
|
||||
let attempts = if close_window { 1 } else { 2 };
|
||||
assert_eq!(provider.contexts.lock().unwrap().len(), attempts);
|
||||
let state = store.state.lock().unwrap();
|
||||
assert_eq!(state.commits, 1);
|
||||
assert_eq!(state.intermediate_failures, attempts - 1);
|
||||
assert_eq!(
|
||||
state.finalizations[0].outcome,
|
||||
if close_window {
|
||||
ExecutionOutcome::Incomplete
|
||||
} else {
|
||||
ExecutionOutcome::Succeeded
|
||||
}
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cancellation_before_pending_delivery_commit_reaches_terminal_state() {
|
||||
let operation = generate_operation();
|
||||
@@ -5726,13 +5807,9 @@ fn configured_rotation_budget_limits_actual_account_changes() {
|
||||
let operation = generate_operation();
|
||||
let route_plan = plan_with_policy(
|
||||
&operation,
|
||||
AccountSelectionPolicy::new(
|
||||
RotationStrategy::Smart,
|
||||
NonZeroU32::new(2).unwrap(),
|
||||
Duration::ZERO,
|
||||
)
|
||||
.with_max_account_rotations(budget)
|
||||
.with_openai_session_affinity_ttl(Duration::from_secs(168 * 3600)),
|
||||
AccountSelectionPolicy::new(RotationStrategy::Smart, NonZeroU32::new(2).unwrap())
|
||||
.with_max_account_rotations(budget)
|
||||
.with_openai_session_affinity_ttl(Duration::from_secs(168 * 3600)),
|
||||
);
|
||||
let (coordinator, store, provider) = coordinator(
|
||||
ROTATION_TEST_ACCOUNTS
|
||||
@@ -5779,12 +5856,8 @@ fn larger_rotation_budget_reaches_a_healthy_account_after_four_failures() {
|
||||
let operation = generate_operation();
|
||||
let route_plan = plan_with_policy(
|
||||
&operation,
|
||||
AccountSelectionPolicy::new(
|
||||
RotationStrategy::Smart,
|
||||
NonZeroU32::new(2).unwrap(),
|
||||
Duration::ZERO,
|
||||
)
|
||||
.with_max_account_rotations(4),
|
||||
AccountSelectionPolicy::new(RotationStrategy::Smart, NonZeroU32::new(2).unwrap())
|
||||
.with_max_account_rotations(4),
|
||||
);
|
||||
let mut scripts = ROTATION_TEST_ACCOUNTS[..4]
|
||||
.iter()
|
||||
@@ -5820,12 +5893,8 @@ fn zero_rotation_budget_still_allows_retrying_the_same_account() {
|
||||
let operation = generate_operation();
|
||||
let route_plan = plan_with_policy(
|
||||
&operation,
|
||||
AccountSelectionPolicy::new(
|
||||
RotationStrategy::Smart,
|
||||
NonZeroU32::new(2).unwrap(),
|
||||
Duration::ZERO,
|
||||
)
|
||||
.with_max_account_rotations(0),
|
||||
AccountSelectionPolicy::new(RotationStrategy::Smart, NonZeroU32::new(2).unwrap())
|
||||
.with_max_account_rotations(0),
|
||||
);
|
||||
let (coordinator, _, provider) = coordinator(vec![
|
||||
Script::Stream {
|
||||
@@ -5854,3 +5923,181 @@ 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]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lease_cancellation_is_preserved_in_the_single_request_terminal() {
|
||||
use gateway_core::lifecycle::{CancellationReason, LeaseFailure, RequestLeaseKind};
|
||||
block_on(async {
|
||||
let operation = generate_operation();
|
||||
let (coordinator, store, _) = coordinator(vec![Script::Pending]);
|
||||
let cancellation = CancellationToken::new();
|
||||
let mut session = coordinator
|
||||
.start(
|
||||
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
|
||||
operation.clone(),
|
||||
plan(&operation),
|
||||
None,
|
||||
None,
|
||||
cancellation.clone(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let reason = CancellationReason::Lease {
|
||||
kind: RequestLeaseKind::ClientAdmission,
|
||||
failure: LeaseFailure::RenewalTimedOut,
|
||||
};
|
||||
let mut next = Box::pin(session.next_event());
|
||||
assert!(futures::poll!(next.as_mut()).is_pending());
|
||||
cancellation.cancel_with_reason(reason);
|
||||
assert!(matches!(next.await, Err(EngineError::Cancelled)));
|
||||
session.cancel_and_finalize().await.unwrap();
|
||||
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].error_message.as_deref(),
|
||||
Some(reason.message())
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
@@ -1119,7 +1119,7 @@ fn request_policy_snapshot(
|
||||
ModelCapabilities::new(BTreeSet::from([OperationKind::Generate]), Some(16_000));
|
||||
RuntimeSnapshot::new(
|
||||
ConfigRevision::new(1).unwrap(),
|
||||
gateway_core::settings::SettingsValues::new(2, 1, "smart", Default::default(), None, None),
|
||||
gateway_core::settings::SettingsValues::new(2, "smart", Default::default(), None, None),
|
||||
providers.to_vec(),
|
||||
providers
|
||||
.iter()
|
||||
@@ -1847,14 +1847,7 @@ fn model_routing_cannot_expand_the_frozen_key_model_scope() {
|
||||
ModelCapabilities::new(BTreeSet::from([OperationKind::Generate]), Some(16_000));
|
||||
let snapshot = RuntimeSnapshot::new(
|
||||
ConfigRevision::new(1).unwrap(),
|
||||
gateway_core::settings::SettingsValues::new(
|
||||
2,
|
||||
1,
|
||||
"smart",
|
||||
Default::default(),
|
||||
None,
|
||||
None,
|
||||
),
|
||||
gateway_core::settings::SettingsValues::new(2, "smart", Default::default(), None, None),
|
||||
vec![kind.clone()],
|
||||
["gpt-start", "gpt-blocked"]
|
||||
.into_iter()
|
||||
@@ -2231,7 +2224,7 @@ fn fallback_observation_snapshot(reference: ExtensionSetReference) -> RuntimeSna
|
||||
ModelCapabilities::new(BTreeSet::from([OperationKind::Generate]), Some(16_000));
|
||||
RuntimeSnapshot::new(
|
||||
ConfigRevision::new(1).unwrap(),
|
||||
gateway_core::settings::SettingsValues::new(2, 1, "smart", Default::default(), None, None),
|
||||
gateway_core::settings::SettingsValues::new(2, "smart", Default::default(), None, None),
|
||||
providers.to_vec(),
|
||||
providers
|
||||
.iter()
|
||||
@@ -3919,7 +3912,7 @@ fn probe_snapshot() -> RuntimeSnapshot {
|
||||
ModelCapabilities::new(BTreeSet::from([OperationKind::Generate]), Some(16_000));
|
||||
RuntimeSnapshot::new(
|
||||
ConfigRevision::new(1).expect("config revision"),
|
||||
gateway_core::settings::SettingsValues::new(1, 1, "smart", Default::default(), None, None),
|
||||
gateway_core::settings::SettingsValues::new(1, "smart", Default::default(), None, None),
|
||||
vec![provider.clone()],
|
||||
vec![ProviderModel::new(
|
||||
provider,
|
||||
@@ -3935,7 +3928,7 @@ fn client_snapshot() -> RuntimeSnapshot {
|
||||
let provider = ProviderKind::new("openai").expect("provider kind");
|
||||
RuntimeSnapshot::new(
|
||||
ConfigRevision::new(1).expect("config revision"),
|
||||
gateway_core::settings::SettingsValues::new(1, 1, "smart", Default::default(), None, None),
|
||||
gateway_core::settings::SettingsValues::new(1, "smart", Default::default(), None, None),
|
||||
vec![provider.clone()],
|
||||
Vec::new(),
|
||||
vec![ClientPolicy::new(
|
||||
@@ -3964,7 +3957,7 @@ fn start_snapshot_with_policy(
|
||||
ModelCapabilities::new(BTreeSet::from([OperationKind::Generate]), Some(16_000));
|
||||
RuntimeSnapshot::new(
|
||||
ConfigRevision::new(revision).expect("config revision"),
|
||||
gateway_core::settings::SettingsValues::new(1, 1, "smart", Default::default(), None, None),
|
||||
gateway_core::settings::SettingsValues::new(1, "smart", Default::default(), None, None),
|
||||
vec![provider.clone()],
|
||||
vec![ProviderModel::new(
|
||||
provider.clone(),
|
||||
@@ -4384,7 +4377,7 @@ fn request_settings_recompute_routing_and_admission_without_changing_sibling_or_
|
||||
let mut settings =
|
||||
serde_json::to_value(modified.request_settings().execution_values().unwrap()).unwrap();
|
||||
settings["runtime"]["model_mappings"] = json!({"request-alias":"gpt-start"});
|
||||
settings["runtime"]["request_interval_ms"] = json!(0);
|
||||
settings["runtime"]["max_account_rotations"] = json!(0);
|
||||
settings["runtime"]["request_profiles"] = json!({"openai":{"identity":"request-local"}});
|
||||
settings["fast_mode"] = json!("default");
|
||||
settings["client_limits"] = json!({"max_concurrency":0,"requests_per_minute":0});
|
||||
@@ -4604,7 +4597,7 @@ fn entry_settings_freeze_authentication_and_rebase_only_explicit_overrides() {
|
||||
let baseline = service.request_settings().unwrap();
|
||||
let mut values = serde_json::to_value(baseline.values()).unwrap();
|
||||
values["responses_max_decompressed_body_bytes"] = json!(1024);
|
||||
values["request_interval_ms"] = json!(0);
|
||||
values["max_account_rotations"] = json!(0);
|
||||
let first = baseline
|
||||
.replace(serde_json::from_value(values).unwrap(), "first")
|
||||
.unwrap();
|
||||
@@ -4621,7 +4614,7 @@ fn entry_settings_freeze_authentication_and_rebase_only_explicit_overrides() {
|
||||
.unwrap();
|
||||
let source = second.inspect();
|
||||
assert_eq!(
|
||||
source["overrides"]["request_interval_ms"]["instance_id"],
|
||||
source["overrides"]["max_account_rotations"]["instance_id"],
|
||||
"first"
|
||||
);
|
||||
assert_eq!(
|
||||
@@ -4819,14 +4812,7 @@ fn child_settings_recompute_key_scope_without_inheriting_parent_defaults() {
|
||||
.collect();
|
||||
let snapshot = RuntimeSnapshot::new(
|
||||
ConfigRevision::new(1).unwrap(),
|
||||
gateway_core::settings::SettingsValues::new(
|
||||
1,
|
||||
0,
|
||||
"smart",
|
||||
Default::default(),
|
||||
None,
|
||||
None,
|
||||
),
|
||||
gateway_core::settings::SettingsValues::new(1, "smart", Default::default(), None, None),
|
||||
vec![provider],
|
||||
vec![],
|
||||
policies,
|
||||
|
||||
@@ -1,11 +1,6 @@
|
||||
//! 验证插件调度策略对软亲和的覆盖与宿主候选范围限制
|
||||
|
||||
use std::{
|
||||
collections::BTreeSet,
|
||||
num::NonZeroU32,
|
||||
sync::Arc,
|
||||
time::{Duration, SystemTime},
|
||||
};
|
||||
use std::{collections::BTreeSet, num::NonZeroU32, sync::Arc, time::SystemTime};
|
||||
|
||||
use futures::{executor::block_on, future::BoxFuture};
|
||||
use gateway_core::{
|
||||
@@ -106,11 +101,7 @@ fn policy(choice: Option<&'static str>) -> RequestPolicyContext {
|
||||
|
||||
fn context() -> AccountSelectionContext {
|
||||
AccountSelectionContext {
|
||||
policy: AccountSelectionPolicy::new(
|
||||
RotationStrategy::Sticky,
|
||||
NonZeroU32::new(2).unwrap(),
|
||||
Duration::ZERO,
|
||||
),
|
||||
policy: AccountSelectionPolicy::new(RotationStrategy::Sticky, NonZeroU32::new(2).unwrap()),
|
||||
now: SystemTime::now(),
|
||||
excluded_accounts: BTreeSet::new(),
|
||||
preferred_account: Some(account_id("acct_a")),
|
||||
@@ -181,15 +172,15 @@ fn smart_switchback_does_not_override_explicit_plugin_choices_or_their_telemetry
|
||||
.clone()
|
||||
.with_scheduling(None, AccountWeight::new(10).unwrap());
|
||||
let mut selection_context = context();
|
||||
selection_context.policy = AccountSelectionPolicy::new(
|
||||
RotationStrategy::Smart,
|
||||
NonZeroU32::new(2).unwrap(),
|
||||
Duration::ZERO,
|
||||
)
|
||||
.with_smart_scheduling(
|
||||
gateway_core::account::SmartSchedulingConfig::new([1.0, 0.8, 1.0, 0.5, 0.0, 0.0], true)
|
||||
.unwrap(),
|
||||
);
|
||||
selection_context.policy =
|
||||
AccountSelectionPolicy::new(RotationStrategy::Smart, NonZeroU32::new(2).unwrap())
|
||||
.with_smart_scheduling(
|
||||
gateway_core::account::SmartSchedulingConfig::new(
|
||||
[1.0, 0.8, 1.0, 0.5, 0.0, 0.0],
|
||||
true,
|
||||
)
|
||||
.unwrap(),
|
||||
);
|
||||
for chosen in [Some("acct_a"), Some("acct_b"), None] {
|
||||
let selected = policy(chosen)
|
||||
.select_account(
|
||||
|
||||
@@ -322,3 +322,25 @@ fn connection_lifecycle_contract_is_object_safe() {
|
||||
|
||||
assert!(!object.is_draining());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cancellation_reason_survives_cleanup_and_propagates_only_to_descendants() {
|
||||
use gateway_core::lifecycle::{CancellationReason, LeaseFailure, RequestLeaseKind};
|
||||
let parent = CancellationToken::new();
|
||||
let child = parent.child_token();
|
||||
let descendant = child.child_token();
|
||||
let reason = CancellationReason::Lease {
|
||||
kind: RequestLeaseKind::Credential,
|
||||
failure: LeaseFailure::Lost,
|
||||
};
|
||||
child.cancel_with_reason(reason);
|
||||
child.cancel();
|
||||
descendant.cancel();
|
||||
assert_eq!(parent.reason(), None);
|
||||
assert_eq!(child.reason(), Some(reason));
|
||||
assert_eq!(descendant.reason(), Some(reason));
|
||||
let requested = CancellationToken::new();
|
||||
requested.cancel();
|
||||
requested.cancel_with_reason(reason);
|
||||
assert_eq!(requested.reason(), Some(CancellationReason::Requested));
|
||||
}
|
||||
|
||||
@@ -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":{
|
||||
|
||||
@@ -23,7 +23,7 @@ mod request_settings;
|
||||
mod snapshot;
|
||||
|
||||
fn settings() -> gateway_core::settings::SettingsValues {
|
||||
gateway_core::settings::SettingsValues::new(3, 50, "smart", BTreeMap::new(), None, None)
|
||||
gateway_core::settings::SettingsValues::new(3, "smart", BTreeMap::new(), None, None)
|
||||
}
|
||||
|
||||
fn capabilities() -> ModelCapabilities {
|
||||
|
||||
@@ -29,7 +29,7 @@ fn snapshot(revision: u64, name: &str) -> Arc<RuntimeSnapshot> {
|
||||
Arc::new(
|
||||
RuntimeSnapshot::new(
|
||||
ConfigRevision::new(revision).unwrap(),
|
||||
SettingsValues::new(3, 50, "smart", BTreeMap::new(), None, None)
|
||||
SettingsValues::new(3, "smart", BTreeMap::new(), None, None)
|
||||
.with_request_profiles(profiles(&[("openai", name), ("xai", name)])),
|
||||
vec![],
|
||||
vec![],
|
||||
|
||||
@@ -136,7 +136,6 @@ fn compile_aliases(
|
||||
revision(1),
|
||||
SettingsValues::new(
|
||||
3,
|
||||
50,
|
||||
"smart",
|
||||
BTreeMap::from([("configured".into(), "upstream-model".into())]),
|
||||
None,
|
||||
@@ -282,7 +281,9 @@ impl ProviderCatalogPort for TestCatalog {
|
||||
}
|
||||
}
|
||||
|
||||
struct MixedCatalog;
|
||||
struct MixedCatalog {
|
||||
xai_available: bool,
|
||||
}
|
||||
|
||||
impl ProviderCatalogPort for MixedCatalog {
|
||||
fn catalog_generations(&self) -> BTreeMap<ProviderKind, ProviderCatalogGeneration> {
|
||||
@@ -307,7 +308,7 @@ impl ProviderCatalogPort for MixedCatalog {
|
||||
) -> BoxFuture<'_, Result<Vec<ProviderModelCapabilities>, ProviderCatalogUnavailable>> {
|
||||
let model = match provider.as_str() {
|
||||
"openai" => "gpt-5.5",
|
||||
"xai" => "grok-4.7",
|
||||
"xai" if self.xai_available => "grok-4.7",
|
||||
_ => return Box::pin(async { Err(ProviderCatalogUnavailable) }),
|
||||
};
|
||||
Box::pin(async move {
|
||||
@@ -325,6 +326,10 @@ impl ProviderCatalogPort for MixedCatalog {
|
||||
}
|
||||
|
||||
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),
|
||||
@@ -346,7 +351,7 @@ fn mixed_catalog_snapshot() -> gateway_core::routing::RuntimeSnapshot {
|
||||
block_on(
|
||||
RuntimeSnapshotCompiler::new(
|
||||
Arc::new(TestSnapshotStore::new(Ok(facts))),
|
||||
Arc::new(MixedCatalog),
|
||||
Arc::new(MixedCatalog { xai_available }),
|
||||
)
|
||||
.compile(),
|
||||
)
|
||||
@@ -405,7 +410,7 @@ fn compiler_accepts_unlimited_default_account_concurrency() {
|
||||
let facts = SnapshotFacts::new(
|
||||
revision(1),
|
||||
revision(1),
|
||||
SettingsValues::new(0, 50, "smart", BTreeMap::new(), None, None),
|
||||
SettingsValues::new(0, "smart", BTreeMap::new(), None, None),
|
||||
Vec::new(),
|
||||
Vec::new(),
|
||||
vec![SnapshotProviderAccountFacts::new(
|
||||
@@ -431,10 +436,6 @@ fn compiler_accepts_unlimited_default_account_concurrency() {
|
||||
plan.account_selection_policy().max_concurrent_per_account(),
|
||||
gateway_core::account::AccountConcurrency::Unlimited
|
||||
);
|
||||
assert_eq!(
|
||||
plan.account_selection_policy().request_interval(),
|
||||
std::time::Duration::from_millis(50)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -449,7 +450,7 @@ fn compiled_plans_keep_their_scheduling_config_after_a_new_snapshot_is_built() {
|
||||
let facts = SnapshotFacts::new(
|
||||
revision(index as u64 + 1),
|
||||
revision(index as u64 + 1),
|
||||
SettingsValues::new(3, 0, "smart", BTreeMap::new(), None, None)
|
||||
SettingsValues::new(3, "smart", BTreeMap::new(), None, None)
|
||||
.with_smart_scheduling(config)
|
||||
.with_openai_guardian_reserved_concurrency(index as u32),
|
||||
vec![],
|
||||
@@ -505,7 +506,6 @@ fn withdrawn_provider_keeps_accounts_without_becoming_a_route_or_expanding_group
|
||||
revision(1),
|
||||
SettingsValues::new(
|
||||
3,
|
||||
0,
|
||||
"smart",
|
||||
BTreeMap::from([("alias".into(), "listed-model".into())]),
|
||||
None,
|
||||
@@ -678,7 +678,7 @@ fn discovery_catalog_should_not_reject_an_unlisted_model() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mixed_catalog_should_prefer_known_model_over_discovery_fallback() {
|
||||
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!(
|
||||
@@ -701,12 +701,12 @@ fn mixed_catalog_should_prefer_known_model_over_discovery_fallback() {
|
||||
.iter()
|
||||
.map(|candidate| candidate.provider().as_str())
|
||||
.collect::<Vec<_>>(),
|
||||
["xai", "openai"]
|
||||
["xai"]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mixed_catalog_should_prioritize_the_mapped_upstream_model() {
|
||||
fn mixed_catalog_should_bind_the_mapped_upstream_model() {
|
||||
let snapshot = mixed_catalog_snapshot();
|
||||
let settings = snapshot
|
||||
.settings()
|
||||
@@ -726,6 +726,7 @@ fn mixed_catalog_should_prioritize_the_mapped_upstream_model() {
|
||||
)
|
||||
.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(),
|
||||
@@ -771,41 +772,106 @@ fn mixed_catalog_should_preserve_provider_order_for_shared_models() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mixed_catalog_should_respect_required_provider_before_model_priority() {
|
||||
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()
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
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_eq!(plan.candidates().len(), 1);
|
||||
assert_eq!(plan.candidates()[0].provider().as_str(), "openai");
|
||||
assert!(matches!(
|
||||
plan,
|
||||
Err(gateway_core::validation::RoutingError::NoCapableProvider { .. })
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mixed_catalog_should_respect_blocked_provider_before_model_priority() {
|
||||
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("grok-4.7").unwrap(),
|
||||
&PublicModelId::new("gpt-5.5").unwrap(),
|
||||
&super::operation(),
|
||||
snapshot.all_account_scope(),
|
||||
&gateway_core::routing::RoutingContext {
|
||||
blocked_providers: std::collections::BTreeSet::from([
|
||||
ProviderKind::new("xai").unwrap()
|
||||
]),
|
||||
..Default::default()
|
||||
},
|
||||
&Default::default(),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(plan.candidates().len(), 1);
|
||||
assert_eq!(plan.candidates()[0].provider().as_str(), "openai");
|
||||
}
|
||||
@@ -991,7 +1057,6 @@ fn facts_with_min_versions(
|
||||
revision(observed_current_revision),
|
||||
SettingsValues::new(
|
||||
3,
|
||||
50,
|
||||
"smart",
|
||||
BTreeMap::from([("public-model".to_owned(), "upstream-model".to_owned())]),
|
||||
desktop,
|
||||
@@ -1033,7 +1098,7 @@ fn global_request_location_should_be_frozen_when_snapshot_is_published() {
|
||||
let facts = SnapshotFacts::new(
|
||||
revision(version),
|
||||
revision(version),
|
||||
SettingsValues::new(3, 0, "smart", BTreeMap::new(), None, None)
|
||||
SettingsValues::new(3, "smart", BTreeMap::new(), None, None)
|
||||
.with_request_location(location, enabled),
|
||||
Vec::new(),
|
||||
Vec::new(),
|
||||
@@ -1110,7 +1175,7 @@ fn decompression_setting_should_validate_and_remain_frozen_across_publication()
|
||||
let facts = SnapshotFacts::new(
|
||||
revision(version),
|
||||
revision(version),
|
||||
SettingsValues::new(3, 0, "smart", BTreeMap::new(), None, None)
|
||||
SettingsValues::new(3, "smart", BTreeMap::new(), None, None)
|
||||
.with_responses_max_decompressed_body_bytes(bytes),
|
||||
Vec::new(),
|
||||
Vec::new(),
|
||||
@@ -1173,7 +1238,7 @@ fn fast_mode_merges_only_bound_groups_without_changing_account_scope() {
|
||||
let facts = SnapshotFacts::new(
|
||||
revision(1),
|
||||
revision(1),
|
||||
SettingsValues::new(3, 0, "smart", BTreeMap::new(), None, None),
|
||||
SettingsValues::new(3, "smart", BTreeMap::new(), None, None),
|
||||
vec![SnapshotClientPolicyFacts::new(
|
||||
ClientApiKeyId::new("key_fast_policy").unwrap(),
|
||||
PlaintextClientApiKey::new("sk_fast_policy").unwrap(),
|
||||
@@ -1254,7 +1319,7 @@ fn key_profiles_replace_whole_global_choice_and_previous_snapshot_stays_frozen()
|
||||
};
|
||||
let build = |global, overridden| {
|
||||
let profiles = BTreeMap::from([(provider.clone(), document(global))]);
|
||||
let settings = SettingsValues::new(3, 0, "smart", BTreeMap::new(), None, None)
|
||||
let settings = SettingsValues::new(3, "smart", BTreeMap::new(), None, None)
|
||||
.with_request_profiles(profiles);
|
||||
let inherited = SnapshotClientPolicyFacts::new(
|
||||
ClientApiKeyId::new("key_inherited").unwrap(),
|
||||
|
||||
@@ -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())
|
||||
})
|
||||
}
|
||||
@@ -740,7 +823,7 @@ fn scoped_facts(value: u64, allow_removed: bool) -> SnapshotFacts {
|
||||
SnapshotFacts::new(
|
||||
revision(value),
|
||||
revision(value),
|
||||
SettingsValues::new(3, 0, "smart", BTreeMap::new(), None, None),
|
||||
SettingsValues::new(3, "smart", BTreeMap::new(), None, None),
|
||||
vec![SnapshotClientPolicyFacts::new(
|
||||
ClientApiKeyId::new("key_audit_synthetic").expect("key ID"),
|
||||
PlaintextClientApiKey::new("sk_audit_synthetic_not_a_real_key").expect("synthetic key"),
|
||||
@@ -775,7 +858,6 @@ fn facts(config_revision: u64, observed_current_revision: u64) -> SnapshotFacts
|
||||
revision(observed_current_revision),
|
||||
SettingsValues::new(
|
||||
3,
|
||||
50,
|
||||
"smart",
|
||||
BTreeMap::from([("public-model".to_owned(), "upstream-model".to_owned())]),
|
||||
None,
|
||||
@@ -800,7 +882,7 @@ fn compiler(store: Arc<dyn SnapshotStorePort>) -> RuntimeSnapshotCompiler {
|
||||
fn empty_snapshot(value: u64) -> RuntimeSnapshot {
|
||||
RuntimeSnapshot::new(
|
||||
revision(value),
|
||||
gateway_core::settings::SettingsValues::new(1, 0, "smart", Default::default(), None, None),
|
||||
gateway_core::settings::SettingsValues::new(1, "smart", Default::default(), None, None),
|
||||
Vec::new(),
|
||||
Vec::new(),
|
||||
Vec::new(),
|
||||
|
||||
@@ -24,6 +24,7 @@ const MAXIMUM_RESPONSE_BYTES: usize = 512 * 1024;
|
||||
const DYNAMIC_CACHE_SECONDS: i64 = 300;
|
||||
const FALLBACK_CACHE_SECONDS: u64 = 60;
|
||||
const MINIMUM_REMAINING_LINK_SECONDS: i64 = 600;
|
||||
const OFFICIAL_METADATA_TIMEOUT: Duration = Duration::from_secs(3);
|
||||
|
||||
const OFFICIAL_X64_URL: &str = "https://persistent.oaistatic.com/codex-app-prod/ChatGPT-x64.msix";
|
||||
const OFFICIAL_ARM64_URL: &str =
|
||||
@@ -53,7 +54,7 @@ impl RgAdguardClientDistribution {
|
||||
pub(crate) fn new() -> Self {
|
||||
let client = Client::builder()
|
||||
.connect_timeout(Duration::from_secs(5))
|
||||
.timeout(Duration::from_secs(12))
|
||||
.timeout(Duration::from_secs(30))
|
||||
.redirect(Policy::none())
|
||||
.build()
|
||||
.inspect_err(|_| {
|
||||
@@ -69,10 +70,11 @@ impl RgAdguardClientDistribution {
|
||||
}
|
||||
|
||||
async fn resolve_uncached(&self) -> (CodexDesktopWindowsDownloads, Duration) {
|
||||
let now = Utc::now();
|
||||
match self.fetch_store_packages(now).await {
|
||||
match self.fetch_store_packages().await {
|
||||
Ok(packages) => {
|
||||
let (packages, warning) = complete_with_official_fallback(packages);
|
||||
let (mut packages, warning) = complete_with_official_fallback(packages);
|
||||
self.fill_official_package_sizes(&mut packages).await;
|
||||
let now = Utc::now();
|
||||
let valid_seconds = packages
|
||||
.iter()
|
||||
.filter_map(|package| package.expires_at)
|
||||
@@ -97,17 +99,33 @@ impl RgAdguardClientDistribution {
|
||||
error_kind = ?error_kind,
|
||||
"RG-Adguard package resolution failed; using official downloads"
|
||||
);
|
||||
(
|
||||
official_fallback(now),
|
||||
Duration::from_secs(FALLBACK_CACHE_SECONDS),
|
||||
)
|
||||
let mut downloads = official_fallback(Utc::now());
|
||||
self.fill_official_package_sizes(&mut downloads.packages)
|
||||
.await;
|
||||
downloads.resolved_at = Utc::now();
|
||||
(downloads, Duration::from_secs(FALLBACK_CACHE_SECONDS))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn fill_official_package_sizes(&self, packages: &mut [ClientDownloadPackage]) {
|
||||
let Some(client) = self.client.as_ref() else {
|
||||
return;
|
||||
};
|
||||
// 只补充回退包的展示大小,超时或缺失长度不能撤销官方稳定下载地址
|
||||
futures::future::join_all(
|
||||
packages
|
||||
.iter_mut()
|
||||
.filter(|package| package.source == ClientDownloadSource::OfficialOpenAi)
|
||||
.map(|package| async {
|
||||
package.size_bytes = official_package_size(client, &package.download_url).await;
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
async fn fetch_store_packages(
|
||||
&self,
|
||||
now: DateTime<Utc>,
|
||||
) -> Result<Vec<ClientDownloadPackage>, ClientDistributionResolutionErrorKind> {
|
||||
let client = self
|
||||
.client
|
||||
@@ -145,7 +163,8 @@ impl RgAdguardClientDistribution {
|
||||
}
|
||||
let html = std::str::from_utf8(&bytes)
|
||||
.map_err(|_| ClientDistributionResolutionErrorKind::InvalidResponse)?;
|
||||
parse_store_packages(html, now)
|
||||
// 上游解析可能耗时几十秒,按收到结果时的时间检查剩余有效期
|
||||
parse_store_packages(html, Utc::now())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -318,11 +337,12 @@ fn parse_size_bytes(value: &str) -> Option<u64> {
|
||||
if parts.next().is_some() {
|
||||
return None;
|
||||
}
|
||||
// RG 页面以 KB/MB/GB 标记二进制大小,换算后才能与 CDN Content-Length 对齐
|
||||
let multiplier = match unit {
|
||||
"B" => 1_u128,
|
||||
"KB" => 1_000,
|
||||
"MB" => 1_000_000,
|
||||
"GB" => 1_000_000_000,
|
||||
"KB" => 1_024,
|
||||
"MB" => 1_048_576,
|
||||
"GB" => 1_073_741_824,
|
||||
_ => return None,
|
||||
};
|
||||
let (whole, fraction) = amount.split_once('.').unwrap_or((amount.as_str(), ""));
|
||||
@@ -351,7 +371,7 @@ fn official_fallback(now: DateTime<Utc>) -> CodexDesktopWindowsDownloads {
|
||||
resolved_at: now,
|
||||
cached: false,
|
||||
warning: Some(
|
||||
"Microsoft Store 离线直链暂不可用,已切换为 OpenAI 官方稳定安装包。".to_owned(),
|
||||
"Microsoft Store 离线直链暂不可用,已切换为 OpenAI 官方稳定安装包".to_owned(),
|
||||
),
|
||||
packages,
|
||||
}
|
||||
@@ -375,7 +395,7 @@ pub fn complete_with_official_fallback(
|
||||
}
|
||||
let warning = (!missing.is_empty()).then(|| {
|
||||
format!(
|
||||
"{} 的 Microsoft Store 临时直链不可用,已切换为 OpenAI 官方稳定安装包。",
|
||||
"{} 的 Microsoft Store 临时直链不可用,已切换为 OpenAI 官方稳定安装包",
|
||||
missing.join("、")
|
||||
)
|
||||
});
|
||||
@@ -397,3 +417,25 @@ fn official_package(architecture: ClientArchitecture) -> ClientDownloadPackage {
|
||||
expires_at: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// 读取官方安装包的响应头大小,不下载包字节
|
||||
pub async fn official_package_size(client: &Client, download_url: &str) -> Option<u64> {
|
||||
let response = client
|
||||
.head(download_url)
|
||||
.timeout(OFFICIAL_METADATA_TIMEOUT)
|
||||
.send()
|
||||
.await
|
||||
.ok()?;
|
||||
if !response.status().is_success() {
|
||||
return None;
|
||||
}
|
||||
// HEAD 没有正文,不能用正文长度推断实际安装包大小
|
||||
response
|
||||
.headers()
|
||||
.get(reqwest::header::CONTENT_LENGTH)?
|
||||
.to_str()
|
||||
.ok()?
|
||||
.parse::<u64>()
|
||||
.ok()
|
||||
.filter(|size| *size > 0)
|
||||
}
|
||||
|
||||
@@ -72,7 +72,7 @@ fn missing_config_field(error: &config::ConfigError) -> Option<String> {
|
||||
}
|
||||
|
||||
/// Host 唯一拥有的进程配置
|
||||
#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
|
||||
#[derive(Debug, Clone, Deserialize, PartialEq)]
|
||||
pub struct HostConfig {
|
||||
#[serde(default)]
|
||||
pub timezone: gateway_core::time::DeploymentTimeZone,
|
||||
@@ -145,7 +145,7 @@ pub struct ListenConfig {
|
||||
}
|
||||
|
||||
/// Host 结构化日志配置
|
||||
#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
|
||||
#[derive(Debug, Clone, Deserialize, PartialEq)]
|
||||
pub struct LoggingConfig {
|
||||
pub level: String,
|
||||
pub stdout: bool,
|
||||
@@ -157,17 +157,17 @@ pub struct LoggingConfig {
|
||||
/// 请求报文的留存窗口由 request_dump_retention_days 独立控制
|
||||
#[serde(default)]
|
||||
pub request_dump: bool,
|
||||
/// 完整报文的最少保留天数;按 UTC 日期整组清理,默认 1 天
|
||||
/// 完整报文的最少保留天数;支持小数,每天按 24 小时换算,默认 1 天
|
||||
#[serde(default = "default_request_dump_retention_days")]
|
||||
pub request_dump_retention_days: usize,
|
||||
pub request_dump_retention_days: f64,
|
||||
}
|
||||
|
||||
fn default_file_retention_days() -> usize {
|
||||
7
|
||||
}
|
||||
|
||||
fn default_request_dump_retention_days() -> usize {
|
||||
1
|
||||
fn default_request_dump_retention_days() -> f64 {
|
||||
1.0
|
||||
}
|
||||
|
||||
impl LoggingConfig {
|
||||
@@ -190,14 +190,19 @@ impl LoggingConfig {
|
||||
"host.logging.file.max_file_size_mb",
|
||||
));
|
||||
}
|
||||
if self.request_dump_retention_days == 0 {
|
||||
return Err(ConfigError::InvalidField(
|
||||
"host.logging.request_dump_retention_days",
|
||||
));
|
||||
}
|
||||
self.request_dump_retention()?;
|
||||
resolve_relative_path(source_dir, &mut self.file.directory);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) fn request_dump_retention(&self) -> Result<Duration, ConfigError> {
|
||||
Duration::try_from_secs_f64(self.request_dump_retention_days * 86_400.0)
|
||||
.ok()
|
||||
.filter(|duration| !duration.is_zero())
|
||||
.ok_or(ConfigError::InvalidField(
|
||||
"host.logging.request_dump_retention_days",
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
/// Host 文件日志配置
|
||||
|
||||
@@ -15,7 +15,7 @@ mod sink;
|
||||
mod writer;
|
||||
|
||||
use sink::{FileLogGuard, FileLogSink, LogHealth};
|
||||
use writer::RotatingLogWriter;
|
||||
use writer::{LogRetention, RotatingLogWriter};
|
||||
|
||||
// 文件前缀统一为 codex-proxy-rs-<kebab-case 类别>;专用 tracing target 使用 snake_case
|
||||
const APPLICATION_LOG_FILE_PREFIX: &str = "codex-proxy-rs-application";
|
||||
@@ -45,9 +45,11 @@ pub enum LogError {
|
||||
AlreadyInitialized,
|
||||
#[error("logging size limit is too large")]
|
||||
SizeOverflow,
|
||||
#[error(transparent)]
|
||||
InvalidConfig(#[from] crate::config::ConfigError),
|
||||
}
|
||||
|
||||
/// 初始化日志;文件按完整日期留存,不以分片数量淘汰窗口内记录
|
||||
/// 初始化日志;按各通道的保留窗口清理,不以分片数量淘汰记录
|
||||
pub fn initialize_logging(
|
||||
config: &LoggingConfig,
|
||||
timezone: gateway_core::time::DeploymentTimeZone,
|
||||
@@ -74,7 +76,7 @@ pub fn initialize_logging(
|
||||
let (writer, guard) = create_file_writer(
|
||||
config,
|
||||
timezone,
|
||||
config.file.retention_days,
|
||||
LogRetention::CalendarDays(config.file.retention_days),
|
||||
APPLICATION_LOG_FILE_PREFIX,
|
||||
"gateway-log-application-file",
|
||||
Arc::clone(&health),
|
||||
@@ -88,7 +90,7 @@ pub fn initialize_logging(
|
||||
let (writer, guard) = create_file_writer(
|
||||
config,
|
||||
timezone,
|
||||
config.file.retention_days,
|
||||
LogRetention::CalendarDays(config.file.retention_days),
|
||||
OAUTH_RECOVERY_LOG_FILE_PREFIX,
|
||||
"gateway-log-oauth-recovery-file",
|
||||
Arc::clone(&health),
|
||||
@@ -102,7 +104,7 @@ pub fn initialize_logging(
|
||||
let (writer, guard) = create_file_writer(
|
||||
config,
|
||||
timezone,
|
||||
config.request_dump_retention_days,
|
||||
LogRetention::Elapsed(config.request_dump_retention()?),
|
||||
REQUEST_DUMP_LOG_FILE_PREFIX,
|
||||
"gateway-log-request-dump-file",
|
||||
Arc::clone(&health),
|
||||
@@ -220,7 +222,7 @@ impl LogGuard {
|
||||
fn create_file_writer(
|
||||
config: &LoggingConfig,
|
||||
timezone: gateway_core::time::DeploymentTimeZone,
|
||||
retention_days: usize,
|
||||
retention: LogRetention,
|
||||
prefix: &'static str,
|
||||
thread_name: &'static str,
|
||||
health: Arc<LogHealth>,
|
||||
@@ -234,7 +236,7 @@ fn create_file_writer(
|
||||
config.file.directory.clone(),
|
||||
prefix,
|
||||
maximum_bytes,
|
||||
retention_days,
|
||||
retention,
|
||||
Arc::clone(&health),
|
||||
timezone,
|
||||
)?;
|
||||
|
||||
@@ -1,22 +1,30 @@
|
||||
//! 按大小轮转、压缩已关闭分片、按部署时区自然日期组清理
|
||||
//! 按大小与时间轮转、压缩已关闭分片,按自然日期或实际时长清理
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
use std::fs::{self, File, OpenOptions};
|
||||
use std::io::{self, BufWriter, Write};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant, SystemTime};
|
||||
|
||||
use chrono::{NaiveDate, Utc};
|
||||
use flate2::{Compression, write::GzEncoder};
|
||||
|
||||
use super::sink::LogHealth;
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
pub(super) enum LogRetention {
|
||||
CalendarDays(usize),
|
||||
Elapsed(Duration),
|
||||
}
|
||||
|
||||
pub(super) struct RotatingLogWriter {
|
||||
timezone: gateway_core::time::DeploymentTimeZone,
|
||||
directory: PathBuf,
|
||||
prefix: &'static str,
|
||||
maximum_bytes: u64,
|
||||
retention_days: usize,
|
||||
retention: LogRetention,
|
||||
opened_at: Instant,
|
||||
date: NaiveDate,
|
||||
segment: usize,
|
||||
bytes_written: u64,
|
||||
@@ -29,19 +37,23 @@ impl RotatingLogWriter {
|
||||
directory: PathBuf,
|
||||
prefix: &'static str,
|
||||
maximum_bytes: u64,
|
||||
retention_days: usize,
|
||||
retention: LogRetention,
|
||||
health: Arc<LogHealth>,
|
||||
timezone: gateway_core::time::DeploymentTimeZone,
|
||||
) -> io::Result<Self> {
|
||||
fs::create_dir_all(&directory)?;
|
||||
let date = timezone.local(Utc::now()).date_naive();
|
||||
cleanup_log_files(&directory, prefix, date, retention_days, timezone)?;
|
||||
cleanup_log_files(&directory, prefix, retention, timezone, None)?;
|
||||
let latest = managed_log_files(&directory, prefix)?
|
||||
.into_iter()
|
||||
.filter(|entry| entry.date == date)
|
||||
.max_by_key(|entry| (entry.segment, entry.compressed));
|
||||
let (segment, bytes_written) = match latest {
|
||||
Some(entry) if !entry.compressed && entry.path.metadata()?.len() < maximum_bytes => {
|
||||
Some(entry)
|
||||
if matches!(retention, LogRetention::CalendarDays(_))
|
||||
&& !entry.compressed
|
||||
&& entry.path.metadata()?.len() < maximum_bytes =>
|
||||
{
|
||||
(entry.segment, entry.path.metadata()?.len())
|
||||
}
|
||||
Some(entry) => (next_segment(entry.segment)?, 0),
|
||||
@@ -54,7 +66,8 @@ impl RotatingLogWriter {
|
||||
directory,
|
||||
prefix,
|
||||
maximum_bytes,
|
||||
retention_days,
|
||||
retention,
|
||||
opened_at: Instant::now(),
|
||||
date,
|
||||
segment,
|
||||
bytes_written,
|
||||
@@ -78,7 +91,14 @@ impl RotatingLogWriter {
|
||||
let day_changed = date != self.date;
|
||||
let size_exceeded = self.bytes_written > 0
|
||||
&& self.bytes_written.saturating_add(incoming_bytes as u64) > self.maximum_bytes;
|
||||
if !day_changed && !size_exceeded {
|
||||
// 小流量转储也需及时关闭分片,避免小数保留期仍要等自然日切换才能生效
|
||||
let age_exceeded = match self.retention {
|
||||
LogRetention::CalendarDays(_) => false,
|
||||
LogRetention::Elapsed(duration) => {
|
||||
self.opened_at.elapsed() >= duration.min(Duration::from_secs(3600))
|
||||
}
|
||||
};
|
||||
if !day_changed && !size_exceeded && !age_exceeded {
|
||||
return Ok(());
|
||||
}
|
||||
let previous = self
|
||||
@@ -104,6 +124,7 @@ impl RotatingLogWriter {
|
||||
self.date = date;
|
||||
self.segment = segment;
|
||||
self.bytes_written = self.file.get_ref().metadata()?.len();
|
||||
self.opened_at = Instant::now();
|
||||
if let Err(error) = compress_log_file(&previous) {
|
||||
self.health.maintenance_failed(error.kind());
|
||||
}
|
||||
@@ -111,9 +132,9 @@ impl RotatingLogWriter {
|
||||
if let Err(error) = cleanup_log_files(
|
||||
&self.directory,
|
||||
self.prefix,
|
||||
date,
|
||||
self.retention_days,
|
||||
self.retention,
|
||||
self.timezone,
|
||||
Some(&path),
|
||||
) {
|
||||
self.health.maintenance_failed(error.kind());
|
||||
}
|
||||
@@ -203,17 +224,38 @@ fn managed_log_files(directory: &Path, prefix: &str) -> io::Result<Vec<ManagedLo
|
||||
fn cleanup_log_files(
|
||||
directory: &Path,
|
||||
prefix: &str,
|
||||
today: NaiveDate,
|
||||
retention_days: usize,
|
||||
retention: LogRetention,
|
||||
timezone: gateway_core::time::DeploymentTimeZone,
|
||||
active_path: Option<&Path>,
|
||||
) -> io::Result<()> {
|
||||
let files = managed_log_files(directory, prefix)?;
|
||||
let retention_days = match retention {
|
||||
LogRetention::CalendarDays(days) => days,
|
||||
LogRetention::Elapsed(duration) => {
|
||||
let now = SystemTime::now();
|
||||
for entry in files {
|
||||
if active_path == Some(entry.path.as_path()) {
|
||||
continue;
|
||||
}
|
||||
let modified = entry.path.metadata()?.modified()?;
|
||||
if now
|
||||
.duration_since(modified)
|
||||
.is_ok_and(|age| age >= duration)
|
||||
{
|
||||
fs::remove_file(entry.path)?;
|
||||
}
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
};
|
||||
let today = timezone.local(Utc::now()).date_naive();
|
||||
let days =
|
||||
u64::try_from(retention_days).map_err(|_| io::Error::other("log retention overflow"))?;
|
||||
let cutoff = today
|
||||
.checked_sub_days(chrono::Days::new(days))
|
||||
.unwrap_or(NaiveDate::MIN);
|
||||
let mut dates = BTreeMap::<NaiveDate, Vec<PathBuf>>::new();
|
||||
for entry in managed_log_files(directory, prefix)? {
|
||||
for entry in files {
|
||||
dates.entry(entry.date).or_default().push(entry.path);
|
||||
}
|
||||
for (date, paths) in dates {
|
||||
|
||||
@@ -11,6 +11,7 @@ mod process;
|
||||
mod release;
|
||||
mod state;
|
||||
mod swap;
|
||||
mod swap_paths;
|
||||
|
||||
pub use config::SystemUpdateConfig;
|
||||
pub use operations::ProcessSystemOperations;
|
||||
|
||||
@@ -245,18 +245,7 @@ fn swap_file(current: &Path, backup: &Path) -> io::Result<()> {
|
||||
if swap.exists() {
|
||||
fs::remove_file(&swap)?;
|
||||
}
|
||||
move_file(current, &swap)?;
|
||||
if let Err(error) = move_file(backup, current) {
|
||||
let _ = move_file(&swap, current);
|
||||
return Err(error);
|
||||
}
|
||||
if let Err(error) = move_file(&swap, backup) {
|
||||
let _ = move_file(current, &swap);
|
||||
let _ = move_file(backup, current);
|
||||
let _ = move_file(&swap, backup);
|
||||
return Err(error);
|
||||
}
|
||||
Ok(())
|
||||
super::swap_paths::swap_paths(current, backup, &swap, move_file)
|
||||
}
|
||||
|
||||
fn swap_dir(current: &Path, backup: &Path) -> io::Result<()> {
|
||||
@@ -270,18 +259,7 @@ fn swap_dir(current: &Path, backup: &Path) -> io::Result<()> {
|
||||
if swap.exists() {
|
||||
remove_dir_all(&swap)?;
|
||||
}
|
||||
move_dir(current, &swap)?;
|
||||
if let Err(error) = move_dir(backup, current) {
|
||||
let _ = move_dir(&swap, current);
|
||||
return Err(error);
|
||||
}
|
||||
if let Err(error) = move_dir(&swap, backup) {
|
||||
let _ = move_dir(current, &swap);
|
||||
let _ = move_dir(backup, current);
|
||||
let _ = move_dir(&swap, backup);
|
||||
return Err(error);
|
||||
}
|
||||
Ok(())
|
||||
super::swap_paths::swap_paths(current, backup, &swap, move_dir)
|
||||
}
|
||||
|
||||
fn restore_dir(current: &Path, backup: &Path) -> io::Result<()> {
|
||||
|
||||
@@ -0,0 +1,59 @@
|
||||
//! 文件和目录交换共用的恢复顺序,保留主失败及补偿失败
|
||||
|
||||
use std::{io, path::Path};
|
||||
|
||||
pub(super) fn swap_paths(
|
||||
current: &Path,
|
||||
backup: &Path,
|
||||
swap: &Path,
|
||||
mut move_path: impl FnMut(&Path, &Path) -> io::Result<()>,
|
||||
) -> io::Result<()> {
|
||||
move_path(current, swap)?;
|
||||
if let Err(error) = move_path(backup, current) {
|
||||
return Err(with_recovery(
|
||||
error,
|
||||
"restore current from swap",
|
||||
move_path(swap, current),
|
||||
[current, backup, swap],
|
||||
));
|
||||
}
|
||||
if let Err(error) = move_path(swap, backup) {
|
||||
// 先把备份版本放回 backup,才能把原版本移回 current
|
||||
// 补偿第一步失败时保留 swap,不能覆盖仍在 current 的备份版本
|
||||
if let Err(recovery) = move_path(current, backup) {
|
||||
return Err(with_recovery(
|
||||
error,
|
||||
"restore backup from current",
|
||||
Err(recovery),
|
||||
[current, backup, swap],
|
||||
));
|
||||
}
|
||||
return Err(with_recovery(
|
||||
error,
|
||||
"restore current from swap",
|
||||
move_path(swap, current),
|
||||
[current, backup, swap],
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn with_recovery(
|
||||
primary: io::Error,
|
||||
action: &'static str,
|
||||
recovery: io::Result<()>,
|
||||
[current, backup, swap]: [&Path; 3],
|
||||
) -> io::Error {
|
||||
match recovery {
|
||||
Ok(()) => primary,
|
||||
Err(recovery) => io::Error::new(
|
||||
primary.kind(),
|
||||
format!(
|
||||
"{primary}; rollback failed: {action}: {recovery}; manual recovery required (current={}, backup={}, swap={})",
|
||||
current.display(),
|
||||
backup.display(),
|
||||
swap.display()
|
||||
),
|
||||
),
|
||||
}
|
||||
}
|
||||
@@ -30,7 +30,7 @@ fn parser_selects_latest_valid_package_for_each_architecture() {
|
||||
assert_eq!(packages[0].architecture, ClientArchitecture::X64);
|
||||
assert_eq!(packages[0].version.as_deref(), Some("26.825.6671.0"));
|
||||
assert!(packages[0].download_url.starts_with("http://"));
|
||||
assert_eq!(packages[0].size_bytes, Some(744_250_000));
|
||||
assert_eq!(packages[0].size_bytes, Some(780_402_688));
|
||||
assert_eq!(packages[1].architecture, ClientArchitecture::Arm64);
|
||||
assert!(packages[1].download_url.starts_with("https://"));
|
||||
}
|
||||
@@ -72,3 +72,74 @@ fn missing_architecture_uses_only_its_official_fallback() {
|
||||
assert_eq!(packages[1].source, ClientDownloadSource::OfficialOpenAi);
|
||||
assert!(warning.is_some_and(|value| value.contains("arm64")));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn official_package_size_reads_head_length_without_downloading_the_body() {
|
||||
use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
|
||||
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let url = format!("http://{}/ChatGPT-x64.msix", listener.local_addr().unwrap());
|
||||
let server = tokio::spawn(async move {
|
||||
let (mut stream, _) = listener.accept().await.unwrap();
|
||||
let mut request = Vec::new();
|
||||
let mut buffer = [0; 1024];
|
||||
while !request.windows(4).any(|bytes| bytes == b"\r\n\r\n") {
|
||||
let read = stream.read(&mut buffer).await.unwrap();
|
||||
assert!(read > 0);
|
||||
request.extend_from_slice(&buffer[..read]);
|
||||
}
|
||||
assert!(request.starts_with(b"HEAD /ChatGPT-x64.msix HTTP/1.1\r\n"));
|
||||
// 只发送响应头,不发送大包正文,证明解析使用文件长度而不是 HEAD 正文长度
|
||||
stream
|
||||
.write_all(b"HTTP/1.1 200 OK\r\ncontent-length: 944929441\r\nconnection: close\r\n\r\n")
|
||||
.await
|
||||
.unwrap();
|
||||
});
|
||||
let client = reqwest::Client::builder().no_proxy().build().unwrap();
|
||||
assert_eq!(
|
||||
gateway_host::client_distribution::official_package_size(&client, &url).await,
|
||||
Some(944_929_441)
|
||||
);
|
||||
server.await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn official_package_size_failure_keeps_metadata_optional() {
|
||||
use wiremock::{
|
||||
Mock, MockServer, ResponseTemplate,
|
||||
matchers::{method, path},
|
||||
};
|
||||
|
||||
let server = MockServer::start().await;
|
||||
for (route, response) in [
|
||||
("/missing.msix", ResponseTemplate::new(404)),
|
||||
("/empty.msix", ResponseTemplate::new(200)),
|
||||
(
|
||||
"/slow.msix",
|
||||
ResponseTemplate::new(200).set_delay(std::time::Duration::from_secs(4)),
|
||||
),
|
||||
] {
|
||||
Mock::given(method("HEAD"))
|
||||
.and(path(route))
|
||||
.respond_with(response)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
}
|
||||
let client = reqwest::Client::builder().no_proxy().build().unwrap();
|
||||
for route in ["/missing.msix", "/empty.msix", "/slow.msix"] {
|
||||
assert_eq!(
|
||||
tokio::time::timeout(
|
||||
std::time::Duration::from_millis(3500),
|
||||
gateway_host::client_distribution::official_package_size(
|
||||
&client,
|
||||
&format!("{}{route}", server.uri()),
|
||||
),
|
||||
)
|
||||
.await
|
||||
.expect("metadata timeout must not wait for the slow response"),
|
||||
None
|
||||
);
|
||||
}
|
||||
server.verify().await;
|
||||
}
|
||||
|
||||
@@ -146,7 +146,7 @@ fn logging_rejects_zero_retention_windows() {
|
||||
for request_dump in [false, true] {
|
||||
let mut config = valid_config();
|
||||
if request_dump {
|
||||
config.logging.request_dump_retention_days = 0;
|
||||
config.logging.request_dump_retention_days = 0.0;
|
||||
} else {
|
||||
config.logging.file.retention_days = 0;
|
||||
}
|
||||
@@ -168,7 +168,7 @@ fn logging_defaults_to_time_retention_and_ignores_removed_count_retention() {
|
||||
"file": { "enabled": true, "directory": "logs", "max_file_size_mb": 20 }
|
||||
});
|
||||
let config: LoggingConfig = serde_json::from_value(value.clone()).unwrap();
|
||||
assert_eq!(config.request_dump_retention_days, 1);
|
||||
assert_eq!(config.request_dump_retention_days, 1.0);
|
||||
assert_eq!(config.file.retention_days, 7);
|
||||
assert!(!config.oauth_recovery);
|
||||
value["file"]["max_files"] = serde_json::json!(20);
|
||||
@@ -178,6 +178,55 @@ fn logging_defaults_to_time_retention_and_ignores_removed_count_retention() {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_dump_retention_accepts_integer_and_fractional_yaml_days() {
|
||||
for (input, expected) in [("1", 1.0), ("0.5", 0.5), ("0.25", 0.25), ("1.5", 1.5)] {
|
||||
let yaml = format!(
|
||||
"level: info\nstdout: true\nrequest_dump_retention_days: {input}\nfile:\n enabled: true\n directory: logs\n max_file_size_mb: 20\n"
|
||||
);
|
||||
let mut config = valid_config();
|
||||
config.logging = config::Config::builder()
|
||||
.add_source(config::File::from_str(&yaml, config::FileFormat::Yaml))
|
||||
.build()
|
||||
.unwrap()
|
||||
.try_deserialize()
|
||||
.unwrap();
|
||||
config
|
||||
.resolve_and_validate(
|
||||
std::path::Path::new("/srv/gateway"),
|
||||
std::path::Path::new("/srv/gateway/web/dist"),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(config.logging.request_dump_retention_days, expected);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_dump_retention_rejects_invalid_durations() {
|
||||
for days in [
|
||||
0.0,
|
||||
-0.5,
|
||||
f64::NAN,
|
||||
f64::INFINITY,
|
||||
f64::NEG_INFINITY,
|
||||
f64::MAX,
|
||||
1e-20,
|
||||
] {
|
||||
let mut config = valid_config();
|
||||
config.logging.request_dump_retention_days = days;
|
||||
let error = config
|
||||
.resolve_and_validate(
|
||||
std::path::Path::new("/srv/gateway"),
|
||||
std::path::Path::new("/srv/gateway/web/dist"),
|
||||
)
|
||||
.unwrap_err();
|
||||
assert!(matches!(
|
||||
error,
|
||||
gateway_host::ConfigError::InvalidField("host.logging.request_dump_retention_days")
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dedicated_file_channels_can_be_enabled_without_application_sinks() {
|
||||
for oauth_recovery in [false, true] {
|
||||
@@ -220,7 +269,7 @@ fn valid_config() -> HostConfig {
|
||||
},
|
||||
oauth_recovery: false,
|
||||
request_dump: false,
|
||||
request_dump_retention_days: 1,
|
||||
request_dump_retention_days: 1.0,
|
||||
},
|
||||
system_update,
|
||||
worker_shutdown_timeout_seconds: 30,
|
||||
|
||||
@@ -48,7 +48,7 @@ fn logging_requires_at_least_one_sink() {
|
||||
},
|
||||
oauth_recovery: false,
|
||||
request_dump: false,
|
||||
request_dump_retention_days: 1,
|
||||
request_dump_retention_days: 1.0,
|
||||
},
|
||||
system_update: SystemUpdateConfig::default(),
|
||||
worker_shutdown_timeout_seconds: 30,
|
||||
@@ -325,7 +325,7 @@ fn logging_config(directory: PathBuf, request_dump: bool) -> HostConfig {
|
||||
},
|
||||
oauth_recovery: true,
|
||||
request_dump,
|
||||
request_dump_retention_days: 1,
|
||||
request_dump_retention_days: 1.0,
|
||||
},
|
||||
system_update: SystemUpdateConfig::default(),
|
||||
worker_shutdown_timeout_seconds: 30,
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
//! 验证日志压缩、轮转与重启后按日期保留完整记录
|
||||
//! 验证日志压缩、轮转,以及重启后的自然日与小数天数保留窗口
|
||||
|
||||
use super::*;
|
||||
use std::time::SystemTime;
|
||||
@@ -10,12 +10,11 @@ fn retention_preserves_complete_dates_across_compression_rotation_and_restart()
|
||||
if env::var_os(CHILD_PROCESS_ENV).is_some() {
|
||||
with_file_logging(
|
||||
PathBuf::from(env::var_os(LOG_DIRECTORY_ENV).unwrap()),
|
||||
true,
|
||||
false,
|
||||
|| {
|
||||
let payload = "x".repeat(1024 * 1024);
|
||||
// 三条记录已跨过 1 MiB 轮转边界;保留日期的 25 个分段由下方单独构造
|
||||
for sequence in 0..ROTATION_RECORDS {
|
||||
tracing::info!(target: REQUEST_DUMP_LOG_TARGET, sequence, payload, "retention record");
|
||||
tracing::info!(target: APPLICATION_LOG_TARGET, sequence, payload, "retention record");
|
||||
tracing::info!(target: OAUTH_RECOVERY_LOG_TARGET, sequence, payload, "retention record");
|
||||
}
|
||||
@@ -32,7 +31,6 @@ fn retention_preserves_complete_dates_across_compression_rotation_and_restart()
|
||||
for (prefix, days) in [
|
||||
(APPLICATION_LOG_FILE_PREFIX, 7),
|
||||
(OAUTH_RECOVERY_LOG_FILE_PREFIX, 7),
|
||||
(REQUEST_DUMP_LOG_FILE_PREFIX, 1),
|
||||
] {
|
||||
let boundary = today - chrono::Days::new(days);
|
||||
for segment in 1..=25 {
|
||||
@@ -90,11 +88,7 @@ fn retention_preserves_complete_dates_across_compression_rotation_and_restart()
|
||||
);
|
||||
assert!(unrelated.exists());
|
||||
|
||||
for prefix in [
|
||||
APPLICATION_LOG_FILE_PREFIX,
|
||||
OAUTH_RECOVERY_LOG_FILE_PREFIX,
|
||||
REQUEST_DUMP_LOG_FILE_PREFIX,
|
||||
] {
|
||||
for prefix in [APPLICATION_LOG_FILE_PREFIX, OAUTH_RECOVERY_LOG_FILE_PREFIX] {
|
||||
let body = read_log_file_set(directory.path(), prefix);
|
||||
let mut sequences = [0; ROTATION_RECORDS];
|
||||
let mut boundaries = std::collections::BTreeSet::new();
|
||||
@@ -141,6 +135,151 @@ fn retention_preserves_complete_dates_across_compression_rotation_and_restart()
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_dump_retention_uses_elapsed_time_across_compression_and_restart() {
|
||||
const RETENTION_ENV: &str = "CPR_LOGGING_TEST_RETENTION_DAYS";
|
||||
if env::var_os(CHILD_PROCESS_ENV).is_some() {
|
||||
let mut config =
|
||||
logging_config(PathBuf::from(env::var_os(LOG_DIRECTORY_ENV).unwrap()), true);
|
||||
config.logging.request_dump_retention_days =
|
||||
env::var(RETENTION_ENV).unwrap().parse().unwrap();
|
||||
with_logging(config, || {
|
||||
let payload = "x".repeat(1024 * 1024);
|
||||
for sequence in 0..ROTATION_RECORDS {
|
||||
tracing::info!(target: REQUEST_DUMP_LOG_TARGET, sequence, payload, "elapsed retention record");
|
||||
}
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
for days in [0.25, 0.5, 1.0, 1.5] {
|
||||
let directory = tempfile::tempdir().unwrap();
|
||||
let now = chrono::Utc::now();
|
||||
let cutoff = now - chrono::Duration::seconds((days * 86_400.0) as i64);
|
||||
let mut retained = Vec::new();
|
||||
let mut expired = Vec::new();
|
||||
// 同一日期内既有过期文件也有有效文件,文件名日期不覆盖实际修改时间
|
||||
for (segment, modified, keep) in [
|
||||
(0, cutoff - chrono::Duration::minutes(1), false),
|
||||
(1, cutoff + chrono::Duration::minutes(1), true),
|
||||
(2, now + chrono::Duration::hours(1), true),
|
||||
] {
|
||||
let path = directory.path().join(format!(
|
||||
"{REQUEST_DUMP_LOG_FILE_PREFIX}{}.{segment}.log",
|
||||
now.date_naive()
|
||||
));
|
||||
fs::write(&path, format!("seed-{segment}\n")).unwrap();
|
||||
fs::File::open(&path)
|
||||
.unwrap()
|
||||
.set_modified(modified.into())
|
||||
.unwrap();
|
||||
if keep {
|
||||
let modified = path.metadata().unwrap().modified().unwrap();
|
||||
retained.push((path, modified));
|
||||
} else {
|
||||
expired.push(path);
|
||||
}
|
||||
}
|
||||
// 归档的修改时间也参与清理,压缩不能重置保留窗口
|
||||
let archive = directory.path().join(format!(
|
||||
"{REQUEST_DUMP_LOG_FILE_PREFIX}{}.3.log.gz",
|
||||
now.date_naive() - chrono::Days::new(3)
|
||||
));
|
||||
let mut encoder = flate2::write::GzEncoder::new(
|
||||
fs::File::create(&archive).unwrap(),
|
||||
flate2::Compression::fast(),
|
||||
);
|
||||
std::io::Write::write_all(&mut encoder, b"expired archive\n").unwrap();
|
||||
encoder
|
||||
.finish()
|
||||
.unwrap()
|
||||
.set_modified((cutoff - chrono::Duration::hours(1)).into())
|
||||
.unwrap();
|
||||
expired.push(archive);
|
||||
let unrelated = directory.path().join("unmanaged.log");
|
||||
fs::write(&unrelated, "unmanaged\n").unwrap();
|
||||
|
||||
for run in 1..=2 {
|
||||
let output = Command::new(env::current_exe().unwrap())
|
||||
.args(["--exact", "logging::writer::request_dump_retention_uses_elapsed_time_across_compression_and_restart"])
|
||||
.env(CHILD_PROCESS_ENV, "1")
|
||||
.env(LOG_DIRECTORY_ENV, directory.path())
|
||||
.env(LOG_TIMEZONE_ENV, "America/New_York")
|
||||
.env(RETENTION_ENV, days.to_string())
|
||||
.env("RUST_LOG", "off")
|
||||
.output().unwrap();
|
||||
assert!(
|
||||
output.status.success(),
|
||||
"{}",
|
||||
String::from_utf8_lossy(&output.stderr)
|
||||
);
|
||||
for (path, modified) in &retained {
|
||||
let archive = path.with_extension("log.gz");
|
||||
assert!(archive.exists());
|
||||
assert_eq!(archive.metadata().unwrap().modified().unwrap(), *modified);
|
||||
}
|
||||
assert!(
|
||||
expired
|
||||
.iter()
|
||||
.all(|path| !path.exists() && !path.with_extension("log.gz").exists())
|
||||
);
|
||||
assert!(unrelated.exists());
|
||||
let body = read_log_file_set(directory.path(), REQUEST_DUMP_LOG_FILE_PREFIX);
|
||||
assert!(body.contains("seed-1\n") && body.contains("seed-2\n"));
|
||||
assert_eq!(
|
||||
body.matches("elapsed retention record").count(),
|
||||
run * ROTATION_RECORDS
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_dump_rotates_small_files_when_the_time_window_elapses() {
|
||||
if env::var_os(CHILD_PROCESS_ENV).is_some() {
|
||||
let directory = PathBuf::from(env::var_os(LOG_DIRECTORY_ENV).unwrap());
|
||||
let mut config = logging_config(directory.clone(), true);
|
||||
config.logging.request_dump_retention_days = 1.0 / 86_400.0;
|
||||
with_logging(config, || {
|
||||
tracing::info!(target: REQUEST_DUMP_LOG_TARGET, "before retention window");
|
||||
let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5);
|
||||
loop {
|
||||
let body = read_log_file_set(&directory, REQUEST_DUMP_LOG_FILE_PREFIX);
|
||||
if body.contains("before retention window") {
|
||||
break;
|
||||
}
|
||||
assert!(
|
||||
std::time::Instant::now() < deadline,
|
||||
"first record was not flushed"
|
||||
);
|
||||
std::thread::sleep(std::time::Duration::from_millis(10));
|
||||
}
|
||||
std::thread::sleep(std::time::Duration::from_millis(1100));
|
||||
tracing::info!(target: REQUEST_DUMP_LOG_TARGET, "after retention window");
|
||||
});
|
||||
return;
|
||||
}
|
||||
let directory = tempfile::tempdir().unwrap();
|
||||
let output = Command::new(env::current_exe().unwrap())
|
||||
.args([
|
||||
"--exact",
|
||||
"logging::writer::request_dump_rotates_small_files_when_the_time_window_elapses",
|
||||
])
|
||||
.env(CHILD_PROCESS_ENV, "1")
|
||||
.env(LOG_DIRECTORY_ENV, directory.path())
|
||||
.env("RUST_LOG", "off")
|
||||
.output()
|
||||
.unwrap();
|
||||
assert!(
|
||||
output.status.success(),
|
||||
"{}",
|
||||
String::from_utf8_lossy(&output.stderr)
|
||||
);
|
||||
let body = read_log_file_set(directory.path(), REQUEST_DUMP_LOG_FILE_PREFIX);
|
||||
assert!(!body.contains("before retention window"));
|
||||
assert!(body.contains("after retention window"));
|
||||
}
|
||||
|
||||
fn seed_log(path: &Path, date: chrono::NaiveDate, body: &str) {
|
||||
fs::write(path, body).unwrap();
|
||||
let modified = gateway_core::time::DeploymentTimeZone::default()
|
||||
|
||||
@@ -56,7 +56,7 @@ pub fn configuration(directory: &Path) -> HostConfig {
|
||||
},
|
||||
oauth_recovery: false,
|
||||
request_dump: false,
|
||||
request_dump_retention_days: 1,
|
||||
request_dump_retention_days: 1.0,
|
||||
},
|
||||
system_update: Default::default(),
|
||||
worker_shutdown_timeout_seconds: 1,
|
||||
|
||||
@@ -7,6 +7,7 @@ mod installation;
|
||||
mod operations;
|
||||
mod release;
|
||||
mod state;
|
||||
mod swap_paths;
|
||||
|
||||
use std::fs;
|
||||
use std::path::PathBuf;
|
||||
|
||||
@@ -0,0 +1,76 @@
|
||||
//! 对真实暂存文件注入移动双故障,验证恢复次序和人工恢复线索
|
||||
|
||||
#[path = "../../src/system_update/swap_paths.rs"]
|
||||
mod implementation;
|
||||
|
||||
use std::{fs, io};
|
||||
|
||||
#[test]
|
||||
fn failed_swap_preserves_both_versions_and_reports_compensation_failures() {
|
||||
for directory in [false, true] {
|
||||
for primary in [2, 3] {
|
||||
let mut recoveries = vec![None, Some(primary + 1)];
|
||||
if primary == 3 {
|
||||
recoveries.push(Some(5));
|
||||
}
|
||||
for recovery in recoveries {
|
||||
let temp = tempfile::tempdir().unwrap();
|
||||
let current = temp.path().join("current");
|
||||
let backup = temp.path().join("backup");
|
||||
let swap = temp.path().join("swap");
|
||||
for (path, content) in [(¤t, "current version"), (&backup, "backup version")]
|
||||
{
|
||||
if directory {
|
||||
fs::create_dir(path).unwrap();
|
||||
fs::write(path.join("asset"), content).unwrap();
|
||||
} else {
|
||||
fs::write(path, content).unwrap();
|
||||
}
|
||||
}
|
||||
let mut moves = 0;
|
||||
let error = implementation::swap_paths(¤t, &backup, &swap, |from, to| {
|
||||
moves += 1;
|
||||
if moves == primary {
|
||||
return Err(io::Error::other("primary injected failure"));
|
||||
}
|
||||
if recovery == Some(moves) {
|
||||
return Err(io::Error::other("recovery injected failure"));
|
||||
}
|
||||
fs::rename(from, to)
|
||||
})
|
||||
.unwrap_err();
|
||||
let message = error.to_string();
|
||||
assert!(message.contains("primary injected failure"));
|
||||
if recovery.is_some() {
|
||||
assert!(message.contains("recovery injected failure"));
|
||||
assert!(message.contains("restore "));
|
||||
assert!(message.contains("manual recovery required"));
|
||||
assert!(message.contains(swap.to_str().unwrap()));
|
||||
} else {
|
||||
assert!(!message.contains("rollback failed"));
|
||||
}
|
||||
let mut contents = [¤t, &backup, &swap]
|
||||
.into_iter()
|
||||
.filter_map(|path| {
|
||||
fs::read_to_string(if directory {
|
||||
path.join("asset")
|
||||
} else {
|
||||
path.clone()
|
||||
})
|
||||
.ok()
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
contents.sort();
|
||||
assert_eq!(contents, ["backup version", "current version"]);
|
||||
if recovery.is_none() {
|
||||
let original = if directory {
|
||||
current.join("asset")
|
||||
} else {
|
||||
current.clone()
|
||||
};
|
||||
assert_eq!(fs::read_to_string(original).unwrap(), "current version");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -87,16 +87,16 @@ pub(super) async fn invoke(
|
||||
let _binding = _binding;
|
||||
tokio::select! {
|
||||
biased;
|
||||
() = context.cancellation.cancelled() => Err(MiddlewareError::Fault),
|
||||
() = context.cancellation.cancelled() => Err(core::upgrade::SessionError::Cancelled),
|
||||
result = async {
|
||||
invocation.session_ready().await?;
|
||||
match stream.next().await {
|
||||
Ok(None) => Ok(()),
|
||||
Ok(Some(_)) => {
|
||||
session.invalid_response(Stage::Http);
|
||||
Err(MiddlewareError::InvalidState)
|
||||
let error = session.invalid_response(Stage::Http);
|
||||
Err(session_error(error, &entry.instance_id))
|
||||
}
|
||||
Err(error) => Err(crate::callback::error::rpc_middleware(error)),
|
||||
Err(error) => Err(session_error(error, &entry.instance_id)),
|
||||
}
|
||||
} => result,
|
||||
}
|
||||
@@ -143,3 +143,32 @@ pub(super) async fn invoke(
|
||||
.await
|
||||
.map_err(|_| invalid())
|
||||
}
|
||||
|
||||
fn session_error(error: crate::RpcError, instance_id: &str) -> core::upgrade::SessionError {
|
||||
use core::upgrade::SessionError;
|
||||
use gateway_plugin_sdk::ErrorCode;
|
||||
|
||||
let kind = match &error {
|
||||
crate::RpcError::Cancelled => return SessionError::Cancelled,
|
||||
// 本地期限已经按调用关联记录,终态消费不能再记同一次超时
|
||||
crate::RpcError::Timeout => return SessionError::Reported,
|
||||
crate::RpcError::Remote(fault) => match fault.code {
|
||||
ErrorCode::Cancelled => return SessionError::Cancelled,
|
||||
ErrorCode::Unsupported => "unsupported",
|
||||
ErrorCode::Rejected => "rejected",
|
||||
ErrorCode::InvalidInput => "invalid_input",
|
||||
ErrorCode::PermissionDenied => "permission_denied",
|
||||
ErrorCode::Upstream => "upstream",
|
||||
ErrorCode::Fault => "fault",
|
||||
ErrorCode::Timeout => "timeout",
|
||||
ErrorCode::Uncertain => "uncertain",
|
||||
ErrorCode::Capacity => "capacity",
|
||||
ErrorCode::Conflict => "conflict",
|
||||
},
|
||||
_ => error.diagnostic().0,
|
||||
};
|
||||
SessionError::Failed {
|
||||
kind,
|
||||
plugin_instance_id: Some(instance_id.to_owned()),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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>>,
|
||||
|
||||
@@ -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)) => {
|
||||
|
||||
@@ -15,14 +15,7 @@ async fn middleware_preserves_three_state_fast_and_sources() {
|
||||
] {
|
||||
let snapshot = gateway_core::routing::RuntimeSnapshot::new(
|
||||
ConfigRevision::new(1).unwrap(),
|
||||
gateway_core::settings::SettingsValues::new(
|
||||
1,
|
||||
0,
|
||||
"smart",
|
||||
Default::default(),
|
||||
None,
|
||||
None,
|
||||
),
|
||||
gateway_core::settings::SettingsValues::new(1, "smart", Default::default(), None, None),
|
||||
vec![],
|
||||
vec![],
|
||||
vec![],
|
||||
|
||||
@@ -395,14 +395,7 @@ async fn active_http_children_inherit_settings_and_preserve_override_sources() {
|
||||
let baseline = RequestSettings::new(Arc::new(
|
||||
RuntimeSnapshot::new(
|
||||
ConfigRevision::new(1).unwrap(),
|
||||
gateway_core::settings::SettingsValues::new(
|
||||
1,
|
||||
10,
|
||||
"smart",
|
||||
Default::default(),
|
||||
None,
|
||||
None,
|
||||
),
|
||||
gateway_core::settings::SettingsValues::new(1, "smart", Default::default(), None, None),
|
||||
vec![],
|
||||
vec![],
|
||||
vec![],
|
||||
@@ -439,7 +432,7 @@ async fn active_http_children_inherit_settings_and_preserve_override_sources() {
|
||||
final_settings["responses_max_decompressed_body_bytes"],
|
||||
8192
|
||||
);
|
||||
assert_eq!(final_settings["request_interval_ms"], 0);
|
||||
assert_eq!(final_settings["max_account_rotations"], 0);
|
||||
assert!(final_settings["min_codex_cli_version"].is_null());
|
||||
let sources = effective.inspect();
|
||||
assert_eq!(
|
||||
@@ -447,14 +440,14 @@ async fn active_http_children_inherit_settings_and_preserve_override_sources() {
|
||||
"dispatch-b"
|
||||
);
|
||||
assert_eq!(
|
||||
sources["overrides"]["request_interval_ms"]["instance_id"],
|
||||
sources["overrides"]["max_account_rotations"]["instance_id"],
|
||||
"dispatch-a"
|
||||
);
|
||||
assert!(
|
||||
sources["overrides"]["min_codex_cli_version"]["order"]
|
||||
.as_u64()
|
||||
.unwrap()
|
||||
> sources["overrides"]["request_interval_ms"]["order"]
|
||||
> sources["overrides"]["max_account_rotations"]["order"]
|
||||
.as_u64()
|
||||
.unwrap()
|
||||
);
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user