mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-12 16:06:16 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
28ad8fb1c4 | ||
|
|
4888681a04 | ||
|
|
9739bb08be | ||
|
|
fac37ec563 | ||
|
|
9f0d712671 | ||
|
|
75be1287ef | ||
|
|
dad91fe9f7 | ||
|
|
b22b0cfa41 | ||
|
|
bd69329a37 | ||
|
|
151a3c6bfd | ||
|
|
e2f194875f | ||
|
|
07b61bdc83 | ||
|
|
c0fdc1bbff | ||
|
|
343c579017 | ||
|
|
914c00785c | ||
|
|
b3de0d637d | ||
|
|
1271fa0186 | ||
|
|
055da40476 | ||
|
|
36ba79ca5e | ||
|
|
1b24029a97 | ||
|
|
431992a7a4 | ||
|
|
d3143406dc | ||
|
|
478f2d1f21 |
Generated
+14
@@ -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",
|
||||
|
||||
@@ -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 != ¤t.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(), ®istry, 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",
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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
|
||||
@@ -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))
|
||||
}
|
||||
@@ -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) {
|
||||
|
||||
@@ -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!(
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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";
|
||||
|
||||
@@ -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()
|
||||
);
|
||||
|
||||
@@ -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
Reference in New Issue
Block a user