Compare commits

..
23 Commits
Author SHA1 Message Date
ZyyandCodex 28ad8fb1c4 chore: bump version to 3.46.0
Co-authored-by: Codex <noreply@openai.com>
2026-11-12 21:12:25 +08:00
ZyyandZyy 4888681a04 feat(scheduling): weight quota resets by window in smart scheduling (#694)
Co-authored-by: Zyy <zyycn@users.noreply.github.com>
2026-11-12 18:56:30 +08:00
9739bb08be fix(egress): honor deployment proxies for unbound outbound requests (#693)
* fix(egress): honor deployment proxies for unbound outbound requests

Prefer the proxy selected for the request. Otherwise apply the deployment proxy and bypass rules consistently to HTTP and WebSocket connections, including background downloads. Keep explicit direct probes separate.

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

* fix(openai): preserve WebSocket dial errors

Recover the OS error or I/O kind from Hyper without retaining the endpoint. Keep refused connections distinguishable for both direct and proxied WebSocket dials.

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

---------

Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
2026-11-12 17:29:35 +08:00
fac37ec563 fix(update): preserve permissions when moving read-only directories (#688)
* fix(update): unlock official plugin dirs before rename on macOS

macOS rejects renaming a 0555 official plugin directory during backup
and rollback, while Linux only needs the parent writable.

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix(update): preserve directory permissions during macOS moves

Fall back to the existing permission-preserving copy path when macOS
rejects a directory rename. Keep chmod limited to source cleanup so
backups and restored releases retain their original directory modes.

Cover read-only staged assets, rollback, failed replacement, and
cancelled or rejected rollback with permission assertions.

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

* fix(update): handle read-only directory moves across Unix platforms

Linux also requires source directory write permission when rename moves
it between parents. Apply the same copy fallback to Unix while retaining
the cross-platform permission regression test.

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

---------

Co-authored-by: CzsGit <czsde@outlook.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
2026-11-12 13:20:16 +08:00
9f0d712671 fix(openai): honor advertised WebSocket message limits (#691)
Use each upstream handshake's message-size limit to select HTTP before
uploading oversized response.create frames. Remove fixed thresholds and
thread size hysteresis while preserving continuation replay and cleanup.

Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
2026-11-12 13:19:56 +08:00
ZyyandCodex 75be1287ef fix(openai): omit output token limits from Responses requests
Co-authored-by: Codex <noreply@openai.com>
2026-11-12 10:24:27 +08:00
ZyyandCodex dad91fe9f7 fix(ui): preserve semantic and neutral colors across theme presets
Co-authored-by: Codex <noreply@openai.com>
2026-11-11 20:48:41 +08:00
github-actions[bot] b22b0cfa41 chore(release): update release channels to 3.45.2 2026-11-11 09:28:26 +00:00
ZyyandCodex bd69329a37 chore: bump version to 3.45.2
Co-authored-by: Codex <noreply@openai.com>
2026-11-11 17:04:39 +08:00
ZyyandCodex 151a3c6bfd fix(openai): negotiate Codex catalogs for Responses upstreams
Co-authored-by: Codex <noreply@openai.com>
2026-11-11 16:49:10 +08:00
ZyyandCodex e2f194875f refactor: clarify protocol ownership and request observations
Separate protocol identifiers from Provider identities and centralize shared
request facts and observations in gateway-protocol. Keep endpoint models
and private continuation state owned by Providers.

Reorganize the architecture documentation and clarify the shutdown guarantees
for buffered charges and observations.

Validated with workspace formatting, strict Clippy, 3013 affected-module
tests and 16 architecture checks. One reference-binary TLS test is ignored.

Co-authored-by: Codex <noreply@openai.com>
2026-11-11 16:04:16 +08:00
ZyyandCodex 07b61bdc83 fix(ui): keep help buttons aligned with form labels and clarify hints
Co-authored-by: Codex <noreply@openai.com>
2026-11-11 15:13:00 +08:00
c0fdc1bbff fix(openai): record explicit Anthropic reasoning effort (#680)
Observe top-level output_config.effort before the OpenAI encoder rejects
Messages payloads. Reuse the protocol parser without inferring an effort
from thinking budgets or changing upstream conversion behavior.

Cover absent and explicit effort values, native Messages forwarding, and
Messages-to-Responses conversion. Document the observation contract.

Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
2026-11-11 15:12:02 +08:00
343c579017 fix(openai): rotate new round-robin sessions without affinity hits (#678)
* fix(openai): rotate new round-robin sessions without affinity hits

Strict affinity keeps a session on one account, but the shared request
cursor advanced on every scheduling load, including affinity hits. A new
session then took whatever slot that request volume happened to leave.

Advance a separate session cursor only when round-robin assigns an
account, and keep the request cursor for other strategies.

* fix(openai): advance round robin when an account assignment succeeds

Commit the session binding and round-robin cursor together after acquiring the account lease. Release a conflicting lease before retrying so concurrent first requests share a binding without consuming extra turns. Count delegated selections and preserve explicit plugin picks.

Verified with OpenAI provider tests, store tests against Redis and PostgreSQL, architecture checks, and workspace Clippy.

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

---------

Co-authored-by: koeria <riennie@qq.com>
Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
2026-11-11 14:40:59 +08:00
koeriaandCursor Agent 914c00785c fix(openai): skip same-account retry on NO_MORE_RETRY (#675)
Treat x-retry-metadata: NO_MORE_RETRY as a hard signal to skip
same-account backoff and transport recovery, even when Retry-After
is within maxRetryAfterSeconds. Capacity failures still rotate under
the existing policy; server_error still returns the upstream error.

Co-authored-by: Cursor Agent <cursoragent@cursor.com>
2026-11-11 14:08:24 +08:00
ZyyandCodex b3de0d637d fix(peer): send routing headers only to discovered nodes
Reuse connection-scoped peer discovery for native HTTP and WebSocket forwarding. Strip the route header for ordinary upstreams and leave explicit plugin forwarding changes to plugins.

Co-authored-by: Codex <noreply@openai.com>
2026-11-11 13:14:25 +08:00
ZyyandCodex 1271fa0186 test(openai): verify platform TLS against official Codex
Compare native HTTP, custom CA HTTP, and WebSocket ClientHello captures with an official Codex executable on the same platform. Document the explicit reference test workflow.

Co-authored-by: Codex <noreply@openai.com>
2026-11-11 12:41:04 +08:00
ZyyandCodex 055da40476 feat(keys): enable async compaction in Codex configs
Co-authored-by: Codex <noreply@openai.com>
2026-11-11 12:37:12 +08:00
ZyyandCodex 36ba79ca5e fix(openai): preserve native Responses service tiers
Limit default tier omission to Messages-to-OAuth conversion so native Responses retain their requested tier and routing hint.

Co-authored-by: Codex <noreply@openai.com>
2026-11-11 11:33:01 +08:00
github-actions[bot] 1b24029a97 chore(release): update release channels to 3.45.1 2026-11-11 00:03:25 +00:00
zyycnandCodex 431992a7a4 chore: bump version to 3.45.1
Co-authored-by: Codex <noreply@openai.com>
2026-11-11 07:41:09 +08:00
ZyyandCodex d3143406dc fix(openai): restore account catalogs without blocking requests
Persist validated directories per credential revision. Restore them in bulk and refresh only missing, expired, or changed accounts with paced background requests. Filter large pools using current per-account evidence while preserving unknown and explicitly allowed candidates.

Co-authored-by: Codex <noreply@openai.com>
2026-11-11 03:46:25 +08:00
github-actions[bot] 478f2d1f21 chore(release): update release channels to 3.45.0 2026-11-10 18:15:24 +00:00
225 changed files with 7309 additions and 4038 deletions
+14
View File
@@ -1807,6 +1807,17 @@ dependencies = [
"uuid",
]
[[package]]
name = "gateway-egress"
version = "0.0.0"
dependencies = [
"gateway-core",
"http",
"hyper-util",
"reqwest",
"url",
]
[[package]]
name = "gateway-host"
version = "0.0.0"
@@ -1822,6 +1833,7 @@ dependencies = [
"futures",
"gateway-admin",
"gateway-core",
"gateway-egress",
"gateway-protocol",
"hex",
"http",
@@ -3353,6 +3365,7 @@ dependencies = [
"futures",
"gateway-admin",
"gateway-core",
"gateway-egress",
"gateway-host",
"gateway-protocol",
"getrandom 0.4.3",
@@ -3410,6 +3423,7 @@ dependencies = [
"futures",
"gateway-admin",
"gateway-core",
"gateway-egress",
"gateway-protocol",
"getrandom 0.4.3",
"jsonwebtoken",
+1
View File
@@ -5,6 +5,7 @@ members = [
"crates/gateway-admin",
"crates/gateway-api",
"crates/gateway-core",
"crates/gateway-egress",
"crates/gateway-host",
"crates/gateway-plugin/runtime",
"crates/gateway-plugin/sdk",
@@ -16,6 +16,7 @@ pub(super) const WORKSPACE_MEMBERS: &[&str] = &[
"crates/gateway-admin",
"crates/gateway-api",
"crates/gateway-core",
"crates/gateway-egress",
"crates/gateway-host",
"crates/gateway-plugin/runtime",
"crates/gateway-plugin/sdk",
@@ -447,6 +448,7 @@ const PACKAGE_TO_MEMBER: &[(&str, &str)] = &[
("gateway-admin", "crates/gateway-admin"),
("gateway-api", "crates/gateway-api"),
("gateway-core", "crates/gateway-core"),
("gateway-egress", "crates/gateway-egress"),
("gateway-host", "crates/gateway-host"),
("gateway-plugin-runtime", "crates/gateway-plugin/runtime"),
("gateway-plugin-sdk", "crates/gateway-plugin/sdk"),
@@ -510,8 +512,10 @@ const ALLOWED_INTERNAL_EDGES: &[(&str, &str)] = &[
("gateway-api", "gateway-core"),
("gateway-api", "gateway-protocol"),
("gateway-core", "gateway-protocol"),
("gateway-egress", "gateway-core"),
("gateway-host", "gateway-admin"),
("gateway-host", "gateway-core"),
("gateway-host", "gateway-egress"),
("gateway-host", "gateway-protocol"),
("gateway-plugin-runtime", "gateway-host"),
("gateway-plugin-runtime", "gateway-core"),
@@ -522,9 +526,11 @@ const ALLOWED_INTERNAL_EDGES: &[(&str, &str)] = &[
("gateway-store", "gateway-core"),
("provider-openai", "gateway-admin"),
("provider-openai", "gateway-core"),
("provider-openai", "gateway-egress"),
("provider-openai", "gateway-protocol"),
("provider-xai", "gateway-admin"),
("provider-xai", "gateway-core"),
("provider-xai", "gateway-egress"),
("provider-xai", "gateway-protocol"),
];
@@ -78,7 +78,12 @@ async fn proxy_probe_should_use_provider_custom_ca_for_https_proxies() {
});
let probe = HttpProxyProbe::new("http://unresolvable.invalid/ip")
.with_client_builder(provider_openai::build_reqwest_client_with_custom_ca);
let result = probe.test(Some(&proxy), false).await;
let result = probe
.test(
gateway_admin::ports::proxy::ProbeRoute::Proxy(&proxy),
false,
)
.await;
assert!(result.success, "{}", result.message);
assert_eq!(result.exit_ip.unwrap().to_string(), "203.0.113.8");
tokio::time::timeout(Duration::from_secs(5), server)
@@ -146,7 +146,7 @@ pub struct AccountRecord {
}
impl AccountRecord {
/// 配置保留绑定,执行请求时禁用代理按直连处理
/// 配置保留绑定,执行请求时禁用代理使用部署默认出口
#[must_use]
pub fn effective_outbound_proxy(&self) -> Option<&gateway_core::account::OutboundProxy> {
self.outbound_proxy
@@ -69,8 +69,14 @@ pub struct ProxyImportReservation {
pub guard: Box<dyn ProxyImportGuard>,
}
#[derive(Clone, Copy)]
pub enum ProbeRoute<'a> {
Direct,
Deployment,
Proxy(&'a OutboundProxy),
}
#[async_trait]
pub trait ProxyProbe: Send + Sync {
/// 未指定代理时探测服务器直连出口,不继承进程代理环境变量
async fn test(&self, proxy: Option<&OutboundProxy>, detect_location: bool) -> ProxyTestResult;
async fn test(&self, proxy: ProbeRoute<'_>, detect_location: bool) -> ProxyTestResult;
}
@@ -52,7 +52,7 @@ impl PeerInfoService {
});
let limits = peer.and_then(|peer| peer.key_limits);
let info = ProviderUpstreamInfo {
discovered: peer.is_some() || stored_failure == Some(UpstreamInfoFailure::Unauthorized),
discovered: snapshot::is_discovered(peer, stored_failure),
observed_at,
stale,
connection_revision: connection.revision.clone(),
@@ -282,6 +282,21 @@ impl PeerInfoService {
}
impl UpstreamInfoSource for PeerInfoService {
fn is_peer(&self, account: &ProviderAccount) -> bool {
if !self.supported(account) {
return false;
}
let connection = UpstreamConnection::for_account(account);
self.observations
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.get(account.id())
.is_some_and(|cached| {
cached.connection == connection
&& snapshot::is_discovered(cached.summary.as_deref(), cached.stored_failure)
})
}
fn signal(
&self,
account: &ProviderAccount,
@@ -16,6 +16,10 @@ use serde::Deserialize as _;
pub(super) const MAX_AGE: Duration = Duration::from_secs(30 * 60);
pub(super) const FRESH_FOR: Duration = Duration::from_secs(5 * 60);
pub(super) fn is_discovered(peer: Option<&PeerLink>, failure: Option<UpstreamInfoFailure>) -> bool {
peer.is_some() || failure == Some(UpstreamInfoFailure::Unauthorized)
}
// 不支持协议是连接的识别结果,复用观测存储与版本保护,不伪造上游额度快照
pub(super) fn unsupported(connection: &UpstreamConnection) -> OpaqueProviderData {
OpaqueProviderData::new(serde_json::Map::from_iter([
@@ -306,7 +306,9 @@ impl PluginsService {
.await
.map_err(|error| map_store_error(error, "proxy"))?;
if !record.enabled {
return Err(AdminError::invalid("下载代理已禁用,请启用或改为直连"));
return Err(AdminError::invalid(
"下载代理已禁用,请启用或改用部署默认出口",
));
}
Ok(Some(PluginDistributionEgress::new(
record.id,
@@ -23,7 +23,7 @@ use crate::{
},
ports::{
provider::ProviderAdminRegistry,
proxy::{PluginProxyAccess, ProxyProbe, ProxyStore},
proxy::{PluginProxyAccess, ProbeRoute, ProxyProbe, ProxyStore},
store::AdminStoreErrorKind,
},
};
@@ -47,7 +47,7 @@ impl PluginProxyAccess for DefaultPluginProxyAccess {
.get(id)
.await
.map_err(|error| map_store_error(error, "proxy"))?;
// 显式选择代理时必须有可用配置,不能沿用账号停用代理后直连的语义
// 显式选择代理时必须有可用配置,不能沿用账号停用代理后使用部署默认出口的语义
if !record.enabled {
return Err(AdminError::invalid("代理已禁用,请启用后再选择"));
}
@@ -114,10 +114,10 @@ pub trait ProxiesService: Send + Sync {
detect_location: bool,
context: &MutationContext,
) -> Result<ProxyRecord, AdminError>;
/// 探测未保存的连接地址或服务器直连出口,不写入配置或修改账号绑定
/// 探测未保存的连接地址、部署默认出口或服务器直连出口,不写入配置或修改账号绑定
async fn probe(
&self,
proxy: Option<&OutboundProxy>,
proxy: ProbeRoute<'_>,
detect_location: bool,
) -> Result<ProxyTestResult, AdminError>;
}
@@ -323,7 +323,7 @@ impl ProxiesService for DefaultProxiesService {
.transpose()
.map_err(|_| AdminError::invalid("代理位置不合法"))?;
command.test = if command.auto_location {
Some(self.probe(Some(&command.proxy), true).await?)
Some(self.probe(ProbeRoute::Proxy(&command.proxy), true).await?)
} else {
None
};
@@ -362,7 +362,7 @@ impl ProxiesService for DefaultProxiesService {
if command.auto_location.unwrap_or(current.auto_location)
&& (!current.auto_location || proxy != &current.proxy)
{
command.test = Some(self.probe(Some(proxy), true).await?);
command.test = Some(self.probe(ProbeRoute::Proxy(proxy), true).await?);
}
}
let result = self
@@ -391,7 +391,7 @@ impl ProxiesService for DefaultProxiesService {
async fn probe(
&self,
proxy: Option<&OutboundProxy>,
proxy: ProbeRoute<'_>,
detect_location: bool,
) -> Result<ProxyTestResult, AdminError> {
let _permit = self
@@ -423,7 +423,10 @@ impl ProxiesService for DefaultProxiesService {
// 手动解析只为本次测试请求位置,不改变自动跟随的持久配置
let result = self
.probe
.test(Some(&record.proxy), record.auto_location || detect_location)
.test(
ProbeRoute::Proxy(&record.proxy),
record.auto_location || detect_location,
)
.await;
let mutation = self
.store
@@ -140,6 +140,10 @@ impl PendingPeer {
}
impl UpstreamInfoSource for PendingPeer {
fn is_peer(&self, _: &ProviderAccount) -> bool {
false
}
fn signal(
&self,
_: &ProviderAccount,
@@ -239,6 +239,7 @@ async fn discovery_uses_provider_capabilities_and_publishes_neutral_signals() {
initialize_peer_upstream(accounts.clone(), &registry, transport.clone()).unwrap();
let initial_accounts = accounts.accounts.lock().unwrap().clone();
let reader = UpstreamInfoReader::default();
assert!(!reader.is_peer(&initial_accounts[0]));
assert!(
reader
.signal(&initial_accounts[0], Some("model"), SystemTime::now())
@@ -246,6 +247,7 @@ async fn discovery_uses_provider_capabilities_and_publishes_neutral_signals() {
);
let source = bundle.scheduling_source();
reader.bind(&source).unwrap();
assert!(!reader.is_peer(&initial_accounts[0]));
assert!(reader.bind(&source).is_err());
drop(source);
let WorkerContribution::Registration(worker) = bundle.take_worker_contributions().remove(0)
@@ -272,6 +274,10 @@ async fn discovery_uses_provider_capabilities_and_publishes_neutral_signals() {
}
);
for account in &initial_accounts[..2] {
assert_eq!(
reader.is_peer(account),
Some(account.provider().as_str()) != failed_provider
);
if Some(account.provider().as_str()) == failed_provider {
assert!(
reader
@@ -304,6 +310,7 @@ async fn discovery_uses_provider_capabilities_and_publishes_neutral_signals() {
);
drop(worker);
drop(bundle);
assert!(!reader.is_peer(&initial_accounts[0]));
assert!(
reader
.signal(&initial_accounts[0], Some("model"), SystemTime::now())
@@ -340,6 +347,7 @@ async fn account_edits_preserve_observation_and_refresh_budget_until_connection_
.is_some()
);
let after = service.refresh_quota(edited.id()).await.unwrap();
assert!(bundle.scheduling_source().is_peer(&edited));
assert_eq!(after.upstream, before.upstream);
assert_eq!(transport.0.lock().unwrap().len(), 1);
@@ -347,6 +355,7 @@ async fn account_edits_preserve_observation_and_refresh_budget_until_connection_
OutboundProxy::parse("http://proxy.invalid:8080").unwrap(),
));
*accounts.accounts.lock().unwrap() = vec![reconnected.clone()];
assert!(!bundle.scheduling_source().is_peer(&reconnected));
service.account_changed(reconnected.id()).await.unwrap();
assert!(
bundle
@@ -422,6 +431,21 @@ fn loaded(account: &ProviderAccount) -> LoadedCredential {
#[async_trait]
impl ProviderAccountStore for TestAccounts {
async fn load_account_catalogs(
&self,
_: &gateway_core::routing::ProviderKind,
) -> Result<Vec<gateway_core::account::StoredAccountCatalog>, gateway_core::error::StoreError>
{
Ok(Vec::new())
}
async fn save_account_catalog(
&self,
_: &gateway_core::account::AccountCatalogObservation,
_: &[u8],
) -> Result<bool, gateway_core::error::StoreError> {
unreachable!("fixture does not discover account catalogs")
}
async fn create_account(&self, _account: NewProviderAccount) -> Result<(), StoreError> {
panic!("unused account operation")
}
@@ -215,6 +215,10 @@ async fn known_peer_identity_changes_only_on_definitive_detection() {
quota.upstream
);
let account = accounts.accounts.lock().unwrap()[0].clone();
assert_eq!(
bundle.scheduling_source().is_peer(&account),
!clears_identity
);
assert_eq!(
bundle
.scheduling_source()
@@ -237,6 +241,7 @@ async fn known_peer_identity_changes_only_on_definitive_detection() {
drop(bundle);
let restarted = initialize(accounts, transport);
let restored = restarted.service().read_quota(&id).await.unwrap().unwrap();
assert!(!restarted.scheduling_source().is_peer(&account));
assert_eq!(
restored.upstream.as_ref().unwrap().discovered,
!clears_identity
@@ -258,6 +263,10 @@ async fn known_peer_identity_changes_only_on_definitive_detection() {
.upstream,
restored.upstream
);
assert_eq!(
restarted.scheduling_source().is_peer(&account),
!clears_identity
);
}
}
@@ -4,7 +4,7 @@ use async_trait::async_trait;
use gateway_admin::{
model::{MutationContext, Revision, proxies::*},
ports::{
proxy::{ProxyProbe, ProxyStore},
proxy::{ProbeRoute, ProxyProbe, ProxyStore},
store::AdminStoreResult,
},
};
@@ -194,7 +194,7 @@ async fn plugin_proxy_resolution_requires_enabled_configuration_without_account_
#[async_trait]
impl ProxyProbe for TestProxies {
async fn test(&self, _: Option<&OutboundProxy>, _: bool) -> ProxyTestResult {
async fn test(&self, _: ProbeRoute<'_>, _: bool) -> ProxyTestResult {
panic!("unexpected proxy probe")
}
}
@@ -456,7 +456,7 @@ impl Drop for ActiveProbe<'_> {
#[async_trait]
impl ProxyProbe for GatedProxyProbe {
async fn test(&self, _: Option<&OutboundProxy>, _: bool) -> ProxyTestResult {
async fn test(&self, _: ProbeRoute<'_>, _: bool) -> ProxyTestResult {
self.active
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let _active = ActiveProbe(&self.active);
@@ -526,8 +526,8 @@ async fn batch_tests_share_capacity_preserve_order_and_release_slots_on_cancella
_ = async { receiver.recv().await.unwrap(); receiver.recv().await.unwrap(); } => {}
}
assert_eq!(probe.active.load(Ordering::Relaxed), 2);
let mut single_a = Box::pin(services.proxies().probe(None, false));
let mut single_b = Box::pin(services.proxies().probe(None, false));
let mut single_a = Box::pin(services.proxies().probe(ProbeRoute::Direct, false));
let mut single_b = Box::pin(services.proxies().probe(ProbeRoute::Direct, false));
tokio::select! {
_ = &mut batch => panic!("batch unexpectedly finished"),
_ = &mut single_a => panic!("single probe unexpectedly finished"),
@@ -538,7 +538,7 @@ async fn batch_tests_share_capacity_preserve_order_and_release_slots_on_cancella
assert_eq!(
services
.proxies()
.probe(None, false)
.probe(ProbeRoute::Direct, false)
.await
.unwrap_err()
.kind(),
@@ -100,7 +100,9 @@ fn contract_type(mut item: Item, name: &str) -> Item {
attrs.retain(|attr| attr.path().is_ident("serde"));
attrs.insert(
0,
if name == "SmartSchedulingConfig" {
if name == "ResetWindowWeights" {
parse_quote!(#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)])
} else if name == "SmartSchedulingConfig" {
parse_quote!(#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)])
} else if comparable {
parse_quote!(#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)])
@@ -256,7 +258,7 @@ fn sdk_settings_contract_matches_host_declarations() {
),
(
"../gateway-core/src/account/smart_scheduling.rs",
&["SmartSchedulingValues"][..],
&["SmartSchedulingValues", "ResetWindowWeights"][..],
),
(
"../gateway-core/src/metering/pricing.rs",
@@ -101,6 +101,8 @@ struct TestRequest {
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct ProbeRequest {
#[serde(default)]
use_deployment_proxy: bool,
#[serde(default)]
detect_location: bool,
proxy_url: AccountProxyUpdate,
@@ -537,10 +539,21 @@ async fn probe<S>(
where
S: SessionState + Send + Sync,
{
use gateway_admin::ports::proxy::ProbeRoute;
let route = match (request.proxy_url.0.as_ref(), request.use_deployment_proxy) {
(Some(_), true) => {
return Err(AdminError::bad_request(
"指定代理地址时不能同时选择部署默认出口",
));
}
(Some(proxy), false) => ProbeRoute::Proxy(proxy),
(None, true) => ProbeRoute::Deployment,
(None, false) => ProbeRoute::Direct,
};
let result = state
.admin_services()
.proxies()
.probe(request.proxy_url.0.as_ref(), request.detect_location)
.probe(route, request.detect_location)
.await
.map_err(map_error)?;
Ok(AdminResponse::new(
@@ -658,7 +658,7 @@ where
let result = state
.admin_services()
.proxies()
.probe(None, true)
.probe(gateway_admin::ports::proxy::ProbeRoute::Direct, true)
.await
.map_err(map_service_error)?;
Ok(AdminResponse::new(
@@ -22,6 +22,7 @@ use gateway_core::{
error::GatewayError,
operation::Operation,
};
use gateway_protocol::openai::OPENAI_PROTOCOL;
use crate::client_error::{ProtocolError, ProtocolErrorBody};
use crate::middleware::headers::encode_headers;
@@ -65,7 +66,7 @@ where
Box::new(move |prepared, request| {
Box::pin(async move {
let (protocol, headers, body) = request_parts(request.clone())?;
if protocol != "openai" {
if protocol != OPENAI_PROTOCOL {
return Err(MiddlewareError::Rejected);
}
let operation = request.apply_capabilities(decode(body, headers).await?)?;
@@ -80,7 +81,8 @@ where
{
Ok(started) => started,
Err(error) => {
return buffered_response("openai", gateway_error_response(&error)).await;
return buffered_response(OPENAI_PROTOCOL, gateway_error_response(&error))
.await;
}
};
collect_raw_json_response(started).await
@@ -89,7 +91,9 @@ where
)
.await;
let response = match result {
Ok(response) => into_http_response(response, ExpectedBody::SingleJson, "openai").await,
Ok(response) => {
into_http_response(response, ExpectedBody::SingleJson, OPENAI_PROTOCOL).await
}
Err(error) => error_response(error),
};
if !execution_entered.load(Ordering::Relaxed) {
@@ -143,7 +147,7 @@ async fn collect_raw_json_response(
.await
.map_err(|_| MiddlewareError::Fault)?;
Ok(pending_execution_response(
"openai".to_owned(),
OPENAI_PROTOCOL.to_owned(),
parts.status.as_u16(),
encode_headers(&parts.headers),
MiddlewareFrame::new(bytes, framing, true).with_transformed(transformed),
@@ -155,7 +159,7 @@ fn raw_json_body(events: Vec<ProviderEvent>) -> Option<Bytes> {
let mut body = None;
for event in events {
let (_, wire) = event.into_parts();
let Some(wire) = wire.filter(|wire| wire.protocol() == "openai") else {
let Some(wire) = wire.filter(|wire| wire.protocol() == OPENAI_PROTOCOL) else {
continue;
};
let Some(raw) = wire.into_raw_json_body() else {
@@ -32,6 +32,7 @@ use gateway_core::engine::{
use gateway_core::error::{GatewayError, GatewayErrorKind};
use gateway_core::event::{ProviderEvent, ProviderResponseHeader};
use gateway_core::lifecycle::{CancellationToken, ConnectionGuard};
use gateway_protocol::openai::OPENAI_PROTOCOL;
use gateway_protocol::openai::sse::response_failed_sse_event_with_id;
use serde_json::Value;
use tokio::time::Instant;
@@ -491,7 +492,9 @@ pub async fn collect_execution_response(session: Box<dyn ExecutionSession>) -> R
)
.await
{
Ok(response) => into_http_response(response, ExpectedBody::SingleJson, "openai").await,
Ok(response) => {
into_http_response(response, ExpectedBody::SingleJson, OPENAI_PROTOCOL).await
}
Err(error) => error_response(error),
}
}
@@ -513,7 +516,9 @@ pub async fn stream_execution_response(
)
.await
{
Ok(response) => into_http_response(response, ExpectedBody::Sse, "openai").await,
Ok(response) => {
into_http_response(response, ExpectedBody::Sse, OPENAI_PROTOCOL).await
}
Err(error) => error_response(error),
}
},
@@ -15,6 +15,8 @@ use gateway_core::{
error::{GatewayError, GatewayErrorKind},
operation::OperationKind,
};
use gateway_protocol::anthropic::ANTHROPIC_PROTOCOL;
use gateway_protocol::openai::OPENAI_PROTOCOL;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum GenerationProtocol {
@@ -27,7 +29,7 @@ impl GenerationProtocol {
pub(crate) fn from_request(request: &DecodedGenerationRequest) -> Self {
if request.operation().kind() == OperationKind::CountTokens {
Self::MessageTokenCount
} else if request.operation().protocol() == "anthropic" {
} else if request.operation().protocol() == ANTHROPIC_PROTOCOL {
Self::Messages
} else {
Self::Responses
@@ -36,8 +38,8 @@ impl GenerationProtocol {
pub(crate) const fn wire(self) -> &'static str {
match self {
Self::Responses => "openai",
Self::Messages | Self::MessageTokenCount => "anthropic",
Self::Responses => OPENAI_PROTOCOL,
Self::Messages | Self::MessageTokenCount => ANTHROPIC_PROTOCOL,
}
}
@@ -2,6 +2,7 @@
use axum::body::Bytes;
use gateway_core::{error::GatewayError, event::ProviderEvent};
use gateway_protocol::anthropic::ANTHROPIC_PROTOCOL;
use gateway_protocol::{
anthropic::MessagesAccumulator,
openai::sse::{
@@ -50,7 +51,7 @@ impl GenerationEncoder {
} => {
let Some(wire) = event
.wire_event()
.filter(|wire| wire.protocol() == "anthropic")
.filter(|wire| wire.protocol() == ANTHROPIC_PROTOCOL)
else {
return Vec::new();
};
@@ -1,6 +1,7 @@
//! 出口校验按协议维护边界事实,不在流式路径保存生成正文
use gateway_core::event::ProviderEvent;
use gateway_protocol::anthropic::ANTHROPIC_PROTOCOL;
use gateway_protocol::openai::sse::parse_sse_events;
use serde_json::Value;
use std::sync::{Arc, Mutex};
@@ -32,7 +33,7 @@ impl GenerationValidationFacts {
Self::Messages(id) => {
if let Some(value) = event
.wire_event()
.filter(|wire| wire.protocol() == "anthropic")
.filter(|wire| wire.protocol() == ANTHROPIC_PROTOCOL)
.map(|wire| wire.data())
&& let Some(value) = value
.pointer("/message/id")
@@ -13,6 +13,7 @@ use gateway_core::error::{GatewayError, GatewayErrorKind};
use gateway_core::operation::{
ImageRequest, ImageRequestKind, Operation, OperationKind, RawJsonPayload,
};
use gateway_protocol::openai::OPENAI_PROTOCOL;
use serde_json::Value;
use crate::ApiState;
@@ -25,7 +26,6 @@ use crate::openai::{
router::{IMAGE_EDITS_PATH, IMAGE_GENERATIONS_PATH},
};
const OPENAI_PROTOCOL: &str = "openai";
const IMAGE_TURN_ID_CONTEXT_KEY: &str = "image_turn_id";
/// `POST /v1/images/generations`
@@ -13,6 +13,7 @@ use gateway_core::engine::execution::StartedExecution;
use gateway_core::event::ProviderEvent;
use gateway_core::live::LiveHangupRequest;
use gateway_core::operation::Operation;
use gateway_protocol::openai::OPENAI_PROTOCOL;
use serde_json::Value;
use super::{
@@ -31,9 +32,6 @@ use crate::{
},
};
/// 引导 wire 事件使用的协议名;与 Provider 侧 `PROVIDER_NAME` 一致。
const LIVE_PROTOCOL: &str = "openai";
pub(crate) async fn live_call(
State(state): State<ApiState>,
connect_info: Option<Extension<ConnectInfo<SocketAddr>>>,
@@ -344,7 +342,7 @@ fn live_call_body(events: Vec<ProviderEvent>) -> Option<Bytes> {
let mut raw_events = 0usize;
for event in events {
let (_, wire) = event.into_parts();
let Some(wire) = wire.filter(|wire| wire.protocol() == LIVE_PROTOCOL) else {
let Some(wire) = wire.filter(|wire| wire.protocol() == OPENAI_PROTOCOL) else {
continue;
};
let raw = wire.into_raw_http_body()?;
@@ -20,6 +20,7 @@ use gateway_core::event::ProviderResponseHeader;
use gateway_core::operation::{
ProviderHttpHeader, ProviderHttpMethod, ProviderHttpRequest, RawHttpPayload,
};
use gateway_protocol::openai::OPENAI_PROTOCOL;
use serde_json::{Value, json};
use crate::openai::error::openai_error_response;
@@ -172,7 +173,7 @@ pub(crate) fn live_call_operation(
"invalid_request",
))
};
let payload = RawHttpPayload::new("openai", body).map_err(|_| invalid())?;
let payload = RawHttpPayload::new(OPENAI_PROTOCOL, body).map_err(|_| invalid())?;
ProviderHttpRequest::new(
LIVE_CALLS_ENDPOINT,
ProviderHttpMethod::Post,
@@ -20,6 +20,7 @@ use gateway_core::{
lifecycle::CancellationToken,
operation::OperationKind,
};
use gateway_protocol::openai::OPENAI_PROTOCOL;
use std::{future::Future, sync::Arc};
use super::http::{ExpectedBody, buffered_response, error_response, into_http_response};
@@ -38,7 +39,7 @@ impl RequestInput {
pub(crate) fn query(endpoint: String, model_hint: Option<String>, headers: &HeaderMap) -> Self {
Self {
endpoint,
protocol: "openai".to_owned(),
protocol: OPENAI_PROTOCOL.to_owned(),
operation: None,
transport: ClientTransport::HttpJson,
model_hint,
@@ -10,6 +10,7 @@ use gateway_core::{
engine::execution::AuthenticatedClient,
routing::{PublicModelDescriptor, PublicModelId, PublicModelProfile},
};
use gateway_protocol::anthropic::ANTHROPIC_PROTOCOL;
use serde::Deserialize;
use serde_json::{Value, json};
@@ -52,7 +53,7 @@ pub(crate) async fn models(
let mut input = RequestInput::query(super::router::MODELS_PATH.to_owned(), None, &headers);
if anthropic {
input.protocol = "anthropic".to_owned();
input.protocol = ANTHROPIC_PROTOCOL.to_owned();
}
let response = middleware::query_response(service.execution(), client, input, move |client| {
models_response(state, client, query, anthropic)
@@ -173,7 +174,7 @@ pub(crate) async fn model_detail(
&headers,
);
if anthropic {
input.protocol = "anthropic".to_owned();
input.protocol = ANTHROPIC_PROTOCOL.to_owned();
}
let response = middleware::query_response(
service.execution(),
@@ -6,6 +6,7 @@ use axum::http::{HeaderMap, HeaderName, HeaderValue};
use base64::{Engine as _, engine::general_purpose::STANDARD};
use bytes::Bytes;
use gateway_core::operation::{ClientTurn, GenerateRequest, Operation, ProtocolPayload};
use gateway_protocol::openai::OPENAI_PROTOCOL;
use gateway_protocol::openai::{
X_OPENAI_INTERNAL_CODEX_RESPONSES_LITE_HEADER, X_OPENAI_MEMGEN_REQUEST_HEADER,
is_transport_managed_request_header,
@@ -17,7 +18,6 @@ use crate::openai::generation::{
ContinuationIntent, DecodedGenerationRequest, GenerationRequestMetadata,
};
const OPENAI_PROTOCOL: &str = "openai";
const OPENAI_SUBAGENT_KEY: &str = "x-openai-subagent";
const CODEX_TURN_METADATA_KEY: &str = "x-codex-turn-metadata";
const PASSTHROUGH_HEADERS_CONTEXT_KEY: &str = "opaque_request_headers";
@@ -4,6 +4,7 @@ use std::collections::BTreeMap;
use bytes::Bytes;
use gateway_core::event::{GatewayEvent, ProtocolWireEvent, ProviderEvent};
use gateway_protocol::openai::OPENAI_PROTOCOL;
use gateway_protocol::openai::sse::{
encode_sse_event_with_metadata, response_failed_sse_data_from_error_event,
};
@@ -12,8 +13,6 @@ use serde_json::Value;
use super::error::ResponseEncodeError;
use crate::client_error::capacity_error_for_client;
const OPENAI_PROTOCOL: &str = "openai";
/// OpenAI Responses wire 的转发器
///
/// OpenAI Provider 与 xAI adapter 都必须交付 OpenAI wire
@@ -6,6 +6,7 @@ use std::{
};
use gateway_core::event::{GatewayEvent, ProtocolWireEvent, ProviderEvent};
use gateway_protocol::openai::OPENAI_PROTOCOL;
use gateway_protocol::openai::sse::{parse_sse_events, sse_frame_is_done};
use serde_json::{Map, Value};
@@ -44,7 +45,7 @@ impl ResponseValidationFacts {
if let Some(wire) = event
.middleware_origin_wire()
.or_else(|| event.wire_event())
.filter(|wire| wire.protocol() == "openai")
.filter(|wire| wire.protocol() == OPENAI_PROTOCOL)
{
facts.observe_wire(wire);
}
@@ -16,6 +16,7 @@ use gateway_core::engine::response_control::{ResponseControl, ResponseControlUna
use gateway_core::error::{GatewayError, GatewayErrorKind};
use gateway_core::event::ProviderResponseHeader;
use gateway_core::operation::ProviderSessionState;
use gateway_protocol::openai::OPENAI_PROTOCOL;
use crate::openai::error::{gateway_error_contract, gateway_error_from_engine};
use crate::openai::middleware::{PendingExecution, pending_execution_response};
@@ -118,7 +119,7 @@ pub(super) async fn execution_response(
.collect();
// 失败也要保留会话到洋葱链返回,否则释放会话会取消仍在等待 next 的插件
return Ok(pending_execution_response(
"openai".to_owned(),
OPENAI_PROTOCOL.to_owned(),
status,
response_headers,
MiddlewareFrame::new(
@@ -142,7 +143,7 @@ pub(super) async fn execution_response(
.map(|header| MiddlewareHeader::new(header.name(), header.value().clone()))
.collect();
Ok(MiddlewareResponse::new(
"openai".to_owned(),
OPENAI_PROTOCOL.to_owned(),
200,
headers,
Box::new(body),
@@ -159,7 +160,7 @@ pub(super) async fn forward_response(
response_control: &ResponseControl,
) -> ForwardOutcome {
let (protocol, status, headers, mut body, _) = response.into_parts();
if protocol != "openai" || StatusCode::from_u16(status).is_err() {
if protocol != OPENAI_PROTOCOL || StatusCode::from_u16(status).is_err() {
return fail_body(connection, body, MiddlewareError::InvalidState, &request_id).await;
}
let response_headers = headers
@@ -26,6 +26,7 @@ use gateway_core::{
lifecycle::{ConnectionGuard, ConnectionLifecycle},
operation::OperationKind,
};
use gateway_protocol::openai::OPENAI_PROTOCOL;
use crate::middleware::headers::encode_headers;
use crate::{
@@ -328,7 +329,7 @@ async fn serve_responses_websocket(socket: Socket, session: ResponsesWebSocketSe
let validation = ResponseValidationFacts::default();
let input = RequestInput {
endpoint: crate::openai::router::RESPONSES_PATH.to_owned(),
protocol: "openai".to_owned(),
protocol: OPENAI_PROTOCOL.to_owned(),
operation: Some(OperationKind::Generate),
transport: ClientTransport::WebSocket,
model_hint,
@@ -349,7 +350,7 @@ async fn serve_responses_websocket(socket: Socket, session: ResponsesWebSocketSe
Box::new(move |prepared, request| {
Box::pin(async move {
let (protocol, headers, body) = request_parts(request.clone())?;
if protocol != "openai" {
if protocol != OPENAI_PROTOCOL {
return Err(MiddlewareError::Rejected);
}
let payload =
@@ -11,6 +11,7 @@ use axum::{
use gateway_core::engine::execution::ClientTransport;
use gateway_core::error::{GatewayError, GatewayErrorKind};
use gateway_core::operation::{Operation, OperationKind, RawJsonPayload, StandaloneSearchRequest};
use gateway_protocol::openai::OPENAI_PROTOCOL;
use crate::ApiState;
use crate::middleware::headers::encode_headers;
@@ -22,8 +23,6 @@ use crate::openai::{
router::SEARCH_PATH,
};
const OPENAI_PROTOCOL: &str = "openai";
/// `POST /v1/alpha/search`
pub(crate) async fn standalone_search(
State(state): State<ApiState>,
@@ -15,6 +15,7 @@ use gateway_core::lifecycle::{ConnectionDraining, ConnectionGuard, ConnectionLif
use gateway_core::routing::{
ProviderCatalogUnavailable, ProviderKind, PublicModelDescriptor, PublicModelId, UpstreamModelId,
};
use gateway_protocol::openai::OPENAI_PROTOCOL;
use uuid::Uuid;
use super::auth::ClientApiKeyAuthError;
@@ -159,7 +160,7 @@ impl OpenAiService {
upstream_model,
operation,
ExecutionRequestMetadata {
protocol: "openai".to_owned(),
protocol: OPENAI_PROTOCOL.to_owned(),
endpoint,
transport: ClientTransport::HttpJson,
stream: false,
@@ -11,11 +11,11 @@ use chrono::Utc;
use gateway_admin::{
model::{MutationContext, Revision, proxies::*},
ports::{
proxy::{ProxyProbe, ProxyStore},
proxy::{ProbeRoute, ProxyProbe, ProxyStore},
store::{AdminStoreError, AdminStoreErrorKind, AdminStoreResult},
},
};
use gateway_core::account::{OutboundProxy, ProviderAccountId};
use gateway_core::account::ProviderAccountId;
use serde_json::{Value, json};
use tower::ServiceExt as _;
@@ -348,17 +348,17 @@ pub(super) struct SuccessfulProbe;
#[async_trait]
impl ProxyProbe for SuccessfulProbe {
async fn test(&self, proxy: Option<&OutboundProxy>, detect_location: bool) -> ProxyTestResult {
if let Some(proxy) = proxy {
async fn test(&self, proxy: ProbeRoute<'_>, detect_location: bool) -> ProxyTestResult {
if let ProbeRoute::Proxy(proxy) = proxy {
assert_eq!(
proxy.expose_url(),
"http://test-user:private-password@proxy.example:8080/"
);
}
let exit_ip = if proxy.is_some() {
"203.0.113.2"
} else {
"203.0.113.3"
let exit_ip = match proxy {
ProbeRoute::Proxy(_) => "203.0.113.2",
ProbeRoute::Direct => "203.0.113.3",
ProbeRoute::Deployment => "203.0.113.4",
};
ProxyTestResult {
location: if detect_location {
@@ -600,6 +600,15 @@ async fn proxy_probe_checks_direct_egress_without_changing_saved_proxies() {
}
);
}
let (status, deployed) = request(
&fixture,
"/api/admin/proxies/probe",
Some(json!({"proxyUrl": "", "useDeploymentProxy": true})),
true,
)
.await;
assert_eq!(status, StatusCode::OK);
assert_eq!(deployed["data"]["exitIpv4"], "203.0.113.4");
let (_, listed) = request(&fixture, "/api/admin/proxies", None, true).await;
assert_eq!(listed["data"]["page"]["total"], 1);
assert_eq!(
@@ -613,6 +622,10 @@ async fn proxy_probe_rejects_missing_or_invalid_addresses() {
let fixture = AdminTestFixture::new().await;
fixture.auth.insert_session("valid-session");
for (body, expected) in [
(
json!({"proxyUrl": "http://proxy.example:8080", "useDeploymentProxy": true}),
StatusCode::BAD_REQUEST,
),
(json!({"proxyUrl": " "}), StatusCode::UNPROCESSABLE_ENTITY),
(json!({"proxyUrl": null}), StatusCode::UNPROCESSABLE_ENTITY),
(json!({}), StatusCode::UNPROCESSABLE_ENTITY),
@@ -268,7 +268,7 @@ async fn smart_settings_round_trip_and_invalid_updates_leave_the_saved_value_int
let fixture = AdminTestFixture::new().await;
fixture.auth.insert_session("valid-session");
let mut body = update_body();
let custom = json!({"loadWeight": 2.0, "quotaWeight": 0.0, "healthWeight": 1.0, "latencyWeight": 0.5, "resetWeight": 1.2, "queueWeight": 2.3});
let custom = json!({"loadWeight": 2.0, "quotaWeight": 0.0, "healthWeight": 1.0, "latencyWeight": 0.5, "resetWeight": 1.2, "queueWeight": 2.3, "resetWindowWeights":{"fiveHour":0.0,"weekly":1.5,"monthly":0.0}});
body["smartScheduling"] = custom.clone();
let response = app(fixture.state())
.oneshot(request(
@@ -301,6 +301,19 @@ async fn smart_settings_round_trip_and_invalid_updates_leave_the_saved_value_int
missing.as_object_mut().unwrap().remove(field);
invalid_values.push(missing);
}
for windows in [
json!({}),
json!({"weekly":1}),
json!(null),
json!({"fiveHour":0,"weekly":1.23,"monthly":0}),
json!({"fiveHour":-1,"weekly":1,"monthly":0}),
json!({"fiveHour":0,"weekly":1,"monthly":10.1}),
json!({"fiveHour":0,"weekly":1,"monthly":0,"daily":1}),
] {
let mut invalid = custom.clone();
invalid["resetWindowWeights"] = windows;
invalid_values.push(invalid);
}
for invalid in invalid_values {
body["smartScheduling"] = invalid;
let response = app(fixture.state())
@@ -10,6 +10,14 @@ use std::time::SystemTime;
pub(super) struct TestPeerInfo(pub Arc<Mutex<Option<ProviderQuota>>>);
impl UpstreamInfoSource for TestPeerInfo {
fn is_peer(&self, _: &ProviderAccount) -> bool {
self.0
.lock()
.unwrap()
.as_ref()
.is_some_and(|quota| quota.upstream.as_ref().is_some_and(|info| info.discovered))
}
fn signal(
&self,
_: &ProviderAccount,
@@ -27,6 +27,8 @@ pub use error::CredentialError;
pub use model::*;
pub(crate) use selection::smart_score;
pub use selection::*;
pub use smart_scheduling::{SmartSchedulingConfig, SmartSchedulingConfigError};
pub use store::ProviderAccountStore;
pub use smart_scheduling::{QuotaResetPeriod, SmartSchedulingConfig, SmartSchedulingConfigError};
pub use store::{
AccountCatalogObservation, AccountCatalogOmissions, ProviderAccountStore, StoredAccountCatalog,
};
pub use upstream::*;
@@ -3,7 +3,7 @@
use std::cmp::Reverse;
use std::collections::{BTreeSet, HashMap};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Mutex, RwLock};
use std::sync::{Arc, Mutex, RwLock};
use std::time::{Duration, Instant, SystemTime};
use futures::FutureExt;
@@ -245,6 +245,7 @@ pub struct AccountRuntimeSignals {
pub in_flight: u32,
pub last_started_at: Option<SystemTime>,
pub quota_reset_at: Option<SystemTime>,
pub quota_reset_windows: Option<Arc<[AccountQuotaResetWindow]>>,
/// Provider 归一化的剩余额度基点:0 耗尽,10_000 全部可用
pub quota_remaining_rank: Option<u64>,
pub cooldown: Option<super::AccountCooldown>,
@@ -459,6 +460,7 @@ impl AccountRuntimeSignals {
pub fn with_provider_quota(mut self, quota: Option<AccountQuotaSignals>) -> Self {
if let Some(quota) = quota {
self.quota_reset_at = quota.reset_at;
self.quota_reset_windows = quota.reset_windows;
self.quota_remaining_rank = quota.remaining_rank;
}
self
@@ -482,9 +484,17 @@ impl AccountRuntimeSignals {
}
}
/// Provider 确认为账号通用且属于受支持周期的重置观测
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountQuotaResetWindow {
pub period: super::QuotaResetPeriod,
pub reset_at: SystemTime,
}
/// Provider 从私有 quota JSON 投影出的中立调度事实
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountQuotaSignals {
reset_windows: Option<Arc<[AccountQuotaResetWindow]>>,
reset_at: Option<SystemTime>,
remaining_rank: Option<u64>,
}
@@ -496,16 +506,23 @@ impl AccountQuotaSignals {
Self {
reset_at,
remaining_rank,
reset_windows: None,
}
}
#[must_use]
pub const fn reset_at(self) -> Option<SystemTime> {
pub fn with_reset_windows(mut self, windows: Vec<AccountQuotaResetWindow>) -> Self {
self.reset_windows = (!windows.is_empty()).then(|| windows.into());
self
}
#[must_use]
pub const fn reset_at(&self) -> Option<SystemTime> {
self.reset_at
}
#[must_use]
pub const fn remaining_rank(self) -> Option<u64> {
pub const fn remaining_rank(&self) -> Option<u64> {
self.remaining_rank
}
}
@@ -940,8 +957,6 @@ impl AccountSelector {
// 首输出 10 秒时延迟得分减半;固定尺度不随其他候选账号变化
const SMART_LATENCY_HALF_SCORE_MS: f64 = 10_000.0;
// 距重置一小时时得分减半;未知和已过期时间不提供重置奖励
const SMART_RESET_HALF_SCORE_SECONDS: f64 = 3_600.0;
fn capacity_utilization(candidate: &AccountCandidate, context: &AccountSelectionContext) -> f64 {
context
@@ -997,15 +1012,6 @@ pub(crate) fn smart_score(candidate: &AccountCandidate, context: &AccountSelecti
SMART_LATENCY_HALF_SCORE_MS / (SMART_LATENCY_HALF_SCORE_MS + latency as f64)
});
let reset = candidate
.signals
.quota_reset_at
.and_then(|reset| reset.duration_since(now).ok())
.filter(|remaining| !remaining.is_zero())
.map_or(0.0, |remaining| {
SMART_RESET_HALF_SCORE_SECONDS
/ (SMART_RESET_HALF_SCORE_SECONDS + remaining.as_secs_f64())
});
let [
load_weight,
quota_weight,
@@ -1018,5 +1024,5 @@ pub(crate) fn smart_score(candidate: &AccountCandidate, context: &AccountSelecti
+ quota_weight * quota
+ health_weight * failure
+ latency_weight * latency
+ reset_weight * reset
+ reset_weight * config.reset_score(&candidate.signals, now)
}
@@ -1,11 +1,40 @@
//! 智能调度参数;固定十分位保证快照可精确比较,浮点数只用于 wire 与评分
use std::time::SystemTime;
use super::AccountRuntimeSignals;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
#[serde(into = "SmartSchedulingValues")]
pub struct SmartSchedulingConfig {
weights: [u8; 6],
reset_windows: [u8; 3],
}
/// 账号通用额度的周期;Provider 按上游周期事实投影,不依赖窗口名称和槽位
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum QuotaResetPeriod {
FiveHour,
Weekly,
Monthly,
}
impl QuotaResetPeriod {
const fn seconds(self) -> u32 {
match self {
Self::FiveHour => 18_000,
Self::Weekly => 604_800,
Self::Monthly => 2_592_000,
}
}
}
fn weight_tenths(value: f64, field: &'static str) -> Result<u8, SmartSchedulingConfigError> {
if !(0.0..=10.0).contains(&value) || (value * 10.0).round() / 10.0 != value {
return Err(SmartSchedulingConfigError::InvalidWeight(field));
}
Ok((value * 10.0).round() as u8)
}
impl<'de> Deserialize<'de> for SmartSchedulingConfig {
@@ -23,6 +52,7 @@ impl Default for SmartSchedulingConfig {
fn default() -> Self {
Self {
weights: [10, 8, 10, 5, 0, 0],
reset_windows: [0; 3],
}
}
}
@@ -50,19 +80,15 @@ impl SmartSchedulingConfig {
.into_iter()
.enumerate()
{
let value = weights[index];
if !value.is_finite()
|| !(0.0..=10.0).contains(&value)
|| (value * 10.0).round() / 10.0 != value
{
return Err(SmartSchedulingConfigError::InvalidWeight(field));
}
tenths[index] = (value * 10.0).round() as u8;
tenths[index] = weight_tenths(weights[index], field)?;
}
if tenths == [0; 6] {
return Err(SmartSchedulingConfigError::EmptyWeights);
}
Ok(Self { weights: tenths })
Ok(Self {
weights: tenths,
..Self::default()
})
}
#[must_use]
@@ -70,11 +96,44 @@ impl SmartSchedulingConfig {
self.weights.map(|weight| f64::from(weight) / 10.0)
}
fn reset_enabled(self) -> bool {
self.weights[4] > 0 && self.reset_windows != [0; 3]
}
pub(crate) fn reset_score(self, signals: &AccountRuntimeSignals, now: SystemTime) -> f64 {
if !self.reset_enabled() {
return 0.0;
}
let mut scores = [0.0_f64; 3];
for window in signals.quota_reset_windows.as_deref().unwrap_or_default() {
let Ok(remaining) = window.reset_at.duration_since(now) else {
continue;
};
if remaining.is_zero() {
continue;
}
let half = f64::from(window.period.seconds()) / 7.0;
let score = half / (half + remaining.as_secs_f64());
// 同一周期只计最近的有效重置,避免重复额度桶叠加得分
let slot = &mut scores[window.period as usize];
*slot = slot.max(score);
}
// 缺失周期得零分,分母固定,避免少报窗口反而提高排名
let weights = self.reset_windows.map(f64::from);
scores
.into_iter()
.zip(weights)
.map(|(score, w)| score * w)
.sum::<f64>()
/ weights.into_iter().sum::<f64>()
}
pub(crate) fn score_tolerance(self) -> f64 {
// 与系数同步缩放,避免仅放大相同比例的系数就改变近似最优候选集合
// 排队系数只参与入队选择,不能放大立即选号的候选容差
// 排队和未启用的重置窗口不能放大立即选号的容差
let count = if self.reset_enabled() { 5 } else { 4 };
0.05 * f64::from(
self.weights[..5]
self.weights[..count]
.iter()
.copied()
.map(u16::from)
@@ -92,20 +151,37 @@ struct SmartSchedulingValues {
latency_weight: f64,
reset_weight: f64,
queue_weight: f64,
#[serde(default)]
reset_window_weights: ResetWindowWeights,
}
#[derive(Default, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct ResetWindowWeights {
five_hour: f64,
weekly: f64,
monthly: f64,
}
impl TryFrom<SmartSchedulingValues> for SmartSchedulingConfig {
type Error = SmartSchedulingConfigError;
fn try_from(value: SmartSchedulingValues) -> Result<Self, Self::Error> {
Self::new([
let mut config = Self::new([
value.load_weight,
value.quota_weight,
value.health_weight,
value.latency_weight,
value.reset_weight,
value.queue_weight,
])
])?;
let windows = value.reset_window_weights;
config.reset_windows = [
weight_tenths(windows.five_hour, "fiveHour")?,
weight_tenths(windows.weekly, "weekly")?,
weight_tenths(windows.monthly, "monthly")?,
];
Ok(config)
}
}
@@ -126,6 +202,11 @@ impl From<SmartSchedulingConfig> for SmartSchedulingValues {
latency_weight,
reset_weight,
queue_weight,
reset_window_weights: ResetWindowWeights {
five_hour: f64::from(value.reset_windows[0]) / 10.0,
weekly: f64::from(value.reset_windows[1]) / 10.0,
monthly: f64::from(value.reset_windows[2]) / 10.0,
},
}
}
}
@@ -12,9 +12,45 @@ use super::{
QuotaObservationTouch, QuotaObservationVersion, QuotaWriteOutcome,
};
/// 一个账号最后成功的目录观测;正文由 Provider 解释,归属不能按相同正文合并
#[derive(Clone)]
pub struct AccountCatalogObservation {
pub account_id: ProviderAccountId,
pub revision: CredentialRevision,
pub plan: Option<String>,
pub etag: Option<String>,
pub observed_at: std::time::SystemTime,
}
/// 批量恢复时每份不同的正文只传输、解析一次
pub struct StoredAccountCatalog {
pub accounts: Vec<AccountCatalogObservation>,
pub document: bytes::Bytes,
}
/// 批量选号的存储优化:只有账号自己的当前版本、有效时间和目录摘要都匹配才可跳过
pub struct AccountCatalogOmissions {
pub documents: Vec<[u8; 32]>,
pub exceptions: Vec<ProviderAccountId>,
pub observed_after: std::time::SystemTime,
pub as_of: std::time::SystemTime,
}
/// `provider_accounts` 的数据库中立端口
#[async_trait]
pub trait ProviderAccountStore: Send + Sync {
async fn load_account_catalogs(
&self,
provider: &ProviderKind,
) -> Result<Vec<StoredAccountCatalog>, StoreError>;
/// 账号删除或凭据已变化时返回 false,不发布迟到的观测
async fn save_account_catalog(
&self,
observation: &AccountCatalogObservation,
document: &[u8],
) -> Result<bool, StoreError>;
async fn create_account(&self, account: NewProviderAccount) -> Result<(), StoreError>;
async fn get_account(
@@ -36,6 +72,16 @@ pub trait ProviderAccountStore: Send + Sync {
account_ids: &[ProviderAccountId],
) -> Result<Vec<ProviderAccount>, StoreError>;
/// 存储可以下推目录筛选,也可以返回完整候选;调用方始终复核最终账号
async fn list_catalog_candidates(
&self,
provider: &ProviderKind,
accounts: &[ProviderAccountId],
_omissions: &AccountCatalogOmissions,
) -> Result<Vec<ProviderAccount>, StoreError> {
self.list_for_provider_accounts(provider, accounts).await
}
/// 一次有界查询返回账号事实和 revision-fenced 明文 credential
async fn list_refresh_candidates(
&self,
@@ -84,6 +84,9 @@ pub struct UpstreamSchedulingSignal {
}
pub trait UpstreamInfoSource: Send + Sync {
/// 当前连接是否已识别为互联节点,不以额度快照的新鲜度代替身份
fn is_peer(&self, account: &ProviderAccount) -> bool;
fn signal(
&self,
account: &ProviderAccount,
@@ -110,6 +113,13 @@ impl UpstreamInfoReader {
}
impl UpstreamInfoSource for UpstreamInfoReader {
fn is_peer(&self, account: &ProviderAccount) -> bool {
self.source
.get()
.and_then(Weak::upgrade)
.is_some_and(|source| source.is_peer(account))
}
fn signal(
&self,
account: &ProviderAccount,
@@ -158,6 +158,7 @@ impl TraceContext {
"recovery": {
"replaySafe": error.replay_is_safe(),
"retryProhibited": error.retry_is_prohibited(),
"sameAccountRetryBlocked": error.same_account_retry_blocked(),
"sameAccountCredentialRetry": error.retries_same_account(),
"intent": error.pre_delivery_retry().map(|intent| format!("{intent:?}")),
"continuationFailure": error.continuation_failure().map(|failure| format!("{failure:?}")),
@@ -346,6 +347,7 @@ impl TraceContext {
for key in [
"replaySafe",
"retryProhibited",
"sameAccountRetryBlocked",
"sameAccountCredentialRetry",
] {
kept[key] = json!(recovery.get(key).and_then(Value::as_bool));
@@ -1672,15 +1672,19 @@ where
error.pre_delivery_retry(),
Some(crate::error::PreDeliveryRetry::SameAccountConnectionRetry { .. })
);
// 服务器建议不可截短;超出请求冻结的上限后沿用安全换号或原错误出口
// 服务器建议不可截短
// 超出上限,或上游拒绝继续重试当前账号时,跳过同账号等待
// 较短的 Retry-After 也直接沿用安全换号或原错误出口
let retry_after_exceeds_limit = error
.retry_after()
.is_some_and(|delay| delay > self.plan.max_retry_after());
let same_account_retry_blocked = error.same_account_retry_blocked();
// Provider 的有界同账号传输恢复是尽力重试策略,不要求把 Sent/Ambiguous
// 改写成安全证明;外部副作用、取消、预算与下游提交仍由 Core 拦截
let transport_recovery = match error.pre_delivery_retry() {
Some(crate::error::PreDeliveryRetry::SameAccountConnectionRetry { transport })
if !execution_effect_observed
if !same_account_retry_blocked
&& !execution_effect_observed
&& !self.connection_budget.exhausted()
&& !self.replay_barrier_reached
&& !self.delivery_pending
@@ -1698,7 +1702,8 @@ where
Some(crate::error::PreDeliveryRetry::SameAccountTransportRetry {
retry_index,
delay,
}) if !retry_after_exceeds_limit
}) if !same_account_retry_blocked
&& !retry_after_exceeds_limit
&& !execution_effect_observed
&& !self.connection_budget.exhausted()
&& !self.replay_barrier_reached
@@ -1707,7 +1712,8 @@ where
Some((AttemptTransport::Retry(retry_index), delay))
}
Some(crate::error::PreDeliveryRetry::SameAccountTransportFallback)
if !retry_after_exceeds_limit
if !same_account_retry_blocked
&& !retry_after_exceeds_limit
&& !execution_effect_observed
&& !self.connection_budget.exhausted()
&& !self.replay_barrier_reached
@@ -1754,7 +1760,11 @@ where
max_delay,
jitter,
..
}) if ordinary_retry && !suppress_transient_retry && !retry_after_exceeds_limit => {
}) if !same_account_retry_blocked
&& ordinary_retry
&& !suppress_transient_retry
&& !retry_after_exceeds_limit =>
{
let retries = self
.transient_retry_counts
.entry(current.metadata.provider_account_id().clone())
@@ -1776,7 +1786,8 @@ where
}
_ => None,
};
let same_account_retry = !execution_effect_observed
let same_account_retry = !same_account_retry_blocked
&& !execution_effect_observed
&& error.retries_same_account()
&& provider_proved_replay_safe
&& !self.replay_barrier_reached
@@ -1834,6 +1845,7 @@ where
"reserveConnectionRotation": reserve_connection_rotation,
"transientRetry": transient_retry.is_some(),
"retryAfterExceeded": retry_after_exceeds_limit,
"sameAccountRetryBlocked": same_account_retry_blocked,
"maxRetryAfterSeconds": self.plan.max_retry_after().as_secs(),
"clientRetrySuppressed": suppress_transient_retry,
"rotationRetry": rotation_retry, "accountRotations": self.account_rotations,
@@ -801,16 +801,38 @@ impl DefaultExecutionService {
active_request,
admission_decision_ms,
} = start_guard;
let observation = plan
.candidates()
.first()
.map_or_else(Default::default, |candidate| {
providers.request_observation(
candidate.provider(),
&operation,
client.policy.key_id(),
)
let observation_provider = match &target {
ExecutionTarget::ProviderEndpoint { provider, .. } => Some(provider),
ExecutionTarget::Model(_) => plan
.candidates()
.iter()
.find(|candidate| {
operation
.provider_session_state(candidate.provider().as_str())
.is_some_and(|state| state.extension_owner().is_none())
})
.or_else(|| plan.candidates().first())
.map(|candidate| candidate.provider()),
};
let provider_observation =
observation_provider.map_or_else(Default::default, |candidate| {
providers.request_observation(candidate, &operation, client.policy.key_id())
});
// 纯协议事实不依赖候选顺序,Provider 仅补充所属会话的旧配置和恢复关联
let observation = match &operation {
Operation::Generate(request) => {
let payload = request.protocol_payload();
gateway_protocol::request::request_observation(
payload.protocol(),
payload.body(),
payload.context(),
provider_observation
.previous_reasoning_effort_override
.as_deref(),
)
}
_ => Default::default(),
};
let (request_kind, subagent_kind) = if let Some(nested) = authorization.nested.as_ref() {
(
Some("plugin_child_model".to_owned()),
@@ -843,7 +865,9 @@ impl DefaultExecutionService {
operation: operation.kind(),
endpoint: metadata.endpoint,
client_transport: metadata.transport.as_str().to_owned(),
requested_model: target.into_public_model().or(observation.requested_model),
requested_model: target
.into_public_model()
.or(provider_observation.requested_model),
client_ip: metadata.client_ip,
user_agent: metadata.user_agent,
reasoning_effort: observation.reasoning_effort,
@@ -851,7 +875,7 @@ impl DefaultExecutionService {
request_kind,
subagent_kind,
compact: observation.compact,
continuation: observation.continuation,
continuation: provider_observation.continuation,
image_generation_requested: operation.image_generation_requested(),
admission_decision_ms,
started_at,
@@ -1088,17 +1088,13 @@ pub struct ProviderRequest {
candidate: ProviderCandidate,
}
/// Provider 对公共观测表可解释的请求语义;未知字段保持空值
/// Provider 私有状态与独立端点的观测补充,通用请求语义由协议层提取
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ProviderRequestObservation {
/// 独立 Provider 端点声明的请求模型,仅用于观测,不参与文本模型目录路由
pub requested_model: Option<PublicModelId>,
/// Provider 从客户端请求与续接配置中解析的推理强度
pub reasoning_effort: Option<String>,
pub reasoning_preset: Option<String>,
pub request_kind: Option<String>,
pub subagent_kind: Option<String>,
pub compact: bool,
/// 已核实归属的上一响应配置;当前请求是否继承由协议层判断
pub previous_reasoning_effort_override: Option<String>,
pub continuation: ContinuationRequestObservation,
}
@@ -516,6 +516,7 @@ pub struct ProviderError {
failure_observation: Option<Box<ProviderErrorFailureObservation>>,
replay_safe: bool,
retry_prohibited: bool,
same_account_retry_blocked: bool,
pre_delivery_retry: Option<Box<PreDeliveryRetry>>,
credential_recovery_required: bool,
retry_same_account: bool,
@@ -559,6 +560,7 @@ impl ProviderError {
failure_observation: None,
replay_safe: false,
retry_prohibited: false,
same_account_retry_blocked: false,
pre_delivery_retry: None,
credential_recovery_required: false,
retry_same_account: false,
@@ -685,6 +687,22 @@ impl ProviderError {
self.retry_prohibited
}
/// 保留已选择的耗尽策略,但禁止执行同账号退避和传输恢复
///
/// Provider 在上游明确拒绝继续使用当前账号时设置
/// Core 不根据错误码自行推断
#[must_use]
pub const fn with_same_account_retry_blocked(mut self) -> Self {
self.same_account_retry_blocked = true;
self
}
/// 返回上游是否拒绝继续重试当前账号
#[must_use]
pub const fn same_account_retry_blocked(&self) -> bool {
self.same_account_retry_blocked
}
/// 允许 Core 仅在客户端尚未收到任何事件时执行一次受预算约束的换号恢复
///
/// 该标记不证明上游未执行,也不等价于 [`Self::with_replay_safe`];它只表达
@@ -1008,6 +1026,7 @@ impl ProviderError {
failure_observation: self.failure_observation.clone(),
replay_safe: self.replay_safe,
retry_prohibited: self.retry_prohibited,
same_account_retry_blocked: self.same_account_retry_blocked,
pre_delivery_retry: self.pre_delivery_retry.clone(),
credential_recovery_required: self.credential_recovery_required,
retry_same_account: self.retry_same_account,
@@ -1036,6 +1055,10 @@ impl fmt::Debug for ProviderError {
)
.field("retry_after", &self.retry_after)
.field("retry_prohibited", &self.retry_prohibited)
.field(
"same_account_retry_blocked",
&self.same_account_retry_blocked,
)
.field("continuation_failure", &self.continuation_failure)
.field(
"continuation_recovery_disposition",
+4 -35
View File
@@ -3,7 +3,7 @@
//! 这里只保留网关需要解释、路由或结算的稳定语义
//! 协议 adapter 无法写入
//! wire body 与连接级事实随 [`ProtocolPayload`] 或 [`RawJsonPayload`] 不透明传递,
//! 由对应 Provider 自己解释
//! 纯协议事实由协议层解析,私有字段与会话由对应 Provider 解释
use std::collections::BTreeSet;
use std::fmt;
@@ -458,8 +458,8 @@ impl ClientTurn {
impl GenerateRequest {
/// 创建携带协议不透明正文的生成请求
///
/// 完整请求只以原始 wire object 保存;Core 仅局部只读已知字段以推导路由能力,
/// Provider 负责解释与转换正文
/// 完整请求以原始 wire object 保存,协议层提取路由与权限所需的纯事实
/// Provider 在选定账号后编排协议转换与私有字段处理
#[must_use]
pub fn from_protocol_payload(protocol_payload: ProtocolPayload) -> Self {
Self {
@@ -521,9 +521,7 @@ impl GenerateRequest {
/// 返回客户端是否请求了图片生成工具
#[must_use]
pub fn image_generation_requested(&self) -> bool {
self.payload.source_image_generation
|| (self.protocol_payload().protocol() == "openai"
&& requests_image_generation(self.body()))
self.payload.source_image_generation || self.payload.wire_facts.image_generation
}
/// 返回指定 Provider 的连接内会话状态
@@ -596,35 +594,6 @@ impl fmt::Debug for GenerateRequest {
}
}
/// 原生工具声明授权上游生成图片,Codex 的被动 namespace 与图片输入不属于该权限
fn requests_image_generation(body: &Map<String, Value>) -> bool {
contains_image_tool(body.get("tools"))
|| body
.get("input")
.and_then(Value::as_array)
.is_some_and(|items| {
items.iter().any(|item| {
item.get("type").and_then(Value::as_str) == Some("additional_tools")
&& contains_image_tool(item.get("tools"))
})
})
|| body.get("tool_choice").is_some_and(|choice| {
is_image_tool(choice)
|| (choice.get("type").and_then(Value::as_str) == Some("allowed_tools")
&& contains_image_tool(choice.get("tools")))
})
}
fn contains_image_tool(value: Option<&Value>) -> bool {
value
.and_then(Value::as_array)
.is_some_and(|tools| tools.iter().any(is_image_tool))
}
fn is_image_tool(value: &Value) -> bool {
value.get("type").and_then(Value::as_str) == Some("image_generation")
}
/// 图像 API 的稳定端点语义
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum ImageRequestKind {
@@ -1,7 +1,7 @@
//! Provider 会话亲和、绑定和排除事实
use super::{ProviderStoreError, ProviderStoreErrorKind};
use crate::{account::ProviderAccountId, identity::ProviderKind};
use crate::{account::ProviderAccountId, identity::ProviderKind, policy::ClientApiKeyId};
use futures::future::BoxFuture;
use std::{collections::BTreeSet, fmt, time::Duration};
@@ -81,6 +81,36 @@ pub struct ProviderSessionAlias {
pub follow_only: bool,
}
/// 发送前需要核对版本的会话绑定;首次认领的 expected 为空
#[derive(Clone, Copy)]
pub struct ProviderSessionClaim<'a> {
pub key: &'a ProviderSessionAffinityKey,
pub expected: Option<&'a ProviderSessionBinding>,
pub account: &'a ProviderAccountId,
pub ttl: Duration,
}
/// 仅成功的内置轮询推进游标;有会话的分配必须同时提交绑定
#[derive(Clone, Copy)]
pub enum ProviderAccountAdmission<'a> {
Session(ProviderSessionClaim<'a>),
RoundRobin {
client_key: &'a ClientApiKeyId,
cursor: u64,
session: Option<ProviderSessionClaim<'a>>,
},
}
impl<'a> ProviderAccountAdmission<'a> {
#[must_use]
pub const fn session(self) -> Option<ProviderSessionClaim<'a>> {
match self {
Self::Session(session) => Some(session),
Self::RoundRobin { session, .. } => session,
}
}
}
/// 客户端作用域内的会话账号绑定,原始会话身份由 Provider 哈希后传入
pub trait ProviderSessionAffinityPort: Send + Sync {
fn load<'a>(
@@ -89,16 +119,19 @@ pub trait ProviderSessionAffinityPort: Send + Sync {
key: &'a ProviderSessionAffinityKey,
) -> BoxFuture<'a, Result<Option<ProviderSessionBinding>, ProviderStoreError>>;
/// 在发送前原子认领、续期或迁移;None 表示快照已变化,调用方必须释放租约重新选择
/// expected 为 None 只允许首次认领,已有快照必须连同版本匹配
fn compare_and_bind<'a>(
/// 读取本 Client Key 和 Provider 的下一轮位置,不推进游标
fn round_robin_cursor<'a>(
&'a self,
provider_kind: &'a ProviderKind,
key: &'a ProviderSessionAffinityKey,
expected: Option<&'a ProviderSessionBinding>,
account_id: &'a ProviderAccountId,
ttl: Duration,
) -> BoxFuture<'a, Result<Option<ProviderSessionBinding>, ProviderStoreError>>;
provider: &'a ProviderKind,
client_key: &'a ClientApiKeyId,
) -> BoxFuture<'a, Result<u64, ProviderStoreError>>;
/// 原子核对并提交会话版本和轮询位置;冲突时均不写入,调用方释放租约重选
fn admit_account<'a>(
&'a self,
provider: &'a ProviderKind,
admission: ProviderAccountAdmission<'a>,
) -> BoxFuture<'a, Result<bool, ProviderStoreError>>;
/// 显式观测到的请求关联只指向会话键,不缓存账号,迁移后仍读取当前绑定
fn load_alias<'a>(
@@ -154,10 +154,16 @@ impl ProviderSchedulingLeaseRequest {
}
}
/// Lease 生命周期由具体 Store guard 管理,Provider 只能持有
pub trait ProviderLeaseGuard: Send + Sync + 'static {}
/// Drop 回收租约;同一请求需要重选时显式等待释放,避免被自己的租约阻塞
pub trait ProviderLeaseGuard: Send + Sync + 'static {
fn release(self: Box<Self>) -> BoxFuture<'static, Result<(), ProviderStoreError>>;
}
impl<T> ProviderLeaseGuard for T where T: Send + Sync + 'static {}
impl ProviderLeaseGuard for () {
fn release(self: Box<Self>) -> BoxFuture<'static, Result<(), ProviderStoreError>> {
Box::pin(async { Ok(()) })
}
}
pub enum ProviderLeaseAcquisition {
Acquired(Box<dyn ProviderLeaseGuard>),
@@ -10,8 +10,9 @@ pub use lease::{
};
mod affinity;
pub use affinity::{
ProviderSessionAffinityKey, ProviderSessionAffinityPort, ProviderSessionAlias,
ProviderSessionBinding, ProviderSessionExclusionPort, ProviderSessionExclusions,
ProviderAccountAdmission, ProviderSessionAffinityKey, ProviderSessionAffinityPort,
ProviderSessionAlias, ProviderSessionBinding, ProviderSessionClaim,
ProviderSessionExclusionPort, ProviderSessionExclusions,
};
mod catalog;
pub use catalog::{
@@ -56,6 +56,7 @@ fn candidate(id: &str, in_flight: u32, remaining: Option<u64>) -> AccountCandida
in_flight,
last_started_at: None,
quota_reset_at: None,
quota_reset_windows: None,
quota_remaining_rank: remaining,
cooldown: None,
failure_rate_basis_points: None,
@@ -228,6 +229,7 @@ fn diagnostic_selection_bypasses_all_local_account_eligibility() {
in_flight: 0,
last_started_at: None,
quota_reset_at: None,
quota_reset_windows: None,
quota_remaining_rank: None,
cooldown: None,
failure_rate_basis_points: None,
@@ -247,6 +249,7 @@ fn diagnostic_selection_bypasses_all_local_account_eligibility() {
in_flight: 0,
last_started_at: None,
quota_reset_at: None,
quota_reset_windows: None,
quota_remaining_rank: None,
cooldown: None,
failure_rate_basis_points: None,
@@ -641,6 +644,7 @@ fn provider_quota_overlay_should_preserve_store_concurrency_facts() {
in_flight: 2,
last_started_at: Some(last_started_at),
quota_reset_at: None,
quota_reset_windows: None,
quota_remaining_rank: None,
cooldown: None,
failure_rate_basis_points: None,
@@ -1,7 +1,8 @@
//! 验证智能调度权重、软亲和与容量约束的组合行为
use gateway_core::account::{
AccountSelector, PreferredAccountSelection, RotationStrategy, SmartSchedulingConfig,
AccountQuotaResetWindow, AccountSelector, PreferredAccountSelection, QuotaResetPeriod,
RotationStrategy, SmartSchedulingConfig,
};
use serde_json::json;
use std::time::{Duration, Instant, SystemTime};
@@ -62,9 +63,10 @@ fn smart_config_rejects_invalid_weights_and_incomplete_wire_values() {
#[test]
fn reset_weight_rewards_only_future_resets_and_preserves_capacity_checks() {
let mut context = context(RotationStrategy::Smart);
context.policy = context
.policy
.with_smart_scheduling(SmartSchedulingConfig::new([0.0, 0.0, 0.0, 0.0, 1.0, 0.0]).unwrap());
context.policy = context.policy.with_smart_scheduling(serde_json::from_value(json!({
"loadWeight":0, "quotaWeight":0, "healthWeight":0, "latencyWeight":0, "resetWeight":1, "queueWeight":0,
"resetWindowWeights":{"fiveHour":0,"weekly":1,"monthly":0}
})).unwrap());
let mut candidates = [
candidate_with_concurrency("acct_unknown", 0, 2),
candidate_with_concurrency("acct_past", 0, 2),
@@ -72,10 +74,20 @@ fn reset_weight_rewards_only_future_resets_and_preserves_capacity_checks() {
candidate_with_concurrency("acct_near", 0, 2),
candidate_with_concurrency("acct_far", 0, 2),
];
candidates[1].signals.quota_reset_at = Some(context.now - Duration::from_secs(1));
candidates[2].signals.quota_reset_at = Some(context.now);
candidates[3].signals.quota_reset_at = Some(context.now + Duration::from_secs(3600));
candidates[4].signals.quota_reset_at = Some(context.now + Duration::from_secs(86_400));
for (candidate, reset_at) in candidates[1..].iter_mut().zip([
context.now - Duration::from_secs(1),
context.now,
context.now + Duration::from_secs(3600),
context.now + Duration::from_secs(86_400),
]) {
candidate.signals.quota_reset_windows = Some(
[AccountQuotaResetWindow {
period: QuotaResetPeriod::Weekly,
reset_at,
}]
.into(),
);
}
assert_eq!(
AccountSelector
.select(&candidates, &context)
@@ -92,8 +104,8 @@ fn reset_weight_rewards_only_future_resets_and_preserves_capacity_checks() {
let scores = snapshot["events"][0]["data"]["candidates"]
.as_array()
.unwrap();
for (candidate, expected) in scores.iter().zip([0.0, 0.0, 0.0, 0.5, 0.04]) {
assert_eq!(candidate["smartScore"].as_f64().unwrap(), expected);
for (candidate, expected) in scores.iter().zip([0.0, 0.0, 0.0, 0.96, 0.5]) {
assert!((candidate["smartScore"].as_f64().unwrap() - expected).abs() < 1e-10);
}
candidates[3].signals.in_flight = 2;
assert_eq!(
@@ -120,6 +132,36 @@ fn reset_weight_rewards_only_future_resets_and_preserves_capacity_checks() {
);
}
#[test]
fn zero_or_absent_window_weights_disable_reset_scores_and_tolerance() {
let mut context = context(RotationStrategy::Smart);
let mut candidate = candidate_with_concurrency("acct_reset", 0, 2);
candidate.signals.quota_reset_windows = Some(
[AccountQuotaResetWindow {
period: QuotaResetPeriod::Monthly,
reset_at: context.now + Duration::from_secs(1),
}]
.into(),
);
let legacy = json!({"loadWeight":0,"quotaWeight":0,"healthWeight":0,"latencyWeight":0,"resetWeight":10,"queueWeight":0});
let mut disabled = legacy.clone();
disabled["resetWindowWeights"] = json!({"fiveHour":0,"weekly":0,"monthly":0});
for wire in [legacy, disabled] {
let config: SmartSchedulingConfig = serde_json::from_value(wire).unwrap();
context.policy = context.policy.with_smart_scheduling(config);
let trace = gateway_core::diagnostics::TraceContext::new("req_zero_windows");
trace.account_selection(std::slice::from_ref(&candidate), &context, None);
let events = trace.snapshot().unwrap();
let data = &events["events"][0]["data"];
assert_eq!(data["candidates"][0]["smartScore"], 0.0);
assert_eq!(data["smartScoreTolerance"], 0.0);
assert_eq!(
serde_json::to_value(config).unwrap()["resetWindowWeights"],
json!({"fiveHour":0.0,"weekly":0.0,"monthly":0.0})
);
}
}
#[test]
fn queue_weight_balances_live_pressure_and_health_only_when_enabled_for_smart() {
for (strategy, queue_weight, expects_shortest) in [
@@ -349,3 +391,41 @@ fn higher_weight_recovery_preserves_usable_soft_affinity() {
);
}
}
#[test]
fn missing_expired_and_duplicate_periods_do_not_inflate_reset_score() {
let mut context = context(RotationStrategy::Smart);
let mut candidate = candidate_with_concurrency("acct_partial", 0, 2);
candidate.signals.quota_reset_at = Some(context.now + Duration::from_secs(1));
candidate.signals.quota_reset_windows = Some(
[
(QuotaResetPeriod::Weekly, 86400),
(QuotaResetPeriod::Weekly, 90000),
(QuotaResetPeriod::Weekly, 0),
(QuotaResetPeriod::FiveHour, 0),
]
.map(|(period, delta)| AccountQuotaResetWindow {
period,
reset_at: context.now + Duration::from_secs(delta),
})
.into(),
);
for scale in [1.0, 0.5] {
let mut config = serde_json::to_value(
SmartSchedulingConfig::new([0.0, 0.0, 0.0, 0.0, 1.0, 0.0]).unwrap(),
)
.unwrap();
config["resetWindowWeights"] = json!({
"weekly": 2.0*scale, "fiveHour": 4.0*scale, "monthly": 4.0*scale
});
context.policy = context
.policy
.with_smart_scheduling(serde_json::from_value(config).unwrap());
let trace = gateway_core::diagnostics::TraceContext::new("req_partial");
trace.account_selection(std::slice::from_ref(&candidate), &context, None);
assert_eq!(
trace.snapshot().unwrap()["events"][0]["data"]["candidates"][0]["smartScore"],
0.1
);
}
}
@@ -4325,6 +4325,87 @@ fn retry_after_above_the_configured_limit_rotates_without_waiting() {
}
}
#[test]
fn no_more_retry_skips_same_account_wait_even_when_retry_after_is_short() {
let advised = Duration::from_secs(30);
for (rejection, rotates) in [
(
transient_rejection()
.with_same_account_retry_blocked()
.with_retry_after(advised),
true,
),
(
transient_rejection()
.with_transient_retry(
NonZeroU32::MIN,
Duration::ZERO,
Duration::ZERO,
Duration::ZERO,
TransientRetryPolicy::RetryThenReturnError,
)
.with_same_account_retry_blocked()
.with_retry_after(advised),
false,
),
] {
let operation = generate_operation();
let route_plan = retry_after_plan(&operation, 240);
let mut scripts = vec![Script::Stream {
account_id: "acct_first",
items: vec![Err(rejection)],
}];
if rotates {
scripts.push(Script::Stream {
account_id: "acct_second",
items: complete_stream(None),
});
}
let (coordinator, _, provider) = coordinator(scripts);
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 trace = session.trace().clone();
let mut collection = Box::pin(session.collect_uncommitted());
let polled = block_on(async { futures::poll!(&mut collection) });
let trace = trace.snapshot().unwrap();
let retry = trace["events"]
.as_array()
.unwrap()
.iter()
.find(|event| event["stage"] == "retry.decided")
.unwrap();
assert_eq!(retry["data"]["sameAccountRetryBlocked"], true);
assert_eq!(retry["data"]["transientRetry"], false);
assert_eq!(retry["data"]["retryAfterExceeded"], false);
if rotates {
assert!(matches!(polled, std::task::Poll::Ready(Ok(_))));
drop(collection);
let contexts = provider.contexts.lock().unwrap();
assert_eq!(contexts.len(), 2);
assert_eq!(contexts[1].required_account(), None);
assert!(
contexts[1]
.excluded_accounts()
.contains(&ProviderAccountId::new("acct_first").unwrap())
);
} else {
assert!(matches!(
polled,
std::task::Poll::Ready(Err(EngineError::Provider(error)))
if error.retry_after() == Some(advised)
));
assert_eq!(provider.contexts.lock().unwrap().len(), 1);
}
}
}
#[test]
fn excessive_retry_after_preserves_the_error_when_rotation_is_unavailable_or_disallowed() {
for (rejection, rotation_allowed) in [
@@ -148,6 +148,209 @@ fn request(service: &DefaultExecutionService, transport: ClientTransport) -> Sta
}
}
#[test]
fn request_observation_is_stable_across_candidate_order_and_respects_session_ownership() {
for order in [["alpha", "beta"], ["beta", "alpha"]] {
for (protocol, state_owner, extension_owned, body, expected) in [
(
"openai",
None,
false,
json!({"reasoning":{"effort":"medium"}}),
"medium",
),
(
"anthropic",
None,
false,
json!({"output_config":{"effort":" Future-Level "}}),
"Future-Level",
),
(
"openai",
Some("beta"),
false,
json!({"previous_response_id":"resp_previous","reasoning":{"effort":"medium"}}),
"high",
),
(
"openai",
Some("unavailable"),
false,
json!({"previous_response_id":"resp_previous","reasoning":{"effort":"medium"}}),
"medium",
),
(
"openai",
Some("beta"),
true,
json!({"previous_response_id":"resp_previous","reasoning":{"effort":"medium"}}),
"medium",
),
(
"openai",
Some("beta"),
false,
json!({"previous_response_id":"resp_previous","reasoning":{"effort":"medium"},"input":[{"type":"configuration_update","reasoning":{}}]}),
"medium",
),
] {
let kinds: Vec<_> = order
.iter()
.map(|name| ProviderKind::new(*name).unwrap())
.collect();
let directory = RuntimeAccountDirectory::new(
kinds
.iter()
.map(|kind| {
(
ProviderAccountId::new(format!("acct_{}", kind.as_str())).unwrap(),
RuntimeAccount::new(kind.clone(), BTreeSet::new()),
)
})
.collect(),
);
let snapshot = RuntimeSnapshot::new(
ConfigRevision::new(1).unwrap(),
gateway_core::settings::SettingsValues::new(
1,
"smart",
Default::default(),
None,
None,
),
kinds.clone(),
kinds
.into_iter()
.map(|kind| {
ProviderModel::new(
kind,
UpstreamModelId::new("gpt-start").unwrap(),
ModelCapabilities::new(BTreeSet::from([OperationKind::Generate]), None)
.with_upstream_feature_validation(),
)
})
.collect(),
vec![ClientPolicy::new(
ClientApiKeyId::new("key_start_test").unwrap(),
PlaintextClientApiKey::new("sk_start_test").unwrap(),
Arc::new(FrozenAccountScope::new(
Arc::new(directory),
ClientRoutingScope::all_accounts(),
)),
true,
RateLimits::unlimited(),
)],
)
.unwrap();
let store = Arc::new(TrackingExecutionStore::default());
let service =
DefaultExecutionService::new(
RuntimeSnapshotHandle::new(snapshot),
store.clone(),
ProviderRegistry::new(order.map(|name| {
Arc::new(SessionObservationProvider(name)) as Arc<dyn Provider>
}))
.unwrap(),
Arc::new(Admissions::default()),
Arc::new(UnusedContinuation),
Arc::new(RecordingClientApiKeyUsage::default()),
Arc::new(crate::RecordingDiagnostics::default()),
);
let mut generate = GenerateRequest::from_protocol_payload(
ProtocolPayload::json_object(protocol, body.as_object().unwrap().clone())
.unwrap()
.with_context(serde_json::Map::from_iter([(
"turn_metadata".to_owned(),
json!("{\"request_kind\":\"review\"}"),
)])),
);
if let Some(owner) = state_owner {
let mut state = gateway_core::operation::ProviderSessionState::new(
owner,
serde_json::Map::from_iter([("effort".to_owned(), json!("high"))]),
)
.unwrap();
if extension_owned {
state = state.with_extension_owner(
gateway_core::operation::ExtensionSessionOwner {
instance_id: "plugin".into(),
contribution_id: "adapter".into(),
adapter_id: "adapter".into(),
generation: 1,
incarnation: "incarnation".into(),
connection_local: true,
},
);
}
generate = generate.with_provider_session_state(state);
}
let mut start = request(&service, ClientTransport::HttpJson);
start.operation = Operation::Generate(generate);
start.metadata.protocol = protocol.to_owned();
let started = block_on(service.start(start)).expect("request observation is recorded");
block_on(started.session.detach_finalize());
let requests = store.requests.lock().unwrap();
let recorded = requests.first().expect("logical request");
assert_eq!(
recorded.reasoning_effort.as_deref(),
Some(expected),
"order={order:?}, state={state_owner:?}, extension={extension_owned}"
);
assert_eq!(
recorded.request_kind.as_deref(),
(protocol == "openai").then_some("review")
);
}
}
}
struct SessionObservationProvider(&'static str);
#[async_trait]
impl Provider for SessionObservationProvider {
fn name(&self) -> &'static str {
self.0
}
fn catalog_generation(&self) -> ProviderCatalogGeneration {
ProviderCatalogGeneration::default()
}
fn request_observation(
&self,
operation: &Operation,
_: &ClientApiKeyId,
) -> ProviderRequestObservation {
ProviderRequestObservation {
previous_reasoning_effort_override: operation
.provider_session_state(self.0)
.filter(|state| state.extension_owner().is_none())
.and_then(|state| state.payload().get("effort"))
.and_then(Value::as_str)
.map(str::to_owned),
..Default::default()
}
}
async fn query_model_capabilities(
&self,
) -> Result<Vec<ProviderModelCapabilities>, ProviderError> {
Ok(Vec::new())
}
async fn execute(
self: Arc<Self>,
_: ProviderRequest,
_: AttemptContext,
) -> Result<ProviderStream, ProviderError> {
Err(ProviderError::new(
ProviderErrorKind::Unavailable,
UpstreamSendState::NotSent,
))
}
}
#[test]
fn default_request_has_no_total_deadline_and_explicit_timeout_can_be_cleared() {
block_on(async {
@@ -2121,6 +2324,7 @@ impl Provider for PolicySelectingProvider {
in_flight: 0,
last_started_at: None,
quota_reset_at: None,
quota_reset_windows: None,
quota_remaining_rank: None,
cooldown: None,
failure_rate_basis_points: None,
@@ -77,6 +77,7 @@ fn candidate(id: &str) -> AccountCandidate {
in_flight: 0,
last_started_at: None,
quota_reset_at: None,
quota_reset_windows: None,
quota_remaining_rank: None,
cooldown: None,
failure_rate_basis_points: None,
+25 -33
View File
@@ -157,40 +157,28 @@ fn prompt_cache_and_image_generation_should_be_read_from_raw_body() {
}
#[test]
fn image_generation_detection_distinguishes_tools_from_image_inputs_and_passive_namespaces() {
for body in [
json!({"tools":[{"type":"image_generation"}]}),
json!({"tool_choice":{"type":"image_generation"}}),
json!({"tool_choice":{"type":"allowed_tools","tools":[{"type":"image_generation"}]}}),
json!({"input":[{"type":"additional_tools","tools":[{"type":"image_generation"}]}]}),
] {
assert!(
generate(body.clone()).image_generation_requested(),
"{body}"
);
}
for body in [
json!({"input":"Please generate an image"}),
json!({"input":[{"role":"user","content":[{"type":"input_image","image_url":"data:image/png;base64,AA=="}]}]}),
json!({"tools":[{"type":"namespace","name":"image_gen","tools":[{"type":"function","name":"imagegen"}]}]}),
json!({"tools":[{"type":"function","name":"image_generation"}]}),
] {
assert!(
!generate(body.clone()).image_generation_requested(),
"{body}"
);
}
let private = GenerateRequest::from_protocol_payload(
ProtocolPayload::json_object(
"private",
json!({"tools":[{"type":"image_generation"}]})
.as_object()
.unwrap()
.clone(),
fn image_generation_fact_tracks_body_replacement_and_ignores_context() {
let operation = Operation::Generate(generate(json!({"input":"hello"})));
let operation = operation
.replace_protocol_wire(
Some(Bytes::from_static(
br#"{"tools":[{"type":"image_generation"}]}"#,
)),
Map::new(),
)
.unwrap(),
);
assert!(!private.image_generation_requested());
.unwrap();
assert!(operation.image_generation_requested());
let operation = operation
.replace_protocol_wire(None, Map::from_iter([("tools".to_owned(), json!([]))]))
.unwrap();
assert!(operation.image_generation_requested());
let operation = operation
.replace_protocol_wire(
Some(Bytes::from_static(br#"{"input":"hello"}"#)),
Map::new(),
)
.unwrap();
assert!(!operation.image_generation_requested());
}
#[test]
@@ -203,6 +191,10 @@ fn image_generation_fact_survives_cross_protocol_translation() {
)
.unwrap();
assert!(translated.image_generation_requested());
let replaced = translated
.replace_protocol_wire(Some(Bytes::from_static(br#"{"messages":[]}"#)), Map::new())
.unwrap();
assert!(replaced.image_generation_requested());
}
#[test]
@@ -200,6 +200,7 @@ fn scheduling_state_preserves_provider_neutral_signals() {
in_flight: 2,
last_started_at: None,
quota_reset_at: None,
quota_reset_windows: None,
quota_remaining_rank: Some(7),
cooldown: None,
failure_rate_basis_points: Some(125),
+17
View File
@@ -0,0 +1,17 @@
[package]
name = "gateway-egress"
version.workspace = true
edition.workspace = true
license.workspace = true
publish.workspace = true
rust-version.workspace = true
[dependencies]
gateway-core = { path = "../gateway-core" }
http = "1.4.0"
hyper-util = { version = "=0.1.21", features = ["client-proxy"] }
reqwest = { version = "=0.12.28", default-features = false, features = ["socks"] }
url = "2.5"
[lints]
workspace = true
+117
View File
@@ -0,0 +1,117 @@
//! 网关 HTTP 与 WebSocket 出口:业务指定的代理优先,否则按目的地址使用部署环境
use std::sync::OnceLock;
use gateway_core::account::{InvalidOutboundProxy, OutboundProxy};
use hyper_util::client::proxy::matcher::Matcher;
use reqwest::{Client, ClientBuilder, NoProxy, Proxy};
use url::Url;
struct Environment {
http: Option<OutboundProxy>,
https: Option<OutboundProxy>,
no_proxy: String,
matcher: Matcher,
}
impl Environment {
fn load() -> Result<Self, InvalidOutboundProxy> {
let var = |name: &str| {
std::env::var(name)
.or_else(|_| std::env::var(name.to_ascii_lowercase()))
.ok()
.filter(|value| !value.is_empty())
};
// CGI 的 HTTP_* 来自入站请求头,不能作为部署配置
let no_proxy = var("NO_PROXY").unwrap_or_default();
let bypass_all = std::env::var_os("REQUEST_METHOD").is_some()
|| no_proxy.split(',').any(|entry| entry.trim() == "*");
let all = var("ALL_PROXY");
let route = |name| {
if bypass_all {
return Ok(None);
}
var(name)
.or_else(|| all.clone())
.map(|value| {
// 部署环境也接受常见的 host:port 写法,服务内保存的代理仍须是完整 URL
if value.contains("://") {
OutboundProxy::parse(&value)
} else {
OutboundProxy::parse(&format!("http://{value}"))
}
})
.transpose()
};
let http = route("HTTP_PROXY")?;
let https = route("HTTPS_PROXY")?;
let matcher = Matcher::builder()
.http(http.as_ref().map_or("", OutboundProxy::expose_url))
.https(https.as_ref().map_or("", OutboundProxy::expose_url))
.no(no_proxy.as_str())
.build();
Ok(Self {
http,
https,
no_proxy,
matcher,
})
}
}
fn environment() -> Result<&'static Environment, InvalidOutboundProxy> {
static ENV: OnceLock<Result<Environment, InvalidOutboundProxy>> = OnceLock::new();
ENV.get_or_init(Environment::load).as_ref().map_err(|e| *e)
}
/// 每个请求(包括重定向)按目标匹配环境代理;显式代理不受 NO_PROXY 影响
pub fn http_client(proxy: Option<&OutboundProxy>) -> Result<ClientBuilder, InvalidOutboundProxy> {
let mut builder = Client::builder().no_proxy();
if let Some(proxy) = proxy {
return Ok(builder.proxy(Proxy::all(proxy.expose_url()).map_err(|_| InvalidOutboundProxy)?));
}
let env = environment()?;
let bypass = NoProxy::from_string(&env.no_proxy);
if let Some(proxy) = &env.http {
builder = builder.proxy(
Proxy::http(proxy.expose_url())
.map_err(|_| InvalidOutboundProxy)?
.no_proxy(bypass.clone()),
);
}
if let Some(proxy) = &env.https {
builder = builder.proxy(
Proxy::https(proxy.expose_url())
.map_err(|_| InvalidOutboundProxy)?
.no_proxy(bypass),
);
}
Ok(builder)
}
/// 自行拨号的传输复用 reqwest 使用的匹配器,但保留原代理的凭据与连接限制
pub fn for_destination<'a>(
proxy: Option<&'a OutboundProxy>,
url: &Url,
) -> Result<Option<&'a OutboundProxy>, InvalidOutboundProxy> {
if proxy.is_some() {
return Ok(proxy);
}
let env = environment()?;
let (scheme, proxy) = match url.scheme() {
"http" | "ws" => ("http", env.http.as_ref()),
"https" | "wss" => ("https", env.https.as_ref()),
_ => return Err(InvalidOutboundProxy),
};
let Some(proxy) = proxy else {
return Ok(None);
};
let authority = &url[url::Position::BeforeHost..url::Position::AfterPort];
let uri = http::Uri::builder()
.scheme(scheme)
.authority(authority)
.path_and_query("/")
.build()
.map_err(|_| InvalidOutboundProxy)?;
Ok(env.matcher.intercept(&uri).map(|_| proxy))
}
+1
View File
@@ -17,6 +17,7 @@ config = { version = "0.15.23", default-features = false, features = ["yaml"] }
flate2 = "1.1.9"
futures = "0.3.32"
gateway-admin = { path = "../gateway-admin" }
gateway-egress = { path = "../gateway-egress" }
gateway-core = { path = "../gateway-core" }
gateway-protocol = { path = "../gateway-protocol" }
hex = "0.4.3"
@@ -52,17 +52,19 @@ struct CachedDownloads {
impl RgAdguardClientDistribution {
pub(crate) fn new() -> Self {
let client = Client::builder()
.connect_timeout(Duration::from_secs(5))
.timeout(Duration::from_secs(30))
.redirect(Policy::none())
.build()
.inspect_err(|_| {
tracing::warn!(
"download resolver HTTP client initialization failed; official downloads will be used"
);
})
.ok();
let client = gateway_egress::http_client(None).ok().and_then(|builder| {
builder
.connect_timeout(Duration::from_secs(5))
.timeout(Duration::from_secs(30))
.redirect(Policy::none())
.build()
.ok()
});
if client.is_none() {
tracing::warn!(
"download resolver HTTP client initialization failed; official downloads will be used"
);
}
Self {
client,
cache: Mutex::new(None),
@@ -216,6 +216,8 @@ impl HttpClient {
{
return Err(HttpError::invalid("URL"));
}
let proxy = gateway_egress::for_destination(proxy, &url)
.map_err(|_| HttpError::invalid("proxy"))?;
let method = Method::from_bytes(request.method.as_bytes())
.map_err(|_| HttpError::invalid("method"))?;
if matches!(method, Method::CONNECT | Method::TRACE) {
+2 -8
View File
@@ -37,17 +37,11 @@ impl PeerTransport for HttpPeerTransport {
account: &ProviderAccount,
authentication: &PeerConnection,
) -> Result<Vec<u8>, PeerFetchFailure> {
let mut builder = reqwest::Client::builder()
.no_proxy()
let builder = gateway_egress::http_client(account.outbound_proxy())
.map_err(|_| UpstreamInfoFailure::Unavailable)?
.redirect(reqwest::redirect::Policy::none())
.connect_timeout(REQUEST_TIMEOUT)
.timeout(REQUEST_TIMEOUT);
if let Some(proxy) = account.outbound_proxy() {
builder = builder.proxy(
reqwest::Proxy::all(proxy.expose_url())
.map_err(|_| UpstreamInfoFailure::Unavailable)?,
);
}
let client = (self.build_client)(builder).map_err(|_| UpstreamInfoFailure::Unavailable)?;
let response = client
.get(format!(
+2 -1
View File
@@ -17,7 +17,8 @@ pub struct ModelsDevPricing;
#[async_trait]
impl PricingSource for ModelsDevPricing {
async fn fetch(&self) -> Result<PricingSyncPreview, AdminError> {
let client = reqwest::Client::builder()
let client = gateway_egress::http_client(None)
.map_err(|_| source_error())?
.https_only(true)
.redirect(reqwest::redirect::Policy::none())
.connect_timeout(Duration::from_secs(10))
+29 -30
View File
@@ -1,4 +1,4 @@
//! 使用与 Provider 请求一致的显式代理或直连路径,执行有超时和响应大小限制的出口测试
//! 按指定出口,执行有超时和响应大小限制的出口测试
use std::{
net::{IpAddr, Ipv4Addr, Ipv6Addr},
@@ -9,11 +9,20 @@ use std::{
use async_trait::async_trait;
use gateway_admin::{
model::proxies::{ProxyLocationDetection, ProxyTestResult},
ports::proxy::ProxyProbe,
ports::proxy::{ProbeRoute, ProxyProbe},
};
use gateway_core::account::{OutboundProxy, RequestLocation};
use gateway_core::account::RequestLocation;
use serde::Deserialize;
fn probe_client(route: ProbeRoute<'_>) -> Result<reqwest::ClientBuilder, &'static str> {
match route {
ProbeRoute::Direct => Ok(reqwest::Client::builder().no_proxy()),
ProbeRoute::Deployment => gateway_egress::http_client(None),
ProbeRoute::Proxy(proxy) => gateway_egress::http_client(Some(proxy)),
}
.map_err(|_| "代理地址不合法")
}
enum ProbeStrategy {
Single(String),
Dual {
@@ -75,18 +84,13 @@ impl HttpProxyProbe {
async fn exit_ip_at(
&self,
proxy: Option<&OutboundProxy>,
proxy: ProbeRoute<'_>,
endpoint: &str,
) -> Result<IpAddr, &'static str> {
let mut builder = reqwest::Client::builder()
.no_proxy()
let builder = probe_client(proxy)?
.connect_timeout(Duration::from_secs(5))
.timeout(Duration::from_secs(12))
.redirect(reqwest::redirect::Policy::none());
if let Some(proxy) = proxy {
builder = builder
.proxy(reqwest::Proxy::all(proxy.expose_url()).map_err(|_| "代理地址不合法")?);
}
let client = (self.build_client)(builder)?;
let mut response = client.get(endpoint).send().await.map_err(|error| {
if error.is_timeout() {
@@ -125,7 +129,7 @@ impl HttpProxyProbe {
impl HttpProxyProbe {
async fn timed_exit_ip_at(
&self,
proxy: Option<&OutboundProxy>,
proxy: ProbeRoute<'_>,
endpoint: &str,
) -> (Result<IpAddr, &'static str>, u64) {
let started = Instant::now();
@@ -134,7 +138,7 @@ impl HttpProxyProbe {
(result, latency_ms)
}
async fn test_connection(&self, proxy: Option<&OutboundProxy>) -> ProxyTestResult {
async fn test_connection(&self, proxy: ProbeRoute<'_>) -> ProxyTestResult {
let started = Instant::now();
let timeout_limit = Duration::from_secs(15);
@@ -276,39 +280,34 @@ impl HttpProxyProbe {
async fn location_for_ip(
&self,
proxy: Option<&OutboundProxy>,
proxy: ProbeRoute<'_>,
ip: IpAddr,
) -> Result<RequestLocation, &'static str> {
// 出口 IP 已由探测确认;按这个固定 IP 从服务端查询位置,避免代理出口屏蔽位置服务
// 部署环境无法直连位置服务时仍尝试原代理路径
match self.lookup_location(ip, None).await {
// 查询的是已经验证的 IP,不要求被测代理能访问位置服务。
// 服务器直连探测保持直连,其余探测先按部署环境查询,必要时再用被测代理。
if !matches!(proxy, ProbeRoute::Proxy(_)) {
return self.lookup_location(ip, proxy).await;
}
match self.lookup_location(ip, ProbeRoute::Deployment).await {
Ok(location) => Ok(location),
Err(direct_error) if proxy.is_some() => {
self.lookup_location(ip, proxy).await.or(Err(direct_error))
}
Err(error) => Err(error),
Err(error) => self.lookup_location(ip, proxy).await.or(Err(error)),
}
}
async fn lookup_location(
&self,
ip: IpAddr,
proxy: Option<&OutboundProxy>,
proxy: ProbeRoute<'_>,
) -> Result<RequestLocation, &'static str> {
let timeout = if proxy.is_some() {
let timeout = if !matches!(proxy, ProbeRoute::Direct) {
Duration::from_secs(4)
} else {
Duration::from_secs(3)
};
let mut builder = reqwest::Client::builder()
.no_proxy()
let builder = probe_client(proxy)?
.connect_timeout(timeout)
.timeout(timeout)
.redirect(reqwest::redirect::Policy::none());
if let Some(proxy) = proxy {
builder = builder
.proxy(reqwest::Proxy::all(proxy.expose_url()).map_err(|_| "代理地址不合法")?);
}
let client = (self.build_client)(builder)?;
// 查询已检测到的具体出口,不能再次查询“我的 IP”,轮换代理可能换到另一个出口
let url = format!("{}/{ip}", self.location_endpoint.trim_end_matches('/'));
@@ -367,7 +366,7 @@ impl HttpProxyProbe {
async fn detect_location(
&self,
proxy: Option<&OutboundProxy>,
proxy: ProbeRoute<'_>,
result: &ProxyTestResult,
) -> ProxyLocationDetection {
let lookup = async |ip: Option<IpAddr>| match ip {
@@ -397,7 +396,7 @@ impl HttpProxyProbe {
#[async_trait]
impl ProxyProbe for HttpProxyProbe {
async fn test(&self, proxy: Option<&OutboundProxy>, detect_location: bool) -> ProxyTestResult {
async fn test(&self, proxy: ProbeRoute<'_>, detect_location: bool) -> ProxyTestResult {
let mut result = self.test_connection(proxy).await;
if detect_location {
result.location = if result.success {
@@ -132,16 +132,17 @@ impl Metadata {
let client = self
.client
.get_or_try_init(|| async {
Client::builder()
gateway_egress::http_client(None)
.map_err(|_| super::error::invalid("部署代理配置不合法"))?
.user_agent("codex-proxy-rs")
.timeout(Duration::from_secs(10))
.connect_timeout(Duration::from_secs(5))
.redirect(reqwest::redirect::Policy::none())
.retry(reqwest::retry::never())
.build()
.map_err(|e| super::error::http("无法初始化发行查询连接", e))
})
.await
.map_err(|e| super::error::http("无法初始化发行查询连接", e))?;
.await?;
let request = client.get(url);
let mut response = if api {
request.header("Accept", "application/vnd.github+json")
@@ -136,7 +136,8 @@ pub(crate) fn download_client(
timeout: Duration,
) -> Result<reqwest::Client, OperationError> {
let api_base = api_base.to_owned();
reqwest::Client::builder()
gateway_egress::http_client(None)
.map_err(|_| invalid("部署代理配置不合法"))?
.timeout(timeout)
.redirect(reqwest::redirect::Policy::custom(move |attempt| {
if attempt.previous().len() < 5 && download_url_allowed(attempt.url(), &api_base) {
@@ -268,7 +268,12 @@ fn move_file(from: &Path, to: &Path) -> io::Result<()> {
fn move_dir(from: &Path, to: &Path) -> io::Result<()> {
match fs::rename(from, to) {
Ok(()) => Ok(()),
Err(error) if error.kind() == io::ErrorKind::CrossesDevices => {
Err(error)
if error.kind() == io::ErrorKind::CrossesDevices
|| (cfg!(unix) && error.kind() == io::ErrorKind::PermissionDenied) =>
{
// Unix 跨父目录移动只读目录可能被拒绝,复用保留权限的复制路径
// 仅清理源目录时放开写权限,备份与恢复后的目录仍保留原权限
if let Err(copy_error) = copy_dir_all(from, to) {
let _ = remove_dir_all(to);
return Err(copy_error);
@@ -354,3 +354,99 @@ async fn managed_http_enforces_the_callers_body_limit_before_sending() {
assert_eq!(requests[0].body.len(), maximum);
assert_eq!(requests[1].body.len(), maximum + 1);
}
#[tokio::test]
async fn deployment_proxy_routes_managed_http_and_account_probes() {
use gateway_admin::ports::proxy::{ProbeRoute, ProxyProbe};
use gateway_host::proxy_probe::HttpProxyProbe;
if let Ok(origin) = std::env::var("CPR_TEST_EGRESS_ORIGIN") {
let selected =
OutboundProxy::parse(&std::env::var("CPR_TEST_EGRESS_SELECTED").unwrap()).unwrap();
let outside = origin.replace("127.0.0.1", "localhost");
let managed = HttpClient::new().unwrap();
let http = |url: String, proxy| {
let managed = &managed;
async move {
managed
.open(
request(url),
proxy,
&local_network(),
Duration::from_secs(2),
)
.await
.unwrap()
.status
}
};
let probe = HttpProxyProbe::new(&outside);
let observed = |result: gateway_admin::model::proxies::ProxyTestResult| {
result.exit_ipv4.unwrap().octets()[3]
};
assert_eq!(observed(probe.test(ProbeRoute::Direct, false).await), 1);
assert_eq!(http(origin.clone(), Some(&selected)).await, 203);
assert_eq!(http(origin, None).await, 200);
let expect_env = std::env::var_os("HTTP_PROXY").is_some();
assert_eq!(
http(outside, None).await,
if expect_env { 202 } else { 200 }
);
assert_eq!(
observed(probe.test(ProbeRoute::Deployment, false).await),
if expect_env { 2 } else { 1 }
);
println!("checked managed egress");
return;
}
let origin = MockServer::start().await;
let env = MockServer::start().await;
let selected = MockServer::start().await;
for (server, code, ip) in [
(&origin, 200, "192.0.2.1"),
(&env, 202, "192.0.2.2"),
(&selected, 203, "192.0.2.3"),
] {
Mock::given(any())
.respond_with(ResponseTemplate::new(code).set_body_json(serde_json::json!({"ip": ip})))
.mount(server)
.await;
}
for configured in [true, false] {
let mut command = tokio::process::Command::new(std::env::current_exe().unwrap());
for name in [
"HTTP_PROXY",
"HTTPS_PROXY",
"ALL_PROXY",
"NO_PROXY",
"REQUEST_METHOD",
] {
command
.env_remove(name)
.env_remove(name.to_ascii_lowercase());
}
if configured {
command
.env("HTTP_PROXY", env.uri())
.env("NO_PROXY", "127.0.0.0/8");
}
let output = command
.args([
"--exact",
"outbound::http::deployment_proxy_routes_managed_http_and_account_probes",
"--nocapture",
])
.env("CPR_TEST_EGRESS_ORIGIN", origin.uri())
.env("CPR_TEST_EGRESS_SELECTED", selected.uri())
.output()
.await
.unwrap();
assert!(
output.status.success(),
"{}",
String::from_utf8_lossy(&output.stderr)
);
assert!(String::from_utf8_lossy(&output.stdout).contains("checked managed egress"));
}
}
@@ -1,6 +1,6 @@
//! 验证代理及直连探测的双栈地址、认证、证书与失败边界
use gateway_admin::ports::proxy::ProxyProbe;
use gateway_admin::ports::proxy::{ProbeRoute, ProxyProbe};
use gateway_core::account::OutboundProxy;
use gateway_host::proxy_probe::HttpProxyProbe;
use serde_json::json;
@@ -28,7 +28,7 @@ async fn proxy_probe_supports_ipv4_and_ipv6_proxies_and_exit_addresses() {
OutboundProxy::parse(&format!("http://user:password@{}", proxy_server.address()))
.unwrap();
let result = HttpProxyProbe::new("http://unresolvable.invalid/ip")
.test(Some(&proxy), false)
.test(ProbeRoute::Proxy(&proxy), false)
.await;
assert!(
result.success,
@@ -55,7 +55,7 @@ async fn proxy_probe_rejects_auth_errors_redirects_and_invalid_or_oversized_resp
.await;
let result = HttpProxyProbe::new("http://unresolvable.invalid/ip")
.test(
Some(&OutboundProxy::parse(&proxy_server.uri()).unwrap()),
ProbeRoute::Proxy(&OutboundProxy::parse(&proxy_server.uri()).unwrap()),
false,
)
.await;
@@ -76,7 +76,7 @@ async fn unavailable_proxy_never_falls_back_to_direct_connection() {
let proxy = OutboundProxy::parse(&format!("http://{}", unused.local_addr().unwrap())).unwrap();
drop(unused);
let result = HttpProxyProbe::new(target.uri())
.test(Some(&proxy), false)
.test(ProbeRoute::Proxy(&proxy), false)
.await;
assert!(!result.success);
}
@@ -91,7 +91,10 @@ async fn invalid_certificate_configuration_should_not_fall_back_or_expose_detail
.await;
let result = HttpProxyProbe::new(target.uri())
.with_client_builder(|_| Err("private-certificate-path"))
.test(Some(&OutboundProxy::parse(&target.uri()).unwrap()), false)
.test(
ProbeRoute::Proxy(&OutboundProxy::parse(&target.uri()).unwrap()),
false,
)
.await;
assert!(!result.success);
assert!(!result.message.contains("private-certificate-path"));
@@ -117,7 +120,7 @@ async fn dual_stack_proxy_probe_reports_both_addresses_when_available() {
format!("{}/v4", proxy_server.uri()),
format!("{}/v6", proxy_server.uri()),
)
.test(Some(&proxy), false)
.test(ProbeRoute::Proxy(&proxy), false)
.await;
assert!(result.success);
@@ -146,7 +149,7 @@ async fn dual_stack_proxy_probe_reports_single_stack_when_only_one_succeeds() {
format!("{}/v4", proxy_server.uri()),
format!("{}/v6", proxy_server.uri()),
)
.test(Some(&proxy), false)
.test(ProbeRoute::Proxy(&proxy), false)
.await;
assert!(result.success);
@@ -178,7 +181,7 @@ async fn dual_stack_latency_measures_the_selected_exit_without_the_other_family_
format!("{}/v4", server.uri()),
format!("{}/v6", server.uri()),
)
.test(None, false)
.test(ProbeRoute::Direct, false)
.await;
assert!(result.success);
assert_eq!(result.exit_ipv4.is_some(), working_family == "v4");
@@ -229,7 +232,7 @@ async fn direct_location_uses_verified_addresses_and_rejects_conflicting_timezon
format!("{}/v6", server.uri()),
)
.with_location_endpoint(format!("{}/geo", server.uri()))
.test(None, true)
.test(ProbeRoute::Direct, true)
.await;
assert!(result.success);
if v6_timezone == Some("America/New_York") {
@@ -261,7 +264,7 @@ async fn direct_location_failure_does_not_retry_the_same_path_or_return_a_locati
.await;
let result = HttpProxyProbe::new(format!("{}/ip", server.uri()))
.with_location_endpoint(format!("{}/geo", server.uri()))
.test(None, true)
.test(ProbeRoute::Direct, true)
.await;
assert_eq!(
result.location,
@@ -294,7 +297,7 @@ async fn location_lookup_uses_verified_exit_ip_without_requiring_proxy_access_to
let result = HttpProxyProbe::new("http://unresolvable.invalid/ip")
.with_location_endpoint(format!("{}/geo", location_server.uri()))
.test(
Some(&OutboundProxy::parse(&proxy_server.uri()).unwrap()),
ProbeRoute::Proxy(&OutboundProxy::parse(&proxy_server.uri()).unwrap()),
true,
)
.await;
@@ -333,7 +336,7 @@ async fn location_lookup_uses_probed_ips_through_the_proxy_and_checks_both_timez
)
.with_location_endpoint("http://unresolvable.invalid/geo/")
.test(
Some(&OutboundProxy::parse(&proxy_server.uri()).unwrap()),
ProbeRoute::Proxy(&OutboundProxy::parse(&proxy_server.uri()).unwrap()),
true,
)
.await;
@@ -379,7 +382,7 @@ async fn geolocation_failure_never_turns_connectivity_into_failure_or_uses_unver
let result = HttpProxyProbe::new("http://unresolvable.invalid/ip")
.with_location_endpoint("http://unresolvable.invalid/geo")
.test(
Some(&OutboundProxy::parse(&proxy_server.uri()).unwrap()),
ProbeRoute::Proxy(&OutboundProxy::parse(&proxy_server.uri()).unwrap()),
true,
)
.await;
@@ -407,7 +410,7 @@ async fn disabled_geolocation_never_calls_the_location_service() {
let result = HttpProxyProbe::new("http://unresolvable.invalid/ip")
.with_location_endpoint("http://unresolvable.invalid/geo")
.test(
Some(&OutboundProxy::parse(&proxy_server.uri()).unwrap()),
ProbeRoute::Proxy(&OutboundProxy::parse(&proxy_server.uri()).unwrap()),
false,
)
.await;
@@ -437,6 +437,13 @@ async fn restart_should_conflict_while_another_system_operation_is_running() {
async fn rollback_should_restore_binary_web_and_version_state() {
let server = MockServer::start().await;
let fixture = Fixture::new();
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt as _;
fs::set_permissions(fixture.official(), fs::Permissions::from_mode(0o555)).unwrap();
fs::set_permissions(fixture.web(), fs::Permissions::from_mode(0o751)).unwrap();
}
fixture
.mount_release(
&server,
@@ -470,6 +477,14 @@ async fn rollback_should_restore_binary_web_and_version_state() {
& 0o777,
0o555
);
assert_eq!(
fs::metadata(fixture.official().with_file_name("official.backup"))
.expect("official backup directory")
.permissions()
.mode()
& 0o777,
0o555
);
for name in ["plugin-release-manifest.json", "new-plugin.tar.gz"] {
assert_eq!(
fs::metadata(fixture.official().join(name))
@@ -487,6 +502,18 @@ async fn rollback_should_restore_binary_web_and_version_state() {
.await
.expect("rollback");
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt as _;
for (path, mode) in [(fixture.official(), 0o555), (fixture.web(), 0o751)] {
assert_eq!(
fs::metadata(path).unwrap().permissions().mode() & 0o777,
mode
);
}
}
assert_eq!(
fs::read(fixture.executable()).expect("binary"),
b"old-binary"
@@ -598,6 +625,8 @@ async fn rollback_should_restore_current_release_when_plugin_revision_changes_du
.expect_err("revision changed after rollback swap");
assert_eq!(error.kind(), SystemOperationErrorKind::Conflict);
#[cfg(unix)]
assert_official_directory_readonly(&fixture);
assert_eq!(
fs::read(fixture.executable()).expect("binary"),
b"new-binary"
@@ -661,6 +690,8 @@ async fn cancelled_rollback_should_restore_current_release_after_the_swap() {
.expect_err("rollback task cancelled")
.is_cancelled()
);
#[cfg(unix)]
assert_official_directory_readonly(&fixture);
assert_eq!(
fs::read(fixture.executable()).expect("binary"),
@@ -1131,10 +1162,94 @@ async fn update_should_replace_web_assets_across_filesystems() {
);
}
#[cfg(unix)]
#[tokio::test]
async fn update_should_preserve_readonly_staged_directory_permissions() {
use std::os::unix::fs::PermissionsExt as _;
struct ReadonlyStagingPreflight(PathBuf);
#[async_trait::async_trait]
impl SystemUpdatePreflight for ReadonlyStagingPreflight {
async fn validate(
&self,
candidate: SystemUpdateCandidate,
) -> Result<Revision, SystemOperationError> {
let staging = fs::read_dir(&self.0)
.unwrap()
.next()
.expect("staging directory")
.unwrap()
.path();
let nested = staging.join("web-dist/assets");
fs::create_dir(&nested).unwrap();
fs::write(nested.join("asset.js"), "asset").unwrap();
for (path, mode) in [
(nested.join("asset.js"), 0o444),
(nested, 0o550),
(staging.join("web-dist"), 0o555),
(staging.join("official-plugins"), 0o555),
] {
fs::set_permissions(path, fs::Permissions::from_mode(mode)).unwrap();
}
AllowingUpdatePreflight.validate(candidate).await
}
async fn confirm_revision(&self, _: Revision) -> Result<(), SystemOperationError> {
Ok(())
}
}
let server = MockServer::start().await;
let fixture = Fixture::new();
fixture
.mount_release(
&server,
TARGET_VERSION,
ArchiveKind::Safe,
ChecksumKind::Valid,
)
.await;
let config = fixture.config(&format!("{}/repos", server.uri()));
let preflight = Arc::new(ReadonlyStagingPreflight(config.update_temp_dir.clone()));
let service = ProcessSystemOperations::new(CancellationToken::new(), config);
service
.perform_update(Some(TARGET_VERSION.to_owned()), None, preflight)
.await
.expect("update accepted");
let status = wait_for_update(&service).await;
assert_eq!(
status.operation.status,
SystemOperationStatus::Succeeded,
"{status:?}"
);
assert_eq!(
fs::read(fixture.web().join("assets/asset.js")).unwrap(),
b"asset"
);
for (path, mode) in [
(fixture.official(), 0o555),
(fixture.web(), 0o555),
(fixture.web().join("assets"), 0o550),
(fixture.web().join("assets/asset.js"), 0o444),
] {
assert_eq!(
fs::metadata(path).unwrap().permissions().mode() & 0o777,
mode
);
}
}
#[tokio::test]
async fn update_should_restore_web_assets_when_binary_backup_fails() {
let server = MockServer::start().await;
let fixture = Fixture::new();
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt as _;
fs::set_permissions(fixture.official(), fs::Permissions::from_mode(0o555)).unwrap();
}
fixture
.mount_release(
&server,
@@ -1162,6 +1277,22 @@ async fn update_should_restore_web_assets_when_binary_backup_fails() {
);
assert!(fixture.official().join("old-plugin.tar.gz").is_file());
assert!(!fixture.official().join("new-plugin.tar.gz").exists());
#[cfg(unix)]
assert_official_directory_readonly(&fixture);
}
#[cfg(unix)]
fn assert_official_directory_readonly(fixture: &Fixture) {
use std::os::unix::fs::PermissionsExt as _;
assert_eq!(
fs::metadata(fixture.official())
.unwrap()
.permissions()
.mode()
& 0o777,
0o555
);
}
#[tokio::test]
@@ -630,6 +630,7 @@ fn candidate(id: &str, weight: u16) -> AccountCandidate {
in_flight: 0,
last_started_at: None,
quota_reset_at: None,
quota_reset_windows: None,
quota_remaining_rank: None,
cooldown: None,
failure_rate_basis_points: None,
@@ -947,7 +947,7 @@ impl PricingSource for UnusedAdminRuntime {
impl ProxyProbe for UnusedAdminRuntime {
async fn test(
&self,
_: Option<&gateway_core::account::OutboundProxy>,
_: gateway_admin::ports::proxy::ProbeRoute<'_>,
_: bool,
) -> ProxyTestResult {
panic!("unexpected proxy probe")
@@ -279,7 +279,7 @@ HTTP 和 request 中间件中的模型回调沿用有效请求设置;attempt
`HostClient::account_http(AccountHttpRequest { account_id, credential_revision, request }, body)` 使用指定账号的当前代理配置,
经 `host.http.account` 返回同一 `HostHttpBody`。每次发送核对账号存在、启用和凭据版本,版本变化须由插件重新读取后决定下一步。
请求 headers 与正文完全由插件提供;宿主不注入账号认证、不申请推理租约,代理不可用时直接失败,不回退直连。
请求 headers 与正文完全由插件提供;宿主不注入账号认证、不申请推理租约。无显式代理时使用部署默认出口,代理不可用时直接失败,不回退直连。
这项能力与上游适配器的账号推理 HTTP 分开,适用于插件自行构造的账号网络请求
`HostClient::proxy_http(ProxyHttpRequest { proxy_id, request }, body)` 经 `host.http.proxy` 使用指定代理的当前配置,
@@ -449,6 +449,9 @@ Ping 的自动 Pong 与关闭握手属于底层 WebSocket 协议事实;插件
或 `WebSocket { uri }`,不能在准备过程中切换传输类型。HTTP 正文是已经序列化、完成内容编码的实际字节,
插件自行处理压缩与 `Content-Encoding` 等头部的一致性;宿主不解码正文、补回字段或修复请求
内置转发在出站准备前,仅为已识别的互联上游组装 `X-CPR-Route`,向普通模型上游剥离。
插件改写目标或自行实现转发时,负责在互联节点间保留并维护路径,在最终模型上游前移除该头;宿主不在插件返回后重新判断或修复互联链路
`OutboundCall.metadata` 提供真实请求、Client Key、分组、Provider、模型、账号与 attempt 序号。
这些归属由宿主冻结,返回值只有请求;修改 URL、认证头或正文不改变账号租约、凭据版本与计费归属。
宿主不从最终字节重建模型、选号或续接事实;需要改变这些业务含义时,插件应在相应的 request / attempt 阶段保持一致。
@@ -85,7 +85,7 @@ let (message, ()) = tokio::try_join!(
```
宿主先注入已选账号认证,插件显式提供的同名头覆盖默认值。路径支持相对基址或完整 HTTP(S) URL;
声明路径用于标记发送用途,不是访问白名单。出站复用账号代理,代理失败不能回退直连,重定向不自动跟随。
声明路径用于标记发送用途,不是访问白名单。出站复用账号代理,未绑定可用代理时使用部署默认出口,代理失败不能回退直连,重定向不自动跟随。
`auxiliary` 路径的出站计入请求副作用水位,防止附件上传等操作在失败后被透明重复执行
## 事件、结算和续接
@@ -256,6 +256,15 @@ pub struct SmartSchedulingConfig {
pub latency_weight: f64,
pub reset_weight: f64,
pub queue_weight: f64,
#[serde(default)]
pub reset_window_weights: ResetWindowWeights,
}
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct ResetWindowWeights {
pub five_hour: f64,
pub weekly: f64,
pub monthly: f64,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
@@ -24,6 +24,7 @@ pub use stream::{
MessagesAccumulator, MessagesToResponsesStream, ResponsesToMessagesStream,
message_to_stream_events,
};
pub use thinking::messages_request_effort;
pub use usage::MessageUsage;
/// Messages 载荷的协议标识
@@ -1,16 +1,22 @@
//! 按目标模型能力映射推理等级与预算;未知兼容上游沿用 CPA 手动预算
//! 提取 Messages 显式 effort,并按目标模型能力映射推理等级与预算
use serde_json::{Map, Value, json};
use super::MessagesModelCapabilities;
pub(super) fn to_responses(request: &Map<String, Value>) -> Value {
let configured = request
/// 返回客户端显式设置的 effort,预算与默认档位由转换器单独处理
#[must_use]
pub fn messages_request_effort(request: &Map<String, Value>) -> Option<&str> {
request
.get("output_config")
.and_then(|value| value.get("effort"))
.and_then(Value::as_str)
.map(|value| value.trim().to_ascii_lowercase())
.filter(|value| !value.is_empty());
.map(str::trim)
.filter(|value| !value.is_empty())
}
pub(super) fn to_responses(request: &Map<String, Value>) -> Value {
let configured = messages_request_effort(request).map(str::to_ascii_lowercase);
let thinking = request.get("thinking");
let effort = match thinking
.and_then(|value| value.get("type"))
@@ -127,7 +127,7 @@ pub fn codex_responses_request_semantics(
) -> CodexResponsesRequestSemantics {
codex_responses_request_semantics_with_turn_metadata(
body,
non_empty_string(context.get("turn_metadata")),
context.get("turn_metadata").and_then(Value::as_str),
None,
)
}
@@ -153,10 +153,12 @@ pub fn codex_responses_request_semantics_with_turn_metadata(
let reasoning_effort = reasoning_effort_override
.clone()
.or_else(|| reasoning_effort(body.get("reasoning")));
// 权威上下文与显式 turnMetadata 即使为空也不退回其他别名
let parsed_turn_metadata = turn_metadata
.or_else(|| body.get("turnMetadata").and_then(Value::as_str))
.or_else(|| request_turn_metadata(body))
.and_then(|value| serde_json::from_str::<Value>(value).ok());
let request_kind = parsed_turn_metadata
let mut request_kind = parsed_turn_metadata
.as_ref()
.and_then(|metadata| non_empty_string(metadata.get("request_kind")))
.map(ToOwned::to_owned);
@@ -174,6 +176,12 @@ pub fn codex_responses_request_semantics_with_turn_metadata(
&& proactive_multi_agent)
.then_some("ultra");
let compact = request_kind.as_deref() == Some("compaction") || compact_trigger;
// generate=false 才是预热,客户端分类提示不能把真实推理排除在用量统计外
if body.get("generate").and_then(Value::as_bool) == Some(false) {
request_kind = Some("prewarm".to_owned());
} else if request_kind.as_deref() == Some("prewarm") {
request_kind = None;
}
CodexResponsesRequestSemantics {
reasoning_effort,
@@ -217,7 +225,10 @@ fn latest_multi_agent_mode<'a>(
input: impl DoubleEndedIterator<Item = &'a Value>,
) -> Option<&'a str> {
input.rev().find_map(|item| {
if non_empty_string(item.get("role")) != Some("developer") {
// 显式 system 消息会由 Codex 编码器转为 developer,观测在编码前保持等价
let system_message = item.get("type").and_then(Value::as_str) == Some("message")
&& item.get("role").and_then(Value::as_str) == Some("system");
if non_empty_string(item.get("role")) != Some("developer") && !system_message {
return None;
}
item.get("content")?
@@ -1,4 +1,4 @@
//! 从 OpenAI 事件与响应头提取用量、额度和重试等待事实
//! 从 OpenAI 事件与响应头提取用量、额度和重试事实
use std::{collections::BTreeMap, time::Duration};
@@ -249,6 +249,37 @@ pub(super) fn retry_after_header(value: &Value) -> Option<(&Value, u64)> {
Some((value, seconds))
}
/// 判断上游错误是否携带 `x-retry-metadata: NO_MORE_RETRY`
///
/// 读取位置与 Retry-After 相同:优先错误对象的 headers,再看事件自身的 headers
/// 只承认这一明确信号,其他取值、数组和非字符串都不构成该事实
pub fn upstream_no_more_retry(value: &Value) -> bool {
let error = value
.pointer("/response/error")
.or_else(|| value.get("error"))
.unwrap_or(value);
no_more_retry_headers(error) || no_more_retry_headers(value)
}
/// 从上游错误正文判断是否声明不再重试当前账号
pub fn upstream_no_more_retry_from_body(body: &str) -> bool {
serde_json::from_str::<Value>(body).is_ok_and(|value| upstream_no_more_retry(&value))
}
fn no_more_retry_headers(value: &Value) -> bool {
value
.get("headers")
.and_then(Value::as_object)
.is_some_and(|headers| {
headers.iter().any(|(name, value)| {
name.eq_ignore_ascii_case("x-retry-metadata")
&& value
.as_str()
.is_some_and(|value| value.trim().eq_ignore_ascii_case("NO_MORE_RETRY"))
})
})
}
/// 标准化的单个限流窗口
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct RateLimitWindow {
@@ -1,5 +1,8 @@
//! 共享的 OpenAI 线协议类型、解析与编码入口
/// OpenAI 线协议载荷与事件的协议标识
pub const OPENAI_PROTOCOL: &str = "openai";
/// Codex Responses Lite 的 HTTP 请求头
pub const X_OPENAI_INTERNAL_CODEX_RESPONSES_LITE_HEADER: &str =
"x-openai-internal-codex-responses-lite";
+84 -3
View File
@@ -1,7 +1,56 @@
//! 从协议正文提取调度需要的纯请求事实
//! 从协议正文提取能力要求与请求观测事实
use serde_json::{Map, Value};
use crate::anthropic::{ANTHROPIC_PROTOCOL, messages_request_effort};
use crate::openai::{OPENAI_PROTOCOL, codex_responses_request_semantics_with_turn_metadata};
/// 与执行 Provider 无关的逻辑请求观测
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct RequestObservation {
/// 当前请求与已核实续接配置共同确定的推理强度
pub reasoning_effort: Option<String>,
/// 协议声明派生的推理预设
pub reasoning_preset: Option<String>,
/// 协议声明的请求类型,可信内部分类由调用方覆盖
pub request_kind: Option<String>,
/// 协议声明的子代理类型
pub subagent_kind: Option<String>,
/// 当前请求是否声明压缩
pub compact: bool,
}
/// 旧配置只由状态所属 Provider 提供,是否继承仍由当前协议的续接语义决定
#[must_use]
pub fn request_observation(
protocol: &str,
body: &Map<String, Value>,
context: &Map<String, Value>,
previous_reasoning_effort_override: Option<&str>,
) -> RequestObservation {
match protocol {
OPENAI_PROTOCOL => {
let semantics = codex_responses_request_semantics_with_turn_metadata(
body,
context.get("turn_metadata").and_then(Value::as_str),
previous_reasoning_effort_override,
);
RequestObservation {
reasoning_effort: semantics.reasoning_effort,
reasoning_preset: semantics.reasoning_preset.map(str::to_owned),
request_kind: semantics.request_kind,
subagent_kind: semantics.subagent_kind,
compact: semantics.compact,
}
}
ANTHROPIC_PROTOCOL => RequestObservation {
reasoning_effort: messages_request_effort(body).map(str::to_owned),
..RequestObservation::default()
},
_ => RequestObservation::default(),
}
}
/// 请求正文表达的能力要求
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct RequestFacts {
@@ -11,6 +60,8 @@ pub struct RequestFacts {
pub tools: bool,
/// 图片输入
pub vision: bool,
/// 图片生成工具声明
pub image_generation: bool,
/// 推理配置
pub reasoning: bool,
/// JSON Schema 输出
@@ -30,12 +81,13 @@ pub fn request_facts(protocol: &str, body: &Map<String, Value>) -> RequestFacts
..RequestFacts::default()
};
match protocol {
"openai" => {
OPENAI_PROTOCOL => {
facts.max_output_tokens = body
.get("max_output_tokens")
.and_then(Value::as_u64)
.filter(|n| *n > 0);
facts.vision = contains_type(body.get("input"), "input_image");
facts.image_generation = requests_image_generation(body);
facts.reasoning = body.get("reasoning").is_some_and(Value::is_object);
facts.json_schema = body
.get("text")
@@ -48,7 +100,7 @@ pub fn request_facts(protocol: &str, body: &Map<String, Value>) -> RequestFacts
.and_then(Value::as_str)
.is_some();
}
"anthropic" => {
ANTHROPIC_PROTOCOL => {
facts.max_output_tokens = body.get("max_tokens").and_then(Value::as_u64);
facts.vision = contains_type(body.get("messages"), "image");
facts.reasoning = body
@@ -68,6 +120,35 @@ pub fn request_facts(protocol: &str, body: &Map<String, Value>) -> RequestFacts
facts
}
/// 原生工具声明授权上游生成图片,Codex 的被动 namespace 与图片输入不属于该权限
fn requests_image_generation(body: &Map<String, Value>) -> bool {
contains_image_tool(body.get("tools"))
|| body
.get("input")
.and_then(Value::as_array)
.is_some_and(|items| {
items.iter().any(|item| {
item.get("type").and_then(Value::as_str) == Some("additional_tools")
&& contains_image_tool(item.get("tools"))
})
})
|| body.get("tool_choice").is_some_and(|choice| {
is_image_tool(choice)
|| (choice.get("type").and_then(Value::as_str) == Some("allowed_tools")
&& contains_image_tool(choice.get("tools")))
})
}
fn contains_image_tool(value: Option<&Value>) -> bool {
value
.and_then(Value::as_array)
.is_some_and(|tools| tools.iter().any(is_image_tool))
}
fn is_image_tool(value: &Value) -> bool {
value.get("type").and_then(Value::as_str) == Some("image_generation")
}
fn contains_type(value: Option<&Value>, expected: &str) -> bool {
match value {
Some(Value::Array(values)) => values.iter().any(|v| contains_type(Some(v), expected)),
@@ -219,6 +219,27 @@ fn codex_semantics_should_detect_ultra_preset_from_proactive_multi_agent() {
assert_eq!(semantics.reasoning_preset, Some("ultra"));
assert!(!semantics.compact);
for (role, message_type, expected) in [
("system", Some("message"), Some("ultra")),
("system", None, None),
("system", Some("function_call"), None),
("user", Some("message"), None),
(" system ", Some("message"), None),
] {
let mut shaped = body.clone();
let item = shaped["input"][0].as_object_mut().unwrap();
item.insert("role".to_owned(), json!(role));
if let Some(message_type) = message_type {
item.insert("type".to_owned(), json!(message_type));
} else {
item.remove("type");
}
assert_eq!(
codex_responses_request_semantics(&shaped, &Map::new()).reasoning_preset,
expected
);
}
for (effort, expected_preset) in [("high", None), ("max", Some("ultra"))] {
body["input"]
.as_array_mut()
@@ -230,6 +251,27 @@ fn codex_semantics_should_detect_ultra_preset_from_proactive_multi_agent() {
}
}
#[test]
fn empty_authoritative_turn_metadata_does_not_fall_back_to_body_aliases() {
for (context, body) in [
(
json!({"turn_metadata":""}),
json!({"turnMetadata":"{\"request_kind\":\"review\"}"}),
),
(
json!({"turn_metadata":" "}),
json!({"turnMetadata":"{\"request_kind\":\"review\"}"}),
),
(
json!({}),
json!({"turnMetadata":"","turn_metadata":"{\"request_kind\":\"review\"}"}),
),
] {
let semantics = codex_responses_request_semantics(&object(body), &object(context));
assert_eq!(semantics.request_kind, None);
}
}
fn object(value: Value) -> Map<String, Value> {
let Value::Object(object) = value else {
panic!("test value must be an object");
@@ -267,3 +309,42 @@ fn account_identity_does_not_infer_a_session_from_turn_or_cache_keys() {
assert_eq!(codex_account_session_id(&object(body), &Map::new()), None);
}
}
#[test]
fn request_observation_classifies_prewarm_from_generate() {
for (generate, request_kind, expected_kind) in [
(Some(json!(false)), None, Some("prewarm")),
(Some(json!(false)), Some("prewarm"), Some("prewarm")),
(Some(json!(false)), Some("review"), Some("prewarm")),
(Some(json!(true)), Some("prewarm"), None),
(None, Some("prewarm"), None),
(Some(json!(null)), Some("prewarm"), None),
(Some(json!("false")), Some("prewarm"), None),
(None, Some("review"), Some("review")),
(None, None, None),
] {
let mut body = Map::from_iter([
("model".to_owned(), json!("gpt-test")),
("input".to_owned(), json!("hello")),
("store".to_owned(), json!(false)),
]);
if let Some(generate) = generate {
body.insert("generate".to_owned(), generate);
}
let context = Map::from_iter([(
"turn_metadata".to_owned(),
Value::String(
json!({"request_kind": request_kind, "subagent_kind": "worker"}).to_string(),
),
)]);
let observation = codex_responses_request_semantics(&body, &context);
assert_eq!(
observation.request_kind.as_deref(),
expected_kind,
"generate={:?}, request_kind={request_kind:?}",
body.get("generate"),
);
assert_eq!(observation.subagent_kind.as_deref(), Some("worker"));
}
}
@@ -4,7 +4,8 @@ use gateway_protocol::openai::events::{
RateLimitKeySource, RateLimitWindow, TokenUsage, billable_usage_is_complete, extract_sse_usage,
extract_usage, is_codex_quota_header_name, is_rate_limit_header_name,
parse_image_rate_limit_headers, parse_rate_limit_headers, parse_rate_limits_event,
rate_limits_to_header_pairs, retry_after_seconds_from_body,
rate_limits_to_header_pairs, retry_after_seconds_from_body, upstream_no_more_retry,
upstream_no_more_retry_from_body,
};
use serde_json::json;
@@ -429,6 +430,51 @@ fn retry_after_should_not_infer_retry_from_unrelated_error_message() {
assert_eq!(retry_after_seconds_from_body(&body), None);
}
#[test]
fn no_more_retry_should_read_the_error_header_without_replacing_retry_after() {
for (body, expected) in [
(
json!({"type":"error","error":{"type":"service_unavailable_error","code":"server_is_overloaded","headers":{"x-retry-metadata":"NO_MORE_RETRY","retry-after":"30","retry-after-ms":"30000"}}}),
true,
),
(
json!({"type":"response.failed","response":{"error":{"headers":{"X-Retry-Metadata":" no_more_retry "}}}}),
true,
),
(
json!({"type":"error","headers":{"x-retry-metadata":"NO_MORE_RETRY"},"error":{"code":"server_is_overloaded"}}),
true,
),
(
json!({"error":{"headers":{"x-retry-metadata":"RETRY","retry-after":"30"}}}),
false,
),
(
json!({"error":{"headers":{"x-retry-metadata":["NO_MORE_RETRY"]}}}),
false,
),
(
json!({"error":{"code":"server_is_overloaded","headers":{"retry-after":"30"}}}),
false,
),
] {
assert_eq!(upstream_no_more_retry(&body), expected, "{body}");
assert_eq!(
upstream_no_more_retry_from_body(&body.to_string()),
expected,
"{body}"
);
}
assert!(!upstream_no_more_retry_from_body("not-json"));
assert_eq!(
retry_after_seconds_from_body(
&json!({"error":{"headers":{"x-retry-metadata":"NO_MORE_RETRY","retry-after":"30"}}})
.to_string()
),
Some(30)
);
}
#[test]
fn retry_after_should_reject_malformed_json() {
assert_eq!(retry_after_seconds_from_body("not-json"), None);
@@ -1,7 +1,7 @@
//! 验证不同协议的共享请求事实
use gateway_protocol::request::request_facts;
use serde_json::json;
use gateway_protocol::request::{request_facts, request_observation};
use serde_json::{Map, json};
#[test]
fn messages_facts_include_images_inside_tool_results() {
@@ -20,3 +20,77 @@ fn responses_facts_keep_native_continuation_separate_from_history() {
assert!(facts.native_continuation);
assert!(!facts.tools && !facts.vision && !facts.reasoning && !facts.json_schema);
}
#[test]
fn image_generation_detection_distinguishes_tools_from_image_inputs_and_passive_namespaces() {
for body in [
json!({"tools":[{"type":"image_generation"}]}),
json!({"tool_choice":{"type":"image_generation"}}),
json!({"tool_choice":{"type":"allowed_tools","tools":[{"type":"image_generation"}]}}),
json!({"input":[{"type":"additional_tools","tools":[{"type":"image_generation"}]}]}),
] {
assert!(
request_facts("openai", body.as_object().expect("请求正文")).image_generation,
"{body}"
);
}
for body in [
json!({"input":"Please generate an image"}),
json!({"input":[{"role":"user","content":[{"type":"input_image","image_url":"data:image/png;base64,AA=="}]}]}),
json!({"tools":[{"type":"namespace","name":"image_gen","tools":[{"type":"function","name":"imagegen"}]}]}),
json!({"tools":[{"type":"function","name":"image_generation"}]}),
] {
assert!(
!request_facts("openai", body.as_object().expect("请求正文")).image_generation,
"{body}"
);
}
let body = json!({"tools":[{"type":"image_generation"}]});
for protocol in ["anthropic", "private"] {
assert!(!request_facts(protocol, body.as_object().expect("请求正文")).image_generation);
}
}
#[test]
fn messages_request_observation_preserves_explicit_effort_without_inventing_defaults() {
for (options, expected) in [
(json!({"output_config":{"effort":"xhigh"}}), Some("xhigh")),
(
json!({"thinking":{"type":"adaptive"},"output_config":{"effort":"low"}}),
Some("low"),
),
(
json!({"thinking":{"type":"disabled"},"output_config":{"effort":"high"}}),
Some("high"),
),
(
json!({"output_config":{"effort":" Future-Level "}}),
Some("Future-Level"),
),
(json!({"output_config":{"effort":" "}}), None),
(json!({"output_config":{"effort":null}}), None),
(json!({"output_config":{"effort":42}}), None),
(json!({"thinking":{"type":"adaptive"}}), None),
(
json!({"thinking":{"type":"enabled","budget_tokens":32768}}),
None,
),
(json!({}), None),
] {
let mut body = json!({"model":"gpt-5.4","max_tokens":128,"messages":[{"role":"user","content":"hello"}]});
body.as_object_mut()
.unwrap()
.extend(options.as_object().unwrap().clone());
let observation = request_observation(
"anthropic",
body.as_object().unwrap(),
&Map::new(),
Some("inapplicable-previous-effort"),
);
assert_eq!(
observation.reasoning_effort.as_deref(),
expected,
"{options}"
);
}
}
@@ -14,6 +14,21 @@ fn loaded_credential_from_record(
#[async_trait]
impl ProviderAccountStore for PgProviderAccountRepository {
async fn load_account_catalogs(
&self,
provider: &ProviderKind,
) -> Result<Vec<gateway_core::account::StoredAccountCatalog>, CoreStoreError> {
self.load_model_catalogs(provider).await
}
async fn save_account_catalog(
&self,
observation: &gateway_core::account::AccountCatalogObservation,
document: &[u8],
) -> Result<bool, CoreStoreError> {
self.save_model_catalog(observation, document).await
}
async fn create_account(&self, account: CoreNewProviderAccount) -> Result<(), CoreStoreError> {
if account.account.revision().get() != 1 {
return Err(CoreStoreError::new(CoreStoreErrorKind::InvalidData));
@@ -88,6 +103,29 @@ impl ProviderAccountStore for PgProviderAccountRepository {
.collect()
}
async fn list_catalog_candidates(
&self,
provider: &ProviderKind,
accounts: &[CoreProviderAccountId],
omissions: &gateway_core::account::AccountCatalogOmissions,
) -> Result<Vec<CoreProviderAccount>, CoreStoreError> {
let ids = accounts
.iter()
.map(CoreProviderAccountId::as_str)
.collect::<Vec<_>>();
self.list_provider_account_summaries(
Some(provider.as_str()),
false,
Some(&ids),
Some(omissions),
)
.await
.map_err(core_store_error)?
.into_iter()
.map(core_account_from_summary)
.collect()
}
async fn list_for_provider_accounts(
&self,
provider: &ProviderKind,
@@ -100,12 +138,17 @@ impl ProviderAccountStore for PgProviderAccountRepository {
.iter()
.map(CoreProviderAccountId::as_str)
.collect::<Vec<_>>();
self.list_provider_account_summaries(Some(provider.as_str()), false, Some(&account_ids))
.await
.map_err(core_store_error)?
.into_iter()
.map(core_account_from_summary)
.collect()
self.list_provider_account_summaries(
Some(provider.as_str()),
false,
Some(&account_ids),
None,
)
.await
.map_err(core_store_error)?
.into_iter()
.map(core_account_from_summary)
.collect()
}
async fn list_refresh_candidates(
@@ -58,6 +58,7 @@ mod admin_queries;
mod authorization;
mod core_adapter;
mod mapping;
mod model_catalogs;
mod quota;
mod quota_forecast;
mod repository;
@@ -0,0 +1,83 @@
//! 账号目录的版本约束写入和按相同正文合并的批量恢复
use super::{
CoreCredentialRevision, CoreProviderAccountId, CoreStoreError, CoreStoreErrorKind, DateTime,
PgProviderAccountRepository, ProviderKind, Utc,
};
use gateway_core::account::{AccountCatalogObservation, StoredAccountCatalog};
impl PgProviderAccountRepository {
pub(super) async fn load_model_catalogs(
&self,
provider: &ProviderKind,
) -> Result<Vec<StoredAccountCatalog>, CoreStoreError> {
// 先聚合轻量归属,再按主键读取代表正文,避免一万个相同账号展开一万份 TOAST
let rows = sqlx::query_as::<_, (Vec<String>, Vec<i64>, Vec<Option<String>>, Vec<Option<String>>, Vec<DateTime<Utc>>, Vec<u8>)>(
"select g.ids, g.revisions, g.plans, g.etags, g.times, d.document from (
select min(c.account_id) representative, array_agg(c.account_id order by c.account_id) ids,
array_agg(c.credential_revision order by c.account_id) revisions,
array_agg(c.plan_type order by c.account_id) plans, array_agg(c.etag order by c.account_id) etags, array_agg(c.observed_at order by c.account_id) times
from provider_account_model_catalogs c join provider_accounts a on a.id=c.account_id
where a.provider_kind=$1 and a.enabled and a.credential_state in ('unknown','ready')
and (a.access_token_expires_at is null or a.access_token_expires_at > now())
and a.credential_revision=c.credential_revision and a.plan_type is not distinct from c.plan_type
group by c.document_hash
) g join provider_account_model_catalogs d on d.account_id=g.representative")
.bind(provider.as_str()).fetch_all(&self.pool).await.map_err(unavailable)?;
rows.into_iter()
.map(|(ids, revisions, plans, etags, times, document)| {
let accounts = ids
.into_iter()
.zip(revisions)
.zip(plans)
.zip(etags)
.zip(times)
.map(|((((id, revision), plan), etag), time)| {
Ok(AccountCatalogObservation {
account_id: CoreProviderAccountId::new(id).map_err(|e| {
CoreStoreError::caused_by(CoreStoreErrorKind::InvalidData, e)
})?,
revision: CoreCredentialRevision::new(
u64::try_from(revision).map_err(|_| {
CoreStoreError::new(CoreStoreErrorKind::InvalidData)
})?,
)
.map_err(|e| {
CoreStoreError::caused_by(CoreStoreErrorKind::InvalidData, e)
})?,
plan,
etag,
observed_at: time.into(),
})
})
.collect::<Result<_, CoreStoreError>>()?;
Ok(StoredAccountCatalog {
accounts,
document: document.into(),
})
})
.collect()
}
pub(super) async fn save_model_catalog(
&self,
observation: &AccountCatalogObservation,
document: &[u8],
) -> Result<bool, CoreStoreError> {
let result = sqlx::query("insert into provider_account_model_catalogs (account_id, credential_revision, plan_type, observed_at, document, etag)
select a.id, a.credential_revision, a.plan_type, $4, $5, $6 from provider_accounts a
where a.id=$1 and a.credential_revision=$2 and a.plan_type is not distinct from $3 for share
on conflict (account_id) do update set credential_revision=excluded.credential_revision, plan_type=excluded.plan_type,
etag=excluded.etag, observed_at=excluded.observed_at, document=excluded.document
where provider_account_model_catalogs.credential_revision < excluded.credential_revision or
(provider_account_model_catalogs.credential_revision=excluded.credential_revision and provider_account_model_catalogs.observed_at <= excluded.observed_at)")
.bind(observation.account_id.as_str()).bind(i64::try_from(observation.revision.get()).map_err(|_|CoreStoreError::new(CoreStoreErrorKind::InvalidData))?)
.bind(observation.plan.as_deref()).bind(DateTime::<Utc>::from(observation.observed_at)).bind(document).bind(observation.etag.as_deref())
.execute(&self.pool).await.map_err(unavailable)?;
Ok(result.rows_affected() == 1)
}
}
fn unavailable(source: sqlx::Error) -> CoreStoreError {
CoreStoreError::caused_by(CoreStoreErrorKind::Unavailable, source)
}
@@ -104,6 +104,7 @@ impl PgProviderAccountRepository {
provider_kind: Option<&str>,
include_disabled: bool,
account_ids: Option<&[&str]>,
omissions: Option<&gateway_core::account::AccountCatalogOmissions>,
) -> StoreResult<Vec<ProviderAccountSummary>> {
let mut query = sqlx::QueryBuilder::<Postgres>::new(
"select coalesce(proxy_enabled, true) as outbound_proxy_enabled, coalesce(max_websocket_connections, 0) as outbound_proxy_max_websocket_connections, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone, outbound_proxy_id, outbound_proxy_url, id, provider_kind, name, notes, email, upstream_user_id,
@@ -129,6 +130,24 @@ impl PgProviderAccountRepository {
.push_bind(account_ids)
.push("::text[])");
}
if let Some(omit) = omissions {
let hashes = omit
.documents
.iter()
.map(<[u8; 32]>::as_slice)
.collect::<Vec<_>>();
let exceptions = omit
.exceptions
.iter()
.map(|id| id.as_str())
.collect::<Vec<_>>();
query.push(" and not exists (select 1 from provider_account_model_catalogs c
where c.account_id=provider_accounts.id and c.credential_revision=provider_accounts.credential_revision
and c.plan_type is not distinct from provider_accounts.plan_type and c.document_hash=any(")
.push_bind(hashes).push("::bytea[]) and c.observed_at > ").push_bind(DateTime::<Utc>::from(omit.observed_after))
.push(" and c.observed_at <= ").push_bind(DateTime::<Utc>::from(omit.as_of))
.push(" and not (c.account_id=any(").push_bind(exceptions).push("::text[])))");
}
query.push(" order by provider_kind, name, id");
let rows = query
.build()
@@ -215,7 +234,7 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
provider_kind: Option<&str>,
include_disabled: bool,
) -> StoreResult<Vec<ProviderAccountSummary>> {
self.list_provider_account_summaries(provider_kind, include_disabled, None)
self.list_provider_account_summaries(provider_kind, include_disabled, None, None)
.await
}
@@ -5,6 +5,9 @@ use std::time::Duration;
use futures::future::BoxFuture;
use gateway_core::lifecycle::CancellationToken;
use gateway_core::provider_ports::{
ProviderLeaseGuard, ProviderStoreError, ProviderStoreErrorKind,
};
use gateway_core::task::{DaemonTask, WorkerDaemonContext, WorkerTaskError};
use tokio::sync::{RwLock, RwLockReadGuard, oneshot};
use tokio::task::JoinSet;
@@ -22,11 +25,32 @@ struct State {
}
pub(crate) struct SchedulingLeaseGuard {
_release: oneshot::Sender<()>,
release: oneshot::Sender<oneshot::Sender<Result<(), ProviderStoreError>>>,
// 活跃请求可能比组合根活得久,不能因端口析构提前中止其续期与租约
_cleanup: SchedulingLeaseCleanup,
}
impl ProviderLeaseGuard for SchedulingLeaseGuard {
fn release(self: Box<Self>) -> BoxFuture<'static, Result<(), ProviderStoreError>> {
Box::pin(async move {
let (ack, released) = oneshot::channel();
self.release.send(ack).map_err(|_| {
ProviderStoreError::new(
ProviderStoreErrorKind::Unavailable,
"request scheduling lease release",
)
})?;
released.await.map_err(|source| {
ProviderStoreError::caused_by(
ProviderStoreErrorKind::Unavailable,
"await scheduling lease release",
source,
)
})?
})
}
}
impl SchedulingLeaseCleanup {
// 关闭必须等已发出的申请完成登记;不能漏掉 Redis 已创建、尚未返回的租约
pub(crate) async fn acquire(&self) -> Option<RwLockReadGuard<'_, ()>> {
@@ -39,7 +63,8 @@ impl SchedulingLeaseCleanup {
cancellation: CancellationToken,
release: impl Future<Output = crate::StoreResult<bool>> + Send + 'static,
) -> SchedulingLeaseGuard {
let (guard, dropped) = oneshot::channel::<()>();
let (guard, dropped) =
oneshot::channel::<oneshot::Sender<Result<(), ProviderStoreError>>>();
let stopping = self.0.stopping.clone();
let mut tasks = self
.0
@@ -50,17 +75,27 @@ impl SchedulingLeaseCleanup {
report_join(result);
}
tasks.spawn(async move {
tokio::select! {
let acknowledgement = tokio::select! {
biased;
() = stopping.cancelled() => cancellation.cancel(),
_ = dropped => {}
}
if let Err(error) = release.await {
() = stopping.cancelled() => { cancellation.cancel(); None },
ack = dropped => ack.ok(),
};
let result = release.await;
if let Err(error) = &result {
tracing::warn!(%error, "账号租约释放失败,依赖 TTL 回收");
}
if let Some(ack) = acknowledgement {
let _ = ack.send(result.map(|_| ()).map_err(|source| {
ProviderStoreError::caused_by(
ProviderStoreErrorKind::Unavailable,
"release scheduling lease",
source,
)
}));
}
});
SchedulingLeaseGuard {
_release: guard,
release: guard,
_cleanup: self.clone(),
}
}
@@ -249,6 +249,20 @@ pub struct CredentialLeaseGuard {
grant: Option<CredentialLeaseGrant>,
}
impl gateway_core::provider_ports::ProviderLeaseGuard for CredentialLeaseGuard {
fn release(
self: Box<Self>,
) -> futures::future::BoxFuture<'static, Result<(), ProviderStoreError>> {
Box::pin(async move {
(*self)
.release()
.await
.map(|_| ())
.map_err(|source| crate::provider_unavailable("release credential lease", source))
})
}
}
impl CredentialLeaseGuard {
/// 自动续期仅属于返回 guard;丢失租约会取消依赖该槽位的请求
pub fn maintain(
@@ -319,6 +333,16 @@ struct RenewingSchedulingLease {
_guard: CredentialLeaseGuard,
}
impl gateway_core::provider_ports::ProviderLeaseGuard for RenewingSchedulingLease {
fn release(
self: Box<Self>,
) -> futures::future::BoxFuture<'static, Result<(), ProviderStoreError>> {
let Self { _renewal, _guard } = *self;
drop(_renewal);
Box::new(_guard).release()
}
}
impl fmt::Debug for CredentialLeaseGuard {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
@@ -430,7 +454,12 @@ impl RedisCredentialLeaseRepository {
client_api_key_id: &ClientApiKeyId,
provider_kind: &ProviderKind,
) -> StoreResult<u64> {
let key = self.scheduling_cursor_key(client_api_key_id, provider_kind)?;
let key = super::scheduling_cursor_key(
&self.namespace,
client_api_key_id,
provider_kind,
"cursor",
)?;
let mut connection = self.connection.clone();
let cursor = Script::new(ADVANCE_SCHEDULING_CURSOR_SCRIPT)
.key(key)
@@ -473,21 +502,6 @@ impl RedisCredentialLeaseRepository {
))
}
fn scheduling_cursor_key(
&self,
client_api_key_id: &ClientApiKeyId,
provider_kind: &ProviderKind,
) -> StoreResult<String> {
let fingerprint = resource_fingerprint(
"provider scheduling cursor",
&format!("{}\0{}", client_api_key_id.as_str(), provider_kind.as_str()),
)?;
Ok(format!(
"{}:scheduler:client-provider:{{{fingerprint}}}:cursor",
self.namespace
))
}
async fn acquire_with_limits(
&self,
request: &CredentialLeaseRequest,
@@ -662,6 +676,7 @@ impl RedisProviderLeaseCoordinator {
in_flight: signal.in_flight,
last_started_at: signal.last_started_at.map(Into::into),
quota_reset_at: None,
quota_reset_windows: None,
quota_remaining_rank: None,
cooldown: None,
failure_rate_basis_points: None,
@@ -43,6 +43,21 @@ pub(crate) fn resource_fingerprint(entity: &'static str, value: &str) -> StoreRe
Ok(hex::encode(Sha256::digest(value.as_bytes())))
}
fn scheduling_cursor_key(
namespace: &str,
client_key: &gateway_core::policy::ClientApiKeyId,
provider: &gateway_core::routing::ProviderKind,
suffix: &str,
) -> StoreResult<String> {
let fingerprint = resource_fingerprint(
"provider scheduling cursor",
&format!("{}\0{}", client_key.as_str(), provider.as_str()),
)?;
Ok(format!(
"{namespace}:scheduler:client-provider:{{{fingerprint}}}:{suffix}"
))
}
pub(crate) fn namespace(value: &str) -> StoreResult<String> {
require_nonempty("Redis namespace", "namespace", value)?;
if value.len() > 128
@@ -3,9 +3,10 @@
use std::time::Duration;
use gateway_core::account::ProviderAccountId;
use gateway_core::policy::ClientApiKeyId;
use gateway_core::provider_ports::{
ProviderSessionAffinityKey, ProviderSessionAffinityPort, ProviderSessionBinding,
ProviderStoreError, ProviderStoreErrorKind,
ProviderAccountAdmission, ProviderSessionAffinityKey, ProviderSessionAffinityPort,
ProviderSessionBinding, ProviderStoreError, ProviderStoreErrorKind,
};
use gateway_core::routing::ProviderKind;
use redis::aio::ConnectionManager;
@@ -17,14 +18,15 @@ use super::{namespace, resource_fingerprint};
const MAX_SESSION_AFFINITY_TTL: Duration =
Duration::from_secs(gateway_core::account::MAX_SESSION_AFFINITY_TTL_HOURS as u64 * 3600);
// 比较完整记录而非账号 ID;冲突时既不改绑定,也不续期
const COMPARE_AND_BIND_SCRIPT: &str = r#"
local current = redis.call('GET', KEYS[1])
if (not current and ARGV[1] == '') or current == ARGV[1] then
redis.call('PSETEX', KEYS[1], tonumber(ARGV[3]), ARGV[2])
return 1
end
return 0
const CURSOR_TTL_MILLIS: u64 = 24 * 60 * 60 * 1000;
// 先比较全部快照再写入;失去首绑或轮询位置的请求不能续期或消耗名额
const ADMIT_ACCOUNT_SCRIPT: &str = r#"
if ARGV[2] ~= '' and (redis.call('GET', KEYS[1]) or '') ~= ARGV[1] then return 0 end
if ARGV[4] ~= '' and (redis.call('GET', KEYS[#KEYS]) or '0') ~= ARGV[4] then return 0 end
if ARGV[2] ~= '' then redis.call('PSETEX', KEYS[1], ARGV[3], ARGV[2]) end
if ARGV[4] ~= '' then redis.call('PSETEX', KEYS[#KEYS], ARGV[6], ARGV[5]) end
return 1
"#;
#[derive(serde::Serialize, serde::Deserialize)]
@@ -75,6 +77,15 @@ impl RedisProviderSessionAffinityRepository {
})
}
fn cursor_key(
&self,
provider: &ProviderKind,
client_key: &ClientApiKeyId,
) -> Result<String, ProviderStoreError> {
super::scheduling_cursor_key(&self.namespace, client_key, provider, "session-cursor")
.map_err(|_| provider_invalid("encode provider round robin cursor key"))
}
fn key(
&self,
provider_kind: &ProviderKind,
@@ -113,42 +124,71 @@ impl ProviderSessionAffinityPort for RedisProviderSessionAffinityRepository {
})
}
fn compare_and_bind<'a>(
fn round_robin_cursor<'a>(
&'a self,
provider_kind: &'a ProviderKind,
key: &'a ProviderSessionAffinityKey,
expected: Option<&'a ProviderSessionBinding>,
account_id: &'a ProviderAccountId,
ttl: Duration,
) -> futures::future::BoxFuture<'a, Result<Option<ProviderSessionBinding>, ProviderStoreError>>
{
provider: &'a ProviderKind,
client_key: &'a ClientApiKeyId,
) -> futures::future::BoxFuture<'a, Result<u64, ProviderStoreError>> {
Box::pin(async move {
let ttl_millis = session_affinity_ttl_millis(ttl)?;
let expected_raw = expected
.map(BindingRecord::encode)
.transpose()?
.unwrap_or_default();
// 同账号命中仅续期;初次认领及换号生成新版本
let binding = match expected.filter(|binding| binding.account_id() == account_id) {
Some(binding) => binding.clone(),
None => ProviderSessionBinding::new(
account_id.clone(),
uuid::Uuid::new_v4().simple().to_string(),
)?,
};
let applied = redis::Script::new(COMPARE_AND_BIND_SCRIPT)
.key(self.key(provider_kind, key)?)
.arg(expected_raw)
.arg(BindingRecord::encode(&binding)?)
.arg(ttl_millis)
.invoke_async::<bool>(&mut self.connection.clone())
let key = self.cursor_key(provider, client_key)?;
redis::cmd("GET")
.arg(key)
.query_async::<Option<u64>>(&mut self.connection.clone())
.await
.map_err(|source| {
crate::provider_unavailable("admit provider session binding", source)
})?;
Ok(applied.then_some(binding))
.map(Option::unwrap_or_default)
.map_err(|source| crate::provider_unavailable("read round robin cursor", source))
})
}
fn admit_account<'a>(
&'a self,
provider: &'a ProviderKind,
admission: ProviderAccountAdmission<'a>,
) -> futures::future::BoxFuture<'a, Result<bool, ProviderStoreError>> {
Box::pin(async move {
let script = redis::Script::new(ADMIT_ACCOUNT_SCRIPT);
let mut invocation = script.prepare_invoke();
if let Some(claim) = admission.session() {
let binding = match claim.expected.filter(|b| b.account_id() == claim.account) {
Some(existing) => existing.clone(),
None => ProviderSessionBinding::new(
claim.account.clone(),
uuid::Uuid::new_v4().simple().to_string(),
)?,
};
invocation
.key(self.key(provider, claim.key)?)
.arg(
claim
.expected
.map(BindingRecord::encode)
.transpose()?
.unwrap_or_default(),
)
.arg(BindingRecord::encode(&binding)?)
.arg(session_affinity_ttl_millis(claim.ttl)?);
} else {
invocation.arg("").arg("").arg(0);
}
if let ProviderAccountAdmission::RoundRobin {
client_key, cursor, ..
} = admission
{
invocation
.key(self.cursor_key(provider, client_key)?)
.arg(cursor)
.arg(cursor.wrapping_add(1))
.arg(CURSOR_TTL_MILLIS);
} else {
invocation.arg("");
}
invocation
.invoke_async(&mut self.connection.clone())
.await
.map_err(|source| crate::provider_unavailable("admit provider account", source))
})
}
fn load_alias<'a>(
&'a self,
provider: &'a ProviderKind,
@@ -83,6 +83,14 @@ async fn shutdown_releases_owned_account_pools_and_preserves_other_instances() {
.unwrap(),
ProviderLeaseAcquisition::Busy { .. }
));
second_guard.release().await.unwrap();
let ProviderLeaseAcquisition::Acquired(second_guard) = second_port
.try_acquire(ProviderLeaseRequest::Scheduling(other.clone()))
.await
.unwrap()
else {
panic!("released slot must be available before reselection")
};
first_requests.push(request);
second_requests.push(other);
first_guards.push(first_guard);
@@ -103,3 +103,101 @@ async fn scoped_account_list_filters_ids_provider_and_live_enabled_state() {
);
database.close().await;
}
#[tokio::test]
async fn bulk_catalog_restore_keeps_ten_thousand_owners_and_fences_late_writes() {
use gateway_core::account::{AccountCatalogObservation, CredentialRevision};
let Some(database) = TestDatabase::create("account_catalog_bulk").await else {
return;
};
let repo = PgProviderAccountRepository::new(database.pool.clone());
sqlx::query("insert into provider_accounts (id,provider_kind,name,authentication_kind,provider_credentials_json,has_refresh_token,credential_state,credential_observed_at,created_at,updated_at,plan_type)
select 'acct_'||i,'openai','source '||i,'oauth','{}',false,'ready',now(),now(),now(),'pro' from generate_series(1,10000) i").execute(&database.pool).await.unwrap();
let observation = AccountCatalogObservation {
account_id: ProviderAccountId::new("acct_1").unwrap(),
revision: CredentialRevision::new(1).unwrap(),
observed_at: std::time::SystemTime::now(),
etag: Some("v1".into()),
plan: Some("pro".into()),
};
let document = vec![b'x'; 64 * 1024];
assert!(
repo.save_account_catalog(&observation, &document)
.await
.unwrap()
);
sqlx::query("insert into provider_account_model_catalogs (account_id,credential_revision,plan_type,observed_at,document)
select a.id,1,'pro',now(),c.document from provider_accounts a cross join provider_account_model_catalogs c where a.id<>'acct_1'").execute(&database.pool).await.unwrap();
let provider = ProviderKind::new("openai").unwrap();
let start = std::time::Instant::now();
let group = repo.load_account_catalogs(&provider).await.unwrap();
eprintln!(
"catalog bulk restore: 10000 accounts, {} documents, {} ms",
group.len(),
start.elapsed().as_millis()
);
assert_eq!(group.len(), 1);
assert_eq!(group[0].accounts.len(), 10000);
assert_eq!(&group[0].document[..], &document);
let older = AccountCatalogObservation {
observed_at: observation.observed_at - std::time::Duration::from_secs(1),
..observation.clone()
};
assert!(!repo.save_account_catalog(&older, b"late").await.unwrap());
sqlx::query("update provider_accounts set credential_revision=2 where id='acct_1'")
.execute(&database.pool)
.await
.unwrap();
assert!(
!repo
.save_account_catalog(&observation, b"old revision")
.await
.unwrap()
);
sqlx::query("update provider_accounts set enabled=false where id='acct_2'")
.execute(&database.pool)
.await
.unwrap();
sqlx::query("delete from provider_accounts where id='acct_3'")
.execute(&database.pool)
.await
.unwrap();
let group = repo.load_account_catalogs(&provider).await.unwrap();
assert_eq!(group[0].accounts.len(), 9997);
assert!(
group[0]
.accounts
.iter()
.all(|a| !matches!(a.account_id.as_str(), "acct_1" | "acct_2" | "acct_3"))
);
// 数据库只能排除有当前版本、未过期的账号证据;显式例外仍进入选号
let hash: Vec<u8> = sqlx::query_scalar("select sha256($1::bytea)")
.bind(&document)
.fetch_one(&database.pool)
.await
.unwrap();
let now = std::time::SystemTime::now();
sqlx::query("update provider_account_model_catalogs set observed_at=now()-interval '1 hour' where account_id='acct_4'").execute(&database.pool).await.unwrap();
let filter = gateway_core::account::AccountCatalogOmissions {
documents: vec![hash.try_into().unwrap()],
exceptions: vec![ProviderAccountId::new("acct_5").unwrap()],
observed_after: now - std::time::Duration::from_secs(1800),
as_of: now,
};
let ids = (1..=10000)
.map(|n| ProviderAccountId::new(format!("acct_{n}")).unwrap())
.collect::<Vec<_>>();
let candidates = repo
.list_catalog_candidates(&provider, &ids, &filter)
.await
.unwrap();
let ids = candidates
.iter()
.map(|a| a.id().as_str())
.collect::<std::collections::BTreeSet<_>>();
assert_eq!(
ids,
std::collections::BTreeSet::from(["acct_1", "acct_4", "acct_5"])
);
database.close().await;
}
@@ -126,21 +126,23 @@ async fn smart_settings_round_trip_publish_and_restore_database_defaults() {
SmartSchedulingConfig::default()
);
let mut update = settings_with_margin(3600);
update.values.smart_scheduling =
SmartSchedulingConfig::new([0.0, 2.0, 1.0, 0.5, 1.2, 2.3]).unwrap();
let expected = update.values.smart_scheduling;
let expected = serde_json::json!({
"loadWeight": 0.0, "quotaWeight": 2.0, "healthWeight": 1.0, "latencyWeight": 0.5, "resetWeight": 1.2, "queueWeight": 2.3,
"resetWindowWeights": {"fiveHour":2.0,"weekly":8.0,"monthly":0.0}
});
update.values.smart_scheduling = serde_json::from_value(expected.clone()).unwrap();
repository.update_runtime_settings(update).await.unwrap();
let reloaded = repository.load_runtime_settings().await.unwrap();
assert_eq!(reloaded.values.smart_scheduling, expected);
assert_eq!(
serde_json::to_value(reloaded.values.smart_scheduling).unwrap(),
expected
);
let snapshot = PgRuntimeSnapshotRepository::new(database.pool.clone())
.load_runtime_snapshot()
.await
.unwrap();
let snapshot_values = serde_json::to_value(&snapshot.settings).unwrap();
assert_eq!(
snapshot_values["smart_scheduling"],
serde_json::to_value(expected).unwrap()
);
assert_eq!(snapshot_values["smart_scheduling"], expected);
assert!(snapshot.config_revision > before.config_revision);
sqlx::query("update runtime_settings set smart_scheduling_json = default where id = 1")
.execute(&database.pool)
@@ -3,12 +3,39 @@
use std::time::Duration;
use gateway_core::account::ProviderAccountId;
use gateway_core::provider_ports::{ProviderSessionAffinityKey, ProviderSessionAffinityPort};
use gateway_core::provider_ports::{
ProviderAccountAdmission, ProviderSessionAffinityKey, ProviderSessionAffinityPort,
ProviderSessionBinding, ProviderSessionClaim, ProviderStoreError,
};
use gateway_core::routing::ProviderKind;
use gateway_store::redis::RedisProviderSessionAffinityRepository;
use redis::aio::ConnectionManager;
use uuid::Uuid;
async fn bind(
repo: &RedisProviderSessionAffinityRepository,
provider: &ProviderKind,
key: &ProviderSessionAffinityKey,
expected: Option<&ProviderSessionBinding>,
account: &ProviderAccountId,
ttl: Duration,
) -> Result<Option<ProviderSessionBinding>, ProviderStoreError> {
let claim = ProviderSessionClaim {
key,
expected,
account,
ttl,
};
if repo
.admit_account(provider, ProviderAccountAdmission::Session(claim))
.await?
{
repo.load(provider, key).await
} else {
Ok(None)
}
}
#[tokio::test]
async fn concurrent_claims_admit_exactly_one_account() {
let Some((repo, mut connection, namespace)) = affinity_repository().await else {
@@ -19,8 +46,22 @@ async fn concurrent_claims_admit_exactly_one_account() {
let first = ProviderAccountId::new("acct_first").unwrap();
let second = ProviderAccountId::new("acct_second").unwrap();
let (a, b) = tokio::join!(
repo.compare_and_bind(&provider, &key, None, &first, Duration::from_secs(60)),
repo.compare_and_bind(&provider, &key, None, &second, Duration::from_secs(60))
bind(
&repo,
&provider,
&key,
None,
&first,
Duration::from_secs(60)
),
bind(
&repo,
&provider,
&key,
None,
&second,
Duration::from_secs(60)
)
);
let winners = [a.unwrap(), b.unwrap()]
.into_iter()
@@ -56,24 +97,21 @@ async fn stale_revision_cannot_overwrite_an_account_that_returned_to_the_session
let a = ProviderAccountId::new("acct_a").unwrap();
let b = ProviderAccountId::new("acct_b").unwrap();
let ttl = Duration::from_secs(60);
let first = repo
.compare_and_bind(&provider, &key, None, &a, ttl)
let first = bind(&repo, &provider, &key, None, &a, ttl)
.await
.unwrap()
.unwrap();
let second = repo
.compare_and_bind(&provider, &key, Some(&first), &b, ttl)
let second = bind(&repo, &provider, &key, Some(&first), &b, ttl)
.await
.unwrap()
.unwrap();
let third = repo
.compare_and_bind(&provider, &key, Some(&second), &a, ttl)
let third = bind(&repo, &provider, &key, Some(&second), &a, ttl)
.await
.unwrap()
.unwrap();
assert_ne!(first.revision(), third.revision());
assert!(
repo.compare_and_bind(&provider, &key, Some(&first), &b, ttl)
bind(&repo, &provider, &key, Some(&first), &b, ttl)
.await
.unwrap()
.is_none()
@@ -90,13 +128,11 @@ async fn renewal_preserves_revision_but_expiration_requires_fresh_claim() {
let key = ProviderSessionAffinityKey::try_new("expired-session").unwrap();
let a = ProviderAccountId::new("acct_a").unwrap();
let ttl = Duration::from_secs(60);
let first = repo
.compare_and_bind(&provider, &key, None, &a, ttl)
let first = bind(&repo, &provider, &key, None, &a, ttl)
.await
.unwrap()
.unwrap();
let renewed = repo
.compare_and_bind(&provider, &key, Some(&first), &a, ttl)
let renewed = bind(&repo, &provider, &key, Some(&first), &a, ttl)
.await
.unwrap()
.unwrap();
@@ -113,14 +149,13 @@ async fn renewal_preserves_revision_but_expiration_requires_fresh_claim() {
.await
.unwrap();
assert!(
repo.compare_and_bind(&provider, &key, Some(&first), &a, ttl)
bind(&repo, &provider, &key, Some(&first), &a, ttl)
.await
.unwrap()
.is_none()
);
assert!(repo.load(&provider, &key).await.unwrap().is_none());
let fresh = repo
.compare_and_bind(&provider, &key, None, &a, ttl)
let fresh = bind(&repo, &provider, &key, None, &a, ttl)
.await
.unwrap()
.unwrap();
@@ -136,7 +171,7 @@ async fn damaged_binding_is_an_error_and_cannot_be_claimed_as_absent() {
let key = ProviderSessionAffinityKey::try_new("damaged-session").unwrap();
let a = ProviderAccountId::new("acct_a").unwrap();
let ttl = Duration::from_secs(60);
repo.compare_and_bind(&provider, &key, None, &a, ttl)
bind(&repo, &provider, &key, None, &a, ttl)
.await
.unwrap()
.unwrap();
@@ -153,13 +188,140 @@ async fn damaged_binding_is_an_error_and_cannot_be_claimed_as_absent() {
.unwrap();
assert!(repo.load(&provider, &key).await.is_err());
assert!(
repo.compare_and_bind(&provider, &key, None, &a, ttl)
bind(&repo, &provider, &key, None, &a, ttl)
.await
.unwrap()
.is_none()
);
}
#[tokio::test]
async fn round_robin_commits_only_with_the_current_cursor_and_session_version() {
use gateway_core::policy::ClientApiKeyId;
let Some((repo, mut connection, namespace)) = affinity_repository().await else {
return;
};
let peer = RedisProviderSessionAffinityRepository::new(connection.clone(), &namespace).unwrap();
let provider = ProviderKind::new("openai").unwrap();
let client = ClientApiKeyId::new("client").unwrap();
let account = ProviderAccountId::new("acct_a").unwrap();
let key = ProviderSessionAffinityKey::try_new("rr-session").unwrap();
let other = ProviderSessionAffinityKey::try_new("rr-other").unwrap();
let claim = ProviderSessionClaim {
key: &key,
expected: None,
account: &account,
ttl: Duration::from_secs(60),
};
let step = |cursor, session| ProviderAccountAdmission::RoundRobin {
client_key: &client,
cursor,
session,
};
let first = step(0, Some(claim));
let (a, b) = tokio::join!(
repo.admit_account(&provider, first),
peer.admit_account(&provider, first)
);
assert_ne!(a.unwrap(), b.unwrap());
let winner = repo.load(&provider, &key).await.unwrap().unwrap();
assert_eq!(winner.account_id(), &account);
assert!(
!repo
.admit_account(&provider, step(1, Some(claim)))
.await
.unwrap()
);
assert_eq!(
peer.round_robin_cursor(&provider, &client).await.unwrap(),
1
);
let other_claim = ProviderSessionClaim {
key: &other,
..claim
};
assert!(
!peer
.admit_account(&provider, step(0, Some(other_claim)))
.await
.unwrap()
);
assert!(repo.load(&provider, &other).await.unwrap().is_none());
assert!(peer.admit_account(&provider, step(1, None)).await.unwrap());
assert_eq!(
repo.round_robin_cursor(&provider, &client).await.unwrap(),
2
);
for (p, c) in [
(&provider, ClientApiKeyId::new("other").unwrap()),
(&ProviderKind::new("xai").unwrap(), client.clone()),
] {
assert_eq!(repo.round_robin_cursor(p, &c).await.unwrap(), 0);
}
let leases =
gateway_store::redis::RedisCredentialLeaseRepository::new(connection.clone(), &namespace)
.unwrap();
assert_eq!(
leases
.advance_scheduling_cursor(&client, &provider)
.await
.unwrap(),
0
);
assert_eq!(
repo.round_robin_cursor(&provider, &client).await.unwrap(),
2
);
let cursor_keys: Vec<String> = redis::cmd("KEYS")
.arg(format!("{namespace}:*:session-cursor"))
.query_async(&mut connection)
.await
.unwrap();
let ttl: i64 = redis::cmd("PTTL")
.arg(&cursor_keys[0])
.query_async(&mut connection)
.await
.unwrap();
assert!((1..=86_400_000).contains(&ttl));
redis::cmd("DEL")
.arg(&cursor_keys)
.query_async::<i64>(&mut connection)
.await
.unwrap();
assert_eq!(
peer.round_robin_cursor(&provider, &client).await.unwrap(),
0
);
assert_eq!(repo.load(&provider, &key).await.unwrap(), Some(winner));
let fresh = ProviderSessionAffinityKey::try_new("rr-fresh").unwrap();
let claims = [
other_claim,
ProviderSessionClaim {
key: &fresh,
..claim
},
];
let (a, b) = tokio::join!(
repo.admit_account(&provider, step(0, Some(claims[0]))),
peer.admit_account(&provider, step(0, Some(claims[1])))
);
let won_first = a.unwrap();
assert_ne!(won_first, b.unwrap());
let loser = claims[usize::from(won_first)];
assert!(repo.load(&provider, loser.key).await.unwrap().is_none());
assert!(
repo.admit_account(&provider, step(1, Some(loser)))
.await
.unwrap()
);
assert_eq!(
peer.round_robin_cursor(&provider, &client).await.unwrap(),
2
);
}
async fn affinity_repository() -> Option<(
RedisProviderSessionAffinityRepository,
ConnectionManager,
@@ -221,12 +383,11 @@ async fn turn_alias_cannot_be_reassigned_and_follows_session_migration() {
.await
.unwrap()
);
let first = repo
.compare_and_bind(&provider, &session, None, &a, ttl)
let first = bind(&repo, &provider, &session, None, &a, ttl)
.await
.unwrap()
.unwrap();
repo.compare_and_bind(&provider, &session, Some(&first), &b, ttl)
bind(&repo, &provider, &session, Some(&first), &b, ttl)
.await
.unwrap()
.unwrap();
@@ -293,8 +454,7 @@ async fn binding_and_turn_alias_can_renew_beyond_one_day_and_shorten_without_sca
let mut binding = None;
for hours in [24, 168, 720, 1] {
let ttl = Duration::from_secs(hours * 3600);
binding = repo
.compare_and_bind(&provider, &session, binding.as_ref(), &account, ttl)
binding = bind(&repo, &provider, &session, binding.as_ref(), &account, ttl)
.await
.unwrap();
assert!(binding.is_some());
@@ -324,7 +484,7 @@ async fn binding_and_turn_alias_can_renew_beyond_one_day_and_shorten_without_sca
}
for ttl in [Duration::ZERO, Duration::from_secs(721 * 3600)] {
assert!(
repo.compare_and_bind(&provider, &session, binding.as_ref(), &account, ttl)
bind(&repo, &provider, &session, binding.as_ref(), &account, ttl)
.await
.is_err()
);
+2 -1
View File
@@ -19,6 +19,7 @@ crc32fast = "1.5.0"
futures = "0.3.32"
flate2 = "1.1.9"
gateway-admin = { path = "../../gateway-admin" }
gateway-egress = { path = "../../gateway-egress" }
gateway-core = { path = "../../gateway-core" }
gateway-protocol = { path = "../../gateway-protocol" }
getrandom = "0.4.3"
@@ -79,7 +80,7 @@ openssl-sys = { version = "=0.9.116", features = ["vendored"] }
gateway-host = { path = "../../gateway-host" }
allocation-counter = "0.8.1"
tempfile = "3.27.0"
tokio = { version = "1.53.3", default-features = false, features = ["rt-multi-thread", "test-util"] }
tokio = { version = "1.53.3", default-features = false, features = ["process", "rt-multi-thread", "test-util"] }
tracing-subscriber = { version = "0.3.23", features = ["json"] }
wiremock = "0.6.5"

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