mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-07 16:05:27 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
14ea5234ce | ||
|
|
f8480f3bdc | ||
|
|
52b8103bc9 | ||
|
|
dacdd3f0e2 | ||
|
|
ac36f5bc13 | ||
|
|
e0a022518c | ||
|
|
37f7a55427 | ||
|
|
f6085206eb | ||
|
|
9813c86055 | ||
|
|
260f775039 | ||
|
|
1424cea3a9 | ||
|
|
787ad28424 | ||
|
|
db23560b0e | ||
|
|
d700df8a4f | ||
|
|
c9adfd9a2b | ||
|
|
654b6baf89 |
@@ -42,8 +42,11 @@ else
|
||||
mkdir -p "$(dirname "$directory")"
|
||||
temporary="$(mktemp -d "${directory}.XXXXXX")"
|
||||
trap 'rm -rf "$temporary"' EXIT
|
||||
mkdir "$temporary/context"
|
||||
"${command[@]}" --output "type=local,dest=$temporary/snapshot" "$temporary/context"
|
||||
mkdir "$temporary"/{context,snapshot}
|
||||
# tar 流避开 local exporter 的逐文件同步阻塞,直接解包以免额外落盘整份归档。
|
||||
# https://github.com/moby/buildkit/issues/2950
|
||||
"${command[@]}" --output type=tar,dest=- "$temporary/context" \
|
||||
| tar -xf - -C "$temporary/snapshot"
|
||||
if usable "$temporary/snapshot"; then
|
||||
rm -rf "$directory"
|
||||
mv "$temporary/snapshot" "$directory"
|
||||
|
||||
Generated
+32
-2
@@ -145,6 +145,18 @@ dependencies = [
|
||||
"serde_json",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "async-compression"
|
||||
version = "0.4.50"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ee19bd99b43e3691acbad4e840420a4881cea6c0b66a208125a824f8fd53f5a1"
|
||||
dependencies = [
|
||||
"compression-codecs",
|
||||
"compression-core",
|
||||
"pin-project-lite",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "async-lock"
|
||||
version = "3.4.2"
|
||||
@@ -931,6 +943,23 @@ dependencies = [
|
||||
"tokio-util",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "compression-codecs"
|
||||
version = "0.4.45"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "98fc98460ba0ad5317075d3632b8dfc45d0be8c4a49347c2a38272019717614a"
|
||||
dependencies = [
|
||||
"compression-core",
|
||||
"flate2",
|
||||
"memchr",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "compression-core"
|
||||
version = "0.4.33"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6e8ccc4ea9f6acc32d102c0f6d471d11d913ad15f20c04de743374861fa1d414"
|
||||
|
||||
[[package]]
|
||||
name = "config"
|
||||
version = "0.15.27"
|
||||
@@ -3750,9 +3779,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "rustls"
|
||||
version = "0.23.45"
|
||||
version = "0.23.46"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0d41d731c7d2f962d1ccc364cec258de3c0e93b38c2fb3ba97ac74513048d634"
|
||||
checksum = "48e13bd8c0e9365c43cfa5c9e8f9ad49d3c8444926c9aac819e0e4dc503c8fdf"
|
||||
dependencies = [
|
||||
"aws-lc-rs",
|
||||
"log",
|
||||
@@ -4743,6 +4772,7 @@ version = "0.7.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "08a05a66a4fdd61cbbe0a1d755ffe0ca6aba159dd4820936a0ff8a8278245b9c"
|
||||
dependencies = [
|
||||
"async-compression",
|
||||
"bitflags",
|
||||
"bytes",
|
||||
"futures-core",
|
||||
|
||||
@@ -4,6 +4,7 @@ use std::{
|
||||
collections::{BTreeMap, BTreeSet},
|
||||
fs,
|
||||
path::{Path, PathBuf},
|
||||
sync::OnceLock,
|
||||
};
|
||||
|
||||
use syn::{Item, visit::Visit};
|
||||
@@ -26,22 +27,24 @@ pub(super) const WORKSPACE_MEMBERS: &[&str] = &[
|
||||
|
||||
#[test]
|
||||
fn workspace_member_list_matches_the_frozen_dag_scope() {
|
||||
let manifest =
|
||||
fs::read_to_string(backend_root().join("Cargo.toml")).expect("read workspace manifest");
|
||||
for member in WORKSPACE_MEMBERS {
|
||||
assert!(
|
||||
manifest.contains(&format!("\"{member}\"")),
|
||||
"workspace must include {member}"
|
||||
);
|
||||
}
|
||||
let member_count = manifest
|
||||
.lines()
|
||||
.map(str::trim)
|
||||
.filter(|line| line.starts_with("\"apps/") || line.starts_with("\"crates/"))
|
||||
.count();
|
||||
let metadata = cargo_metadata_json();
|
||||
let root = backend_root();
|
||||
let actual: BTreeSet<_> = metadata["packages"]
|
||||
.as_array()
|
||||
.expect("Cargo packages")
|
||||
.iter()
|
||||
.map(|package| {
|
||||
Path::new(package["manifest_path"].as_str().expect("package manifest"))
|
||||
.parent()
|
||||
.expect("package directory")
|
||||
.strip_prefix(&root)
|
||||
.expect("workspace member path")
|
||||
.to_path_buf()
|
||||
})
|
||||
.collect();
|
||||
assert_eq!(
|
||||
member_count,
|
||||
WORKSPACE_MEMBERS.len(),
|
||||
actual,
|
||||
WORKSPACE_MEMBERS.iter().map(PathBuf::from).collect(),
|
||||
"new workspace members must extend the dependency DAG rules"
|
||||
);
|
||||
}
|
||||
@@ -378,7 +381,7 @@ const ALLOWED_INTERNAL_EDGES: &[(&str, &str)] = &[
|
||||
#[test]
|
||||
fn workspace_internal_dependency_edges_match_frozen_dag() {
|
||||
let metadata = cargo_metadata_json();
|
||||
let actual = workspace_runtime_dependency_edges(&metadata);
|
||||
let actual = workspace_runtime_dependency_edges(metadata);
|
||||
let mut expected: Vec<(String, String)> = ALLOWED_INTERNAL_EDGES
|
||||
.iter()
|
||||
.map(|(from, to)| ((*from).to_owned(), (*to).to_owned()))
|
||||
@@ -474,24 +477,6 @@ fn adapter_public_module_surfaces_match_allowlist() {
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn production_sources_do_not_host_tests() {
|
||||
for member in WORKSPACE_MEMBERS {
|
||||
let src = backend_root().join(member).join("src");
|
||||
let sources = super::rust_files(&src);
|
||||
assert!(!sources.is_empty(), "{member} has no production sources");
|
||||
for relative in sources {
|
||||
let path = src.join(relative);
|
||||
let source = fs::read_to_string(&path).expect("read production source");
|
||||
assert!(
|
||||
!source.contains("#[cfg(test)]"),
|
||||
"{} hosts tests in production src",
|
||||
path.display()
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn workspace_modules_follow_conventional_file_layout() {
|
||||
let targets = test_source_roots();
|
||||
@@ -748,40 +733,46 @@ pub(super) fn backend_root() -> PathBuf {
|
||||
.expect("backend workspace root")
|
||||
}
|
||||
|
||||
fn cargo_metadata_json() -> serde_json::Value {
|
||||
let output = std::process::Command::new("cargo")
|
||||
.args(["metadata", "--format-version", "1", "--no-deps"])
|
||||
.current_dir(backend_root())
|
||||
.output()
|
||||
.expect("run cargo metadata");
|
||||
assert!(
|
||||
output.status.success(),
|
||||
"cargo metadata failed: {}",
|
||||
String::from_utf8_lossy(&output.stderr)
|
||||
);
|
||||
serde_json::from_slice(&output.stdout).expect("parse cargo metadata")
|
||||
fn cargo_metadata_json() -> &'static serde_json::Value {
|
||||
static METADATA: OnceLock<serde_json::Value> = OnceLock::new();
|
||||
METADATA.get_or_init(|| {
|
||||
let output = std::process::Command::new("cargo")
|
||||
.args(["metadata", "--format-version", "1", "--no-deps", "--locked"])
|
||||
.current_dir(backend_root())
|
||||
.output()
|
||||
.expect("run cargo metadata");
|
||||
assert!(
|
||||
output.status.success(),
|
||||
"cargo metadata failed: {}",
|
||||
String::from_utf8_lossy(&output.stderr)
|
||||
);
|
||||
serde_json::from_slice(&output.stdout).expect("parse cargo metadata")
|
||||
})
|
||||
}
|
||||
|
||||
/// 提取成员 `[dependencies]` 段内声明的依赖名;段落以下一个 `[` 表头结束
|
||||
/// 复用 Cargo 解析后的实际包名,覆盖重命名依赖与平台条件,排除测试和构建依赖
|
||||
fn dependency_names(member: &str) -> Vec<String> {
|
||||
let manifest = fs::read_to_string(backend_root().join(member).join("Cargo.toml"))
|
||||
.expect("read member manifest");
|
||||
let mut in_dependencies = false;
|
||||
let mut names = Vec::new();
|
||||
for line in manifest.lines() {
|
||||
let line = line.trim();
|
||||
if line.starts_with('[') {
|
||||
in_dependencies = line == "[dependencies]";
|
||||
continue;
|
||||
}
|
||||
if !in_dependencies || line.is_empty() || line.starts_with('#') {
|
||||
continue;
|
||||
}
|
||||
if let Some((name, _)) = line.split_once('=') {
|
||||
names.push(name.trim().to_owned());
|
||||
}
|
||||
}
|
||||
names
|
||||
let manifest = backend_root().join(member).join("Cargo.toml");
|
||||
let package = cargo_metadata_json()["packages"]
|
||||
.as_array()
|
||||
.expect("Cargo packages")
|
||||
.iter()
|
||||
.find(|package| {
|
||||
Path::new(package["manifest_path"].as_str().expect("package manifest")) == manifest
|
||||
})
|
||||
.expect("workspace member");
|
||||
package["dependencies"]
|
||||
.as_array()
|
||||
.expect("package dependencies")
|
||||
.iter()
|
||||
.filter(|dependency| dependency["kind"].is_null())
|
||||
.map(|dependency| {
|
||||
dependency["name"]
|
||||
.as_str()
|
||||
.expect("dependency name")
|
||||
.to_owned()
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn assert_no_dependency(member: &str, forbidden: &[&str]) {
|
||||
|
||||
@@ -59,9 +59,16 @@ fn cargo_library_root_is_conventional_lib() {
|
||||
fn workspace_production_files_have_no_hidden_modules_or_test_hooks() {
|
||||
for member in architecture::WORKSPACE_MEMBERS {
|
||||
let root = architecture::backend_root().join(member).join("src");
|
||||
for relative in rust_files(&root) {
|
||||
let sources = rust_files(&root);
|
||||
assert!(!sources.is_empty(), "{member} has no production sources");
|
||||
for relative in sources {
|
||||
let path = root.join(relative);
|
||||
let source = fs::read_to_string(&path).expect("read production source");
|
||||
assert!(
|
||||
!source.contains("#[cfg(test)]"),
|
||||
"{} hosts tests in production src",
|
||||
path.display()
|
||||
);
|
||||
assert!(
|
||||
!source.contains("include!("),
|
||||
"{} uses include!",
|
||||
|
||||
@@ -124,6 +124,7 @@ pub struct AccountRecord {
|
||||
pub weight: AccountWeight,
|
||||
pub model_access: gateway_core::account::AccountModelAccess,
|
||||
pub outbound_proxy: Option<gateway_core::account::OutboundProxy>,
|
||||
pub outbound_proxy_enabled: bool,
|
||||
pub credential_state: CredentialState,
|
||||
pub credential_observed_at: DateTime<Utc>,
|
||||
pub quota: QuotaState,
|
||||
@@ -133,6 +134,16 @@ pub struct AccountRecord {
|
||||
pub updated_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
impl AccountRecord {
|
||||
/// 配置保留绑定,执行请求时禁用代理按直连处理
|
||||
#[must_use]
|
||||
pub fn effective_outbound_proxy(&self) -> Option<&gateway_core::account::OutboundProxy> {
|
||||
self.outbound_proxy
|
||||
.as_ref()
|
||||
.filter(|_| self.outbound_proxy_enabled)
|
||||
}
|
||||
}
|
||||
|
||||
pub use super::observability::{
|
||||
AccountModelUsage, AccountPoolMetrics as AccountSummary, AccountRequestBucket,
|
||||
CurrencyCost as AccountCost,
|
||||
|
||||
@@ -69,6 +69,7 @@ pub struct ProxyTestResult {
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ProxyRecord {
|
||||
pub enabled: bool,
|
||||
pub auto_location: bool,
|
||||
pub detected_location: Option<DetectedProxyLocation>,
|
||||
pub location: Option<gateway_core::account::RequestLocation>,
|
||||
@@ -103,6 +104,7 @@ pub struct NewProxy {
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct UpdateProxy {
|
||||
pub enabled: Option<bool>,
|
||||
pub auto_location: Option<bool>,
|
||||
pub test: Option<ProxyTestResult>,
|
||||
/// 外层为空保留配置,内层为空恢复全局继承
|
||||
|
||||
@@ -106,7 +106,8 @@ pub(super) async fn pending_authorization(
|
||||
proxy = required_credential(accounts, provider_kind, account_id, resource)
|
||||
.await?
|
||||
.credential
|
||||
.outbound_proxy;
|
||||
.effective_outbound_proxy()
|
||||
.cloned();
|
||||
AuthorizationMutationTarget::Reauthorize {
|
||||
account_id: account_id.clone(),
|
||||
}
|
||||
@@ -121,6 +122,9 @@ pub(super) async fn pending_authorization(
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "proxy"))?;
|
||||
proxy_id = Some(record.id);
|
||||
if !record.enabled {
|
||||
return Err(AdminError::invalid("代理已禁用,请启用后再选择"));
|
||||
}
|
||||
Some(record.proxy)
|
||||
}
|
||||
Some(AccountProxySelection::Direct) | None => None,
|
||||
|
||||
@@ -305,6 +305,9 @@ impl PluginsService {
|
||||
.get(id)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "proxy"))?;
|
||||
if !record.enabled {
|
||||
return Err(AdminError::invalid("下载代理已禁用,请启用或改为直连"));
|
||||
}
|
||||
Ok(Some(PluginDistributionEgress::new(
|
||||
record.id,
|
||||
record.revision,
|
||||
|
||||
@@ -3137,6 +3137,7 @@ fn quota_local_usage(account_id: &str, total_tokens: u64) -> AccountUsage {
|
||||
pub(super) fn account_record(kind: &str) -> AccountRecord {
|
||||
let now = Utc::now();
|
||||
AccountRecord {
|
||||
outbound_proxy_enabled: true,
|
||||
notes: None,
|
||||
model_access: Default::default(),
|
||||
outbound_proxy: None,
|
||||
|
||||
@@ -534,6 +534,7 @@ fn preview_fixture() -> Fixture {
|
||||
fn source_proxy() -> ProxyRecord {
|
||||
let now = chrono::Utc::now();
|
||||
ProxyRecord {
|
||||
enabled: true,
|
||||
auto_location: false,
|
||||
detected_location: None,
|
||||
location: None,
|
||||
@@ -608,6 +609,38 @@ async fn saved_source_proxy_is_resolved_for_queries_and_frozen_into_verified_sou
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn disabled_source_proxy_rejects_queries_and_downloads_before_network_access() {
|
||||
let mut fixture = preview_fixture();
|
||||
fixture.binding.get_mut().unwrap().outbound_proxy_id = Some("source-proxy".into());
|
||||
let mut proxy = source_proxy();
|
||||
proxy.enabled = false;
|
||||
*fixture.proxy.get_mut().unwrap() = Some(proxy);
|
||||
let fixture = Arc::new(fixture);
|
||||
assert_eq!(
|
||||
fixture
|
||||
.service()
|
||||
.check_update("test.example", &[])
|
||||
.await
|
||||
.unwrap_err()
|
||||
.kind(),
|
||||
AdminErrorKind::Invalid
|
||||
);
|
||||
let mut request = preview_request();
|
||||
request.outbound_proxy_id = Some("source-proxy".into());
|
||||
assert_eq!(
|
||||
fixture
|
||||
.service()
|
||||
.verify_remote(request)
|
||||
.await
|
||||
.unwrap_err()
|
||||
.kind(),
|
||||
AdminErrorKind::Invalid
|
||||
);
|
||||
assert!(fixture.egresses.lock().unwrap().is_empty());
|
||||
assert_eq!(*fixture.downloads.lock().unwrap(), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn proxy_revision_change_during_download_rejects_the_verified_candidate() {
|
||||
let mut fixture = preview_fixture();
|
||||
|
||||
@@ -126,6 +126,7 @@ async fn authorization_uses_selected_proxy_regardless_of_probe_status() {
|
||||
.accounts(FakeAccountStore::new(kind, events.clone()))
|
||||
.proxies(Arc::new(TestProxies {
|
||||
record: Some(ProxyRecord {
|
||||
enabled: true,
|
||||
auto_location: false,
|
||||
detected_location: None,
|
||||
location: None,
|
||||
|
||||
@@ -82,6 +82,7 @@ fn contract_type(mut item: Item, name: &str) -> Item {
|
||||
| "RegeneratedAdminApiKey"
|
||||
| "CodexPrivacyPolicy"
|
||||
| "PrivacyRule"
|
||||
| "PrivacyStep"
|
||||
| "PrivacyPreviewRequest"
|
||||
| "PrivacyPreviewResult"
|
||||
);
|
||||
@@ -137,6 +138,7 @@ fn sdk_settings_contract_matches_host_declarations() {
|
||||
"PrivacyScope",
|
||||
"PrivacyAction",
|
||||
"PrivacyRule",
|
||||
"PrivacyStep",
|
||||
"PrivacyRuleOutcome",
|
||||
"PrivacyPreviewRequest",
|
||||
"PrivacyPreviewResult",
|
||||
|
||||
@@ -32,7 +32,7 @@ sha2 = "0.11.0"
|
||||
thiserror = "2.0.18"
|
||||
tokio = { version = "1.52.3", default-features = false, features = ["macros", "rt", "sync", "time"] }
|
||||
tower = { version = "0.5.3", features = ["util"] }
|
||||
tower-http = { version = "0.7.0", features = ["cors", "fs", "request-id", "timeout", "trace"] }
|
||||
tower-http = { version = "0.7.0", features = ["compression-gzip", "cors", "fs", "request-id", "timeout", "trace"] }
|
||||
tracing = "0.1.44"
|
||||
url = "2.5.8"
|
||||
uuid = { version = "1.23.3", features = ["v7"] }
|
||||
|
||||
@@ -0,0 +1,25 @@
|
||||
//! 管理端普通 JSON 响应的压缩范围与大小门槛
|
||||
|
||||
use http::{Response, header};
|
||||
use http_body::Body;
|
||||
use tower_http::compression::predicate::Predicate;
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
pub(super) struct AdminJson;
|
||||
|
||||
impl Predicate for AdminJson {
|
||||
fn should_compress<B: Body>(&self, response: &Response<B>) -> bool {
|
||||
// 只处理长度已知的 JSON,不为判断门槛收集流式正文
|
||||
response
|
||||
.body()
|
||||
.size_hint()
|
||||
.exact()
|
||||
.is_some_and(|size| size >= 1024)
|
||||
&& response
|
||||
.headers()
|
||||
.get(header::CONTENT_TYPE)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.and_then(|value| value.split(';').next())
|
||||
.is_some_and(|value| value.trim().eq_ignore_ascii_case("application/json"))
|
||||
}
|
||||
}
|
||||
@@ -9,12 +9,14 @@ use axum::{
|
||||
response::Response,
|
||||
routing::any,
|
||||
};
|
||||
use tower_http::compression::{CompressionLayer, CompressionLevel};
|
||||
|
||||
pub mod account_groups;
|
||||
pub mod accounts;
|
||||
pub mod auth;
|
||||
pub mod backups;
|
||||
pub mod client_keys;
|
||||
mod compression;
|
||||
pub(crate) mod diagnostics;
|
||||
mod extract;
|
||||
pub mod observability;
|
||||
@@ -54,6 +56,12 @@ where
|
||||
.method_not_allowed_fallback(method_not_allowed)
|
||||
.route("/api/admin", any(admin_not_found))
|
||||
.route("/api/admin/{*path}", any(admin_not_found))
|
||||
.layer(
|
||||
CompressionLayer::new()
|
||||
.quality(CompressionLevel::Fastest)
|
||||
.compress_when(compression::AdminJson),
|
||||
)
|
||||
.merge(plugins::public_router().method_not_allowed_fallback(method_not_allowed))
|
||||
.layer(middleware::map_response(no_store))
|
||||
}
|
||||
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
//! 观测领域事实到管理页面的展示投影
|
||||
|
||||
use super::*;
|
||||
use gateway_core::metering::Usage;
|
||||
|
||||
pub(crate) fn display_duration(value: Option<u64>) -> String {
|
||||
let Some(value) = value.and_then(|value| i64::try_from(value).ok()) else {
|
||||
@@ -100,7 +101,7 @@ pub(crate) fn attempt_metrics_view(metrics: &domain::AttemptMetrics) -> AttemptM
|
||||
}
|
||||
|
||||
pub(crate) fn token_details(record: &domain::UsageRecord) -> TokenDetailsView {
|
||||
TokenDetailsView {
|
||||
token_details_view(&Usage {
|
||||
input_tokens: record.input_tokens,
|
||||
output_tokens: record.output_tokens,
|
||||
cached_input_tokens: record.cached_input_tokens,
|
||||
@@ -109,35 +110,11 @@ pub(crate) fn token_details(record: &domain::UsageRecord) -> TokenDetailsView {
|
||||
image_input_tokens: record.image_input_tokens,
|
||||
image_output_tokens: record.image_output_tokens,
|
||||
total_tokens: record.total_tokens,
|
||||
input_tokens_display: record
|
||||
.input_tokens
|
||||
.map_or_else(|| "-".to_owned(), format_number),
|
||||
output_tokens_display: record
|
||||
.output_tokens
|
||||
.map_or_else(|| "-".to_owned(), format_number),
|
||||
cached_input_tokens_display: record
|
||||
.cached_input_tokens
|
||||
.map_or_else(|| "-".to_owned(), format_compact_number),
|
||||
cache_write_input_tokens_display: record
|
||||
.cache_write_input_tokens
|
||||
.map_or_else(|| "-".to_owned(), format_compact_number),
|
||||
reasoning_output_tokens_display: record
|
||||
.reasoning_output_tokens
|
||||
.map_or_else(|| "-".to_owned(), format_number),
|
||||
image_input_tokens_display: record
|
||||
.image_input_tokens
|
||||
.map_or_else(|| "-".to_owned(), format_number),
|
||||
image_output_tokens_display: record
|
||||
.image_output_tokens
|
||||
.map_or_else(|| "-".to_owned(), format_number),
|
||||
total_tokens_display: record
|
||||
.total_tokens
|
||||
.map_or_else(|| "-".to_owned(), format_number),
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn usage_list_token_details(record: &domain::UsageListRecord) -> TokenDetailsView {
|
||||
TokenDetailsView {
|
||||
token_details_view(&Usage {
|
||||
input_tokens: record.input_tokens,
|
||||
output_tokens: record.output_tokens,
|
||||
cached_input_tokens: record.cached_input_tokens,
|
||||
@@ -146,28 +123,41 @@ pub(crate) fn usage_list_token_details(record: &domain::UsageListRecord) -> Toke
|
||||
image_input_tokens: record.image_input_tokens,
|
||||
image_output_tokens: record.image_output_tokens,
|
||||
total_tokens: record.total_tokens,
|
||||
input_tokens_display: record
|
||||
})
|
||||
}
|
||||
|
||||
fn token_details_view(usage: &Usage) -> TokenDetailsView {
|
||||
TokenDetailsView {
|
||||
input_tokens: usage.input_tokens,
|
||||
output_tokens: usage.output_tokens,
|
||||
cached_input_tokens: usage.cached_input_tokens,
|
||||
cache_write_input_tokens: usage.cache_write_input_tokens,
|
||||
reasoning_output_tokens: usage.reasoning_output_tokens,
|
||||
image_input_tokens: usage.image_input_tokens,
|
||||
image_output_tokens: usage.image_output_tokens,
|
||||
total_tokens: usage.total_tokens,
|
||||
input_tokens_display: usage
|
||||
.input_tokens
|
||||
.map_or_else(|| "-".to_owned(), format_number),
|
||||
output_tokens_display: record
|
||||
output_tokens_display: usage
|
||||
.output_tokens
|
||||
.map_or_else(|| "-".to_owned(), format_number),
|
||||
cached_input_tokens_display: record
|
||||
cached_input_tokens_display: usage
|
||||
.cached_input_tokens
|
||||
.map_or_else(|| "-".to_owned(), format_compact_number),
|
||||
cache_write_input_tokens_display: record
|
||||
cache_write_input_tokens_display: usage
|
||||
.cache_write_input_tokens
|
||||
.map_or_else(|| "-".to_owned(), format_compact_number),
|
||||
reasoning_output_tokens_display: record
|
||||
reasoning_output_tokens_display: usage
|
||||
.reasoning_output_tokens
|
||||
.map_or_else(|| "-".to_owned(), format_number),
|
||||
image_input_tokens_display: record
|
||||
image_input_tokens_display: usage
|
||||
.image_input_tokens
|
||||
.map_or_else(|| "-".to_owned(), format_number),
|
||||
image_output_tokens_display: record
|
||||
image_output_tokens_display: usage
|
||||
.image_output_tokens
|
||||
.map_or_else(|| "-".to_owned(), format_number),
|
||||
total_tokens_display: record
|
||||
total_tokens_display: usage
|
||||
.total_tokens
|
||||
.map_or_else(|| "-".to_owned(), format_number),
|
||||
}
|
||||
|
||||
@@ -53,8 +53,19 @@ pub(super) fn router<S: SessionState + Clone + Send + Sync + 'static>() -> Route
|
||||
.route("/api/admin/plugins/extensions/{instance_id}/{artifact_sha256}/{revision}/callback-tickets", post(start_callback::<S>))
|
||||
.route("/api/admin/plugins/extensions/{instance_id}/{artifact_sha256}/{revision}/api/{*path}", any(handle::<S>))
|
||||
.route("/api/admin/plugins/extensions/{instance_id}/{artifact_sha256}/{revision}/resources/{*path}", get(resource::<S>))
|
||||
.route("/plugins/resources/{instance_id}/{artifact_sha256}/{revision}/{*path}", get(public_resource::<S>))
|
||||
.route("/plugins/callbacks/{instance_id}/{artifact_sha256}/{revision}/{*path}", get(callback::<S>))
|
||||
}
|
||||
|
||||
pub(in crate::admin) fn public_router<S: SessionState + Clone + Send + Sync + 'static>() -> Router<S>
|
||||
{
|
||||
Router::new()
|
||||
.route(
|
||||
"/plugins/resources/{instance_id}/{artifact_sha256}/{revision}/{*path}",
|
||||
get(public_resource::<S>),
|
||||
)
|
||||
.route(
|
||||
"/plugins/callbacks/{instance_id}/{artifact_sha256}/{revision}/{*path}",
|
||||
get(callback::<S>),
|
||||
)
|
||||
}
|
||||
|
||||
pub(in crate::admin) fn model_router() -> Router<crate::ApiState> {
|
||||
|
||||
@@ -5,7 +5,7 @@ mod distribution;
|
||||
mod instances;
|
||||
mod management;
|
||||
|
||||
pub(super) use management::model_router;
|
||||
pub(super) use management::{model_router, public_router};
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
|
||||
@@ -60,6 +60,7 @@ struct CreateRequest {
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
struct UpdateRequest {
|
||||
enabled: Option<bool>,
|
||||
auto_location: Option<bool>,
|
||||
#[serde(default, deserialize_with = "deserialize_location_update")]
|
||||
location: Option<Option<gateway_core::account::RequestLocation>>,
|
||||
@@ -145,6 +146,7 @@ struct ProxyAccountView {
|
||||
#[derive(Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct ProxyView {
|
||||
enabled: bool,
|
||||
auto_location: bool,
|
||||
detected_location: Option<gateway_admin::model::proxies::DetectedProxyLocation>,
|
||||
location: Option<gateway_core::account::RequestLocation>,
|
||||
@@ -168,6 +170,7 @@ impl From<(ProxyRecord, crate::time::TimePresenter)> for ProxyView {
|
||||
let endpoint = record.proxy.endpoint();
|
||||
Self {
|
||||
auto_location: record.auto_location,
|
||||
enabled: record.enabled,
|
||||
detected_location: record.detected_location,
|
||||
location: record.location,
|
||||
id: record.id,
|
||||
@@ -419,6 +422,7 @@ where
|
||||
.proxies()
|
||||
.update(
|
||||
UpdateProxy {
|
||||
enabled: request.enabled,
|
||||
auto_location: request.auto_location,
|
||||
test: None,
|
||||
location: request.location,
|
||||
|
||||
@@ -17,9 +17,10 @@ use axum::{
|
||||
http::{HeaderName, HeaderValue, Request, StatusCode},
|
||||
response::{IntoResponse as _, Response},
|
||||
};
|
||||
use gateway_core::engine::middleware::FrozenMiddlewarePlan;
|
||||
use gateway_core::engine::middleware::http as contract;
|
||||
use gateway_core::engine::middleware::http::Settings;
|
||||
use gateway_core::engine::middleware::{FrozenMiddlewarePlan, MiddlewareError};
|
||||
use gateway_core::error::GatewayErrorKind;
|
||||
use gateway_core::lifecycle::CancellationToken;
|
||||
use gateway_core::middleware::compose;
|
||||
use http_body::{Body as HttpBody, Frame, SizeHint};
|
||||
@@ -270,6 +271,11 @@ async fn handle(
|
||||
}
|
||||
response.map(|body| Body::new(ResponseBody::new(body, guard)))
|
||||
}
|
||||
Err(MiddlewareError::Gateway(error))
|
||||
if error.kind() == GatewayErrorKind::MiddlewareCapacityUnavailable =>
|
||||
{
|
||||
crate::openai::error::gateway_error_response(&error)
|
||||
}
|
||||
Err(error) => {
|
||||
let status = if error.is_rejected() {
|
||||
StatusCode::FORBIDDEN
|
||||
|
||||
@@ -13,6 +13,8 @@ use gateway_core::{
|
||||
};
|
||||
use serde_json::{Value, json};
|
||||
|
||||
pub use gateway_core::engine::execution::gateway_error_from_engine;
|
||||
|
||||
use super::responses::ProtocolErrorBody;
|
||||
|
||||
/// OpenAI 风格错误响应
|
||||
@@ -148,39 +150,6 @@ pub fn model_not_found_response() -> (StatusCode, Json<Value>) {
|
||||
)
|
||||
}
|
||||
|
||||
/// 将 core Engine 错误收敛为稳定 Gateway 分类
|
||||
#[must_use]
|
||||
pub fn gateway_error_from_engine(error: &EngineError) -> GatewayError {
|
||||
match error {
|
||||
EngineError::Provider(error) => GatewayError::from_provider(error),
|
||||
EngineError::Cancelled => {
|
||||
GatewayError::new(GatewayErrorKind::Cancelled, "request was cancelled")
|
||||
}
|
||||
EngineError::Deadline => {
|
||||
GatewayError::new(GatewayErrorKind::Timeout, "gateway request timed out")
|
||||
}
|
||||
EngineError::ProviderNotRegistered { .. } | EngineError::EmptyRoutingPlan => {
|
||||
GatewayError::new(
|
||||
GatewayErrorKind::NoAvailableProvider,
|
||||
"no upstream provider is currently available for this request",
|
||||
)
|
||||
}
|
||||
EngineError::Store(source) => {
|
||||
GatewayError::new(GatewayErrorKind::Internal, "gateway execution failed")
|
||||
.with_source(source.clone())
|
||||
}
|
||||
EngineError::ProviderMetadataMismatch
|
||||
| EngineError::ContinuationPinMismatch
|
||||
| EngineError::RequiredAccountMismatch
|
||||
| EngineError::AccountOutsideClientScope
|
||||
| EngineError::DownstreamCommitRequired
|
||||
| EngineError::NoActiveAttempt
|
||||
| EngineError::InvalidDeliveryState => {
|
||||
GatewayError::new(GatewayErrorKind::Internal, "gateway execution failed")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Gateway 错误的 OpenAI HTTP 表达
|
||||
pub fn gateway_error_response(error: &GatewayError) -> Response {
|
||||
let (status, default_type, default_code) = gateway_error_contract(error.kind());
|
||||
@@ -310,6 +279,11 @@ pub const fn gateway_error_contract(
|
||||
"server_error",
|
||||
"provider_infrastructure_unavailable",
|
||||
),
|
||||
GatewayErrorKind::MiddlewareCapacityUnavailable => (
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
"server_error",
|
||||
"middleware_capacity_unavailable",
|
||||
),
|
||||
GatewayErrorKind::ConcurrencyQueueFull => (
|
||||
StatusCode::TOO_MANY_REQUESTS,
|
||||
"rate_limit_error",
|
||||
|
||||
@@ -5,7 +5,9 @@ use std::{borrow::Cow, fmt, net::IpAddr};
|
||||
use axum::http::HeaderMap;
|
||||
use base64::{Engine as _, engine::general_purpose::STANDARD};
|
||||
use bytes::Bytes;
|
||||
use gateway_core::operation::{GenerateRequest, Operation, ProtocolPayload, ProviderSessionState};
|
||||
use gateway_core::operation::{
|
||||
ClientTurn, GenerateRequest, Operation, ProtocolPayload, ProviderSessionState,
|
||||
};
|
||||
use gateway_protocol::openai::{
|
||||
X_OPENAI_INTERNAL_CODEX_RESPONSES_LITE_HEADER, X_OPENAI_MEMGEN_REQUEST_HEADER,
|
||||
is_transport_managed_request_header,
|
||||
@@ -497,12 +499,30 @@ pub(super) fn decode_request_object(
|
||||
};
|
||||
let protocol_context = request_headers.protocol_context(use_websocket, frame_turn_metadata);
|
||||
|
||||
let semantics =
|
||||
gateway_protocol::openai::codex_responses_request_semantics(&object, &protocol_context);
|
||||
let client_turn = if object.get("generate").and_then(Value::as_bool) != Some(false)
|
||||
&& semantics.request_kind.as_deref() == Some("turn")
|
||||
&& !semantics.compact
|
||||
{
|
||||
gateway_protocol::openai::codex_account_thread_id(&object, &protocol_context)
|
||||
.zip(gateway_protocol::openai::codex_turn_id(
|
||||
&object,
|
||||
&protocol_context,
|
||||
))
|
||||
.and_then(|(thread, turn)| ClientTurn::new(thread, turn))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let payload = ProtocolPayload::json_object(OPENAI_PROTOCOL, object)
|
||||
.map_err(|_| RequestDecodeError::CanonicalContract {
|
||||
field: "request".to_owned(),
|
||||
})?
|
||||
.with_context(protocol_context);
|
||||
let operation = Operation::Generate(GenerateRequest::from_protocol_payload(payload));
|
||||
let operation = Operation::Generate(
|
||||
GenerateRequest::from_protocol_payload(payload).with_client_turn(client_turn),
|
||||
);
|
||||
Ok(DecodedResponsesRequest {
|
||||
operation,
|
||||
metadata: ResponsesRequestMetadata {
|
||||
|
||||
@@ -112,6 +112,7 @@ fn account_page_item(
|
||||
|
||||
AccountPageItem {
|
||||
account: AccountRecord {
|
||||
outbound_proxy_enabled: true,
|
||||
id: "acct_cooldown".to_owned(),
|
||||
provider_kind: gateway_core::routing::ProviderKind::new("openai").unwrap(),
|
||||
groups: Vec::new(),
|
||||
|
||||
@@ -0,0 +1,111 @@
|
||||
//! 管理端 JSON 压缩的协商、响应合同与重复编码回归测试
|
||||
|
||||
use std::io::Read as _;
|
||||
|
||||
use axum::{
|
||||
Router,
|
||||
body::{Body, to_bytes},
|
||||
http::{Request, StatusCode, header},
|
||||
response::Response,
|
||||
};
|
||||
use flate2::read::GzDecoder;
|
||||
use tower::ServiceExt as _;
|
||||
use tower_http::compression::CompressionLayer;
|
||||
|
||||
use super::AdminTestFixture;
|
||||
|
||||
async fn app() -> Router {
|
||||
let fixture = AdminTestFixture::new().await;
|
||||
fixture.auth.insert_session("valid-session");
|
||||
crate::openai::api_router_with_admin(fixture.services)
|
||||
}
|
||||
|
||||
async fn get(router: Router, uri: &str, encoding: Option<&str>) -> Response {
|
||||
let mut request = Request::builder()
|
||||
.uri(uri)
|
||||
.header(header::COOKIE, "cpr_session=valid-session")
|
||||
.header("x-request-id", "req_admin_compression");
|
||||
if let Some(encoding) = encoding {
|
||||
request = request.header(header::ACCEPT_ENCODING, encoding);
|
||||
}
|
||||
router
|
||||
.oneshot(request.body(Body::empty()).expect("request"))
|
||||
.await
|
||||
.expect("response")
|
||||
}
|
||||
|
||||
async fn bytes(response: Response) -> Vec<u8> {
|
||||
to_bytes(response.into_body(), 1024 * 1024)
|
||||
.await
|
||||
.expect("response body")
|
||||
.to_vec()
|
||||
}
|
||||
|
||||
fn decode_gzip(body: &[u8]) -> Vec<u8> {
|
||||
let mut decoded = Vec::new();
|
||||
GzDecoder::new(body)
|
||||
.read_to_end(&mut decoded)
|
||||
.expect("valid gzip body");
|
||||
decoded
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gzip_should_preserve_admin_json_and_response_headers() {
|
||||
let router = app().await;
|
||||
let expected = bytes(get(router.clone(), "/api/admin/settings", None).await).await;
|
||||
assert!(expected.len() >= 1024);
|
||||
|
||||
for encoding in ["gzip", "br, gzip;q=0.8, deflate;q=0.5", "*;q=0.5"] {
|
||||
let response = get(router.clone(), "/api/admin/settings", Some(encoding)).await;
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(response.headers()[header::CONTENT_ENCODING], "gzip");
|
||||
assert_eq!(response.headers()[header::CONTENT_TYPE], "application/json");
|
||||
assert_eq!(response.headers()[header::CACHE_CONTROL], "no-store");
|
||||
assert_eq!(response.headers()["x-request-id"], "req_admin_compression");
|
||||
assert_eq!(response.headers()[header::VARY], "accept-encoding");
|
||||
assert!(!response.headers().contains_key(header::CONTENT_LENGTH));
|
||||
let compressed = bytes(response).await;
|
||||
assert!(compressed.len() < expected.len());
|
||||
assert_eq!(decode_gzip(&compressed), expected, "{encoding}");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn clients_without_gzip_support_should_receive_original_json() {
|
||||
let router = app().await;
|
||||
let expected = bytes(get(router.clone(), "/api/admin/settings", None).await).await;
|
||||
for encoding in [None, Some("identity"), Some("br"), Some("gzip;q=0")] {
|
||||
let response = get(router.clone(), "/api/admin/settings", encoding).await;
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert!(!response.headers().contains_key(header::CONTENT_ENCODING));
|
||||
assert_eq!(response.headers()[header::VARY], "accept-encoding");
|
||||
assert_eq!(bytes(response).await, expected, "{encoding:?}");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn small_admin_json_should_remain_uncompressed() {
|
||||
let router = app().await;
|
||||
let uri = "/api/admin/settings/admin-api-key";
|
||||
let expected = bytes(get(router.clone(), uri, None).await).await;
|
||||
assert!(expected.len() < 1024);
|
||||
let response = get(router, uri, Some("gzip")).await;
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert!(!response.headers().contains_key(header::CONTENT_ENCODING));
|
||||
assert_eq!(bytes(response).await, expected);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_second_compression_layer_should_not_encode_gzip_again() {
|
||||
let router = app().await;
|
||||
let expected = bytes(get(router.clone(), "/api/admin/settings", None).await).await;
|
||||
let response = get(
|
||||
router.layer(CompressionLayer::new()),
|
||||
"/api/admin/settings",
|
||||
Some("gzip"),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(response.headers()[header::CONTENT_ENCODING], "gzip");
|
||||
assert_eq!(response.headers()[header::VARY], "accept-encoding");
|
||||
assert_eq!(decode_gzip(&bytes(response).await), expected);
|
||||
}
|
||||
@@ -511,6 +511,7 @@ mod provider {
|
||||
total_slots: None,
|
||||
},
|
||||
account: AccountRecord {
|
||||
outbound_proxy_enabled: true,
|
||||
notes: None,
|
||||
model_access: Default::default(),
|
||||
id: "acct_error_test".to_owned(),
|
||||
|
||||
@@ -88,6 +88,7 @@ mod account_groups;
|
||||
mod accounts;
|
||||
mod auth;
|
||||
mod client_keys;
|
||||
mod compression;
|
||||
mod error_diagnostics;
|
||||
mod errors;
|
||||
mod observability;
|
||||
|
||||
@@ -435,6 +435,35 @@ async fn resource_request(
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn large_html_resource_should_remain_uncompressed() {
|
||||
let fixture = published_fixture().await;
|
||||
let path = format!("ui/{}.html", "a".repeat(2048));
|
||||
let response = crate::openai::api_router_with_admin(fixture.services)
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.uri(format!(
|
||||
"/api/admin/plugins/extensions/plugin-instance-diagnostic/{}/7/resources/{path}",
|
||||
"a".repeat(64)
|
||||
))
|
||||
.header("cookie", "cpr_session=management-session")
|
||||
.header("accept-encoding", "gzip")
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(response.headers()["content-type"], "text/html");
|
||||
assert!(!response.headers().contains_key("content-encoding"));
|
||||
assert_eq!(
|
||||
axum::body::to_bytes(response.into_body(), 4096)
|
||||
.await
|
||||
.unwrap(),
|
||||
path
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn static_resources_revalidate_privately_after_authorization() {
|
||||
use axum::body::to_bytes;
|
||||
|
||||
@@ -89,6 +89,43 @@ async fn proxy_location_rejects_invalid_and_incomplete_input() {
|
||||
#[derive(Default)]
|
||||
pub(super) struct MemoryProxies(Mutex<Option<ProxyRecord>>);
|
||||
|
||||
#[tokio::test]
|
||||
async fn proxy_enabled_updates_preserve_connection_and_omitted_state() {
|
||||
let fixture = AdminTestFixture::new().await;
|
||||
fixture.auth.insert_session("valid-session");
|
||||
let (status, created) = request(
|
||||
&fixture,
|
||||
"/api/admin/proxies/create",
|
||||
Some(json!({
|
||||
"name":"Office", "proxyUrl":"http://proxy.example:8080"
|
||||
})),
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, StatusCode::CREATED);
|
||||
assert_eq!(created["data"]["record"]["enabled"], true);
|
||||
for (revision, enabled, expected) in [
|
||||
(1, Some(false), false),
|
||||
(2, None, false),
|
||||
(3, Some(true), true),
|
||||
] {
|
||||
let mut body = json!({"id":"proxy_test", "revision":revision, "name":"Office"});
|
||||
if let Some(enabled) = enabled {
|
||||
body["enabled"] = json!(enabled);
|
||||
}
|
||||
let (status, changed) =
|
||||
request(&fixture, "/api/admin/proxies/update", Some(body), true).await;
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
assert_eq!(changed["data"]["record"]["enabled"], expected);
|
||||
assert_eq!(
|
||||
changed["data"]["record"]["endpoint"],
|
||||
created["data"]["record"]["endpoint"]
|
||||
);
|
||||
let (_, listed) = request(&fixture, "/api/admin/proxies", None, true).await;
|
||||
assert_eq!(listed["data"]["items"][0]["enabled"], expected);
|
||||
}
|
||||
}
|
||||
|
||||
fn missing() -> AdminStoreError {
|
||||
AdminStoreError::new(AdminStoreErrorKind::NotFound, "proxy", "missing proxy")
|
||||
}
|
||||
@@ -165,6 +202,7 @@ impl ProxyStore for MemoryProxies {
|
||||
_: &MutationContext,
|
||||
) -> AdminStoreResult<ProxyMutation> {
|
||||
let mut record = ProxyRecord {
|
||||
enabled: true,
|
||||
auto_location: command.auto_location,
|
||||
detected_location: None,
|
||||
location: command.location,
|
||||
@@ -197,6 +235,9 @@ impl ProxyStore for MemoryProxies {
|
||||
let mut stored = self.0.lock().unwrap();
|
||||
let record = stored.as_mut().ok_or_else(missing)?;
|
||||
record.name = command.name;
|
||||
if let Some(enabled) = command.enabled {
|
||||
record.enabled = enabled;
|
||||
}
|
||||
if let Some(auto_location) = command.auto_location {
|
||||
if auto_location != record.auto_location {
|
||||
record.detected_location = None;
|
||||
|
||||
@@ -84,9 +84,10 @@ async fn privacy_save_requires_provider_validation_and_preview_requires_admin_au
|
||||
fixture.auth.insert_session("valid-session");
|
||||
let mut body = update_body();
|
||||
let policy = json!({"enabled":true,"onError":"reject_request","rules":[{
|
||||
"id":"auth","name":"自定义认证头","enabled":true,"scope":"request_header",
|
||||
"selector":"authorization","action":"remove_field","pattern":null,
|
||||
"replacement":"","value":null,"replaceAll":true,"caseInsensitive":false,"multiLine":false
|
||||
"id":"auth","name":"自定义认证头","enabled":true,"steps":[{
|
||||
"scope":"request_header","selector":"authorization","action":"remove_field","pattern":null,
|
||||
"replacement":"","value":null,"replaceAll":true,"caseInsensitive":false,"multiLine":false
|
||||
}]
|
||||
}]});
|
||||
body["codexPrivacyPolicy"] = policy.clone();
|
||||
let response = app(fixture.state())
|
||||
|
||||
@@ -9,7 +9,7 @@ use axum::{
|
||||
http::{Request, StatusCode, header},
|
||||
};
|
||||
use chrono::Utc;
|
||||
use futures::stream;
|
||||
use futures::{StreamExt as _, stream};
|
||||
use gateway_admin::{
|
||||
model::system::{
|
||||
SystemOperationAccepted, SystemOperationKind, SystemOperationState, SystemOperationStatus,
|
||||
@@ -19,6 +19,7 @@ use gateway_admin::{
|
||||
ports::system::{SystemOperationError, SystemOperations, SystemUpdateEventStream},
|
||||
};
|
||||
use gateway_api::admin::system::{self, UpdateDetailQuery, UpdateRequest};
|
||||
use http_body_util::BodyExt as _;
|
||||
use tower::ServiceExt as _;
|
||||
|
||||
use super::{AdminTestFixture, AdminTestState, unavailable_system};
|
||||
@@ -40,25 +41,39 @@ fn update_request_should_preserve_target_version_for_domain_validation() {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn update_event_stream_should_preserve_event_id_and_download_progress_percent() {
|
||||
let fixture = AdminTestFixture::with_system(Arc::new(ProgressSystem)).await;
|
||||
let fixture = AdminTestFixture::with_system(Arc::new(ProgressSystem::default())).await;
|
||||
fixture.auth.insert_session("valid-session");
|
||||
let response = app(fixture.state())
|
||||
let response = gateway_api::admin::router::<AdminTestState>()
|
||||
.with_state(fixture.state())
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.uri("/api/admin/system/update/events")
|
||||
.header(header::COOKIE, "cpr_session=valid-session")
|
||||
.header(header::ACCEPT_ENCODING, "gzip")
|
||||
.header("x-request-id", "req_system_update_events")
|
||||
.body(Body::empty())
|
||||
.expect("update event request"),
|
||||
)
|
||||
);
|
||||
let response = tokio::time::timeout(std::time::Duration::from_secs(1), response)
|
||||
.await
|
||||
.expect("SSE response must not wait for stream completion")
|
||||
.expect("update event response");
|
||||
let body = to_bytes(response.into_body(), 1024 * 1024)
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response.headers()[header::CONTENT_TYPE],
|
||||
"text/event-stream"
|
||||
);
|
||||
assert!(!response.headers().contains_key(header::CONTENT_ENCODING));
|
||||
let mut stream = response.into_body();
|
||||
let frame = tokio::time::timeout(std::time::Duration::from_secs(1), stream.frame())
|
||||
.await
|
||||
.expect("update event body");
|
||||
let body = String::from_utf8(body.to_vec()).expect("UTF-8 event stream");
|
||||
.expect("first SSE event must arrive before stream completion")
|
||||
.expect("first event")
|
||||
.expect("event body");
|
||||
let body = String::from_utf8(frame.into_data().expect("event data").to_vec())
|
||||
.expect("UTF-8 event stream");
|
||||
|
||||
assert!(
|
||||
body.contains(r#""id":"update-progress-1""#),
|
||||
@@ -72,7 +87,7 @@ async fn update_event_stream_should_preserve_event_id_and_download_progress_perc
|
||||
|
||||
#[tokio::test]
|
||||
async fn update_should_return_accepted_and_report_restart_only_through_status() {
|
||||
let fixture = AdminTestFixture::with_system(Arc::new(ProgressSystem)).await;
|
||||
let fixture = AdminTestFixture::with_system(Arc::new(ProgressSystem::default())).await;
|
||||
fixture.auth.insert_session("valid-session");
|
||||
let router = app(fixture.state());
|
||||
let response = router
|
||||
@@ -119,7 +134,55 @@ fn app(state: AdminTestState) -> Router {
|
||||
system::router::<AdminTestState>().with_state(state)
|
||||
}
|
||||
|
||||
struct ProgressSystem;
|
||||
#[tokio::test]
|
||||
async fn update_detail_should_compress_json_starting_at_one_kib() {
|
||||
async fn detail(notes: String, encoding: &str) -> axum::response::Response {
|
||||
let fixture =
|
||||
AdminTestFixture::with_system(Arc::new(ProgressSystem { notes: Some(notes) })).await;
|
||||
fixture.auth.insert_session("valid-session");
|
||||
crate::openai::api_router_with_admin(fixture.services)
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.uri("/api/admin/system/update/detail")
|
||||
.header(header::COOKIE, "cpr_session=valid-session")
|
||||
.header(header::ACCEPT_ENCODING, encoding)
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
let response = detail(String::new(), "identity").await;
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let base = to_bytes(response.into_body(), 4096).await.unwrap();
|
||||
assert!(base.len() < 1023);
|
||||
for size in [1023, 1024] {
|
||||
let response = detail("x".repeat(size - base.len()), "gzip").await;
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let compressed = response.headers().get(header::CONTENT_ENCODING).cloned();
|
||||
let bytes = to_bytes(response.into_body(), 4096).await.unwrap();
|
||||
if size == 1023 {
|
||||
assert!(compressed.is_none());
|
||||
assert_eq!(bytes.len(), size);
|
||||
} else {
|
||||
use std::io::Read as _;
|
||||
assert_eq!(compressed.unwrap(), "gzip");
|
||||
let mut decoded = Vec::new();
|
||||
flate2::read::GzDecoder::new(bytes.as_ref())
|
||||
.read_to_end(&mut decoded)
|
||||
.unwrap();
|
||||
assert_eq!(decoded.len(), size);
|
||||
let body: serde_json::Value = serde_json::from_slice(&decoded).unwrap();
|
||||
assert_eq!(body["data"]["notes"], "x".repeat(size - base.len()));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct ProgressSystem {
|
||||
notes: Option<String>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl SystemOperations for ProgressSystem {
|
||||
@@ -144,7 +207,7 @@ impl SystemOperations for ProgressSystem {
|
||||
deployment_mode: "docker".into(),
|
||||
build_type: "release".into(),
|
||||
release_url: None,
|
||||
notes: None,
|
||||
notes: self.notes.clone(),
|
||||
cached: false,
|
||||
update_supported: true,
|
||||
unsupported_reason: None,
|
||||
@@ -153,18 +216,21 @@ impl SystemOperations for ProgressSystem {
|
||||
}
|
||||
|
||||
fn update_events(&self) -> SystemUpdateEventStream {
|
||||
Box::pin(stream::once(async {
|
||||
SystemUpdateEvent {
|
||||
id: "update-progress-1".to_owned(),
|
||||
operation_id: Some("update-1".to_owned()),
|
||||
level: SystemUpdateEventLevel::Info,
|
||||
step: Some("download".to_owned()),
|
||||
message: "已下载 4.0 MiB / 10.0 MiB (40%)".to_owned(),
|
||||
terminal: true,
|
||||
progress_percent: Some(40),
|
||||
occurred_at: Utc::now(),
|
||||
}
|
||||
}))
|
||||
Box::pin(
|
||||
stream::once(async {
|
||||
SystemUpdateEvent {
|
||||
id: "update-progress-1".to_owned(),
|
||||
operation_id: Some("update-1".to_owned()),
|
||||
level: SystemUpdateEventLevel::Info,
|
||||
step: Some("download".to_owned()),
|
||||
message: "已下载 4.0 MiB / 10.0 MiB (40%)".to_owned(),
|
||||
terminal: false,
|
||||
progress_percent: Some(40),
|
||||
occurred_at: Utc::now(),
|
||||
}
|
||||
})
|
||||
.chain(stream::pending()),
|
||||
)
|
||||
}
|
||||
|
||||
async fn perform_update(
|
||||
@@ -216,7 +282,7 @@ impl SystemOperations for ProgressSystem {
|
||||
|
||||
#[tokio::test]
|
||||
async fn update_detail_should_validate_and_forward_the_temporary_channel() {
|
||||
let fixture = AdminTestFixture::with_system(Arc::new(ProgressSystem)).await;
|
||||
let fixture = AdminTestFixture::with_system(Arc::new(ProgressSystem::default())).await;
|
||||
fixture.auth.insert_session("valid-session");
|
||||
let router = app(fixture.state());
|
||||
for (authenticated, query, expected) in [
|
||||
|
||||
@@ -22,13 +22,6 @@ fn manifest_should_depend_only_on_core_protocol_and_http_adapter_layers() {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn workspace_should_include_gateway_api() {
|
||||
let workspace = include_str!("../../../Cargo.toml");
|
||||
|
||||
assert!(workspace.contains("\"crates/gateway-api\""));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn admin_routes_should_use_only_get_post_and_static_paths() {
|
||||
let root = Path::new(env!("CARGO_MANIFEST_DIR")).join("src/admin");
|
||||
|
||||
@@ -35,6 +35,7 @@ use tower::ServiceExt as _;
|
||||
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
enum Mode {
|
||||
Overloaded,
|
||||
Rewrite,
|
||||
Upgrade {
|
||||
reject: bool,
|
||||
@@ -142,6 +143,12 @@ impl MiddlewarePlan for Plan {
|
||||
let mode = self.mode;
|
||||
Box::pin(async move {
|
||||
match mode {
|
||||
Mode::Overloaded => Err(MiddlewareError::Gateway(
|
||||
gateway_core::error::GatewayError::new(
|
||||
gateway_core::error::GatewayErrorKind::MiddlewareCapacityUnavailable,
|
||||
"Request middleware is temporarily busy",
|
||||
),
|
||||
)),
|
||||
Mode::Service => unreachable!("service-only plan bypasses HTTP layers"),
|
||||
Mode::Upgrade { reject, idle, .. } => {
|
||||
let upgrade = request
|
||||
@@ -659,3 +666,25 @@ async fn plugin_http_dispatch_uses_host_identity_without_client_credentials() {
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn http_middleware_overload_is_an_explicit_503_before_native_routing() {
|
||||
let (router, plan) = router(Mode::Overloaded, None).await;
|
||||
for path in [
|
||||
"/healthz",
|
||||
"/assets/index.js",
|
||||
"/api/settings",
|
||||
"/v1/responses",
|
||||
] {
|
||||
let response = router
|
||||
.clone()
|
||||
.oneshot(Request::builder().uri(path).body(Body::empty()).unwrap())
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
|
||||
let body = response.into_body().collect().await.unwrap().to_bytes();
|
||||
let body: serde_json::Value = serde_json::from_slice(&body).unwrap();
|
||||
assert_eq!(body["error"]["code"], "middleware_capacity_unavailable");
|
||||
}
|
||||
assert_eq!(plan.calls.load(Ordering::Acquire), 4);
|
||||
}
|
||||
|
||||
@@ -263,6 +263,18 @@ fn rate_limited_error_should_map_to_openai_retryable_status() {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn middleware_capacity_is_service_unavailable() {
|
||||
assert_eq!(
|
||||
gateway_error_contract(GatewayErrorKind::MiddlewareCapacityUnavailable),
|
||||
(
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
"server_error",
|
||||
"middleware_capacity_unavailable"
|
||||
),
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn queue_rejections_keep_distinct_retryable_http_codes() {
|
||||
for (kind, code) in [
|
||||
@@ -548,7 +560,7 @@ fn engine_store_error_should_collapse_to_safe_internal_error() {
|
||||
|
||||
let projected = gateway_error_from_engine(&error);
|
||||
assert_eq!(projected.kind(), GatewayErrorKind::Internal);
|
||||
assert_eq!(projected.safe_message(), "gateway execution failed");
|
||||
assert_eq!(projected.safe_message(), "request execution failed");
|
||||
}
|
||||
#[test]
|
||||
fn openai_error_response_should_preserve_only_safe_contract_fields() {
|
||||
|
||||
@@ -34,6 +34,49 @@ fn generate_request(body: Value) -> DecodedResponsesRequest {
|
||||
decode_request(body.to_string().as_bytes()).expect("test request should decode")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decoder_should_attach_only_identified_codex_inference_turns() {
|
||||
let body = json!({"model":"smart-code", "input":"hello", "client_metadata": {
|
||||
"thread_id":"thread", "turn_id":"turn", "x-codex-turn-metadata": r#"{"request_kind":"turn"}"#
|
||||
}});
|
||||
let decoded = generate_request(body.clone());
|
||||
let identity = generate_operation(&decoded).client_turn().unwrap();
|
||||
assert_eq!(identity.thread_id(), "thread");
|
||||
assert_eq!(identity.turn_id(), "turn");
|
||||
assert_eq!(Value::Object(openai_wire_body(&decoded).clone()), body);
|
||||
for variant in [
|
||||
"warmup",
|
||||
"compaction",
|
||||
"missing_thread",
|
||||
"missing_turn",
|
||||
"unknown_kind",
|
||||
"oversized_id",
|
||||
] {
|
||||
let mut request = body.clone();
|
||||
match variant {
|
||||
"warmup" => request["generate"] = json!(false),
|
||||
"compaction" => {
|
||||
request["client_metadata"]["x-codex-turn-metadata"] =
|
||||
json!(r#"{"request_kind":"compaction"}"#)
|
||||
}
|
||||
"missing_thread" => request["client_metadata"]["thread_id"] = json!(""),
|
||||
"missing_turn" => request["client_metadata"]["turn_id"] = json!(""),
|
||||
"unknown_kind" => {
|
||||
request["client_metadata"]["x-codex-turn-metadata"] =
|
||||
json!(r#"{"request_kind":"other"}"#)
|
||||
}
|
||||
"oversized_id" => request["client_metadata"]["turn_id"] = json!("x".repeat(257)),
|
||||
_ => unreachable!(),
|
||||
}
|
||||
assert!(
|
||||
generate_operation(&generate_request(request))
|
||||
.client_turn()
|
||||
.is_none(),
|
||||
"{variant}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
fn generate_operation(
|
||||
decoded: &DecodedResponsesRequest,
|
||||
) -> &gateway_core::operation::GenerateRequest {
|
||||
|
||||
@@ -1,10 +1,192 @@
|
||||
//! 验证未知控制消息、迟到中断、结算期间收发和创建请求的串行准入
|
||||
|
||||
use super::*;
|
||||
use gateway_core::engine::client_retry::{ClientRetryKey, ClientRetryStore};
|
||||
use gateway_core::engine::response_control::{
|
||||
ResponseControlTransport, ResponseControlUnavailable,
|
||||
};
|
||||
|
||||
#[derive(Default)]
|
||||
struct RetryMarker(Mutex<Option<ModelRequestId>>);
|
||||
|
||||
impl ClientRetryStore for RetryMarker {
|
||||
fn load<'a>(
|
||||
&'a self,
|
||||
_: &'a ClientRetryKey,
|
||||
) -> BoxFuture<'a, Result<Option<ModelRequestId>, StoreError>> {
|
||||
Box::pin(async { Ok(self.0.lock().unwrap().clone()) })
|
||||
}
|
||||
fn record_failure<'a>(
|
||||
&'a self,
|
||||
_: &'a ClientRetryKey,
|
||||
version: &'a ModelRequestId,
|
||||
_: Duration,
|
||||
) -> BoxFuture<'a, Result<(), StoreError>> {
|
||||
Box::pin(async move {
|
||||
*self.0.lock().unwrap() = Some(version.clone());
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
fn clear_failure<'a>(
|
||||
&'a self,
|
||||
_: &'a ClientRetryKey,
|
||||
expected: &'a ModelRequestId,
|
||||
) -> BoxFuture<'a, Result<(), StoreError>> {
|
||||
Box::pin(async move {
|
||||
let mut marker = self.0.lock().unwrap();
|
||||
if marker.as_ref() == Some(expected) {
|
||||
*marker = None;
|
||||
}
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct RetryBudgetProvider(AtomicUsize);
|
||||
|
||||
#[async_trait]
|
||||
impl Provider for RetryBudgetProvider {
|
||||
fn name(&self) -> &'static str {
|
||||
"openai"
|
||||
}
|
||||
fn catalog_generation(&self) -> ProviderCatalogGeneration {
|
||||
ProviderCatalogGeneration::default()
|
||||
}
|
||||
async fn query_model_capabilities(
|
||||
&self,
|
||||
) -> Result<Vec<ProviderModelCapabilities>, ProviderError> {
|
||||
Ok(Vec::new())
|
||||
}
|
||||
async fn execute(
|
||||
self: Arc<Self>,
|
||||
request: ProviderRequest,
|
||||
context: AttemptContext,
|
||||
) -> Result<ProviderStream, ProviderError> {
|
||||
let number = self.0.fetch_add(1, Ordering::SeqCst) + 1;
|
||||
if number <= 7 || matches!(number, 9 | 10) {
|
||||
let error = ProviderError::new(ProviderErrorKind::Unavailable, UpstreamSendState::Sent)
|
||||
.with_status(503)
|
||||
.with_replay_safe()
|
||||
.with_client_visible_upstream_response(gateway_core::error::ClientVisibleUpstreamResponse::new(
|
||||
503, Some(b"application/json".to_vec()),
|
||||
Bytes::from_static(br#"{"error":{"code":"server_error","message":"synthetic unavailable"}}"#),
|
||||
))
|
||||
.with_transient_retry(
|
||||
std::num::NonZeroU32::new(5).unwrap(),
|
||||
Duration::ZERO,
|
||||
Duration::ZERO,
|
||||
Duration::ZERO,
|
||||
gateway_core::error::TransientRetryPolicy::RetryThenReturnError,
|
||||
);
|
||||
return Ok(ProviderStream::new(
|
||||
ProviderCallMetadata::new(
|
||||
request.candidate().provider().clone(),
|
||||
request.candidate().upstream_model().unwrap().clone(),
|
||||
ProviderAccountId::new("acct_api_test").unwrap(),
|
||||
UpstreamTransport::new("websocket").unwrap(),
|
||||
),
|
||||
futures::stream::iter([Err(error)]),
|
||||
(),
|
||||
));
|
||||
}
|
||||
Arc::new(ChargedWebSocketProvider)
|
||||
.execute(request, context)
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn client_retry_budget_should_flow_through_http_and_live_websocket() {
|
||||
for websocket in [false, true] {
|
||||
let provider = Arc::new(RetryBudgetProvider::default());
|
||||
let marker = Arc::new(RetryMarker::default());
|
||||
let execution = Arc::new(
|
||||
DefaultExecutionService::new(
|
||||
RuntimeSnapshotHandle::new(snapshot("sk_retry_probe", "openai")),
|
||||
Arc::new(SettlementPorts::default()),
|
||||
ProviderRegistry::new([provider.clone() as Arc<dyn Provider>]).unwrap(),
|
||||
Arc::new(InterruptAdmissions::default()),
|
||||
Arc::new(UnusedContinuation),
|
||||
Arc::new(IgnoredClientApiKeyUsage),
|
||||
Arc::new(crate::support::RecordingDiagnostics::default()),
|
||||
)
|
||||
.with_client_retry_store(marker.clone()),
|
||||
);
|
||||
let app = api_router(execution).await;
|
||||
let mut body = json!({"model":"model-a", "input":"hello", "stream":true, "client_metadata": {
|
||||
"thread_id":"thread", "turn_id":"turn", "x-codex-turn-metadata":r#"{"request_kind":"turn"}"#
|
||||
}});
|
||||
let mut socket = None;
|
||||
let mut server = None;
|
||||
if websocket {
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
let app = app.clone();
|
||||
server = Some(Server(tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.unwrap()
|
||||
})));
|
||||
let mut request = format!("ws://{address}/v1/responses")
|
||||
.into_client_request()
|
||||
.unwrap();
|
||||
request.headers_mut().insert(
|
||||
AUTHORIZATION,
|
||||
HeaderValue::from_static("Bearer sk_retry_probe"),
|
||||
);
|
||||
socket = Some(connect_async(request).await.unwrap().0);
|
||||
body["type"] = json!("response.create");
|
||||
}
|
||||
for (expected_attempts, success) in [(6, false), (7, false), (8, true), (11, true)] {
|
||||
if let Some(socket) = &mut socket {
|
||||
send(socket, body.clone()).await;
|
||||
loop {
|
||||
let event = next(socket).await;
|
||||
match event["type"].as_str() {
|
||||
Some("response.completed") => {
|
||||
assert!(success);
|
||||
break;
|
||||
}
|
||||
Some("response.failed" | "error") => {
|
||||
assert!(!success);
|
||||
break;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
let request = Request::builder()
|
||||
.method("POST")
|
||||
.uri("/v1/responses")
|
||||
.header(AUTHORIZATION, "Bearer sk_retry_probe")
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(body.to_string()))
|
||||
.unwrap();
|
||||
let response = app.clone().oneshot(request).await.unwrap();
|
||||
assert_eq!(
|
||||
response.status(),
|
||||
if success {
|
||||
StatusCode::OK
|
||||
} else {
|
||||
StatusCode::SERVICE_UNAVAILABLE
|
||||
}
|
||||
);
|
||||
let payload = axum::body::to_bytes(response.into_body(), usize::MAX)
|
||||
.await
|
||||
.unwrap();
|
||||
if success {
|
||||
assert!(String::from_utf8_lossy(&payload).contains("response.completed"));
|
||||
}
|
||||
}
|
||||
assert_eq!(provider.0.load(Ordering::SeqCst), expected_attempts);
|
||||
assert_eq!(marker.0.lock().unwrap().is_none(), success);
|
||||
}
|
||||
if let Some(mut socket) = socket {
|
||||
socket.close(None).await.unwrap();
|
||||
}
|
||||
drop(server);
|
||||
}
|
||||
}
|
||||
|
||||
struct ControlTransport {
|
||||
id: String,
|
||||
active: std::sync::atomic::AtomicBool,
|
||||
|
||||
@@ -8,6 +8,30 @@ use serde_json::json;
|
||||
|
||||
use super::decode_response_create;
|
||||
|
||||
#[test]
|
||||
fn client_retry_identity_should_follow_each_frame_instead_of_handshake_metadata() {
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert(
|
||||
"x-codex-turn-metadata",
|
||||
HeaderValue::from_static(
|
||||
r#"{"thread_id":"old-thread","turn_id":"old-turn","request_kind":"compaction"}"#,
|
||||
),
|
||||
);
|
||||
let headers = OpenAiRequestHeaders::from_headers(&headers);
|
||||
for turn in ["turn-one", "turn-two"] {
|
||||
let decoded = decode_response_create_with_context(&json!({
|
||||
"type":"response.create", "model":"smart-code", "input":"hello",
|
||||
"client_metadata": {"thread_id":"thread", "turn_id":turn, "x-codex-turn-metadata":r#"{"request_kind":"turn"}"#}
|
||||
}).to_string(), &headers).unwrap();
|
||||
let Operation::Generate(request) = decoded.operation() else {
|
||||
panic!("generate");
|
||||
};
|
||||
let identity = request.client_turn().unwrap();
|
||||
assert_eq!(identity.thread_id(), "thread");
|
||||
assert_eq!(identity.turn_id(), turn);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn response_create_should_carry_source_headers_to_provider_on_every_frame() {
|
||||
let mut headers = HeaderMap::new();
|
||||
|
||||
@@ -0,0 +1,124 @@
|
||||
//! 客户端失败轮次的短期重试抑制与终态前状态更新
|
||||
|
||||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use futures::{FutureExt, future::BoxFuture, pin_mut, select_biased};
|
||||
use futures_timer::Delay;
|
||||
|
||||
use super::ModelRequestId;
|
||||
use crate::{error::StoreError, operation::ClientTurn, policy::ClientApiKeyId};
|
||||
|
||||
const FAILURE_TTL: Duration = Duration::from_secs(300);
|
||||
const STORE_TIMEOUT: Duration = Duration::from_millis(100);
|
||||
|
||||
/// 不同客户端不能共享轮次失败状态
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ClientRetryKey {
|
||||
pub client_key_id: ClientApiKeyId,
|
||||
pub turn: ClientTurn,
|
||||
}
|
||||
|
||||
/// 可丢失的失败标记,版本使用创建该标记的模型请求 ID
|
||||
pub trait ClientRetryStore: Send + Sync {
|
||||
fn load<'a>(
|
||||
&'a self,
|
||||
key: &'a ClientRetryKey,
|
||||
) -> BoxFuture<'a, Result<Option<ModelRequestId>, StoreError>>;
|
||||
|
||||
fn record_failure<'a>(
|
||||
&'a self,
|
||||
key: &'a ClientRetryKey,
|
||||
version: &'a ModelRequestId,
|
||||
ttl: Duration,
|
||||
) -> BoxFuture<'a, Result<(), StoreError>>;
|
||||
|
||||
/// 比较与删除必须原子执行,迟到成功不能删除较新的失败
|
||||
fn clear_failure<'a>(
|
||||
&'a self,
|
||||
key: &'a ClientRetryKey,
|
||||
expected: &'a ModelRequestId,
|
||||
) -> BoxFuture<'a, Result<(), StoreError>>;
|
||||
}
|
||||
|
||||
pub(super) struct ClientRetryGuard {
|
||||
store: Arc<dyn ClientRetryStore>,
|
||||
key: ClientRetryKey,
|
||||
request_id: ModelRequestId,
|
||||
observed: Option<ModelRequestId>,
|
||||
settled: bool,
|
||||
}
|
||||
|
||||
impl ClientRetryGuard {
|
||||
pub(super) async fn load(
|
||||
store: Arc<dyn ClientRetryStore>,
|
||||
key: ClientRetryKey,
|
||||
request_id: ModelRequestId,
|
||||
) -> Self {
|
||||
let observed = bounded("load", &request_id, store.load(&key))
|
||||
.await
|
||||
.flatten();
|
||||
Self {
|
||||
store,
|
||||
key,
|
||||
request_id,
|
||||
observed,
|
||||
settled: false,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn has_failure(&self) -> bool {
|
||||
self.observed.is_some()
|
||||
}
|
||||
|
||||
pub(super) async fn failure(&mut self) {
|
||||
if self.settled {
|
||||
return;
|
||||
}
|
||||
bounded(
|
||||
"record",
|
||||
&self.request_id,
|
||||
self.store
|
||||
.record_failure(&self.key, &self.request_id, FAILURE_TTL),
|
||||
)
|
||||
.await;
|
||||
self.settled = true;
|
||||
}
|
||||
|
||||
pub(super) async fn success(&mut self) {
|
||||
if self.settled {
|
||||
return;
|
||||
}
|
||||
if let Some(observed) = &self.observed {
|
||||
bounded(
|
||||
"clear",
|
||||
&self.request_id,
|
||||
self.store.clear_failure(&self.key, observed),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
self.settled = true;
|
||||
}
|
||||
}
|
||||
|
||||
async fn bounded<T>(
|
||||
action: &'static str,
|
||||
request_id: &ModelRequestId,
|
||||
operation: BoxFuture<'_, Result<T, StoreError>>,
|
||||
) -> Option<T> {
|
||||
let operation = operation.fuse();
|
||||
let timeout = Delay::new(STORE_TIMEOUT).fuse();
|
||||
pin_mut!(operation, timeout);
|
||||
select_biased! {
|
||||
result = operation => match result {
|
||||
Ok(value) => Some(value),
|
||||
Err(error) => {
|
||||
tracing::warn!(action, request_id = request_id.as_str(), error_kind = ?error.kind(), "客户端重试标记不可用,保留原请求结果与重试策略");
|
||||
None
|
||||
}
|
||||
},
|
||||
() = timeout => {
|
||||
tracing::warn!(action, request_id = request_id.as_str(), "客户端重试标记操作超时,保留原请求结果与重试策略");
|
||||
None
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -9,6 +9,7 @@ use std::num::NonZeroU32;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant, SystemTime};
|
||||
|
||||
use super::client_retry::{ClientRetryGuard, ClientRetryKey};
|
||||
use super::nested::ExecutionEffects;
|
||||
use super::observation::{
|
||||
RequestObservationDispatch, ResponseObservation, WebSocketResponseAttempt,
|
||||
@@ -358,6 +359,7 @@ where
|
||||
last_retryable_failure_events: Vec::new(),
|
||||
last_observed_provider: None,
|
||||
pending_terminal_failure: None,
|
||||
client_retry: None,
|
||||
};
|
||||
|
||||
if session.cancellation.is_cancelled() {
|
||||
@@ -368,6 +370,26 @@ where
|
||||
session.finish_interruption(&EngineError::Deadline).await?;
|
||||
return Err(EngineError::Deadline);
|
||||
}
|
||||
if !session.account_selection.is_diagnostic()
|
||||
&& let Some(store) = &self.engine.client_retry
|
||||
&& let Operation::Generate(generate) = &session.operation
|
||||
&& let Some(turn) = generate.client_turn()
|
||||
{
|
||||
let guard = ClientRetryGuard::load(
|
||||
Arc::clone(store),
|
||||
ClientRetryKey {
|
||||
client_key_id: session.client_api_key_ref.clone(),
|
||||
turn: turn.clone(),
|
||||
},
|
||||
session.request_id.clone(),
|
||||
)
|
||||
.await;
|
||||
session.trace.record(
|
||||
"client_retry.loaded",
|
||||
json!({"previousFailure": guard.has_failure()}),
|
||||
);
|
||||
session.client_retry = Some(guard);
|
||||
}
|
||||
Ok(session)
|
||||
}
|
||||
}
|
||||
@@ -502,6 +524,7 @@ pub struct ResponseExecutionSession<S: ?Sized> {
|
||||
last_observed_provider: Option<crate::identity::ProviderKind>,
|
||||
/// 原子失败批次已交给协议层、但尚待下游提交后收敛的原 Provider 错误
|
||||
pending_terminal_failure: Option<PendingTerminalFailure>,
|
||||
client_retry: Option<ClientRetryGuard>,
|
||||
}
|
||||
|
||||
struct PendingTerminalFailure {
|
||||
@@ -575,6 +598,7 @@ where
|
||||
error,
|
||||
send_state,
|
||||
} => {
|
||||
self.record_client_retry_failure(&error).await;
|
||||
let requirement = if self.delivery_committed {
|
||||
CommitRequirement::AlreadyCommitted
|
||||
} else {
|
||||
@@ -912,7 +936,13 @@ where
|
||||
.translate_native_response(event, terminal);
|
||||
match translated {
|
||||
Ok(events) if events.is_empty() => continue,
|
||||
Ok(events) => return Ok(PullOutcome::Events(events)),
|
||||
Ok(events) => {
|
||||
// 成功终态交付后客户端即可续发工具结果,不能等流 EOF 或异步观测再清理
|
||||
if terminal && let Some(guard) = &mut self.client_retry {
|
||||
guard.success().await;
|
||||
}
|
||||
return Ok(PullOutcome::Events(events));
|
||||
}
|
||||
Err(error) => {
|
||||
self.finish_provider_error(&error).await?;
|
||||
return Err(provider_engine_error(error));
|
||||
@@ -1582,22 +1612,32 @@ where
|
||||
&& attempt_send_state != UpstreamSendState::Ambiguous
|
||||
&& provider_proved_replay_safe
|
||||
&& self.routing_attempts < self.plan.max_attempts().get();
|
||||
let suppress_transient_retry = is_return_error_retry(&error)
|
||||
&& self
|
||||
.client_retry
|
||||
.as_ref()
|
||||
.is_some_and(ClientRetryGuard::has_failure);
|
||||
let transient_retry = match error.pre_delivery_retry() {
|
||||
Some(crate::error::PreDeliveryRetry::SameAccountTransientRetry {
|
||||
max_retries,
|
||||
initial_delay,
|
||||
max_delay,
|
||||
}) if ordinary_retry => {
|
||||
jitter,
|
||||
..
|
||||
}) if ordinary_retry && !suppress_transient_retry => {
|
||||
let retries = self
|
||||
.transient_retry_counts
|
||||
.entry(current.metadata.provider_account_id().clone())
|
||||
.or_default();
|
||||
if *retries < max_retries.get() {
|
||||
let multiplier = 1_u32.checked_shl(*retries).unwrap_or(u32::MAX);
|
||||
// 服务器建议优先于本地退避,不能被本地上限缩短或再次指数放大
|
||||
let delay = error
|
||||
.retry_after()
|
||||
.unwrap_or_else(|| initial_delay.saturating_mul(multiplier).min(max_delay));
|
||||
// 服务器建议优先于本地退避,不截断、不放大,也不叠加随机等待
|
||||
let delay = error.retry_after().unwrap_or_else(|| {
|
||||
initial_delay
|
||||
.saturating_mul(multiplier)
|
||||
.min(max_delay)
|
||||
.saturating_add(jitter)
|
||||
});
|
||||
*retries = retries.saturating_add(1);
|
||||
Some(delay)
|
||||
} else {
|
||||
@@ -1620,9 +1660,11 @@ where
|
||||
// 同一排除分支。换号预算只在这里消耗;同账号分支(瞬态退避、传输恢复、
|
||||
// 凭据恢复重放、continuation 精确重连)不消耗。不能直接把预算门写进
|
||||
// ordinary_retry:它是同账号瞬态退避的前置条件,会被连带误伤
|
||||
// RetryThenReturnError 耗尽后直接交付错误,不能落入 ordinary_retry 的换号兜底
|
||||
let rotation_retry = !continuation_retry
|
||||
&& !same_account_retry
|
||||
&& transient_retry.is_none()
|
||||
&& !is_return_error_retry(&error)
|
||||
&& transport_recovery.is_none()
|
||||
&& (ordinary_retry || account_rotation_retry)
|
||||
&& self.account_rotations
|
||||
@@ -1668,6 +1710,7 @@ where
|
||||
"sameAccountRetry": same_account_retry, "accountRotationRetry": account_rotation_retry,
|
||||
"ordinaryRetry": ordinary_retry, "transportRecovery": transport_recovery.is_some(),
|
||||
"transientRetry": transient_retry.is_some(),
|
||||
"clientRetrySuppressed": suppress_transient_retry,
|
||||
"rotationRetry": rotation_retry, "accountRotations": self.account_rotations,
|
||||
"connectionRetry": connection_retry_requested,
|
||||
"connectionRetries": self.connection_retries,
|
||||
@@ -1948,6 +1991,9 @@ where
|
||||
if self.is_finalized() {
|
||||
return Ok(());
|
||||
}
|
||||
if let Some(guard) = &mut self.client_retry {
|
||||
guard.success().await;
|
||||
}
|
||||
let completed_at = SystemTime::now();
|
||||
self.observation.finish();
|
||||
let upstream_request_id = self
|
||||
@@ -2022,6 +2068,7 @@ where
|
||||
error: &ProviderError,
|
||||
send_state: UpstreamSendState,
|
||||
) -> Result<(), EngineError> {
|
||||
self.record_client_retry_failure(error).await;
|
||||
let send_state = self.raise_send_watermark(send_state);
|
||||
let outcome = if error.kind() == ProviderErrorKind::Cancelled {
|
||||
ExecutionOutcome::Cancelled
|
||||
@@ -2049,6 +2096,16 @@ where
|
||||
.await
|
||||
}
|
||||
|
||||
async fn record_client_retry_failure(&mut self, error: &ProviderError) {
|
||||
// 只记录本策略的终态失败,不让其他错误改变客户端轮次的退避预算
|
||||
if self.attempts > 0
|
||||
&& is_return_error_retry(error)
|
||||
&& let Some(guard) = &mut self.client_retry
|
||||
{
|
||||
guard.failure().await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn finish_interruption(&mut self, error: &EngineError) -> Result<(), EngineError> {
|
||||
// 结算 I/O 可能继续等待,先冻结交付事实以保持终态与实际 HTTP 状态一致
|
||||
self.response_commit.finish();
|
||||
@@ -2343,6 +2400,16 @@ where
|
||||
}
|
||||
}
|
||||
|
||||
fn is_return_error_retry(error: &ProviderError) -> bool {
|
||||
matches!(
|
||||
error.pre_delivery_retry(),
|
||||
Some(crate::error::PreDeliveryRetry::SameAccountTransientRetry {
|
||||
policy: crate::error::TransientRetryPolicy::RetryThenReturnError,
|
||||
..
|
||||
})
|
||||
)
|
||||
}
|
||||
|
||||
async fn best_effort_store_write<T>(
|
||||
operation: &'static str,
|
||||
request_id: &ModelRequestId,
|
||||
|
||||
@@ -92,6 +92,7 @@ pub struct DefaultExecutionService {
|
||||
execution_extensions: Option<ExecutionExtensionIndex>,
|
||||
frontend_authentication: Option<FrontendAuthenticationExtensionIndex>,
|
||||
session_affinity: Option<Arc<dyn ProviderSessionAffinityPort>>,
|
||||
client_retry: Option<Arc<dyn crate::engine::client_retry::ClientRetryStore>>,
|
||||
snapshot_refresh: Option<Arc<RuntimeSnapshotPublisher>>,
|
||||
active_requests: ActiveRequestRegistry,
|
||||
}
|
||||
@@ -122,6 +123,7 @@ impl DefaultExecutionService {
|
||||
execution_extensions: None,
|
||||
frontend_authentication: None,
|
||||
session_affinity: None,
|
||||
client_retry: None,
|
||||
snapshot_refresh: None,
|
||||
active_requests: Arc::default(),
|
||||
}
|
||||
@@ -133,6 +135,15 @@ impl DefaultExecutionService {
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_client_retry_store(
|
||||
mut self,
|
||||
store: Arc<dyn crate::engine::client_retry::ClientRetryStore>,
|
||||
) -> Self {
|
||||
self.client_retry = Some(store);
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_request_observers(mut self, observers: RequestObserverExtensionIndex) -> Self {
|
||||
self.request_observers = Some(observers);
|
||||
@@ -829,8 +840,15 @@ impl DefaultExecutionService {
|
||||
started_at,
|
||||
deadline_at: authorization.deadline_at,
|
||||
};
|
||||
let coordinator =
|
||||
AttemptCoordinator::new(GatewayEngine::new(self.observations.clone(), providers));
|
||||
let mut engine = GatewayEngine::new(self.observations.clone(), providers);
|
||||
// 插件内部调用不共享外部客户端的失败轮次,探测另走 transient engine
|
||||
if authorization.nested.is_none()
|
||||
&& authorization.bound.is_none()
|
||||
&& let Some(store) = &self.client_retry
|
||||
{
|
||||
engine = engine.with_client_retry_store(Arc::clone(store));
|
||||
}
|
||||
let coordinator = AttemptCoordinator::new(engine);
|
||||
let core = match coordinator
|
||||
.start_observed(
|
||||
new_request,
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
pub mod admission;
|
||||
pub mod authentication;
|
||||
pub mod budget;
|
||||
pub mod client_retry;
|
||||
pub mod connection;
|
||||
pub mod continuation;
|
||||
pub mod coordinator;
|
||||
@@ -1121,12 +1122,26 @@ pub enum EngineError {
|
||||
pub struct GatewayEngine<S: ?Sized> {
|
||||
store: Arc<S>,
|
||||
providers: provider::ProviderRegistry,
|
||||
client_retry: Option<Arc<dyn client_retry::ClientRetryStore>>,
|
||||
}
|
||||
|
||||
impl<S: ?Sized> GatewayEngine<S> {
|
||||
#[must_use]
|
||||
pub const fn new(store: Arc<S>, providers: provider::ProviderRegistry) -> Self {
|
||||
Self { store, providers }
|
||||
Self {
|
||||
store,
|
||||
providers,
|
||||
client_retry: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_client_retry_store(
|
||||
mut self,
|
||||
store: Arc<dyn client_retry::ClientRetryStore>,
|
||||
) -> Self {
|
||||
self.client_retry = Some(store);
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
|
||||
@@ -638,7 +638,8 @@ fn provider_event_to_middleware_frame(
|
||||
|
||||
fn encode_sse_wire_event(wire: &ProtocolWireEvent) -> Result<Bytes, ProviderError> {
|
||||
let mut encoded = Vec::new();
|
||||
if let Some(id) = wire.sse_id() {
|
||||
// 源事件元数据保持不透明;仅编码可表示为单行 SSE 字段的值,避免生成额外帧
|
||||
if let Some(id) = wire.sse_id().filter(|id| !id.contains(['\r', '\n', '\0'])) {
|
||||
encoded.extend_from_slice(b"id: ");
|
||||
encoded.extend_from_slice(id.as_bytes());
|
||||
encoded.push(b'\n');
|
||||
@@ -646,7 +647,10 @@ fn encode_sse_wire_event(wire: &ProtocolWireEvent) -> Result<Bytes, ProviderErro
|
||||
if let Some(retry) = wire.sse_retry() {
|
||||
encoded.extend_from_slice(format!("retry: {retry}\n").as_bytes());
|
||||
}
|
||||
if let Some(event_type) = wire.event_type() {
|
||||
if let Some(event_type) = wire
|
||||
.event_type()
|
||||
.filter(|event_type| !event_type.is_empty() && !event_type.contains(['\r', '\n']))
|
||||
{
|
||||
encoded.extend_from_slice(b"event: ");
|
||||
encoded.extend_from_slice(event_type.as_bytes());
|
||||
encoded.push(b'\n');
|
||||
|
||||
@@ -120,19 +120,34 @@ pub enum ContinuationRecoveryDisposition {
|
||||
ProviderReplayAllowed,
|
||||
}
|
||||
|
||||
/// 同账号瞬态错误的恢复策略,明确重试耗尽后的去向
|
||||
///
|
||||
/// Provider 按上游错误场景选择,Core 共用同一套退避与安全重放检查
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
||||
pub enum TransientRetryPolicy {
|
||||
/// 同账号退避重试,耗尽后允许在安全边界与换号预算内选择其他账号
|
||||
RetryThenRotateAccount,
|
||||
/// 同账号退避重试,耗尽后返回最后一次上游错误,不因该错误换号
|
||||
RetryThenReturnError,
|
||||
}
|
||||
|
||||
/// Provider 在下游尚未提交事件时请求的隐藏恢复方式
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
||||
pub enum PreDeliveryRetry {
|
||||
/// 排除本次账号,按普通调度策略重新选号
|
||||
AccountRotation,
|
||||
/// 上游明确拒绝后的同账号退避;耗尽后仍由普通安全重放规则决定是否换号
|
||||
/// 上游明确拒绝后的同账号退避;耗尽动作由 Provider 显式选择
|
||||
SameAccountTransientRetry {
|
||||
/// 每个账号在本请求中的重试上限,同时受总路由预算约束
|
||||
max_retries: NonZeroU32,
|
||||
/// 指数退避的起始间隔
|
||||
initial_delay: Duration,
|
||||
/// 单次等待上限,避免上游提示把请求拖入长时间等待
|
||||
/// 本地指数退避的基准上限,不含附加抖动,也不缩短上游 Retry-After
|
||||
max_delay: Duration,
|
||||
/// Provider 为本次失败采样的附加等待,仅在没有 Retry-After 时使用
|
||||
jitter: Duration,
|
||||
/// 同账号恢复策略,不提供额外安全重放资格
|
||||
policy: TransientRetryPolicy,
|
||||
},
|
||||
/// 固定本次账号,并按 Provider 给出的序号重试当前传输
|
||||
/// 可用于发送结果不明确时的有界尽力恢复,不等于上游未执行的证明
|
||||
@@ -682,18 +697,22 @@ impl ProviderError {
|
||||
self
|
||||
}
|
||||
|
||||
/// 请求有界同账号退避;该意图本身不提供安全重放证据
|
||||
/// 请求有界同账号退避并显式选择耗尽去向;该意图本身不提供安全重放证据
|
||||
#[must_use]
|
||||
pub fn with_transient_retry(
|
||||
mut self,
|
||||
max_retries: NonZeroU32,
|
||||
initial_delay: Duration,
|
||||
max_delay: Duration,
|
||||
jitter: Duration,
|
||||
policy: TransientRetryPolicy,
|
||||
) -> Self {
|
||||
self.pre_delivery_retry = Some(Box::new(PreDeliveryRetry::SameAccountTransientRetry {
|
||||
max_retries,
|
||||
initial_delay,
|
||||
max_delay,
|
||||
jitter,
|
||||
policy,
|
||||
}));
|
||||
self
|
||||
}
|
||||
@@ -1048,6 +1067,8 @@ pub enum GatewayErrorKind {
|
||||
NoAvailableProvider,
|
||||
/// 符合条件的上游账号暂时没有调度容量
|
||||
AccountCapacityUnavailable,
|
||||
/// 中间件在途资源或执行等待队列已满
|
||||
MiddlewareCapacityUnavailable,
|
||||
ConcurrencyQueueFull,
|
||||
ConcurrencyQueueTimeout,
|
||||
/// Provider 的本地账号基础设施不可用
|
||||
@@ -1079,6 +1100,7 @@ impl GatewayErrorKind {
|
||||
Self::ModelNotFound => "model_not_found",
|
||||
Self::NoAvailableProvider => "no_available_provider",
|
||||
Self::AccountCapacityUnavailable => "account_capacity_unavailable",
|
||||
Self::MiddlewareCapacityUnavailable => "middleware_capacity_unavailable",
|
||||
Self::ConcurrencyQueueFull => "concurrency_queue_full",
|
||||
Self::ConcurrencyQueueTimeout => "concurrency_queue_timeout",
|
||||
Self::ProviderInfrastructureUnavailable => "provider_infrastructure_unavailable",
|
||||
|
||||
@@ -62,6 +62,7 @@ pub struct CoreStorePorts {
|
||||
client_api_key_usage: Arc<dyn ClientApiKeyUsageSink>,
|
||||
budget: Option<Arc<dyn engine::budget::ClientBudgetPort>>,
|
||||
session_affinity: Option<Arc<dyn ProviderSessionAffinityPort>>,
|
||||
client_retry: Option<Arc<dyn engine::client_retry::ClientRetryStore>>,
|
||||
}
|
||||
|
||||
impl CoreStorePorts {
|
||||
@@ -91,6 +92,7 @@ impl CoreStorePorts {
|
||||
client_api_key_usage,
|
||||
budget: None,
|
||||
session_affinity: None,
|
||||
client_retry: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -100,6 +102,15 @@ impl CoreStorePorts {
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_client_retry_store(
|
||||
mut self,
|
||||
store: Arc<dyn engine::client_retry::ClientRetryStore>,
|
||||
) -> Self {
|
||||
self.client_retry = Some(store);
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_session_affinity(
|
||||
mut self,
|
||||
@@ -209,6 +220,9 @@ fn build_execution_service(
|
||||
if let Some(budget) = &ports.budget {
|
||||
service = service.with_budget(Arc::clone(budget));
|
||||
}
|
||||
if let Some(store) = &ports.client_retry {
|
||||
service = service.with_client_retry_store(Arc::clone(store));
|
||||
}
|
||||
if let Some(request_observers) = request_observers {
|
||||
service = service.with_request_observers(request_observers);
|
||||
}
|
||||
|
||||
@@ -404,11 +404,38 @@ pub struct GenerateRequest {
|
||||
#[derive(Clone, PartialEq)]
|
||||
struct GeneratePayload {
|
||||
protocol_payload: ProtocolPayload,
|
||||
client_turn: Option<ClientTurn>,
|
||||
provider_session_state: Option<ProviderSessionState>,
|
||||
source_requirements: Option<CapabilityRequirements>,
|
||||
inherited_capability_requirements: Option<CapabilityRequirements>,
|
||||
}
|
||||
|
||||
/// 协议入口确认的客户端推理轮次,不使用连接路由键猜测身份
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ClientTurn {
|
||||
thread_id: String,
|
||||
turn_id: String,
|
||||
}
|
||||
|
||||
impl ClientTurn {
|
||||
pub fn new(thread_id: String, turn_id: String) -> Option<Self> {
|
||||
if [thread_id.as_str(), turn_id.as_str()].iter().any(|value| {
|
||||
value.trim().is_empty() || value.len() > 256 || value.chars().any(char::is_control)
|
||||
}) {
|
||||
return None;
|
||||
}
|
||||
Some(Self { thread_id, turn_id })
|
||||
}
|
||||
|
||||
pub fn thread_id(&self) -> &str {
|
||||
&self.thread_id
|
||||
}
|
||||
|
||||
pub fn turn_id(&self) -> &str {
|
||||
&self.turn_id
|
||||
}
|
||||
}
|
||||
|
||||
impl GenerateRequest {
|
||||
/// 创建携带协议不透明正文的生成请求
|
||||
///
|
||||
@@ -419,6 +446,7 @@ impl GenerateRequest {
|
||||
Self {
|
||||
payload: Arc::new(GeneratePayload {
|
||||
protocol_payload,
|
||||
client_turn: None,
|
||||
provider_session_state: None,
|
||||
source_requirements: None,
|
||||
inherited_capability_requirements: None,
|
||||
@@ -426,27 +454,14 @@ impl GenerateRequest {
|
||||
}
|
||||
}
|
||||
|
||||
/// 写回原生 Provider 编码结果,保留已准入的能力、连接上下文和会话状态
|
||||
///
|
||||
/// 原生编码不消耗中间件的一次跨协议改写额度:可以保留协议名,也可以在一次
|
||||
/// 中间件改写后编码到 Provider wire
|
||||
/// 调用方仍须验证保护字段没有被改动
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// 目标协议名无效时返回错误
|
||||
pub fn with_native_encoded_body(
|
||||
mut self,
|
||||
protocol: impl Into<String>,
|
||||
body: Map<String, Value>,
|
||||
) -> Result<Self, OperationError> {
|
||||
let encoded = ProtocolPayload::json_object(protocol, body)?
|
||||
.with_context(self.protocol_payload().context().clone());
|
||||
let requirements = self.requirements();
|
||||
let payload = Arc::make_mut(&mut self.payload);
|
||||
payload.protocol_payload = encoded;
|
||||
payload.source_requirements = Some(requirements);
|
||||
Ok(self)
|
||||
#[must_use]
|
||||
pub fn with_client_turn(mut self, client_turn: Option<ClientTurn>) -> Self {
|
||||
Arc::make_mut(&mut self.payload).client_turn = client_turn;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn client_turn(&self) -> Option<&ClientTurn> {
|
||||
self.payload.client_turn.as_ref()
|
||||
}
|
||||
|
||||
/// 附着同一客户端连接上一轮由 Provider 返回的不透明状态
|
||||
@@ -951,41 +966,6 @@ impl Operation {
|
||||
self.replace_protocol_wire_inner(None, body, context)
|
||||
}
|
||||
|
||||
/// 替换 Provider HTTP 正文、非 wire 上下文和已解析 headers,同时保持 endpoint、
|
||||
/// method 与 query 不变
|
||||
/// headers 继续由 [`ProviderHttpRequest::new`] 统一验证;
|
||||
/// Core 不解析插件传输使用的 base64 表示
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// 当前 operation 不是 Provider HTTP,或替换后的 headers 超出请求边界时返回错误
|
||||
pub fn replace_provider_http_wire(
|
||||
self,
|
||||
body: Option<Bytes>,
|
||||
context: Map<String, Value>,
|
||||
headers: Vec<ProviderHttpHeader>,
|
||||
) -> Result<Self, OperationError> {
|
||||
let Self::ProviderHttp(request) = self else {
|
||||
return Err(OperationError::EmptyField {
|
||||
field: "provider_http_request",
|
||||
});
|
||||
};
|
||||
let payload = request.payload;
|
||||
ProviderHttpRequest::new(
|
||||
request.endpoint,
|
||||
request.method,
|
||||
request.query,
|
||||
headers,
|
||||
RawHttpPayload {
|
||||
protocol: payload.protocol,
|
||||
body: body.unwrap_or(payload.body),
|
||||
context,
|
||||
translated: payload.translated,
|
||||
},
|
||||
)
|
||||
.map(Self::ProviderHttp)
|
||||
}
|
||||
|
||||
/// 把一次已选择的协议转换结果写回 operation,同时保留非 wire 上下文和
|
||||
/// Provider 会话状态
|
||||
/// 该方法只验证公共载荷形状,不解释目标协议字段
|
||||
|
||||
@@ -56,6 +56,12 @@ pub struct PrivacyRule {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub enabled: bool,
|
||||
pub steps: Vec<PrivacyStep>,
|
||||
}
|
||||
|
||||
#[derive(Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
pub struct PrivacyStep {
|
||||
pub scope: PrivacyScope,
|
||||
pub selector: String,
|
||||
pub action: PrivacyAction,
|
||||
@@ -71,9 +77,17 @@ pub struct PrivacyRule {
|
||||
impl fmt::Debug for PrivacyRule {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("PrivacyRule")
|
||||
.field("enabled", &self.enabled)
|
||||
.field("step_count", &self.steps.len())
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for PrivacyStep {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("PrivacyStep")
|
||||
.field("scope", &self.scope)
|
||||
.field("action", &self.action)
|
||||
.field("enabled", &self.enabled)
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -16,6 +16,7 @@ use futures::executor::block_on;
|
||||
|
||||
use gateway_core::account::FastMode;
|
||||
use gateway_core::account::{AccountSelectionPolicy, ProviderAccountId, RotationStrategy};
|
||||
use gateway_core::engine::client_retry::{ClientRetryKey, ClientRetryStore};
|
||||
use gateway_core::engine::continuation::{
|
||||
ContinuationBinding, NativeContinuationPin, PreviousResponseId,
|
||||
};
|
||||
@@ -31,7 +32,7 @@ use gateway_core::engine::{
|
||||
use gateway_core::error::{
|
||||
ClientVisibleUpstreamError, ClientVisibleUpstreamResponse, ContinuationFailure,
|
||||
ContinuationRecoveryDisposition, GatewayErrorKind, OpaqueUpstreamValue, ProviderError,
|
||||
ProviderErrorKind, StoreError, StoreErrorKind,
|
||||
ProviderErrorKind, StoreError, StoreErrorKind, TransientRetryPolicy,
|
||||
};
|
||||
use gateway_core::event::{
|
||||
ContentItem, ContentKind, GatewayEvent, ProtocolWireEvent, ProviderEvent,
|
||||
@@ -40,6 +41,7 @@ use gateway_core::event::{
|
||||
};
|
||||
use gateway_core::lifecycle::CancellationToken;
|
||||
use gateway_core::metering::{CalculatedCost, CostSource, ProviderReportedCost, Usage};
|
||||
use gateway_core::operation::ClientTurn;
|
||||
use gateway_core::operation::{GenerateRequest, Operation, ProtocolPayload, ProviderSessionState};
|
||||
use gateway_core::policy::ClientApiKeyId;
|
||||
use gateway_core::routing::{
|
||||
@@ -51,6 +53,407 @@ use gateway_core::routing::{
|
||||
use gateway_core::upstream::{UpstreamSendState, UpstreamTransport};
|
||||
use serde_json::{Map, Value, json};
|
||||
|
||||
#[derive(Default)]
|
||||
struct RetryMarkers {
|
||||
entries: Mutex<BTreeMap<String, ModelRequestId>>,
|
||||
unavailable: bool,
|
||||
stalled: bool,
|
||||
}
|
||||
|
||||
fn retry_marker_key(key: &ClientRetryKey) -> String {
|
||||
format!(
|
||||
"{}/{}/{}",
|
||||
key.client_key_id.as_str(),
|
||||
key.turn.thread_id(),
|
||||
key.turn.turn_id()
|
||||
)
|
||||
}
|
||||
|
||||
impl ClientRetryStore for RetryMarkers {
|
||||
fn load<'a>(
|
||||
&'a self,
|
||||
key: &'a ClientRetryKey,
|
||||
) -> futures::future::BoxFuture<'a, Result<Option<ModelRequestId>, StoreError>> {
|
||||
Box::pin(async move {
|
||||
if self.stalled {
|
||||
return futures::future::pending().await;
|
||||
}
|
||||
if self.unavailable {
|
||||
return Err(StoreError::new(StoreErrorKind::Unavailable));
|
||||
}
|
||||
Ok(self
|
||||
.entries
|
||||
.lock()
|
||||
.unwrap()
|
||||
.get(&retry_marker_key(key))
|
||||
.cloned())
|
||||
})
|
||||
}
|
||||
|
||||
fn record_failure<'a>(
|
||||
&'a self,
|
||||
key: &'a ClientRetryKey,
|
||||
version: &'a ModelRequestId,
|
||||
ttl: Duration,
|
||||
) -> futures::future::BoxFuture<'a, Result<(), StoreError>> {
|
||||
Box::pin(async move {
|
||||
assert_eq!(ttl, Duration::from_secs(300));
|
||||
if self.stalled {
|
||||
return futures::future::pending().await;
|
||||
}
|
||||
if self.unavailable {
|
||||
return Err(StoreError::new(StoreErrorKind::Unavailable));
|
||||
}
|
||||
self.entries
|
||||
.lock()
|
||||
.unwrap()
|
||||
.insert(retry_marker_key(key), version.clone());
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
|
||||
fn clear_failure<'a>(
|
||||
&'a self,
|
||||
key: &'a ClientRetryKey,
|
||||
expected: &'a ModelRequestId,
|
||||
) -> futures::future::BoxFuture<'a, Result<(), StoreError>> {
|
||||
Box::pin(async move {
|
||||
if self.stalled {
|
||||
return futures::future::pending().await;
|
||||
}
|
||||
if self.unavailable {
|
||||
return Err(StoreError::new(StoreErrorKind::Unavailable));
|
||||
}
|
||||
let mut entries = self.entries.lock().unwrap();
|
||||
if entries.get(&retry_marker_key(key)) == Some(expected) {
|
||||
entries.remove(&retry_marker_key(key));
|
||||
}
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn retry_turn_operation(thread: &str, turn: &str) -> Operation {
|
||||
let Operation::Generate(request) = generate_operation() else {
|
||||
unreachable!()
|
||||
};
|
||||
Operation::Generate(
|
||||
request.with_client_turn(ClientTurn::new(thread.to_owned(), turn.to_owned())),
|
||||
)
|
||||
}
|
||||
|
||||
fn retry_failure() -> Script {
|
||||
Script::Stream {
|
||||
account_id: "acct_first",
|
||||
items: vec![Err(ProviderError::new(
|
||||
ProviderErrorKind::Unavailable,
|
||||
UpstreamSendState::Sent,
|
||||
)
|
||||
.with_status(503)
|
||||
.with_replay_safe()
|
||||
.with_transient_retry(
|
||||
NonZeroU32::new(5).unwrap(),
|
||||
Duration::ZERO,
|
||||
Duration::ZERO,
|
||||
Duration::ZERO,
|
||||
TransientRetryPolicy::RetryThenReturnError,
|
||||
))],
|
||||
}
|
||||
}
|
||||
|
||||
fn retry_coordinator(
|
||||
scripts: Vec<Script>,
|
||||
markers: Arc<RetryMarkers>,
|
||||
) -> (AttemptCoordinator<FakeStore>, Arc<ScriptedProvider>) {
|
||||
let provider = Arc::new(ScriptedProvider::new(scripts));
|
||||
let mut registry = ProviderRegistry::builder();
|
||||
registry.register(provider.clone()).unwrap();
|
||||
let engine = GatewayEngine::new(Arc::new(FakeStore::default()), registry.build())
|
||||
.with_client_retry_store(markers);
|
||||
(AttemptCoordinator::new(engine), provider)
|
||||
}
|
||||
|
||||
fn retry_session(
|
||||
coordinator: &AttemptCoordinator<FakeStore>,
|
||||
id: &str,
|
||||
operation: Operation,
|
||||
client: &str,
|
||||
) -> gateway_core::engine::ResponseExecutionSession<FakeStore> {
|
||||
let mut request = model_request(&operation, SystemTime::now() + Duration::from_secs(30));
|
||||
request.id = ModelRequestId::new(id).unwrap();
|
||||
request.client_api_key_ref = ClientApiKeyId::new(client).unwrap();
|
||||
block_on(coordinator.start(
|
||||
request,
|
||||
operation.clone(),
|
||||
plan(&operation),
|
||||
None,
|
||||
None,
|
||||
CancellationToken::new(),
|
||||
))
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn client_retry_should_skip_repeated_budget_and_clear_before_completed_delivery() {
|
||||
let markers = Arc::new(RetryMarkers::default());
|
||||
let mut scripts = (0..7).map(|_| retry_failure()).collect::<Vec<_>>();
|
||||
scripts.push(Script::HangingStream {
|
||||
account_id: "acct_first",
|
||||
items: complete_stream(None),
|
||||
});
|
||||
scripts.extend([
|
||||
retry_failure(),
|
||||
retry_failure(),
|
||||
Script::Stream {
|
||||
account_id: "acct_first",
|
||||
items: complete_stream(None),
|
||||
},
|
||||
]);
|
||||
let (coordinator, provider) = retry_coordinator(scripts, markers.clone());
|
||||
let operation = retry_turn_operation("thread", "turn");
|
||||
for (id, expected) in [("req_initial", 6), ("req_retry", 7)] {
|
||||
let mut session = retry_session(&coordinator, id, operation.clone(), "key_client_1");
|
||||
assert!(block_on(session.collect_uncommitted()).is_err());
|
||||
assert_eq!(provider.contexts.lock().unwrap().len(), expected);
|
||||
}
|
||||
let mut recovered = retry_session(
|
||||
&coordinator,
|
||||
"req_recovered",
|
||||
operation.clone(),
|
||||
"key_client_1",
|
||||
);
|
||||
block_on(recovered.next_event()).unwrap().unwrap();
|
||||
block_on(recovered.commit_downstream(Some(200))).unwrap();
|
||||
let completed = block_on(recovered.next_event()).unwrap().unwrap();
|
||||
assert!(completed.into_provider_events().iter().any(|event| {
|
||||
event
|
||||
.canonical_facts()
|
||||
.iter()
|
||||
.any(|fact| matches!(fact, GatewayEvent::Completed(_)))
|
||||
}));
|
||||
assert!(markers.entries.lock().unwrap().is_empty());
|
||||
assert!(!recovered.is_finalized());
|
||||
// 输入可以随工具结果改变,正常续请求沿用同轮次但重新获得网关预算
|
||||
let mut followup = retry_session(&coordinator, "req_followup", operation, "key_client_1");
|
||||
block_on(followup.collect_uncommitted()).unwrap();
|
||||
block_on(followup.commit_downstream(Some(200))).unwrap();
|
||||
assert_eq!(provider.contexts.lock().unwrap().len(), 11);
|
||||
block_on(recovered.cancel_and_finalize()).unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn client_retry_should_record_atomic_failure_before_delivering_its_events() {
|
||||
let markers = Arc::new(RetryMarkers::default());
|
||||
let (coordinator, provider) = retry_coordinator(
|
||||
vec![
|
||||
Script::Stream {
|
||||
account_id: "acct_first",
|
||||
items: vec![
|
||||
Ok(GatewayEvent::Started(ResponseMeta::new("resp", "gpt-5"))),
|
||||
Err(
|
||||
atomic_response_failed("resp", "failed").with_transient_retry(
|
||||
NonZeroU32::new(5).unwrap(),
|
||||
Duration::ZERO,
|
||||
Duration::ZERO,
|
||||
Duration::ZERO,
|
||||
TransientRetryPolicy::RetryThenReturnError,
|
||||
),
|
||||
),
|
||||
],
|
||||
},
|
||||
retry_failure(),
|
||||
],
|
||||
markers.clone(),
|
||||
);
|
||||
let operation = retry_turn_operation("thread", "turn");
|
||||
let mut first = retry_session(
|
||||
&coordinator,
|
||||
"req_atomic",
|
||||
operation.clone(),
|
||||
"key_client_1",
|
||||
);
|
||||
block_on(first.next_event()).unwrap().unwrap();
|
||||
block_on(first.commit_downstream(Some(200))).unwrap();
|
||||
block_on(first.next_event()).unwrap().unwrap();
|
||||
assert_eq!(
|
||||
markers
|
||||
.entries
|
||||
.lock()
|
||||
.unwrap()
|
||||
.values()
|
||||
.next()
|
||||
.unwrap()
|
||||
.as_str(),
|
||||
"req_atomic"
|
||||
);
|
||||
// 不驱动旧请求 EOF,新请求此时已经必须看到失败标记
|
||||
let mut retry = retry_session(&coordinator, "req_atomic_retry", operation, "key_client_1");
|
||||
assert!(block_on(retry.collect_uncommitted()).is_err());
|
||||
assert_eq!(provider.contexts.lock().unwrap().len(), 2);
|
||||
block_on(first.cancel_and_finalize()).unwrap();
|
||||
assert_eq!(
|
||||
markers
|
||||
.entries
|
||||
.lock()
|
||||
.unwrap()
|
||||
.values()
|
||||
.next()
|
||||
.unwrap()
|
||||
.as_str(),
|
||||
"req_atomic_retry"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn client_retry_late_success_should_keep_a_newer_failure() {
|
||||
let markers = Arc::new(RetryMarkers::default());
|
||||
markers.entries.lock().unwrap().insert(
|
||||
"key_client_1/thread/turn".to_owned(),
|
||||
ModelRequestId::new("req_old").unwrap(),
|
||||
);
|
||||
let (coordinator, _) = retry_coordinator(
|
||||
vec![
|
||||
Script::HangingStream {
|
||||
account_id: "acct_first",
|
||||
items: complete_stream(None),
|
||||
},
|
||||
retry_failure(),
|
||||
],
|
||||
markers.clone(),
|
||||
);
|
||||
let operation = retry_turn_operation("thread", "turn");
|
||||
let mut slow = retry_session(&coordinator, "req_slow", operation.clone(), "key_client_1");
|
||||
block_on(slow.next_event()).unwrap().unwrap();
|
||||
block_on(slow.commit_downstream(Some(200))).unwrap();
|
||||
let mut newer = retry_session(&coordinator, "req_newer", operation, "key_client_1");
|
||||
assert!(block_on(newer.collect_uncommitted()).is_err());
|
||||
block_on(slow.next_event()).unwrap().unwrap();
|
||||
assert_eq!(
|
||||
markers
|
||||
.entries
|
||||
.lock()
|
||||
.unwrap()
|
||||
.values()
|
||||
.next()
|
||||
.unwrap()
|
||||
.as_str(),
|
||||
"req_newer"
|
||||
);
|
||||
block_on(slow.cancel_and_finalize()).unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn client_retry_should_isolate_client_thread_and_turn_and_ignore_missing_identity() {
|
||||
for (client, operation) in [
|
||||
("key_other", retry_turn_operation("thread", "turn")),
|
||||
("key_client_1", retry_turn_operation("other", "turn")),
|
||||
("key_client_1", retry_turn_operation("thread", "other")),
|
||||
("key_client_1", generate_operation()),
|
||||
] {
|
||||
let markers = Arc::new(RetryMarkers::default());
|
||||
markers.entries.lock().unwrap().insert(
|
||||
"key_client_1/thread/turn".to_owned(),
|
||||
ModelRequestId::new("req_old").unwrap(),
|
||||
);
|
||||
let (coordinator, provider) = retry_coordinator(
|
||||
vec![
|
||||
retry_failure(),
|
||||
Script::Stream {
|
||||
account_id: "acct_first",
|
||||
items: complete_stream(None),
|
||||
},
|
||||
],
|
||||
markers.clone(),
|
||||
);
|
||||
let mut session = retry_session(&coordinator, "req_distinct", operation, client);
|
||||
block_on(session.collect_uncommitted()).unwrap();
|
||||
block_on(session.commit_downstream(Some(200))).unwrap();
|
||||
assert_eq!(provider.contexts.lock().unwrap().len(), 2);
|
||||
assert_eq!(markers.entries.lock().unwrap().len(), 1);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn client_retry_store_failure_or_timeout_should_preserve_original_retries_and_error() {
|
||||
for stalled in [false, true] {
|
||||
let markers = Arc::new(RetryMarkers {
|
||||
unavailable: !stalled,
|
||||
stalled,
|
||||
..Default::default()
|
||||
});
|
||||
let (coordinator, provider) =
|
||||
retry_coordinator((0..6).map(|_| retry_failure()).collect(), markers);
|
||||
let started = std::time::Instant::now();
|
||||
let mut session = retry_session(
|
||||
&coordinator,
|
||||
"req_unavailable",
|
||||
retry_turn_operation("thread", "turn"),
|
||||
"key_client_1",
|
||||
);
|
||||
let Err(EngineError::Provider(error)) = block_on(session.collect_uncommitted()) else {
|
||||
panic!("original upstream failure");
|
||||
};
|
||||
assert_eq!(error.upstream_status(), Some(503));
|
||||
assert_eq!(provider.contexts.lock().unwrap().len(), 6);
|
||||
assert!(started.elapsed() < Duration::from_secs(2));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn client_retry_should_preserve_other_retry_policies_and_not_mark_their_failures() {
|
||||
for seeded in [false, true] {
|
||||
let markers = Arc::new(RetryMarkers::default());
|
||||
if seeded {
|
||||
markers.entries.lock().unwrap().insert(
|
||||
"key_client_1/thread/turn".to_owned(),
|
||||
ModelRequestId::new("req_old").unwrap(),
|
||||
);
|
||||
}
|
||||
let failure = || {
|
||||
ProviderError::new(ProviderErrorKind::RateLimited, UpstreamSendState::Sent)
|
||||
.with_status(429)
|
||||
.with_replay_safe()
|
||||
.with_transient_retry(
|
||||
NonZeroU32::new(3).unwrap(),
|
||||
Duration::ZERO,
|
||||
Duration::ZERO,
|
||||
Duration::ZERO,
|
||||
TransientRetryPolicy::RetryThenRotateAccount,
|
||||
)
|
||||
};
|
||||
let scripts = vec![
|
||||
Script::Stream {
|
||||
account_id: "acct_first",
|
||||
items: vec![Err(failure())],
|
||||
},
|
||||
// 已提交输出后的失败直接交回客户端,不应创建或刷新本策略标记
|
||||
Script::Stream {
|
||||
account_id: "acct_first",
|
||||
items: vec![
|
||||
Ok(GatewayEvent::Started(ResponseMeta::new("resp", "gpt-5"))),
|
||||
Err(failure()),
|
||||
],
|
||||
},
|
||||
];
|
||||
let (coordinator, provider) = retry_coordinator(scripts, markers.clone());
|
||||
let mut session = retry_session(
|
||||
&coordinator,
|
||||
"req_other_policy",
|
||||
retry_turn_operation("thread", "turn"),
|
||||
"key_client_1",
|
||||
);
|
||||
block_on(session.next_event()).unwrap().unwrap();
|
||||
block_on(session.commit_downstream(Some(200))).unwrap();
|
||||
assert!(block_on(session.next_event()).is_err());
|
||||
assert_eq!(provider.contexts.lock().unwrap().len(), 2);
|
||||
let entries = markers.entries.lock().unwrap();
|
||||
assert_eq!(entries.len(), usize::from(seeded));
|
||||
if seeded {
|
||||
assert_eq!(entries.values().next().unwrap().as_str(), "req_old");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
struct FinalState {
|
||||
error_message: Option<String>,
|
||||
@@ -3668,6 +4071,8 @@ fn transient_rejection() -> ProviderError {
|
||||
NonZeroU32::new(3).expect("retry budget"),
|
||||
Duration::ZERO,
|
||||
Duration::ZERO,
|
||||
Duration::ZERO,
|
||||
TransientRetryPolicy::RetryThenRotateAccount,
|
||||
)
|
||||
}
|
||||
|
||||
@@ -3679,6 +4084,15 @@ fn server_retry_advice_should_override_local_backoff_and_transport_fallback() {
|
||||
NonZeroU32::MIN,
|
||||
Duration::from_millis(1),
|
||||
Duration::from_millis(1),
|
||||
Duration::from_secs(1),
|
||||
TransientRetryPolicy::RetryThenRotateAccount,
|
||||
),
|
||||
transient_rejection().with_transient_retry(
|
||||
NonZeroU32::MIN,
|
||||
Duration::from_millis(1),
|
||||
Duration::from_millis(1),
|
||||
Duration::from_secs(1),
|
||||
TransientRetryPolicy::RetryThenReturnError,
|
||||
),
|
||||
ProviderError::new(ProviderErrorKind::Transport, UpstreamSendState::NotSent)
|
||||
.with_pre_delivery_transport_fallback(),
|
||||
@@ -3710,6 +4124,14 @@ fn server_retry_advice_should_override_local_backoff_and_transport_fallback() {
|
||||
assert_eq!(provider.contexts.lock().unwrap().len(), 1);
|
||||
block_on(collection).unwrap();
|
||||
assert!(started.elapsed() >= advised_delay);
|
||||
let trace = session.trace().snapshot().unwrap();
|
||||
let retry = trace["events"]
|
||||
.as_array()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.find(|event| event["stage"] == "retry.decided")
|
||||
.unwrap();
|
||||
assert_eq!(retry["data"]["delayMs"], 120);
|
||||
let contexts = provider.contexts.lock().unwrap();
|
||||
assert_eq!(contexts.len(), 2);
|
||||
assert_eq!(
|
||||
@@ -3721,38 +4143,45 @@ fn server_retry_advice_should_override_local_backoff_and_transport_fallback() {
|
||||
|
||||
#[test]
|
||||
fn zero_server_retry_advice_should_override_nonzero_local_backoff() {
|
||||
let operation = generate_operation();
|
||||
let route_plan = plan(&operation);
|
||||
let (coordinator, _, _) = coordinator(vec![
|
||||
Script::Stream {
|
||||
account_id: "acct_first",
|
||||
items: vec![Err(transient_rejection()
|
||||
.with_transient_retry(
|
||||
NonZeroU32::MIN,
|
||||
Duration::from_secs(60),
|
||||
Duration::from_secs(60),
|
||||
)
|
||||
.with_retry_after(Duration::ZERO))],
|
||||
},
|
||||
Script::Stream {
|
||||
account_id: "acct_first",
|
||||
items: complete_stream(None),
|
||||
},
|
||||
]);
|
||||
let mut session = block_on(coordinator.start(
|
||||
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
|
||||
operation,
|
||||
route_plan,
|
||||
None,
|
||||
None,
|
||||
CancellationToken::new(),
|
||||
))
|
||||
.unwrap();
|
||||
let mut collection = Box::pin(session.collect_uncommitted());
|
||||
assert!(matches!(
|
||||
block_on(async { futures::poll!(&mut collection) }),
|
||||
std::task::Poll::Ready(Ok(_))
|
||||
));
|
||||
for policy in [
|
||||
TransientRetryPolicy::RetryThenRotateAccount,
|
||||
TransientRetryPolicy::RetryThenReturnError,
|
||||
] {
|
||||
let operation = generate_operation();
|
||||
let route_plan = plan(&operation);
|
||||
let (coordinator, _, _) = coordinator(vec![
|
||||
Script::Stream {
|
||||
account_id: "acct_first",
|
||||
items: vec![Err(transient_rejection()
|
||||
.with_transient_retry(
|
||||
NonZeroU32::MIN,
|
||||
Duration::from_secs(60),
|
||||
Duration::from_secs(60),
|
||||
Duration::from_secs(60),
|
||||
policy,
|
||||
)
|
||||
.with_retry_after(Duration::ZERO))],
|
||||
},
|
||||
Script::Stream {
|
||||
account_id: "acct_first",
|
||||
items: complete_stream(None),
|
||||
},
|
||||
]);
|
||||
let mut session = block_on(coordinator.start(
|
||||
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
|
||||
operation,
|
||||
route_plan,
|
||||
None,
|
||||
None,
|
||||
CancellationToken::new(),
|
||||
))
|
||||
.unwrap();
|
||||
let mut collection = Box::pin(session.collect_uncommitted());
|
||||
assert!(matches!(
|
||||
block_on(async { futures::poll!(&mut collection) }),
|
||||
std::task::Poll::Ready(Ok(_))
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -3842,6 +4271,113 @@ fn transient_retries_back_off_on_the_same_account_then_rotate_with_a_fresh_budge
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transient_return_to_client_policy_keeps_the_account_and_preserves_the_final_error() {
|
||||
for succeeds in [false, true] {
|
||||
let operation = generate_operation();
|
||||
let route_plan = plan(&operation);
|
||||
let jitter_millis = [3, 7, 2, 6, 4, 0];
|
||||
let failure = |attempt: usize| {
|
||||
ProviderError::new(ProviderErrorKind::Unavailable, UpstreamSendState::Sent)
|
||||
.with_status(503)
|
||||
.with_replay_safe()
|
||||
.with_client_visible_upstream_response(ClientVisibleUpstreamResponse::new(
|
||||
503,
|
||||
Some(b"application/json".to_vec()),
|
||||
format!(
|
||||
r#"{{"error":{{"code":"server_error","message":"attempt {attempt}"}}}}"#
|
||||
)
|
||||
.into(),
|
||||
))
|
||||
.with_transient_retry(
|
||||
NonZeroU32::new(5).expect("retry budget"),
|
||||
Duration::from_millis(5),
|
||||
Duration::from_millis(20),
|
||||
Duration::from_millis(jitter_millis[attempt - 1]),
|
||||
TransientRetryPolicy::RetryThenReturnError,
|
||||
)
|
||||
};
|
||||
let mut scripts = (1..=5)
|
||||
.map(|attempt| Script::Stream {
|
||||
account_id: "acct_first",
|
||||
items: vec![Err(failure(attempt))],
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
scripts.push(Script::Stream {
|
||||
account_id: "acct_first",
|
||||
items: if succeeds {
|
||||
complete_stream(None)
|
||||
} else {
|
||||
vec![Err(failure(6))]
|
||||
},
|
||||
});
|
||||
scripts.push(Script::Stream {
|
||||
account_id: "acct_second",
|
||||
items: complete_stream(None),
|
||||
});
|
||||
let (coordinator, store, 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(),
|
||||
))
|
||||
.expect("start execution");
|
||||
let result = block_on(session.collect_uncommitted());
|
||||
if succeeds {
|
||||
result.expect("same account recovers");
|
||||
block_on(session.commit_downstream(Some(200))).expect("commit");
|
||||
} else {
|
||||
let Err(EngineError::Provider(error)) = result else {
|
||||
panic!("retry exhaustion must return the last failure without rotating");
|
||||
};
|
||||
assert_eq!(error.upstream_status(), Some(503));
|
||||
assert_eq!(
|
||||
error
|
||||
.client_visible_upstream_response()
|
||||
.unwrap()
|
||||
.body()
|
||||
.as_ref(),
|
||||
failure(6)
|
||||
.client_visible_upstream_response()
|
||||
.unwrap()
|
||||
.body()
|
||||
.as_ref(),
|
||||
);
|
||||
}
|
||||
let original = ProviderAccountId::new("acct_first").unwrap();
|
||||
let trace = session.trace().snapshot().unwrap();
|
||||
let delays: Vec<_> = trace["events"]
|
||||
.as_array()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.filter(|event| event["stage"] == "retry.decided")
|
||||
.filter_map(|event| event["data"]["delayMs"].as_u64())
|
||||
.collect();
|
||||
// 每次使用该次失败的抖动,指数基准封顶后再相加,不放大或截断抖动
|
||||
assert_eq!(delays, [8, 17, 22, 26, 24]);
|
||||
let contexts = provider.contexts.lock().unwrap();
|
||||
assert_eq!(contexts.len(), 6);
|
||||
for context in &contexts[1..] {
|
||||
assert_eq!(context.required_account(), Some(&original));
|
||||
assert!(!context.excluded_accounts().contains(&original));
|
||||
}
|
||||
let state = store.state.lock().unwrap();
|
||||
assert_eq!(state.intermediate_failures, 5);
|
||||
assert_eq!(state.finalizations[0].attempt_count, 6);
|
||||
assert_eq!(
|
||||
state.finalizations[0].outcome,
|
||||
if succeeds {
|
||||
ExecutionOutcome::Succeeded
|
||||
} else {
|
||||
ExecutionOutcome::Failed
|
||||
}
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transient_retries_consume_the_total_routing_budget() {
|
||||
let operation = generate_operation();
|
||||
@@ -3854,6 +4390,8 @@ fn transient_retries_consume_the_total_routing_budget() {
|
||||
max_attempts,
|
||||
Duration::ZERO,
|
||||
Duration::ZERO,
|
||||
Duration::ZERO,
|
||||
TransientRetryPolicy::RetryThenRotateAccount,
|
||||
))],
|
||||
})
|
||||
.collect();
|
||||
@@ -3890,7 +4428,13 @@ fn transient_retry_intent_does_not_supply_replay_proof_or_override_ambiguous_sen
|
||||
let operation = generate_operation();
|
||||
let route_plan = plan(&operation);
|
||||
let mut error = ProviderError::new(ProviderErrorKind::Unavailable, send_state)
|
||||
.with_transient_retry(NonZeroU32::MIN, Duration::ZERO, Duration::ZERO);
|
||||
.with_transient_retry(
|
||||
NonZeroU32::MIN,
|
||||
Duration::ZERO,
|
||||
Duration::ZERO,
|
||||
Duration::ZERO,
|
||||
TransientRetryPolicy::RetryThenRotateAccount,
|
||||
);
|
||||
if prove_safe {
|
||||
error = error.with_replay_safe();
|
||||
}
|
||||
@@ -3917,30 +4461,37 @@ fn transient_retry_intent_does_not_supply_replay_proof_or_override_ambiguous_sen
|
||||
|
||||
#[test]
|
||||
fn transient_backoff_obeys_request_deadline() {
|
||||
let operation = generate_operation();
|
||||
let route_plan = plan(&operation);
|
||||
let (coordinator, _, provider) = coordinator(vec![Script::Stream {
|
||||
account_id: "acct_first",
|
||||
items: vec![Err(transient_rejection().with_transient_retry(
|
||||
NonZeroU32::MIN,
|
||||
Duration::from_secs(60),
|
||||
Duration::from_secs(60),
|
||||
))],
|
||||
}]);
|
||||
let mut session = block_on(coordinator.start(
|
||||
model_request(&operation, SystemTime::now() + Duration::from_millis(100)),
|
||||
operation,
|
||||
route_plan,
|
||||
None,
|
||||
None,
|
||||
CancellationToken::new(),
|
||||
))
|
||||
.expect("start execution");
|
||||
assert!(matches!(
|
||||
block_on(session.collect_uncommitted()),
|
||||
Err(EngineError::Deadline)
|
||||
));
|
||||
assert_eq!(provider.contexts.lock().expect("contexts").len(), 1);
|
||||
for policy in [
|
||||
TransientRetryPolicy::RetryThenRotateAccount,
|
||||
TransientRetryPolicy::RetryThenReturnError,
|
||||
] {
|
||||
let operation = generate_operation();
|
||||
let route_plan = plan(&operation);
|
||||
let (coordinator, _, provider) = coordinator(vec![Script::Stream {
|
||||
account_id: "acct_first",
|
||||
items: vec![Err(transient_rejection().with_transient_retry(
|
||||
NonZeroU32::MIN,
|
||||
Duration::from_secs(60),
|
||||
Duration::from_secs(60),
|
||||
Duration::ZERO,
|
||||
policy,
|
||||
))],
|
||||
}]);
|
||||
let mut session = block_on(coordinator.start(
|
||||
model_request(&operation, SystemTime::now() + Duration::from_millis(100)),
|
||||
operation,
|
||||
route_plan,
|
||||
None,
|
||||
None,
|
||||
CancellationToken::new(),
|
||||
))
|
||||
.expect("start execution");
|
||||
assert!(matches!(
|
||||
block_on(session.collect_uncommitted()),
|
||||
Err(EngineError::Deadline)
|
||||
));
|
||||
assert_eq!(provider.contexts.lock().expect("contexts").len(), 1);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -4022,36 +4573,43 @@ fn transient_retry_is_suppressed_for_required_accounts_and_committed_output() {
|
||||
|
||||
#[test]
|
||||
fn transient_backoff_can_be_cancelled_before_the_next_attempt() {
|
||||
let operation = generate_operation();
|
||||
let route_plan = plan(&operation);
|
||||
let (coordinator, _, provider) = coordinator(vec![Script::Stream {
|
||||
account_id: "acct_first",
|
||||
items: vec![Err(transient_rejection().with_transient_retry(
|
||||
NonZeroU32::MIN,
|
||||
Duration::from_secs(60),
|
||||
Duration::from_secs(60),
|
||||
))],
|
||||
}]);
|
||||
let cancellation = CancellationToken::new();
|
||||
let mut session = block_on(coordinator.start(
|
||||
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
|
||||
operation,
|
||||
route_plan,
|
||||
None,
|
||||
None,
|
||||
cancellation.clone(),
|
||||
))
|
||||
.expect("start execution");
|
||||
let cancel = std::thread::spawn(move || {
|
||||
std::thread::sleep(Duration::from_millis(50));
|
||||
cancellation.cancel();
|
||||
});
|
||||
assert!(matches!(
|
||||
block_on(session.collect_uncommitted()),
|
||||
Err(EngineError::Cancelled)
|
||||
));
|
||||
cancel.join().expect("cancel task");
|
||||
assert_eq!(provider.contexts.lock().expect("contexts").len(), 1);
|
||||
for policy in [
|
||||
TransientRetryPolicy::RetryThenRotateAccount,
|
||||
TransientRetryPolicy::RetryThenReturnError,
|
||||
] {
|
||||
let operation = generate_operation();
|
||||
let route_plan = plan(&operation);
|
||||
let (coordinator, _, provider) = coordinator(vec![Script::Stream {
|
||||
account_id: "acct_first",
|
||||
items: vec![Err(transient_rejection().with_transient_retry(
|
||||
NonZeroU32::MIN,
|
||||
Duration::from_secs(60),
|
||||
Duration::from_secs(60),
|
||||
Duration::ZERO,
|
||||
policy,
|
||||
))],
|
||||
}]);
|
||||
let cancellation = CancellationToken::new();
|
||||
let mut session = block_on(coordinator.start(
|
||||
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
|
||||
operation,
|
||||
route_plan,
|
||||
None,
|
||||
None,
|
||||
cancellation.clone(),
|
||||
))
|
||||
.expect("start execution");
|
||||
let cancel = std::thread::spawn(move || {
|
||||
std::thread::sleep(Duration::from_millis(50));
|
||||
cancellation.cancel();
|
||||
});
|
||||
assert!(matches!(
|
||||
block_on(session.collect_uncommitted()),
|
||||
Err(EngineError::Cancelled)
|
||||
));
|
||||
cancel.join().expect("cancel task");
|
||||
assert_eq!(provider.contexts.lock().expect("contexts").len(), 1);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -5711,6 +6269,8 @@ fn provider_candidate_advance_stops_at_the_account_rotation_budget() {
|
||||
NonZeroU32::new(3).expect("retry budget"),
|
||||
Duration::ZERO,
|
||||
Duration::ZERO,
|
||||
Duration::ZERO,
|
||||
TransientRetryPolicy::RetryThenRotateAccount,
|
||||
))],
|
||||
},
|
||||
Script::Error(ProviderError::new(
|
||||
|
||||
@@ -509,6 +509,84 @@ fn attempt_passthrough_preserves_parsed_wire_for_every_client_transport() {
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn attempt_sse_encoding_keeps_opaque_metadata_from_creating_extra_frames() {
|
||||
#[derive(Debug)]
|
||||
struct InspectingPlan {
|
||||
expected: &'static str,
|
||||
}
|
||||
|
||||
impl MiddlewarePlan for InspectingPlan {
|
||||
fn handle(
|
||||
&self,
|
||||
_: MiddlewareContext,
|
||||
request: MiddlewareRequest,
|
||||
next: MiddlewareNext,
|
||||
) -> BoxFuture<'static, Result<MiddlewareResponse, MiddlewareError>> {
|
||||
let expected = self.expected;
|
||||
Box::pin(async move {
|
||||
let (protocol, status, headers, mut body, envelope) =
|
||||
next.run(request).await?.into_parts();
|
||||
let frame = body.next_frame().await?.expect("provider event");
|
||||
assert_eq!(frame.bytes().as_ref(), expected.as_bytes());
|
||||
Ok(MiddlewareResponse::new(protocol, status, headers, body)
|
||||
.with_envelope(envelope.expect("provider metadata")))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
futures::executor::block_on(async {
|
||||
for (event_type, id, expected) in [
|
||||
(
|
||||
"future.event",
|
||||
"event-1",
|
||||
"id: event-1\nretry: 10\nevent: future.event\ndata: {\"value\":1}\n\n",
|
||||
),
|
||||
(
|
||||
"future.event\ndata: injected\n\n",
|
||||
"event-1\r\ndata: injected\r\n\r\n",
|
||||
"retry: 10\ndata: {\"value\":1}\n\n",
|
||||
),
|
||||
("", "event\0id", "retry: 10\ndata: {\"value\":1}\n\n"),
|
||||
] {
|
||||
let plan = FrozenMiddlewarePlan::new(
|
||||
Arc::new(InspectingPlan { expected }),
|
||||
ExtensionSetReference::new(
|
||||
ExtensionSetId::new("sse-metadata".to_owned()).unwrap(),
|
||||
Arc::new(ReadyGeneration),
|
||||
),
|
||||
);
|
||||
let wire = ProtocolWireEvent::json_with_sse_metadata(
|
||||
"openai",
|
||||
Some(event_type.to_owned()),
|
||||
json!({"value":1}),
|
||||
Some(id.to_owned()),
|
||||
Some(10),
|
||||
)
|
||||
.unwrap();
|
||||
let terminal: ProviderMiddlewareTerminal = Box::new(move |_, _| {
|
||||
Box::pin(async move {
|
||||
Ok(ProviderStream::new(
|
||||
empty_provider_stream().metadata().clone(),
|
||||
stream::iter([Ok(ProviderEvent::wire(wire))]),
|
||||
(),
|
||||
))
|
||||
})
|
||||
});
|
||||
|
||||
execute_attempt_middleware(
|
||||
Some(&plan),
|
||||
context(),
|
||||
operation(),
|
||||
ClientTransport::HttpSse,
|
||||
terminal,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn attempt_middleware_preserves_multi_value_and_authentication_headers() {
|
||||
futures::executor::block_on(async {
|
||||
|
||||
@@ -8,7 +8,6 @@ mod execution;
|
||||
mod extensions;
|
||||
mod middleware;
|
||||
mod policy;
|
||||
mod probe;
|
||||
mod provider;
|
||||
mod response_control;
|
||||
|
||||
|
||||
@@ -1,9 +0,0 @@
|
||||
//! 验证账号探测可通过对象安全的执行链端口调用
|
||||
|
||||
use gateway_core::engine::probe::AccountProbe;
|
||||
|
||||
#[test]
|
||||
fn account_probe_is_an_object_safe_execution_chain_port() {
|
||||
fn accepts(_: Option<&dyn AccountProbe>) {}
|
||||
accepts(None);
|
||||
}
|
||||
@@ -1,24 +1,10 @@
|
||||
//! 验证健康探针的对象安全接口与 Worker 身份投影
|
||||
//! 验证 Worker 健康状态的身份投影
|
||||
|
||||
use futures::future::BoxFuture;
|
||||
use gateway_core::health::{
|
||||
HealthProbe, HealthState, WorkerHealthKey, WorkerHealthSnapshot, WorkerHealthSource,
|
||||
WorkerRuntimeState,
|
||||
WorkerHealthKey, WorkerHealthSnapshot, WorkerHealthSource, WorkerRuntimeState,
|
||||
};
|
||||
use gateway_core::task::{WorkerId, WorkerKind};
|
||||
|
||||
struct HealthyProbe;
|
||||
|
||||
impl HealthProbe for HealthyProbe {
|
||||
fn name(&self) -> &'static str {
|
||||
"test"
|
||||
}
|
||||
|
||||
fn check(&self) -> BoxFuture<'_, HealthState> {
|
||||
Box::pin(async { HealthState::Healthy })
|
||||
}
|
||||
}
|
||||
|
||||
struct SingleSnapshot;
|
||||
|
||||
impl WorkerHealthSource for SingleSnapshot {
|
||||
@@ -38,13 +24,6 @@ impl WorkerHealthSource for SingleSnapshot {
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn health_probe_contract_is_object_safe() {
|
||||
let probe: &dyn HealthProbe = &HealthyProbe;
|
||||
|
||||
assert_eq!(probe.name(), "test");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn worker_health_source_preserves_worker_identity() {
|
||||
let source: &dyn WorkerHealthSource = &SingleSnapshot;
|
||||
|
||||
@@ -315,14 +315,6 @@ fn connection_registration_rejects_after_drain_linearization() {
|
||||
assert!(matches!(lifecycle.try_register(), Err(ConnectionDraining)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn connection_lifecycle_contract_is_object_safe() {
|
||||
let lifecycle = TestLifecycle::new();
|
||||
let object: &dyn ConnectionLifecycle = &lifecycle;
|
||||
|
||||
assert!(!object.is_draining());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cancellation_reason_survives_cleanup_and_propagates_only_to_descendants() {
|
||||
use gateway_core::lifecycle::{CancellationReason, LeaseFailure, RequestLeaseKind};
|
||||
|
||||
@@ -109,68 +109,6 @@ fn middleware_protocol_change_keeps_source_capability_facts_without_interpreting
|
||||
assert_eq!(replaced.capability_requirements(), source_requirements);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn native_encoding_keeps_admission_context_and_session_after_optional_translation() {
|
||||
let context = Map::from_iter([("conversation_id".to_owned(), json!("conversation"))]);
|
||||
let session = ProviderSessionState::new(
|
||||
"xai",
|
||||
Map::from_iter([("session_id".to_owned(), json!("session"))]),
|
||||
)
|
||||
.unwrap();
|
||||
let request = GenerateRequest::from_protocol_payload(
|
||||
ProtocolPayload::json_object(
|
||||
"openai",
|
||||
json!({"input":"hello", "max_output_tokens":64, "reasoning":{"effort":"high"}})
|
||||
.as_object()
|
||||
.unwrap()
|
||||
.clone(),
|
||||
)
|
||||
.unwrap()
|
||||
.with_context(context.clone()),
|
||||
)
|
||||
.with_provider_session_state(session.clone());
|
||||
let original = Operation::Generate(request.clone());
|
||||
let translated = original
|
||||
.clone()
|
||||
.replace_middleware_wire("native-input", Bytes::from_static(br#"{"prompt":"hello"}"#))
|
||||
.unwrap();
|
||||
let Operation::Generate(translated) = translated else {
|
||||
panic!("generate request");
|
||||
};
|
||||
|
||||
for (request, target) in [(request, "openai"), (translated, "xai")] {
|
||||
let body = Map::from_iter([("native_prompt".to_owned(), json!("encoded"))]);
|
||||
let encoded = request
|
||||
.with_native_encoded_body(target, body.clone())
|
||||
.unwrap();
|
||||
assert_eq!(encoded.protocol_payload().protocol(), target);
|
||||
assert_eq!(encoded.protocol_payload().body(), &body);
|
||||
assert_eq!(encoded.protocol_payload().context(), &context);
|
||||
assert_eq!(encoded.provider_session_state("xai"), Some(&session));
|
||||
let operation = Operation::Generate(encoded);
|
||||
assert_eq!(
|
||||
operation.capability_requirements(),
|
||||
original.capability_requirements()
|
||||
);
|
||||
assert!(
|
||||
operation
|
||||
.replace_middleware_wire("another", Bytes::from_static(b"{}"))
|
||||
.is_err()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn native_encoding_rejects_invalid_protocol_names() {
|
||||
for name in ["", "invalid\nprotocol"] {
|
||||
assert!(
|
||||
generate(json!({"input":"hello"}))
|
||||
.with_native_encoded_body(name, Map::new())
|
||||
.is_err()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn middleware_processing_recomputes_requirements_but_protocol_change_is_direct_once() {
|
||||
let operation = Operation::Generate(generate(json!({
|
||||
@@ -336,68 +274,6 @@ fn provider_http_should_reject_urls_controls_and_unbounded_headers() {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_http_wire_replacement_keeps_frozen_target_and_validates_typed_headers() {
|
||||
let operation = Operation::ProviderHttp(
|
||||
ProviderHttpRequest::new(
|
||||
"inspect",
|
||||
ProviderHttpMethod::Post,
|
||||
Some("mode=stable".to_owned()),
|
||||
vec![ProviderHttpHeader::new(
|
||||
"x-original",
|
||||
Bytes::from_static(b"old"),
|
||||
)],
|
||||
RawHttpPayload::new("provider-http", Bytes::from_static(b"old-body"))
|
||||
.expect("provider HTTP payload"),
|
||||
)
|
||||
.expect("provider HTTP request"),
|
||||
);
|
||||
assert!(
|
||||
operation
|
||||
.clone()
|
||||
.replace_provider_http_wire(
|
||||
None,
|
||||
Map::new(),
|
||||
vec![ProviderHttpHeader::new(
|
||||
"x-invalid",
|
||||
Bytes::from_static(b"line\r\nbreak"),
|
||||
)],
|
||||
)
|
||||
.is_err()
|
||||
);
|
||||
let replacement = operation
|
||||
.replace_provider_http_wire(
|
||||
Some(Bytes::from_static(&[0x00, 0xff, b'n', b'e', b'w'])),
|
||||
Map::from_iter([("connection_fact".to_owned(), json!("kept"))]),
|
||||
vec![
|
||||
ProviderHttpHeader::new("x-repeat", Bytes::from_static(&[0xff, 0xfe])),
|
||||
ProviderHttpHeader::new("x-repeat", Bytes::from_static(b"two")),
|
||||
],
|
||||
)
|
||||
.expect("typed Provider HTTP replacement");
|
||||
let Operation::ProviderHttp(request) = replacement else {
|
||||
panic!("operation kind changed")
|
||||
};
|
||||
assert_eq!(request.endpoint(), "inspect");
|
||||
assert_eq!(request.method(), ProviderHttpMethod::Post);
|
||||
assert_eq!(request.query(), Some("mode=stable"));
|
||||
assert_eq!(
|
||||
request.payload().body().as_ref(),
|
||||
&[0x00, 0xff, b'n', b'e', b'w']
|
||||
);
|
||||
assert_eq!(request.payload().context()["connection_fact"], "kept");
|
||||
assert_eq!(request.headers().len(), 2);
|
||||
assert_eq!(request.headers()[0].value().as_ref(), &[0xff, 0xfe]);
|
||||
assert_eq!(request.headers()[1].value().as_ref(), b"two");
|
||||
|
||||
let generate = Operation::Generate(generate(json!({"model":"gpt-test"})));
|
||||
assert!(
|
||||
generate
|
||||
.replace_provider_http_wire(None, Map::new(), Vec::new())
|
||||
.is_err()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn standalone_search_should_keep_raw_body_and_context_opaque() {
|
||||
let body = Bytes::from_static(
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
//! 隐私策略在快照与请求计划中的冻结、复用及关闭行为
|
||||
//! 隐私策略的数据合同,以及快照与请求计划中的冻结、复用及关闭行为
|
||||
|
||||
use super::*;
|
||||
use gateway_core::settings::privacy::*;
|
||||
@@ -16,7 +16,7 @@ impl PrivacyPolicyCompiler for Compiler {
|
||||
policy: &CodexPrivacyPolicy,
|
||||
) -> Result<Arc<dyn CompiledPrivacyPolicy>, PrivacyError> {
|
||||
self.0.fetch_add(1, Ordering::SeqCst);
|
||||
Ok(Arc::new(Compiled(policy.rules[0].value.clone())))
|
||||
Ok(Arc::new(Compiled(policy.rules[0].steps[0].value.clone())))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -41,19 +41,41 @@ fn policy(value: &str) -> CodexPrivacyPolicy {
|
||||
id: "rule".into(),
|
||||
name: "规则".into(),
|
||||
enabled: true,
|
||||
scope: PrivacyScope::RequestBody,
|
||||
selector: "$.metadata".into(),
|
||||
action: PrivacyAction::SetValue,
|
||||
pattern: None,
|
||||
replacement: String::new(),
|
||||
value: serde_json::json!(value),
|
||||
replace_all: false,
|
||||
case_insensitive: false,
|
||||
multi_line: false,
|
||||
steps: vec![PrivacyStep {
|
||||
scope: PrivacyScope::RequestBody,
|
||||
selector: "$.metadata".into(),
|
||||
action: PrivacyAction::SetValue,
|
||||
pattern: None,
|
||||
replacement: String::new(),
|
||||
value: serde_json::json!(value),
|
||||
replace_all: false,
|
||||
case_insensitive: false,
|
||||
multi_line: false,
|
||||
}],
|
||||
}],
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn privacy_rule_requires_steps_and_rejects_unknown_fields() {
|
||||
let expected = policy("user").rules.remove(0);
|
||||
let value = serde_json::to_value(&expected).unwrap();
|
||||
assert_eq!(
|
||||
serde_json::from_value::<PrivacyRule>(value.clone()).unwrap(),
|
||||
expected
|
||||
);
|
||||
|
||||
let mut missing_steps = value.clone();
|
||||
missing_steps.as_object_mut().unwrap().remove("steps");
|
||||
assert!(serde_json::from_value::<PrivacyRule>(missing_steps).is_err());
|
||||
let mut unknown_rule_field = value.clone();
|
||||
unknown_rule_field["scope"] = serde_json::json!("request_body");
|
||||
assert!(serde_json::from_value::<PrivacyRule>(unknown_rule_field).is_err());
|
||||
let mut unknown_step_field = value;
|
||||
unknown_step_field["steps"][0]["unexpected"] = serde_json::json!(true);
|
||||
assert!(serde_json::from_value::<PrivacyRule>(unknown_step_field).is_err());
|
||||
}
|
||||
|
||||
fn compiled(snapshot: &RuntimeSnapshot) -> Option<Arc<dyn CompiledPrivacyPolicy>> {
|
||||
snapshot
|
||||
.plan_provider_endpoint(
|
||||
|
||||
@@ -324,25 +324,6 @@ fn publisher_should_suspend_but_still_notify_after_committed_refresh_failure() {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn snapshot_ports_and_control_should_remain_object_safe() {
|
||||
fn accept_store(_: &dyn SnapshotStorePort) {}
|
||||
fn accept_subscriptions(_: &dyn SnapshotSubscriptionPort) {}
|
||||
fn accept_control(_: &dyn SnapshotControl) {}
|
||||
|
||||
let store = Arc::new(TestSnapshotStore::new(Ok(facts(1, 1))));
|
||||
let subscriptions = Arc::new(TestSnapshotSubscriptions::default());
|
||||
let publisher = RuntimeSnapshotPublisher::new(
|
||||
Arc::new(compiler(store.clone())),
|
||||
RuntimeSnapshotHandle::new(empty_snapshot(1)),
|
||||
subscriptions.clone(),
|
||||
);
|
||||
|
||||
accept_store(store.as_ref());
|
||||
accept_subscriptions(subscriptions.as_ref());
|
||||
accept_control(&publisher);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn publisher_should_contribute_reconciliation_and_subscription_workers() {
|
||||
let store = Arc::new(TestSnapshotStore::new(Ok(facts(1, 1))));
|
||||
|
||||
@@ -6,7 +6,7 @@ use std::time::Duration;
|
||||
use futures::future::BoxFuture;
|
||||
use gateway_core::lifecycle::CancellationToken;
|
||||
use gateway_core::task::{
|
||||
DaemonRestartPolicy, DaemonTask, ScheduledTask, WorkerCycleContext, WorkerDefinitionError,
|
||||
DaemonRestartPolicy, ScheduledTask, WorkerCycleContext, WorkerDefinitionError,
|
||||
WorkerFencingToken, WorkerId, WorkerKind, WorkerLeaseRequest, WorkerRegistration,
|
||||
WorkerRunnable, WorkerSchedule, WorkerTaskError,
|
||||
};
|
||||
@@ -22,14 +22,6 @@ impl ScheduledTask for NoopScheduledTask {
|
||||
}
|
||||
}
|
||||
|
||||
struct NoopDaemonTask;
|
||||
|
||||
impl DaemonTask for NoopDaemonTask {
|
||||
fn run(&self, _cancellation: CancellationToken) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
|
||||
Box::pin(async { Ok(()) })
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn schedule_rejects_every_invalid_duration_boundary() {
|
||||
let one = Duration::from_nanos(1);
|
||||
@@ -135,15 +127,6 @@ fn worker_cycle_context_preserves_optional_fencing_token() {
|
||||
assert_eq!(context.fencing_token(), Some(token));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn scheduled_and_daemon_tasks_are_object_safe() {
|
||||
fn accept_scheduled(_task: &dyn ScheduledTask) {}
|
||||
fn accept_daemon(_task: &dyn DaemonTask) {}
|
||||
|
||||
accept_scheduled(&NoopScheduledTask);
|
||||
accept_daemon(&NoopDaemonTask);
|
||||
}
|
||||
|
||||
fn worker_id(owner: &str) -> WorkerId {
|
||||
WorkerId::try_new(WorkerKind::OAuthRefresh, owner).expect("valid worker identity")
|
||||
}
|
||||
|
||||
@@ -27,7 +27,7 @@ hyper-util = { version = "0.1.20", features = ["client-legacy", "client-proxy",
|
||||
ipnet = "2.11.0"
|
||||
percent-encoding = "2.3.2"
|
||||
reqwest = { version = "=0.12.28", default-features = false, features = ["json", "stream", "rustls-tls-native-roots", "socks"] }
|
||||
rustls = { version = "=0.23.45", default-features = false, features = ["aws_lc_rs", "std", "tls12"] }
|
||||
rustls = { version = "=0.23.46", default-features = false, features = ["aws_lc_rs", "std", "tls12"] }
|
||||
rustls-native-certs = "=0.8.3"
|
||||
scraper = "0.27.0"
|
||||
semver = "1.0.27"
|
||||
|
||||
@@ -39,7 +39,11 @@ pub struct HttpPluginDistribution {
|
||||
|
||||
impl HttpPluginDistribution {
|
||||
pub fn new(http: Arc<HttpClient>) -> Result<Self, AdminError> {
|
||||
Self::with_transport("https://api.github.com/", http, source_network()?)
|
||||
Self::with_transport(
|
||||
"https://api.github.com/",
|
||||
http,
|
||||
NetworkPolicy::unrestricted(),
|
||||
)
|
||||
}
|
||||
|
||||
/// 测试可注入受限网络;生产来源由管理员显式配置,因此保持既有内外网可达范围
|
||||
@@ -173,10 +177,6 @@ impl PluginDistribution for HttpPluginDistribution {
|
||||
}
|
||||
}
|
||||
|
||||
fn source_network() -> Result<NetworkPolicy, AdminError> {
|
||||
NetworkPolicy::new(&["0.0.0.0/0".into(), "::/0".into()]).map_err(|_| source_error())
|
||||
}
|
||||
|
||||
fn credential_ids(credentials: &[SourceCredential]) -> Vec<String> {
|
||||
credentials
|
||||
.iter()
|
||||
|
||||
@@ -139,7 +139,7 @@ fn rpc_failure(error: RpcError) -> &'static str {
|
||||
RpcError::Handshake(_) | RpcError::Protocol(_) | RpcError::InvalidResponse(..) => {
|
||||
"插件协议错误"
|
||||
}
|
||||
RpcError::Capacity => "插件调用容量不足",
|
||||
RpcError::Capacity | RpcError::Overloaded => "插件调用容量不足",
|
||||
RpcError::Context(_) => "调用上下文无效",
|
||||
RpcError::Remote(_) => "插件拒绝命令",
|
||||
}
|
||||
|
||||
@@ -221,6 +221,7 @@ async fn recover_rpc(
|
||||
RpcError::Remote(fault) if fault.code == ErrorCode::Rejected
|
||||
);
|
||||
let delegate_safe = !rejected
|
||||
&& !error.is_overloaded()
|
||||
&& !matches!(
|
||||
&error,
|
||||
RpcError::Remote(fault) if fault.send_state != SendState::NotSent
|
||||
|
||||
@@ -4,7 +4,7 @@ use gateway_core::{
|
||||
engine::{EngineError, execution::gateway_error_from_engine, middleware::MiddlewareError},
|
||||
error::{
|
||||
ContinuationFailure, ContinuationRecoveryDisposition, GatewayError, GatewayErrorKind,
|
||||
PreDeliveryRetry, ProviderError,
|
||||
PreDeliveryRetry, ProviderError, TransientRetryPolicy,
|
||||
},
|
||||
upstream::UpstreamSendState,
|
||||
};
|
||||
@@ -20,6 +20,7 @@ pub(super) fn gateway(error: GatewayError) -> PluginFault {
|
||||
GatewayErrorKind::Unauthorized => ErrorCode::PermissionDenied,
|
||||
GatewayErrorKind::PolicyDenied => ErrorCode::Rejected,
|
||||
GatewayErrorKind::AccountCapacityUnavailable
|
||||
| GatewayErrorKind::MiddlewareCapacityUnavailable
|
||||
| GatewayErrorKind::ConcurrencyQueueFull
|
||||
| GatewayErrorKind::ConcurrencyQueueTimeout
|
||||
| GatewayErrorKind::RateLimited
|
||||
@@ -85,7 +86,15 @@ fn provider(error: &ProviderError) -> PluginFault {
|
||||
gateway_core::engine::AttemptTransport::Fallback => json!({"kind":"fallback"}),
|
||||
}}),
|
||||
PreDeliveryRetry::SameAccountTransportRetry {retry_index,delay} => json!({"kind":"same_account_transport_retry","retry_index":retry_index.get(),"delay_ms":delay.as_millis()}),
|
||||
PreDeliveryRetry::SameAccountTransientRetry {max_retries,initial_delay,max_delay} => json!({"kind":"same_account_transient_retry","max_retries":max_retries.get(),"initial_delay_ms":initial_delay.as_millis(),"max_delay_ms":max_delay.as_millis()}),
|
||||
PreDeliveryRetry::SameAccountTransientRetry {max_retries,initial_delay,max_delay,jitter,policy} => json!({
|
||||
"kind":"same_account_transient_retry","max_retries":max_retries.get(),
|
||||
"initial_delay_ms":initial_delay.as_millis(),"max_delay_ms":max_delay.as_millis(),
|
||||
"jitter_ms":jitter.as_millis(),
|
||||
"policy": match policy {
|
||||
TransientRetryPolicy::RetryThenRotateAccount => "retry_then_rotate_account",
|
||||
TransientRetryPolicy::RetryThenReturnError => "retry_then_return_error",
|
||||
},
|
||||
}),
|
||||
}),
|
||||
"replay_safe": error.replay_is_safe(),
|
||||
"retry_same_account": error.retries_same_account(),
|
||||
@@ -134,13 +143,22 @@ pub(crate) fn rpc(error: &crate::RpcError) -> PluginFault {
|
||||
RpcError::Remote(fault) => return fault.clone(),
|
||||
RpcError::Timeout => ErrorCode::Timeout,
|
||||
RpcError::Cancelled => ErrorCode::Cancelled,
|
||||
RpcError::Capacity => ErrorCode::Capacity,
|
||||
RpcError::Capacity | RpcError::Overloaded => ErrorCode::Capacity,
|
||||
_ => ErrorCode::Fault,
|
||||
};
|
||||
PluginFault::new(code, error.to_string())
|
||||
}
|
||||
|
||||
pub(crate) fn rpc_middleware(error: crate::RpcError) -> MiddlewareError {
|
||||
if error.is_overloaded() {
|
||||
return MiddlewareError::Gateway(
|
||||
GatewayError::new(
|
||||
GatewayErrorKind::MiddlewareCapacityUnavailable,
|
||||
"Request middleware is temporarily busy",
|
||||
)
|
||||
.with_source(gateway_core::error::ErrorSource::new(error)),
|
||||
);
|
||||
}
|
||||
let rejected =
|
||||
matches!(&error, crate::RpcError::Remote(fault) if fault.code == ErrorCode::Rejected);
|
||||
MiddlewareError::Remote {
|
||||
|
||||
@@ -194,7 +194,10 @@ impl PluginRuntime {
|
||||
.map_err(|_| AdminError::internal("插件集合 ID 无效"))?;
|
||||
let observers = crate::adapter::observer::PluginObserverPlan::compile(
|
||||
observer_entries,
|
||||
self.config.rpc_limits.maximum_calls,
|
||||
self.config
|
||||
.rpc_limits
|
||||
.maximum_calls
|
||||
.min(self.config.rpc_limits.maximum_executing_calls),
|
||||
self.config.rpc_limits.maximum_call_timeout,
|
||||
self.config.rpc_limits.maximum_buffered_body_bytes,
|
||||
)
|
||||
|
||||
@@ -0,0 +1,50 @@
|
||||
//! 区分插件执行名额与初次准入队列,恢复中的调用沿用在途资源预算
|
||||
|
||||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
|
||||
|
||||
use super::RpcError;
|
||||
|
||||
pub(super) struct ExecutionCapacity {
|
||||
active: Arc<Semaphore>,
|
||||
queued: Arc<Semaphore>,
|
||||
wait_timeout: Duration,
|
||||
}
|
||||
|
||||
impl ExecutionCapacity {
|
||||
pub(super) fn new(active: usize, queued: usize, wait_timeout: Duration) -> Self {
|
||||
Self {
|
||||
active: Arc::new(Semaphore::new(active)),
|
||||
queued: Arc::new(Semaphore::new(queued)),
|
||||
wait_timeout,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn enter(&self) -> Result<OwnedSemaphorePermit, RpcError> {
|
||||
if let Ok(permit) = self.active.clone().try_acquire_owned() {
|
||||
return Ok(permit);
|
||||
}
|
||||
let _queued = self.queued.try_acquire().map_err(|error| {
|
||||
if matches!(error, tokio::sync::TryAcquireError::Closed) {
|
||||
RpcError::Closed(None)
|
||||
} else {
|
||||
RpcError::Overloaded
|
||||
}
|
||||
})?;
|
||||
self.resume().await
|
||||
}
|
||||
|
||||
/// 恢复的等待者已持有在途名额,不再竞争初次准入队列
|
||||
pub(super) async fn resume(&self) -> Result<OwnedSemaphorePermit, RpcError> {
|
||||
tokio::time::timeout(self.wait_timeout, self.active.clone().acquire_owned())
|
||||
.await
|
||||
.map_err(|_| RpcError::Overloaded)?
|
||||
.map_err(|_| RpcError::Closed(None))
|
||||
}
|
||||
|
||||
pub(super) fn close(&self) {
|
||||
self.active.close();
|
||||
self.queued.close();
|
||||
}
|
||||
}
|
||||
@@ -125,7 +125,7 @@ async fn read_loop<R: AsyncRead + Unpin>(
|
||||
Err(RpcError::Protocol(None))
|
||||
} else if let Some(context) = shared.context(parent_id) {
|
||||
last_callback = id;
|
||||
if let Some(permit) = shared.try_callback_slot() {
|
||||
if let Some(permit) = shared.try_callback_slot(&method) {
|
||||
let handler = Arc::clone(&callbacks);
|
||||
let response = Arc::clone(&shared);
|
||||
let (ready, started) = oneshot::channel();
|
||||
@@ -138,6 +138,12 @@ async fn read_loop<R: AsyncRead + Unpin>(
|
||||
return;
|
||||
}
|
||||
let result = handler.call(context, method, params, frame.payload).await;
|
||||
if pause.is_some()
|
||||
&& let Err(error) = response.resume_next(parent_id).await
|
||||
{
|
||||
response.cancel(parent_id, error);
|
||||
return;
|
||||
}
|
||||
drop(pause);
|
||||
let mut reply = match result {
|
||||
Ok(reply) => Frame {
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
//! 插件双向 RPC 会话、调用结果与生命周期的公共入口
|
||||
|
||||
mod admission;
|
||||
mod dispatch;
|
||||
mod session;
|
||||
|
||||
|
||||
@@ -15,17 +15,21 @@ use gateway_plugin_sdk::{
|
||||
CallContext, Frame, Handshake, Message, PROTOCOL_VERSION, PluginFault, Stage,
|
||||
client::{CallDeadline, CallDeadlinePause, read_frame, validate_frame, write_frame},
|
||||
};
|
||||
use tokio::sync::{Mutex as AsyncMutex, Semaphore, mpsc, oneshot, watch};
|
||||
use tokio::sync::{Mutex as AsyncMutex, OwnedSemaphorePermit, Semaphore, mpsc, oneshot, watch};
|
||||
|
||||
use crate::{PreparedPackage, RpcStream, stream::StreamIngress};
|
||||
|
||||
use super::dispatch;
|
||||
use super::{admission::ExecutionCapacity, dispatch};
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct RpcLimits {
|
||||
/// 受管回调与观察队列的缓冲预算,不是请求正文或 IPC 消息总量上限
|
||||
pub maximum_buffered_body_bytes: usize,
|
||||
/// 包含初次排队、下游等待、取消确认和未消费流的总在途上限
|
||||
pub maximum_calls: usize,
|
||||
pub maximum_executing_calls: usize,
|
||||
pub maximum_queued_calls: usize,
|
||||
pub maximum_queue_timeout: Duration,
|
||||
pub maximum_callbacks: usize,
|
||||
pub handshake_timeout: Duration,
|
||||
pub maximum_call_timeout: Duration,
|
||||
@@ -36,8 +40,11 @@ impl Default for RpcLimits {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
maximum_buffered_body_bytes: 1024 * 1024,
|
||||
maximum_calls: 16,
|
||||
maximum_callbacks: 16,
|
||||
maximum_calls: 256,
|
||||
maximum_executing_calls: 64,
|
||||
maximum_queued_calls: 64,
|
||||
maximum_queue_timeout: Duration::from_secs(30),
|
||||
maximum_callbacks: 64,
|
||||
handshake_timeout: Duration::from_secs(5),
|
||||
maximum_call_timeout: Duration::from_secs(120),
|
||||
maximum_stderr_bytes: 64 * 1024,
|
||||
@@ -63,6 +70,8 @@ pub enum RpcError {
|
||||
Cancelled,
|
||||
#[error("plugin call capacity is exhausted")]
|
||||
Capacity,
|
||||
#[error("plugin admission capacity is exhausted")]
|
||||
Overloaded,
|
||||
#[error("plugin call context is invalid")]
|
||||
Context(#[source] Option<gateway_core::error::ErrorSource>),
|
||||
#[error("plugin returned an error")]
|
||||
@@ -70,6 +79,11 @@ pub enum RpcError {
|
||||
}
|
||||
|
||||
impl RpcError {
|
||||
pub(crate) fn is_overloaded(&self) -> bool {
|
||||
matches!(self, Self::Overloaded)
|
||||
|| matches!(self, Self::Remote(fault) if fault.code == gateway_plugin_sdk::ErrorCode::Capacity)
|
||||
}
|
||||
|
||||
/// 会话状态与请求错误共用安全分类,不展开插件返回正文
|
||||
pub(crate) fn diagnostic(&self) -> (&'static str, &'static str) {
|
||||
match self {
|
||||
@@ -101,6 +115,7 @@ impl RpcError {
|
||||
RpcError::Timeout => ("timeout", "插件调用超时"),
|
||||
RpcError::Cancelled => ("cancelled", "插件会话已取消"),
|
||||
RpcError::Capacity => ("capacity", "插件调用容量已耗尽"),
|
||||
RpcError::Overloaded => ("overloaded", "插件暂时繁忙"),
|
||||
RpcError::Context(_) => ("context", "插件调用上下文无效"),
|
||||
RpcError::Remote(_) => ("remote", "插件返回错误"),
|
||||
}
|
||||
@@ -216,6 +231,8 @@ pub(crate) struct Shared {
|
||||
pub process: ProcessControl,
|
||||
pub stopped: watch::Sender<bool>,
|
||||
callback_slots: Arc<Semaphore>,
|
||||
next_slots: Arc<Semaphore>,
|
||||
execution: ExecutionCapacity,
|
||||
}
|
||||
|
||||
struct State {
|
||||
@@ -224,6 +241,8 @@ struct State {
|
||||
}
|
||||
|
||||
struct PendingCall {
|
||||
_slot: Arc<OwnedSemaphorePermit>,
|
||||
execution: Option<OwnedSemaphorePermit>,
|
||||
deadline: CallDeadline,
|
||||
handler: Arc<dyn CallbackHandler>,
|
||||
transmitted: bool,
|
||||
@@ -274,6 +293,7 @@ impl Shared {
|
||||
.flatten();
|
||||
(std::mem::take(&mut state.calls), observer)
|
||||
};
|
||||
self.execution.close();
|
||||
for call in calls.into_values() {
|
||||
call.finish(Err(error.clone()));
|
||||
}
|
||||
@@ -303,6 +323,7 @@ impl Shared {
|
||||
call.deadline.clear();
|
||||
}
|
||||
let initial = call.result.take().ok_or(RpcError::Protocol(None))?;
|
||||
call.execution.take();
|
||||
call.handler.streaming(&call.context);
|
||||
let _ = initial.send(result);
|
||||
} else {
|
||||
@@ -384,11 +405,40 @@ impl Shared {
|
||||
}
|
||||
|
||||
pub fn pause_next(&self, id: u64) -> Option<CallDeadlinePause> {
|
||||
self.state()
|
||||
let mut state = self.state();
|
||||
let call = state
|
||||
.calls
|
||||
.get(&id)
|
||||
.get_mut(&id)
|
||||
.filter(|call| call.cancellation.is_none())?;
|
||||
let pause = call.deadline.pause()?;
|
||||
call.execution.take();
|
||||
Some(pause)
|
||||
}
|
||||
|
||||
pub async fn resume_next(&self, id: u64) -> Result<(), RpcError> {
|
||||
{
|
||||
let state = self.state();
|
||||
let call = state
|
||||
.calls
|
||||
.get(&id)
|
||||
.filter(|call| call.cancellation.is_none())
|
||||
.ok_or(RpcError::Cancelled)?;
|
||||
// 已交付响应头的流只占在途资源,不重新占用首结果执行名额
|
||||
if call.result.is_none() {
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
let permit = self.execution.resume().await?;
|
||||
let mut state = self.state();
|
||||
let call = state
|
||||
.calls
|
||||
.get_mut(&id)
|
||||
.filter(|call| call.cancellation.is_none())
|
||||
.and_then(|call| call.deadline.pause())
|
||||
.ok_or(RpcError::Cancelled)?;
|
||||
if call.result.is_some() {
|
||||
call.execution = Some(permit);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn is_cancelling(&self, id: u64) -> bool {
|
||||
@@ -495,21 +545,23 @@ impl Shared {
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn try_callback_slot(&self) -> Option<tokio::sync::OwnedSemaphorePermit> {
|
||||
self.callback_slots.clone().try_acquire_owned().ok()
|
||||
pub(super) fn try_callback_slot(&self, method: &str) -> Option<OwnedSemaphorePermit> {
|
||||
let slots = if method == gateway_plugin_sdk::call::middleware::NEXT_METHOD {
|
||||
&self.next_slots
|
||||
} else {
|
||||
&self.callback_slots
|
||||
};
|
||||
slots.clone().try_acquire_owned().ok()
|
||||
}
|
||||
|
||||
async fn wait_for_callbacks(&self, maximum_callbacks: usize) {
|
||||
let Ok(maximum_callbacks) = u32::try_from(maximum_callbacks) else {
|
||||
return;
|
||||
};
|
||||
if let Ok(permits) = self
|
||||
.callback_slots
|
||||
.clone()
|
||||
.acquire_many_owned(maximum_callbacks)
|
||||
.await
|
||||
{
|
||||
drop(permits);
|
||||
async fn wait_for_callbacks(&self, limits: RpcLimits) {
|
||||
for (slots, maximum) in [
|
||||
(&self.callback_slots, limits.maximum_callbacks),
|
||||
(&self.next_slots, limits.maximum_calls),
|
||||
] {
|
||||
if let Ok(permits) = slots.clone().acquire_many_owned(maximum as u32).await {
|
||||
drop(permits);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -558,6 +610,10 @@ impl RpcSession {
|
||||
) -> Result<Self, RpcError> {
|
||||
if limits.maximum_calls == 0
|
||||
|| limits.maximum_calls > 256
|
||||
|| limits.maximum_executing_calls == 0
|
||||
|| limits.maximum_executing_calls > 256
|
||||
|| limits.maximum_queued_calls > 256
|
||||
|| limits.maximum_queue_timeout.is_zero()
|
||||
|| limits.maximum_callbacks == 0
|
||||
|| limits.maximum_callbacks > 256
|
||||
|| limits.maximum_buffered_body_bytes < 1024
|
||||
@@ -657,7 +713,7 @@ impl RpcSession {
|
||||
}
|
||||
let (data, data_receiver) = mpsc::channel(limits.maximum_calls);
|
||||
let (control, control_receiver) =
|
||||
mpsc::channel(limits.maximum_callbacks + limits.maximum_calls + 8);
|
||||
mpsc::channel(limits.maximum_callbacks + 2 * limits.maximum_calls + 8);
|
||||
let (stopped, _) = watch::channel(false);
|
||||
let callback_slots = Arc::new(Semaphore::new(limits.maximum_callbacks));
|
||||
let shared = Arc::new(Shared {
|
||||
@@ -674,6 +730,12 @@ impl RpcSession {
|
||||
process,
|
||||
stopped,
|
||||
callback_slots,
|
||||
next_slots: Arc::new(Semaphore::new(limits.maximum_calls)),
|
||||
execution: ExecutionCapacity::new(
|
||||
limits.maximum_executing_calls,
|
||||
limits.maximum_queued_calls,
|
||||
limits.maximum_queue_timeout,
|
||||
),
|
||||
});
|
||||
dispatch::start(
|
||||
input,
|
||||
@@ -754,11 +816,7 @@ impl RpcSession {
|
||||
}
|
||||
}
|
||||
|
||||
fn call_deadline(
|
||||
&self,
|
||||
method: &str,
|
||||
context: &CallContext,
|
||||
) -> Result<tokio::time::Instant, RpcError> {
|
||||
fn call_timeout(&self, method: &str, context: &CallContext) -> Result<Duration, RpcError> {
|
||||
if context.instance_id != self.handshake.instance_id
|
||||
|| context.generation != self.handshake.generation
|
||||
|| context.incarnation != self.handshake.incarnation
|
||||
@@ -769,7 +827,7 @@ impl RpcSession {
|
||||
{
|
||||
return Err(RpcError::Context(None));
|
||||
}
|
||||
Ok(tokio::time::Instant::now() + Duration::from_millis(context.timeout_ms))
|
||||
Ok(Duration::from_millis(context.timeout_ms))
|
||||
}
|
||||
|
||||
fn begin_call(
|
||||
@@ -777,6 +835,8 @@ impl RpcSession {
|
||||
context: CallContext,
|
||||
deadline: tokio::time::Instant,
|
||||
stream: Option<StreamIngress>,
|
||||
slot: Arc<OwnedSemaphorePermit>,
|
||||
execution: OwnedSemaphorePermit,
|
||||
) -> Result<StartedCall, RpcError> {
|
||||
let id = context.call_id;
|
||||
let (result, received) = oneshot::channel();
|
||||
@@ -784,14 +844,13 @@ impl RpcSession {
|
||||
if let Some(error) = &state.failure {
|
||||
return Err(error.clone());
|
||||
}
|
||||
if state.calls.len() >= self.limits.maximum_calls {
|
||||
return Err(RpcError::Capacity);
|
||||
}
|
||||
let deadline = CallDeadline::new(deadline);
|
||||
self.callbacks.begin(&context, &deadline);
|
||||
state.calls.insert(
|
||||
id,
|
||||
PendingCall {
|
||||
_slot: slot,
|
||||
execution: Some(execution),
|
||||
deadline: deadline.clone(),
|
||||
handler: self.callbacks.clone(),
|
||||
transmitted: false,
|
||||
@@ -818,7 +877,16 @@ impl RpcSession {
|
||||
stream: Option<StreamIngress>,
|
||||
credit: Option<(u32, u32)>,
|
||||
) -> Result<CompletedCall, RpcError> {
|
||||
let deadline = self.call_deadline(method, &context)?;
|
||||
let timeout = self.call_timeout(method, &context)?;
|
||||
let slot = Arc::new(self.slots.clone().try_acquire_owned().map_err(|error| {
|
||||
if matches!(error, tokio::sync::TryAcquireError::Closed) {
|
||||
RpcError::Closed(None)
|
||||
} else {
|
||||
RpcError::Overloaded
|
||||
}
|
||||
})?);
|
||||
let execution = self.shared.execution.enter().await?;
|
||||
let deadline = tokio::time::Instant::now() + timeout;
|
||||
// FIFO 锁同时拥有 ID 与 data 入队顺序;Call/初始 Credit 入队后必须立即释放,
|
||||
// 不能把插件回复或重入回调纳入串行区,也不能为等待锁重置调用期限
|
||||
let mut next_call = tokio::time::timeout_at(deadline, self.next_call.lock())
|
||||
@@ -841,7 +909,7 @@ impl RpcSession {
|
||||
validate_frame(&frame).map_err(|source| {
|
||||
RpcError::Context(Some(gateway_core::error::ErrorSource::new(source)))
|
||||
})?;
|
||||
let started = self.begin_call(context, deadline, stream)?;
|
||||
let started = self.begin_call(context, deadline, stream, slot.clone(), execution)?;
|
||||
let mut guard = CallGuard {
|
||||
shared: Arc::clone(&self.shared),
|
||||
id: started.id,
|
||||
@@ -906,6 +974,7 @@ impl RpcSession {
|
||||
}
|
||||
};
|
||||
Ok(CompletedCall {
|
||||
slot,
|
||||
id: started.id,
|
||||
deadline: started.deadline,
|
||||
reply,
|
||||
@@ -919,7 +988,6 @@ impl RpcSession {
|
||||
params: serde_json::Value,
|
||||
payload: Vec<u8>,
|
||||
) -> Result<RpcReply, RpcError> {
|
||||
let _slot = self.slots.try_acquire().map_err(|_| RpcError::Capacity)?;
|
||||
self.send_call(method, context, params, payload, None, None)
|
||||
.await
|
||||
.map(|completed| completed.reply)
|
||||
@@ -932,9 +1000,6 @@ impl RpcSession {
|
||||
params: serde_json::Value,
|
||||
payload: Vec<u8>,
|
||||
) -> Result<RpcStream, RpcError> {
|
||||
let slot = Arc::clone(&self.slots)
|
||||
.try_acquire_owned()
|
||||
.map_err(|_| RpcError::Capacity)?;
|
||||
let bytes = self.limits.maximum_buffered_body_bytes.min(256 * 1024) as u32;
|
||||
let frames = 32;
|
||||
let (stream, chunks, terminal) = StreamIngress::new(bytes, frames);
|
||||
@@ -955,7 +1020,7 @@ impl RpcSession {
|
||||
deadline: completed.deadline,
|
||||
id: completed.id,
|
||||
shared: Arc::clone(&self.shared),
|
||||
_slot: slot,
|
||||
_slot: completed.slot,
|
||||
_session: Arc::clone(self),
|
||||
})
|
||||
}
|
||||
@@ -988,9 +1053,7 @@ impl RpcSession {
|
||||
self.shared.process.exited().await;
|
||||
// 不等待监控任务调度;先撤销父调用并中止其回调,再确认回调 future 已析构
|
||||
self.shared.fail(RpcError::Closed(None));
|
||||
self.shared
|
||||
.wait_for_callbacks(self.limits.maximum_callbacks)
|
||||
.await;
|
||||
self.shared.wait_for_callbacks(self.limits).await;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1022,6 +1085,7 @@ struct StartedCall {
|
||||
}
|
||||
|
||||
struct CompletedCall {
|
||||
slot: Arc<OwnedSemaphorePermit>,
|
||||
id: u64,
|
||||
deadline: CallDeadline,
|
||||
reply: RpcReply,
|
||||
|
||||
@@ -62,7 +62,7 @@ pub struct RpcStream {
|
||||
pub(crate) deadline: CallDeadline,
|
||||
pub(crate) id: u64,
|
||||
pub(crate) shared: Arc<Shared>,
|
||||
pub(crate) _slot: OwnedSemaphorePermit,
|
||||
pub(crate) _slot: Arc<OwnedSemaphorePermit>,
|
||||
pub(crate) _session: Arc<RpcSession>,
|
||||
}
|
||||
|
||||
|
||||
@@ -99,6 +99,8 @@ impl tracing::Subscriber for Capture {
|
||||
|
||||
#[tokio::test]
|
||||
async fn real_worker_logs_are_correlated_sanitized_validated_and_bounded() {
|
||||
// 低于每秒日志上限,单独验证阻塞写入的槽位及跨代保留行为
|
||||
const MAXIMUM_CALLBACKS: usize = 16;
|
||||
let capture = Capture::default();
|
||||
tracing::subscriber::set_global_default(capture.clone()).unwrap();
|
||||
let directory = tempfile::tempdir().unwrap();
|
||||
@@ -155,7 +157,10 @@ async fn real_worker_logs_are_correlated_sanitized_validated_and_bounded() {
|
||||
cache_directory: directory.path().join("cache"),
|
||||
host_version: "1.0.0".parse().unwrap(),
|
||||
package_limits: PackageLimits::default(),
|
||||
rpc_limits: RpcLimits::default(),
|
||||
rpc_limits: RpcLimits {
|
||||
maximum_callbacks: MAXIMUM_CALLBACKS,
|
||||
..RpcLimits::default()
|
||||
},
|
||||
restart_circuit: Default::default(),
|
||||
},
|
||||
Arc::new(gateway_host::outbound::HttpClient::new().unwrap()),
|
||||
@@ -270,7 +275,7 @@ async fn real_worker_logs_are_correlated_sanitized_validated_and_bounded() {
|
||||
blocked.config_revision = Revision::new(9).unwrap();
|
||||
blocked.instances[0].id = "log-blocked".into();
|
||||
blocked.instances[0].configuration = json!({
|
||||
"log_method":"plugin.register", "log_entries":[{"params":{"event":"fixture.blocked"},"repeat":32}],
|
||||
"log_method":"plugin.register", "log_entries":[{"params":{"event":"fixture.blocked"},"repeat":MAXIMUM_CALLBACKS * 2}],
|
||||
"log_marker":blocked_marker,
|
||||
});
|
||||
// 日志后端停住时,注册的业务结果和父调用结束仍须及时完成
|
||||
@@ -288,7 +293,7 @@ async fn real_worker_logs_are_correlated_sanitized_validated_and_bounded() {
|
||||
.iter()
|
||||
.filter(|result| result["recorded"] == true)
|
||||
.count(),
|
||||
RpcLimits::default().maximum_callbacks
|
||||
MAXIMUM_CALLBACKS
|
||||
);
|
||||
assert!(
|
||||
blocked_results
|
||||
|
||||
@@ -27,6 +27,7 @@ struct Facts(AtomicUsize, AtomicUsize, Option<String>);
|
||||
fn account() -> AccountRecord {
|
||||
let at = "2026-01-01T00:00:00Z".parse().unwrap();
|
||||
AccountRecord {
|
||||
outbound_proxy_enabled: true,
|
||||
id: "acct_facts".into(),
|
||||
provider_kind: ProviderKind::new("openai").unwrap(),
|
||||
groups: vec![gateway_admin::model::account_groups::AccountGroupRef {
|
||||
|
||||
@@ -0,0 +1,162 @@
|
||||
//! 验证 next 等待不占执行名额,硬上限拒绝请求且不跳过插件
|
||||
|
||||
use super::*;
|
||||
use gateway_core::error::GatewayErrorKind;
|
||||
|
||||
fn context() -> core::Context {
|
||||
core::Context {
|
||||
plugin_instance_id: None,
|
||||
call_id: uuid::Uuid::new_v4().to_string(),
|
||||
parent_call_id: None,
|
||||
extensions: Default::default(),
|
||||
plan: None,
|
||||
request_id: uuid::Uuid::new_v4().to_string(),
|
||||
cancellation: CancellationToken::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn plugin_request() -> core::Request {
|
||||
::http::Request::builder()
|
||||
.uri("/v1/responses")
|
||||
.header("authorization", "Bearer fixture-secret")
|
||||
.body(core::empty_body())
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn http_waiting_next_releases_execution_for_delegate_plugin() {
|
||||
assert_capacity(PluginFailurePolicy::Delegate, RpcLimits::default(), false).await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn http_waiting_next_releases_execution_for_reject_plugin() {
|
||||
assert_capacity(PluginFailurePolicy::Reject, RpcLimits::default(), false).await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn http_in_flight_limit_never_delegates_and_recovers_after_cancellation() {
|
||||
for policy in [PluginFailurePolicy::Delegate, PluginFailurePolicy::Reject] {
|
||||
assert_capacity(
|
||||
policy,
|
||||
RpcLimits {
|
||||
maximum_calls: 64,
|
||||
..Default::default()
|
||||
},
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn assert_capacity(policy: PluginFailurePolicy, limits: RpcLimits, saturated: bool) {
|
||||
const CONCURRENT_CALLS: usize = 64;
|
||||
let (cache, runtime, generation, plan) =
|
||||
http_plugin_with_policy("passthrough", limits, policy).await;
|
||||
let release = Arc::new(tokio::sync::Semaphore::new(0));
|
||||
let (started, mut entered) = tokio::sync::mpsc::unbounded_channel();
|
||||
let mut slow_calls = tokio::task::JoinSet::new();
|
||||
for index in 0..CONCURRENT_CALLS {
|
||||
let plan = plan.clone();
|
||||
let release = release.clone();
|
||||
let started = started.clone();
|
||||
slow_calls.spawn(async move {
|
||||
plan.handle_http(
|
||||
context(),
|
||||
plugin_request(),
|
||||
compose(Vec::new(), move |request: core::Request| {
|
||||
assert_eq!(request.uri(), "/rewritten?encoded=%2F");
|
||||
let release = release.clone();
|
||||
Box::pin(async move {
|
||||
started.send(()).unwrap();
|
||||
release.acquire().await.unwrap().forget();
|
||||
Ok(core::Response::new(core::empty_body()))
|
||||
})
|
||||
}),
|
||||
)
|
||||
.await
|
||||
});
|
||||
tokio::time::timeout(Duration::from_secs(5), entered.recv())
|
||||
.await
|
||||
.unwrap_or_else(|_| panic!("第 {} 个调用未能进入下游", index + 1))
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
for uri in [
|
||||
"/healthz",
|
||||
"/assets/index.js",
|
||||
"/api/settings",
|
||||
"/v1/responses",
|
||||
] {
|
||||
let request = ::http::Request::builder()
|
||||
.uri(uri)
|
||||
.header("authorization", "Bearer fixture-secret")
|
||||
.body(core::empty_body())
|
||||
.unwrap();
|
||||
let result = tokio::time::timeout(
|
||||
Duration::from_secs(3),
|
||||
plan.handle_http(
|
||||
context(),
|
||||
request,
|
||||
compose(Vec::new(), move |request: core::Request| {
|
||||
assert!(!saturated, "容量耗尽时不得执行下游");
|
||||
assert_eq!(request.uri(), "/rewritten?encoded=%2F");
|
||||
Box::pin(async { Ok(core::Response::new(core::empty_body())) })
|
||||
}),
|
||||
),
|
||||
)
|
||||
.await
|
||||
.expect("下游等待不应阻塞其它请求");
|
||||
if saturated {
|
||||
assert!(matches!(result, Err(MiddlewareError::Gateway(error))
|
||||
if error.kind() == GatewayErrorKind::MiddlewareCapacityUnavailable));
|
||||
} else {
|
||||
assert_eq!(result.unwrap().headers()["x-http-plugin"], "active");
|
||||
}
|
||||
}
|
||||
|
||||
if saturated {
|
||||
slow_calls.abort_all();
|
||||
while let Some(result) = slow_calls.join_next().await {
|
||||
assert!(matches!(result, Err(error) if error.is_cancelled()));
|
||||
}
|
||||
}
|
||||
release.add_permits(CONCURRENT_CALLS);
|
||||
while let Some(result) = slow_calls.join_next().await {
|
||||
let response = result.unwrap().unwrap();
|
||||
assert_eq!(response.headers()["x-http-plugin"], "active");
|
||||
}
|
||||
// 取消确认到达前仍持有在途名额,探测只允许遇到明确的暂时过载
|
||||
let response = tokio::time::timeout(Duration::from_secs(2), async {
|
||||
loop {
|
||||
let result = plan
|
||||
.handle_http(
|
||||
context(),
|
||||
plugin_request(),
|
||||
compose(Vec::new(), |request: core::Request| {
|
||||
Box::pin(async move {
|
||||
assert_eq!(request.uri(), "/rewritten?encoded=%2F");
|
||||
Ok(core::Response::new(core::empty_body()))
|
||||
})
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
match result {
|
||||
Ok(response) => break response,
|
||||
Err(MiddlewareError::Gateway(error))
|
||||
if error.kind() == GatewayErrorKind::MiddlewareCapacityUnavailable =>
|
||||
{
|
||||
tokio::task::yield_now().await
|
||||
}
|
||||
Err(error) => panic!("unexpected error: {error}"),
|
||||
}
|
||||
}
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.headers()["x-http-plugin"], "active");
|
||||
drop(response);
|
||||
drop(plan);
|
||||
drop(generation);
|
||||
runtime.shutdown().await;
|
||||
crate::generation::wait_until_empty(cache.path()).await;
|
||||
}
|
||||
@@ -1,5 +1,6 @@
|
||||
//! 验证真实插件 HTTP 中间件的请求改写、流式转换与短路响应
|
||||
|
||||
mod capacity;
|
||||
mod dispatch;
|
||||
mod websocket;
|
||||
|
||||
@@ -27,6 +28,19 @@ async fn http_plugin_with_limits(
|
||||
PluginRuntime,
|
||||
ExtensionSetReference,
|
||||
gateway_core::engine::middleware::FrozenMiddlewarePlan,
|
||||
) {
|
||||
http_plugin_with_policy(mode, limits, PluginFailurePolicy::Reject).await
|
||||
}
|
||||
|
||||
async fn http_plugin_with_policy(
|
||||
mode: &str,
|
||||
limits: RpcLimits,
|
||||
failure_policy: PluginFailurePolicy,
|
||||
) -> (
|
||||
tempfile::TempDir,
|
||||
PluginRuntime,
|
||||
ExtensionSetReference,
|
||||
gateway_core::engine::middleware::FrozenMiddlewarePlan,
|
||||
) {
|
||||
let worker = std::fs::read(env!("CARGO_BIN_EXE_gateway-plugin-test-middleware")).unwrap();
|
||||
let package = crate::support::package_with_contributions(
|
||||
@@ -42,12 +56,7 @@ async fn http_plugin_with_limits(
|
||||
vec![InstanceFixture {
|
||||
id: "http-plugin",
|
||||
configuration: serde_json::json!({"http":true,"mode":mode}),
|
||||
bindings: vec![binding(
|
||||
MIDDLEWARE_CONTRIBUTION,
|
||||
"http",
|
||||
0,
|
||||
PluginFailurePolicy::Reject,
|
||||
)],
|
||||
bindings: vec![binding(MIDDLEWARE_CONTRIBUTION, "http", 0, failure_policy)],
|
||||
}],
|
||||
package,
|
||||
limits,
|
||||
|
||||
@@ -327,3 +327,79 @@ async fn sdk_waiting_next_obeys_request_deadline_and_client_cancellation() {
|
||||
super::super::wait_until_empty(cache.path()).await;
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn request_middleware_never_delegates_when_total_in_flight_capacity_is_full() {
|
||||
for policy in [PluginFailurePolicy::Delegate, PluginFailurePolicy::Reject] {
|
||||
let package = crate::support::package_with_contributions(
|
||||
crate::support::worker(),
|
||||
Contributions::from([crate::support::contribution(
|
||||
Capability::Middleware,
|
||||
vec![Stage::Request],
|
||||
vec!["openai".into()],
|
||||
vec!["openai".into()],
|
||||
)]),
|
||||
);
|
||||
let (cache, runtime) = setup_package_with_limits(
|
||||
vec![InstanceFixture {
|
||||
id: "busy-middleware",
|
||||
configuration: serde_json::json!({}),
|
||||
bindings: vec![binding(MIDDLEWARE_CONTRIBUTION, "request", 0, policy)],
|
||||
}],
|
||||
package,
|
||||
RpcLimits {
|
||||
maximum_calls: 1,
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await;
|
||||
let generation = prepare(&runtime).await;
|
||||
let plan = runtime
|
||||
.execution_registry()
|
||||
.middleware(&generation)
|
||||
.unwrap();
|
||||
let (entered, waiting) = tokio::sync::oneshot::channel();
|
||||
let first = tokio::spawn({
|
||||
let plan = plan.clone();
|
||||
async move {
|
||||
plan.handle(
|
||||
middleware_context(ClientTransport::HttpJson),
|
||||
MiddlewareRequest::new(
|
||||
"openai",
|
||||
middleware_headers(),
|
||||
Bytes::from_static(b"{}"),
|
||||
),
|
||||
gateway_core::middleware::compose(Vec::new(), move |_| {
|
||||
Box::pin(async move {
|
||||
entered.send(()).unwrap();
|
||||
std::future::pending().await
|
||||
})
|
||||
}),
|
||||
)
|
||||
.await
|
||||
}
|
||||
});
|
||||
tokio::time::timeout(Duration::from_secs(3), waiting)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
let result = plan
|
||||
.handle(
|
||||
middleware_context(ClientTransport::HttpJson),
|
||||
MiddlewareRequest::new("openai", middleware_headers(), Bytes::from_static(b"{}")),
|
||||
gateway_core::middleware::compose(Vec::new(), |_| {
|
||||
Box::pin(async { panic!("RPC 过载不能跳过请求中间件") })
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
assert!(
|
||||
matches!(result, Err(MiddlewareError::Gateway(error)) if error.kind() == gateway_core::error::GatewayErrorKind::MiddlewareCapacityUnavailable)
|
||||
);
|
||||
first.abort();
|
||||
assert!(matches!(first.await, Err(error) if error.is_cancelled()));
|
||||
drop(plan);
|
||||
drop(generation);
|
||||
runtime.shutdown().await;
|
||||
super::super::wait_until_empty(cache.path()).await;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,328 @@
|
||||
//! 验证执行排队、next 恢复、流式在途预算与取消回收
|
||||
|
||||
use std::{
|
||||
sync::{Arc, atomic::Ordering},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use super::session::{Callbacks, LifecycleCallbacks, prepared_session};
|
||||
use futures::future::BoxFuture;
|
||||
use gateway_plugin_runtime::{CallbackHandler, RpcError, RpcLimits, RpcReply, RpcSession};
|
||||
use gateway_plugin_sdk::{CallContext, PluginFault, Stage, call::middleware::NEXT_METHOD};
|
||||
use serde_json::json;
|
||||
use tokio::sync::{Semaphore, mpsc};
|
||||
|
||||
async fn limited<C: CallbackHandler + 'static>(
|
||||
callbacks: Arc<C>,
|
||||
limits: RpcLimits,
|
||||
) -> (tempfile::TempDir, Arc<RpcSession>) {
|
||||
let (cache, prepared, handshake, processes) = prepared_session();
|
||||
let session = RpcSession::start(prepared, handshake, limits, &processes, callbacks)
|
||||
.await
|
||||
.unwrap();
|
||||
(cache, Arc::new(session))
|
||||
}
|
||||
|
||||
fn limits() -> RpcLimits {
|
||||
RpcLimits {
|
||||
maximum_calls: 4,
|
||||
maximum_executing_calls: 1,
|
||||
maximum_queued_calls: 1,
|
||||
maximum_queue_timeout: Duration::from_millis(200),
|
||||
maximum_callbacks: 1,
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
async fn wait_begun(callbacks: &LifecycleCallbacks, count: usize) {
|
||||
tokio::time::timeout(Duration::from_secs(2), async {
|
||||
while callbacks.begun.load(Ordering::Acquire) < count {
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
fn call(session: &RpcSession, method: &str) -> impl Future<Output = Result<RpcReply, RpcError>> {
|
||||
session.call(
|
||||
method,
|
||||
session.context(Stage::Request, Duration::from_secs(2)),
|
||||
json!({}),
|
||||
Vec::new(),
|
||||
)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bounded_initial_queue_waits_and_releases_cancelled_or_expired_waiters() {
|
||||
let callbacks = Arc::new(LifecycleCallbacks::default());
|
||||
let (_cache, session) = limited(callbacks.clone(), limits()).await;
|
||||
let active = tokio::spawn({
|
||||
let session = session.clone();
|
||||
async move { call(&session, "hang").await }
|
||||
});
|
||||
wait_begun(&callbacks, 1).await;
|
||||
|
||||
let mut queued = Box::pin(call(&session, "echo"));
|
||||
assert!(futures::poll!(&mut queued).is_pending());
|
||||
assert!(matches!(
|
||||
call(&session, "echo").await,
|
||||
Err(RpcError::Overloaded)
|
||||
));
|
||||
assert_eq!(callbacks.begun.load(Ordering::Acquire), 1);
|
||||
drop(queued);
|
||||
|
||||
let expired = session
|
||||
.call(
|
||||
"echo",
|
||||
session.context(Stage::Request, Duration::from_millis(30)),
|
||||
json!({}),
|
||||
Vec::new(),
|
||||
)
|
||||
.await;
|
||||
assert!(matches!(expired, Err(RpcError::Overloaded)));
|
||||
assert_eq!(callbacks.begun.load(Ordering::Acquire), 1);
|
||||
|
||||
let mut queued = Box::pin(call(&session, "echo"));
|
||||
assert!(futures::poll!(&mut queued).is_pending());
|
||||
active.abort();
|
||||
assert!(matches!(active.await, Err(error) if error.is_cancelled()));
|
||||
assert!(queued.await.is_ok());
|
||||
assert_eq!(callbacks.begun.load(Ordering::Acquire), 2);
|
||||
assert!(session.is_ready());
|
||||
session.shutdown(Duration::from_secs(1)).await;
|
||||
}
|
||||
|
||||
struct NextCallbacks {
|
||||
entered: mpsc::UnboundedSender<()>,
|
||||
release: Arc<Semaphore>,
|
||||
}
|
||||
|
||||
impl CallbackHandler for NextCallbacks {
|
||||
fn call(
|
||||
&self,
|
||||
_context: CallContext,
|
||||
method: String,
|
||||
params: serde_json::Value,
|
||||
payload: Vec<u8>,
|
||||
) -> BoxFuture<'static, Result<RpcReply, PluginFault>> {
|
||||
let release = self.release.clone();
|
||||
self.entered.send(()).unwrap();
|
||||
Box::pin(async move {
|
||||
if method == NEXT_METHOD {
|
||||
release.acquire().await.unwrap().forget();
|
||||
}
|
||||
Ok(RpcReply {
|
||||
result: params,
|
||||
payload,
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn next_resumption_does_not_need_a_vacant_initial_queue_slot() {
|
||||
let (entered, mut started) = mpsc::unbounded_channel();
|
||||
let release = Arc::new(Semaphore::new(0));
|
||||
let (_cache, session) = limited(
|
||||
Arc::new(NextCallbacks {
|
||||
entered,
|
||||
release: release.clone(),
|
||||
}),
|
||||
limits(),
|
||||
)
|
||||
.await;
|
||||
let waiting = tokio::spawn({
|
||||
let session = session.clone();
|
||||
async move {
|
||||
session
|
||||
.call(
|
||||
"callback_method",
|
||||
session.context(Stage::Request, Duration::from_secs(2)),
|
||||
json!({"method": NEXT_METHOD}),
|
||||
Vec::new(),
|
||||
)
|
||||
.await
|
||||
}
|
||||
});
|
||||
started.recv().await.unwrap();
|
||||
// 等待 next 不挤占唯一的普通回调额度
|
||||
let mut occupying = Box::pin(call(&session, "callback"));
|
||||
assert!(futures::poll!(&mut occupying).is_pending());
|
||||
started.recv().await.unwrap();
|
||||
occupying.await.unwrap();
|
||||
let mut active = Box::pin(call(&session, "hang"));
|
||||
assert!(futures::poll!(&mut active).is_pending());
|
||||
let mut queued = Box::pin(call(&session, "echo"));
|
||||
assert!(futures::poll!(&mut queued).is_pending());
|
||||
assert!(matches!(
|
||||
call(&session, "echo").await,
|
||||
Err(RpcError::Overloaded)
|
||||
));
|
||||
release.add_permits(1);
|
||||
// 恢复等待不能立即因初次队列已满而拒绝,也不能绕过执行名额完成
|
||||
tokio::time::sleep(Duration::from_millis(30)).await;
|
||||
assert!(!waiting.is_finished());
|
||||
drop(active);
|
||||
assert!(queued.await.is_ok());
|
||||
assert!(waiting.await.unwrap().is_ok());
|
||||
assert!(session.is_ready());
|
||||
session.shutdown(Duration::from_secs(1)).await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn response_stream_releases_execution_but_retains_total_capacity_until_drop() {
|
||||
let callbacks = Arc::new(Callbacks::default());
|
||||
let (_cache, session) = limited(
|
||||
callbacks,
|
||||
RpcLimits {
|
||||
maximum_calls: 2,
|
||||
..limits()
|
||||
},
|
||||
)
|
||||
.await;
|
||||
let mut streams = Vec::new();
|
||||
for _ in 0..2 {
|
||||
streams.push(
|
||||
session
|
||||
.call_stream(
|
||||
"stream",
|
||||
session.context(Stage::Request, Duration::from_secs(5)),
|
||||
json!({}),
|
||||
Vec::new(),
|
||||
)
|
||||
.await
|
||||
.unwrap(),
|
||||
);
|
||||
}
|
||||
assert!(matches!(
|
||||
call(&session, "echo").await,
|
||||
Err(RpcError::Overloaded)
|
||||
));
|
||||
drop(streams.pop());
|
||||
tokio::time::timeout(Duration::from_secs(1), async {
|
||||
loop {
|
||||
match call(&session, "echo").await {
|
||||
Ok(_) => break,
|
||||
Err(RpcError::Overloaded) => tokio::task::yield_now().await,
|
||||
Err(error) => panic!("unexpected error: {error}"),
|
||||
}
|
||||
}
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(session.is_ready());
|
||||
drop(streams);
|
||||
session.shutdown(Duration::from_secs(1)).await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn queued_resumption_is_bounded_and_cancellable_without_failing_session() {
|
||||
for cancelled in [false, true] {
|
||||
let (entered, mut started) = mpsc::unbounded_channel();
|
||||
let release = Arc::new(Semaphore::new(0));
|
||||
let (_cache, session) = limited(
|
||||
Arc::new(NextCallbacks {
|
||||
entered,
|
||||
release: release.clone(),
|
||||
}),
|
||||
limits(),
|
||||
)
|
||||
.await;
|
||||
let waiting = tokio::spawn({
|
||||
let session = session.clone();
|
||||
async move {
|
||||
session
|
||||
.call(
|
||||
"callback_method",
|
||||
session.context(Stage::Request, Duration::from_millis(200)),
|
||||
json!({"method": NEXT_METHOD}),
|
||||
Vec::new(),
|
||||
)
|
||||
.await
|
||||
}
|
||||
});
|
||||
started.recv().await.unwrap();
|
||||
let mut active = Box::pin(call(&session, "hang"));
|
||||
assert!(futures::poll!(&mut active).is_pending());
|
||||
release.add_permits(1);
|
||||
if cancelled {
|
||||
tokio::time::sleep(Duration::from_millis(30)).await;
|
||||
assert!(!waiting.is_finished());
|
||||
waiting.abort();
|
||||
assert!(matches!(waiting.await, Err(error) if error.is_cancelled()));
|
||||
} else {
|
||||
assert!(matches!(waiting.await.unwrap(), Err(RpcError::Overloaded)));
|
||||
}
|
||||
drop(active);
|
||||
assert!(call(&session, "echo").await.is_ok());
|
||||
assert!(session.is_ready());
|
||||
session.shutdown(Duration::from_secs(1)).await;
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn initial_queue_wait_does_not_consume_the_plugin_processing_budget() {
|
||||
let (_cache, session) = limited(
|
||||
Arc::new(Callbacks::default()),
|
||||
RpcLimits {
|
||||
maximum_queue_timeout: Duration::from_secs(2),
|
||||
..limits()
|
||||
},
|
||||
)
|
||||
.await;
|
||||
let mut active = Box::pin(call(&session, "hang"));
|
||||
assert!(futures::poll!(&mut active).is_pending());
|
||||
let mut queued = Box::pin(session.call(
|
||||
"echo",
|
||||
session.context(Stage::Request, Duration::from_millis(200)),
|
||||
json!({}),
|
||||
Vec::new(),
|
||||
));
|
||||
assert!(futures::poll!(&mut queued).is_pending());
|
||||
tokio::time::sleep(Duration::from_millis(350)).await;
|
||||
drop(active);
|
||||
assert!(queued.await.is_ok());
|
||||
assert!(session.is_ready());
|
||||
session.shutdown(Duration::from_secs(1)).await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resumption_queue_wait_preserves_the_remaining_plugin_processing_budget() {
|
||||
let (entered, mut started) = mpsc::unbounded_channel();
|
||||
let release = Arc::new(Semaphore::new(0));
|
||||
let (_cache, session) = limited(
|
||||
Arc::new(NextCallbacks {
|
||||
entered,
|
||||
release: release.clone(),
|
||||
}),
|
||||
RpcLimits {
|
||||
maximum_queue_timeout: Duration::from_secs(2),
|
||||
..limits()
|
||||
},
|
||||
)
|
||||
.await;
|
||||
let waiting = tokio::spawn({
|
||||
let session = session.clone();
|
||||
async move {
|
||||
session
|
||||
.call(
|
||||
"callback_method",
|
||||
session.context(Stage::Request, Duration::from_millis(200)),
|
||||
json!({"method": NEXT_METHOD}),
|
||||
Vec::new(),
|
||||
)
|
||||
.await
|
||||
}
|
||||
});
|
||||
started.recv().await.unwrap();
|
||||
let mut active = Box::pin(call(&session, "hang"));
|
||||
assert!(futures::poll!(&mut active).is_pending());
|
||||
release.add_permits(1);
|
||||
tokio::time::sleep(Duration::from_millis(350)).await;
|
||||
assert!(!waiting.is_finished());
|
||||
drop(active);
|
||||
assert!(waiting.await.unwrap().is_ok());
|
||||
assert!(session.is_ready());
|
||||
session.shutdown(Duration::from_secs(1)).await;
|
||||
}
|
||||
@@ -1,4 +1,5 @@
|
||||
//! 插件双向 RPC 分派与会话生命周期的测试入口
|
||||
|
||||
mod admission;
|
||||
mod dispatch;
|
||||
mod session;
|
||||
|
||||
@@ -17,7 +17,7 @@ use serde_json::json;
|
||||
use tokio::sync::{Barrier, Notify};
|
||||
|
||||
#[derive(Default)]
|
||||
struct Callbacks {
|
||||
pub(super) struct Callbacks {
|
||||
called: AtomicUsize,
|
||||
nested: Mutex<Option<Weak<RpcSession>>>,
|
||||
}
|
||||
@@ -29,8 +29,8 @@ struct BlockingCallbacks {
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct LifecycleCallbacks {
|
||||
begun: AtomicUsize,
|
||||
pub(super) struct LifecycleCallbacks {
|
||||
pub(super) begun: AtomicUsize,
|
||||
finished: AtomicUsize,
|
||||
}
|
||||
|
||||
@@ -193,7 +193,7 @@ where
|
||||
(cache, Arc::new(session))
|
||||
}
|
||||
|
||||
fn prepared_session() -> (
|
||||
pub(super) fn prepared_session() -> (
|
||||
tempfile::TempDir,
|
||||
Arc<gateway_plugin_runtime::PreparedPackage>,
|
||||
Handshake,
|
||||
|
||||
@@ -52,6 +52,12 @@ use gateway_plugin_sdk::client::{read_frame, write_frame};
|
||||
`PluginSession::accept/run` 和 `PluginHandler`。组合插件从作者清单建立类型化方法,SDK 自动生成
|
||||
`plugin.register`;作者不需要复制贡献项或手写方法名、阶段和敏感载荷编解码。
|
||||
会话管理握手、有界并发、父调用关联、回调、流信用、取消、排空与关闭。
|
||||
|
||||
`SessionConfig.maximum_calls` 默认 256,限制包含下游等待和响应流在内的在途调用;`maximum_callbacks` 默认 64,
|
||||
限制普通宿主回调。`next` 使用独立回调预算,上限为 `maximum_calls`,不挤占普通回调容量。
|
||||
宿主负责执行与准入调度,等待与处理独立计时;默认预算与过载行为见[插件使用](../../../../docs/plugins.md#故障处理)。
|
||||
插件显式覆盖配置时需与宿主协调在途预算;较小的插件侧上限会提前拒绝新调用
|
||||
|
||||
`ResponseStream::channel` 提供有界流生产入口,`ResponseStream::pull` 按消费进度驱动生产者;`PluginCall` 的宿主客户端只代表当前父调用,
|
||||
不能把一个调用的资源误用到另一个调用。插件程序自行选择运行时和启动入口,不会启动网关。
|
||||
`SessionConfig.maximum_stream_chunk_bytes` 默认 1 MiB,只约束单个业务流分块;
|
||||
|
||||
@@ -243,7 +243,9 @@ SDK 不自动重试。超时或断连不能证明写入未提交;重试上限
|
||||
`next.run(input)` 消费自身,最多执行一次;结果按相反顺序经过外层。重试或独立子请求通过宿主执行服务发起
|
||||
|
||||
调用上下文的 `timeout_ms` 限制插件自身处理时间,默认上限为 120 秒。等待宿主 `next` 时暂停这份预算,
|
||||
收到结果或错误后恢复剩余时间;入站处理与返回处理合计使用同一预算,普通宿主回调不会暂停计时。
|
||||
宿主等待期间归还执行名额,返回后先重新取得名额再交付结果或错误,插件继续同一次调用。
|
||||
初次准入和恢复排队各有独立等待期限,默认 30 秒,不扣插件处理预算;恢复等待受总在途上限约束,不占初次准入队列。
|
||||
入站处理与返回处理合计使用同一预算,普通宿主回调不会暂停计时。
|
||||
请求总期限、下游操作超时和取消信号独立生效,客户端取消或插件退出仍会结束等待并回收资源。
|
||||
宿主与插件 SDK 必须同时支持此语义;仍使用固定绝对期限的 SDK 构建的插件需重新构建,单独升级宿主不能解除插件侧期限
|
||||
|
||||
|
||||
@@ -51,6 +51,11 @@ pub struct PrivacyRule {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub enabled: bool,
|
||||
pub steps: Vec<PrivacyStep>,
|
||||
}
|
||||
#[derive(Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
pub struct PrivacyStep {
|
||||
pub scope: PrivacyScope,
|
||||
pub selector: String,
|
||||
pub action: PrivacyAction,
|
||||
|
||||
@@ -25,8 +25,8 @@ impl Default for SessionConfig {
|
||||
Self {
|
||||
// 流分块还需匹配宿主授予的信用窗口
|
||||
maximum_stream_chunk_bytes: 1024 * 1024,
|
||||
maximum_calls: 32,
|
||||
maximum_callbacks: 32,
|
||||
maximum_calls: 256,
|
||||
maximum_callbacks: 64,
|
||||
maximum_buffered_stream_chunks: 1_024,
|
||||
handshake_timeout: Duration::from_secs(5),
|
||||
maximum_call_timeout: Duration::from_secs(120),
|
||||
|
||||
@@ -91,6 +91,7 @@ where
|
||||
.config
|
||||
.maximum_calls
|
||||
.checked_add(self.config.maximum_callbacks)
|
||||
.and_then(|capacity| capacity.checked_add(self.config.maximum_calls))
|
||||
.ok_or(SessionError::Configuration)?;
|
||||
let control_reserve = self
|
||||
.config
|
||||
@@ -105,7 +106,10 @@ where
|
||||
sender: outbound,
|
||||
data_slots: Arc::new(Semaphore::new(data_capacity)),
|
||||
};
|
||||
let callbacks = Arc::new(CallbackRegistry::new(self.config.maximum_callbacks));
|
||||
let callbacks = Arc::new(CallbackRegistry::new(
|
||||
self.config.maximum_callbacks,
|
||||
self.config.maximum_calls,
|
||||
));
|
||||
let handler: Arc<dyn PluginHandler> = Arc::new(handler);
|
||||
let lifecycle = LifecycleHooks::new(Arc::clone(&handler), control_reserve)?;
|
||||
write_frame(
|
||||
|
||||
@@ -67,7 +67,7 @@ impl HostClient {
|
||||
ordered = self.callbacks.order.lock() => ordered,
|
||||
};
|
||||
let managed_next = method == NEXT_METHOD;
|
||||
let (id, received) = self.callbacks.begin(self.parent_id)?;
|
||||
let (id, received) = self.callbacks.begin(self.parent_id, managed_next)?;
|
||||
let frame = Frame {
|
||||
message: Message::Callback {
|
||||
id,
|
||||
@@ -114,6 +114,7 @@ pub(super) struct CallbackRegistry {
|
||||
state: Mutex<CallbackState>,
|
||||
order: tokio::sync::Mutex<()>,
|
||||
maximum: usize,
|
||||
maximum_next: usize,
|
||||
maximum_retired: usize,
|
||||
}
|
||||
|
||||
@@ -126,12 +127,13 @@ struct CallbackState {
|
||||
|
||||
struct PendingCallback {
|
||||
parent_id: u64,
|
||||
managed_next: bool,
|
||||
pause: Option<CallDeadlinePause>,
|
||||
result: oneshot::Sender<Result<HostReply, SessionError>>,
|
||||
}
|
||||
|
||||
impl CallbackRegistry {
|
||||
pub(super) fn new(maximum: usize) -> Self {
|
||||
pub(super) fn new(maximum: usize, maximum_next: usize) -> Self {
|
||||
Self {
|
||||
state: Mutex::new(CallbackState {
|
||||
next_id: 2,
|
||||
@@ -141,19 +143,32 @@ impl CallbackRegistry {
|
||||
}),
|
||||
order: tokio::sync::Mutex::new(()),
|
||||
maximum,
|
||||
maximum_retired: maximum.saturating_mul(2),
|
||||
maximum_next,
|
||||
maximum_retired: (maximum + maximum_next).saturating_mul(2),
|
||||
}
|
||||
}
|
||||
|
||||
fn begin(
|
||||
&self,
|
||||
parent_id: u64,
|
||||
managed_next: bool,
|
||||
) -> Result<(u64, oneshot::Receiver<Result<HostReply, SessionError>>), SessionError> {
|
||||
let mut state = self
|
||||
.state
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
if state.pending.len() >= self.maximum {
|
||||
let maximum = if managed_next {
|
||||
self.maximum_next
|
||||
} else {
|
||||
self.maximum
|
||||
};
|
||||
if state
|
||||
.pending
|
||||
.values()
|
||||
.filter(|pending| pending.managed_next == managed_next)
|
||||
.count()
|
||||
>= maximum
|
||||
{
|
||||
return Err(SessionError::Capacity);
|
||||
}
|
||||
let id = state.next_id;
|
||||
@@ -163,6 +178,7 @@ impl CallbackRegistry {
|
||||
id,
|
||||
PendingCallback {
|
||||
parent_id,
|
||||
managed_next,
|
||||
result,
|
||||
pause: None,
|
||||
},
|
||||
|
||||
@@ -140,10 +140,15 @@ impl PluginHandler for TestHandler {
|
||||
tokio::time::sleep(Duration::from_millis(80)).await;
|
||||
Ok(CallReply::unary(call.params, call.payload))
|
||||
}
|
||||
"callback" => {
|
||||
"callback" | "next_callback" => {
|
||||
let method = if call.method == "next_callback" {
|
||||
NEXT_METHOD
|
||||
} else {
|
||||
"host.log"
|
||||
};
|
||||
let reply = call
|
||||
.host
|
||||
.call("host.log", call.params, call.payload)
|
||||
.call(method, call.params, call.payload)
|
||||
.await
|
||||
.map_err(SessionError::into_plugin_fault)?;
|
||||
Ok(CallReply::unary(reply.result, reply.payload))
|
||||
@@ -248,23 +253,31 @@ pub(super) async fn start_session<H: PluginHandler>(
|
||||
async fn start_session_with_capacity<H: PluginHandler>(
|
||||
handler: H,
|
||||
capacity: usize,
|
||||
) -> (HostPeer, JoinHandle<Result<(), SessionError>>) {
|
||||
start_configured_session(
|
||||
handler,
|
||||
capacity,
|
||||
SessionConfig {
|
||||
maximum_stream_chunk_bytes: MAXIMUM_STREAM_CHUNK_BYTES,
|
||||
maximum_calls: 4,
|
||||
maximum_callbacks: 4,
|
||||
maximum_buffered_stream_chunks: 16,
|
||||
handshake_timeout: Duration::from_secs(1),
|
||||
maximum_call_timeout: Duration::from_secs(2),
|
||||
},
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn start_configured_session<H: PluginHandler>(
|
||||
handler: H,
|
||||
capacity: usize,
|
||||
config: SessionConfig,
|
||||
) -> (HostPeer, JoinHandle<Result<(), SessionError>>) {
|
||||
let (host, plugin) = tokio::io::duplex(capacity);
|
||||
let (plugin_reader, plugin_writer) = tokio::io::split(plugin);
|
||||
let task = tokio::spawn(async move {
|
||||
let session = PluginSession::accept(
|
||||
plugin_reader,
|
||||
plugin_writer,
|
||||
SessionConfig {
|
||||
maximum_stream_chunk_bytes: MAXIMUM_STREAM_CHUNK_BYTES,
|
||||
maximum_calls: 4,
|
||||
maximum_callbacks: 4,
|
||||
maximum_buffered_stream_chunks: 16,
|
||||
handshake_timeout: Duration::from_secs(1),
|
||||
maximum_call_timeout: Duration::from_secs(2),
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
let session = PluginSession::accept(plugin_reader, plugin_writer, config).await?;
|
||||
session.run(handler).await
|
||||
});
|
||||
let (mut reader, mut writer) = tokio::io::split(host);
|
||||
|
||||
@@ -297,3 +297,57 @@ async fn typed_key_and_quota_calls_keep_payloads_and_do_not_retry_failures() {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn managed_next_does_not_consume_the_bounded_ordinary_callback_budget() {
|
||||
let (mut host, task) = start_configured_session(
|
||||
TestHandler::default(),
|
||||
MAXIMUM_STREAM_CHUNK_BYTES * 2,
|
||||
SessionConfig {
|
||||
maximum_calls: 4,
|
||||
maximum_callbacks: 1,
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await;
|
||||
send_call(&mut host, 1, "next_callback", json!({}), Vec::new()).await;
|
||||
let Message::Callback {
|
||||
id: next,
|
||||
parent_id: 1,
|
||||
method,
|
||||
..
|
||||
} = receive(&mut host).await.message
|
||||
else {
|
||||
panic!("expected next callback")
|
||||
};
|
||||
assert_eq!(method, NEXT_METHOD);
|
||||
send_call(&mut host, 3, "callback", json!({}), Vec::new()).await;
|
||||
let Message::Callback {
|
||||
id: ordinary,
|
||||
parent_id: 3,
|
||||
method,
|
||||
..
|
||||
} = receive(&mut host).await.message
|
||||
else {
|
||||
panic!("next must not occupy ordinary callback capacity")
|
||||
};
|
||||
assert_eq!(method, "host.log");
|
||||
send_call(&mut host, 5, "callback", json!({}), Vec::new()).await;
|
||||
assert!(
|
||||
matches!(receive(&mut host).await.message, Message::Error { id: 5, error } if error.code == ErrorCode::Capacity)
|
||||
);
|
||||
for (callback, parent) in [(ordinary, 3), (next, 1)] {
|
||||
send_control(
|
||||
&mut host,
|
||||
Message::Result {
|
||||
id: callback,
|
||||
result: json!({}),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
assert!(
|
||||
matches!(receive(&mut host).await.message, Message::Result { id, .. } if id == parent)
|
||||
);
|
||||
}
|
||||
shutdown(&mut host, task).await;
|
||||
}
|
||||
|
||||
@@ -599,7 +599,7 @@ fn retry_after_seconds_from_resets_at(error: &Value) -> Option<u64> {
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.ok()?
|
||||
.as_secs();
|
||||
(resets_at > now).then_some(resets_at - now)
|
||||
(resets_at > now).then(|| resets_at - now)
|
||||
}
|
||||
|
||||
fn positive_u64(value: &Value) -> Option<u64> {
|
||||
|
||||
@@ -394,6 +394,19 @@ fn retry_after_should_read_string_reset_timestamp() {
|
||||
assert!((59..=60).contains(&retry_after));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn retry_after_should_ignore_elapsed_reset_timestamps_and_keep_later_advice() {
|
||||
for reset_at in [json!(1), json!("1")] {
|
||||
let mut body = json!({
|
||||
"error": {"type": "usage_limit_reached", "resets_at": reset_at}
|
||||
});
|
||||
assert_eq!(retry_after_seconds_from_body(&body.to_string()), None);
|
||||
|
||||
body["retry_after_seconds"] = json!(7);
|
||||
assert_eq!(retry_after_seconds_from_body(&body.to_string()), Some(7));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn retry_after_should_read_nested_error_delay() {
|
||||
let body = json!({
|
||||
|
||||
@@ -255,7 +255,11 @@ async fn connect(
|
||||
.with_budget(Arc::new(
|
||||
postgres::PgClientBudgetStore::new(pool.clone()).with_timezone(config.timezone),
|
||||
))
|
||||
.with_session_affinity(Arc::clone(&provider_session_affinity));
|
||||
.with_session_affinity(Arc::clone(&provider_session_affinity))
|
||||
.with_client_retry_store(Arc::new(redis::RedisClientRetryRepository::new(
|
||||
redis_connection.clone(),
|
||||
REDIS_NAMESPACE,
|
||||
)?));
|
||||
|
||||
let provider_ports = ProviderStorePorts::new(
|
||||
account_store,
|
||||
|
||||
@@ -226,7 +226,7 @@ impl AccountGroupStore for PgAccountGroupRepository {
|
||||
.map(|group_id| group_id.as_str().to_owned())
|
||||
.collect::<Vec<_>>();
|
||||
let rows = sqlx::query(
|
||||
"select proxy.auto_location, proxy.detected_location_json, proxy.location_country, proxy.location_region, proxy.location_city, proxy.location_timezone, membership.account_group_id,
|
||||
"select coalesce(proxy.enabled, true) as outbound_proxy_enabled, proxy.auto_location, proxy.detected_location_json, proxy.location_country, proxy.location_region, proxy.location_city, proxy.location_timezone, membership.account_group_id,
|
||||
account.id, account.provider_kind, account.name, account.notes, account.email,
|
||||
account.upstream_user_id, account.upstream_account_id, account.plan_type,
|
||||
account.authentication_kind, account.credential_revision, account.outbound_proxy_url,
|
||||
|
||||
@@ -106,7 +106,7 @@ pub(crate) async fn load_admin_account_page(
|
||||
from runtime_settings
|
||||
where id = 1
|
||||
)
|
||||
select p.auto_location, p.detected_location_json, p.location_country, p.location_region, p.location_city, p.location_timezone, a.outbound_proxy_url, a.id, a.provider_kind, a.name, a.notes, a.email, a.upstream_user_id,
|
||||
select coalesce(p.enabled, true) as outbound_proxy_enabled, p.auto_location, p.detected_location_json, p.location_country, p.location_region, p.location_city, p.location_timezone, a.outbound_proxy_url, a.id, a.provider_kind, a.name, a.notes, a.email, a.upstream_user_id,
|
||||
a.upstream_account_id, a.plan_type, a.authentication_kind,
|
||||
a.credential_revision, a.has_refresh_token, a.access_token_expires_at,
|
||||
a.next_refresh_at, a.enabled, a.concurrency_limit, a.weight, a.model_access_json,
|
||||
|
||||
@@ -10,6 +10,7 @@ pub(crate) fn admin_account_record(
|
||||
summary: ProviderAccountSummary,
|
||||
) -> AdminStoreResult<AccountRecord> {
|
||||
Ok(AccountRecord {
|
||||
outbound_proxy_enabled: summary.outbound_proxy_enabled,
|
||||
outbound_proxy: summary.outbound_proxy,
|
||||
id: summary.id,
|
||||
provider_kind: ProviderKind::new(summary.provider_kind).map_err(|_| {
|
||||
|
||||
@@ -122,14 +122,14 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
|
||||
return Err(invalid("plugin account page limit must be positive"));
|
||||
}
|
||||
let rows = sqlx::query(
|
||||
"select auto_location, detected_location_json, location_country, location_region, location_city, location_timezone, outbound_proxy_url, id, provider_kind, name, notes, email, upstream_user_id,
|
||||
"select coalesce(proxy_enabled, true) as outbound_proxy_enabled, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone, outbound_proxy_url, id, provider_kind, name, notes, email, upstream_user_id,
|
||||
upstream_account_id, plan_type, authentication_kind, credential_revision, has_refresh_token,
|
||||
access_token_expires_at, next_refresh_at, enabled, concurrency_limit, weight, model_access_json, credential_state,
|
||||
credential_observed_at, quota_access_state, quota_evidence,
|
||||
quota_access_observed_at, quota_reset_at,
|
||||
quota_observed_at, last_error_reason, last_error_message, created_at, updated_at
|
||||
from provider_accounts
|
||||
left join (select id as location_proxy_id, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone from outbound_proxies) proxy_location
|
||||
left join (select id as location_proxy_id, enabled as proxy_enabled, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone from outbound_proxies) proxy_location
|
||||
on outbound_proxy_id = location_proxy_id
|
||||
where ($1::text is null or provider_kind = $1)
|
||||
and ($2::text is null or id > $2)
|
||||
@@ -151,14 +151,14 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
|
||||
include_disabled: bool,
|
||||
) -> StoreResult<Vec<ProviderAccountSummary>> {
|
||||
let rows = sqlx::query(
|
||||
"select auto_location, detected_location_json, location_country, location_region, location_city, location_timezone, outbound_proxy_url, id, provider_kind, name, notes, email, upstream_user_id,
|
||||
"select coalesce(proxy_enabled, true) as outbound_proxy_enabled, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone, outbound_proxy_url, id, provider_kind, name, notes, email, upstream_user_id,
|
||||
upstream_account_id, plan_type, authentication_kind, credential_revision, has_refresh_token,
|
||||
access_token_expires_at, next_refresh_at, enabled, concurrency_limit, weight, model_access_json, credential_state,
|
||||
credential_observed_at, quota_access_state, quota_evidence,
|
||||
quota_access_observed_at, quota_reset_at,
|
||||
quota_observed_at, last_error_reason, last_error_message, created_at, updated_at
|
||||
from provider_accounts
|
||||
left join (select id as location_proxy_id, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone from outbound_proxies) proxy_location
|
||||
left join (select id as location_proxy_id, enabled as proxy_enabled, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone from outbound_proxies) proxy_location
|
||||
on outbound_proxy_id = location_proxy_id
|
||||
where ($1::text is null or provider_kind = $1) and ($2 or enabled)
|
||||
order by provider_kind, name, id",
|
||||
|
||||
@@ -51,6 +51,7 @@ pub(crate) fn parse_error_reason(value: Option<String>) -> StoreResult<Option<Ac
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ProviderAccountSummary {
|
||||
pub outbound_proxy_enabled: bool,
|
||||
pub request_location: Option<gateway_core::account::RequestLocation>,
|
||||
pub outbound_proxy: Option<gateway_core::account::OutboundProxy>,
|
||||
pub id: String,
|
||||
@@ -323,37 +324,37 @@ impl ProviderAccountStateUpdate {
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) const ACCOUNT_SELECT: &str = "select auto_location, detected_location_json, location_country, location_region, location_city, location_timezone, outbound_proxy_url, id, provider_kind, name, notes, email, upstream_user_id,
|
||||
pub(crate) const ACCOUNT_SELECT: &str = "select coalesce(proxy_enabled, true) as outbound_proxy_enabled, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone, outbound_proxy_url, id, provider_kind, name, notes, email, upstream_user_id,
|
||||
upstream_account_id, plan_type, authentication_kind, provider_credentials_json, credential_revision,
|
||||
has_refresh_token, access_token_expires_at, next_refresh_at, enabled, concurrency_limit, weight, model_access_json, credential_state,
|
||||
provider_quota_json, quota_access_state, quota_evidence, quota_access_observed_at, quota_reset_at,
|
||||
last_error_reason, last_error_message,
|
||||
credential_observed_at, quota_observed_at, created_at, updated_at
|
||||
from provider_accounts
|
||||
left join (select id as location_proxy_id, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone from outbound_proxies) proxy_location
|
||||
left join (select id as location_proxy_id, enabled as proxy_enabled, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone from outbound_proxies) proxy_location
|
||||
on outbound_proxy_id = location_proxy_id
|
||||
where id = $1";
|
||||
|
||||
pub(crate) const ACCOUNT_SELECT_BY_IDS: &str = "select auto_location, detected_location_json, location_country, location_region, location_city, location_timezone, outbound_proxy_url, id, provider_kind, name, notes, email, upstream_user_id,
|
||||
pub(crate) const ACCOUNT_SELECT_BY_IDS: &str = "select coalesce(proxy_enabled, true) as outbound_proxy_enabled, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone, outbound_proxy_url, id, provider_kind, name, notes, email, upstream_user_id,
|
||||
upstream_account_id, plan_type, authentication_kind, provider_credentials_json, credential_revision,
|
||||
has_refresh_token, access_token_expires_at, next_refresh_at, enabled, concurrency_limit, weight, model_access_json, credential_state,
|
||||
provider_quota_json, quota_access_state, quota_evidence, quota_access_observed_at, quota_reset_at,
|
||||
last_error_reason, last_error_message,
|
||||
credential_observed_at, quota_observed_at, created_at, updated_at
|
||||
from provider_accounts
|
||||
left join (select id as location_proxy_id, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone from outbound_proxies) proxy_location
|
||||
left join (select id as location_proxy_id, enabled as proxy_enabled, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone from outbound_proxies) proxy_location
|
||||
on outbound_proxy_id = location_proxy_id
|
||||
where id = any($1::text[]) and provider_kind = $2
|
||||
order by id";
|
||||
|
||||
pub(crate) const REFRESH_CANDIDATES_SELECT: &str = "select auto_location, detected_location_json, location_country, location_region, location_city, location_timezone, outbound_proxy_url, id, provider_kind, name, notes, email, upstream_user_id,
|
||||
pub(crate) const REFRESH_CANDIDATES_SELECT: &str = "select coalesce(proxy_enabled, true) as outbound_proxy_enabled, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone, outbound_proxy_url, id, provider_kind, name, notes, email, upstream_user_id,
|
||||
upstream_account_id, plan_type, authentication_kind, provider_credentials_json, credential_revision,
|
||||
has_refresh_token, access_token_expires_at, next_refresh_at, enabled, concurrency_limit, weight, model_access_json, credential_state,
|
||||
provider_quota_json, quota_access_state, quota_evidence, quota_access_observed_at, quota_reset_at,
|
||||
last_error_reason, last_error_message,
|
||||
credential_observed_at, quota_observed_at, created_at, updated_at
|
||||
from provider_accounts
|
||||
left join (select id as location_proxy_id, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone from outbound_proxies) proxy_location
|
||||
left join (select id as location_proxy_id, enabled as proxy_enabled, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone from outbound_proxies) proxy_location
|
||||
on outbound_proxy_id = location_proxy_id
|
||||
where provider_kind = $1
|
||||
and has_refresh_token
|
||||
@@ -424,7 +425,11 @@ pub(crate) fn core_account_from_summary(
|
||||
)
|
||||
.with_scheduling(summary.concurrency_limit, summary.weight)
|
||||
.with_model_access(summary.model_access)
|
||||
.with_outbound_proxy(summary.outbound_proxy)
|
||||
.with_outbound_proxy(
|
||||
summary
|
||||
.outbound_proxy
|
||||
.filter(|_| summary.outbound_proxy_enabled),
|
||||
)
|
||||
.with_request_location(summary.request_location)
|
||||
.with_refresh_schedule(
|
||||
summary.has_refresh_token,
|
||||
@@ -460,6 +465,7 @@ pub(crate) fn account_summary_from_row(
|
||||
.and_then(AccountWeight::new)
|
||||
.ok_or_else(|| invalid("invalid weight"))?;
|
||||
Ok(ProviderAccountSummary {
|
||||
outbound_proxy_enabled: get(&row, "outbound_proxy_enabled")?,
|
||||
request_location: super::super::proxies::effective_location_from_row(&row)?,
|
||||
outbound_proxy: get::<Option<String>>(&row, "outbound_proxy_url")?
|
||||
.map(|url| {
|
||||
|
||||
@@ -113,6 +113,9 @@ fn record(row: PgRow) -> StoreResult<ProxyRecord> {
|
||||
.try_get("last_test_latency_ms")
|
||||
.map_err(|source| invalid().with_source(source))?;
|
||||
Ok(ProxyRecord {
|
||||
enabled: row
|
||||
.try_get("enabled")
|
||||
.map_err(|source| invalid().with_source(source))?,
|
||||
auto_location: row
|
||||
.try_get("auto_location")
|
||||
.map_err(|source| invalid().with_source(source))?,
|
||||
@@ -345,8 +348,8 @@ pub(crate) async fn resolve_proxy_selection(
|
||||
Ok((Some(id), Some(proxy.clone())))
|
||||
}
|
||||
AccountProxySelection::Saved(id) => {
|
||||
let value: String = sqlx::query_scalar(
|
||||
"select proxy_url from outbound_proxies where id = $1 for share",
|
||||
let (value, enabled): (String, bool) = sqlx::query_as(
|
||||
"select proxy_url, enabled from outbound_proxies where id = $1 for share",
|
||||
)
|
||||
.bind(id)
|
||||
.fetch_optional(&mut **transaction)
|
||||
@@ -357,6 +360,9 @@ pub(crate) async fn resolve_proxy_selection(
|
||||
entity: ENTITY,
|
||||
id: id.clone(),
|
||||
})?;
|
||||
if !enabled {
|
||||
return Err(conflict(id));
|
||||
}
|
||||
Ok((
|
||||
Some(id.clone()),
|
||||
Some(OutboundProxy::parse(&value).map_err(|source| invalid().with_source(source))?),
|
||||
@@ -561,7 +567,13 @@ impl ProxyStore for PgProxyRepository {
|
||||
if !acquired {
|
||||
return Err(store_error(conflict(id)));
|
||||
}
|
||||
let record = match self.get(id).await {
|
||||
let record = match self.get(id).await.and_then(|record| {
|
||||
if record.enabled {
|
||||
Ok(record)
|
||||
} else {
|
||||
Err(store_error(conflict(id)))
|
||||
}
|
||||
}) {
|
||||
Ok(record) => record,
|
||||
Err(error) => {
|
||||
// 拒绝预留时先等待数据库释放锁,避免连接关闭尚未生效就误挡后续代理操作
|
||||
@@ -730,7 +742,7 @@ impl ProxyStore for PgProxyRepository {
|
||||
}
|
||||
let changed = sqlx::query(
|
||||
"update outbound_proxies set name = $3, proxy_url = coalesce($4, proxy_url), revision = revision + 1, updated_at = now(),
|
||||
auto_location = coalesce($5, auto_location),
|
||||
auto_location = coalesce($5, auto_location), enabled = coalesce($6, enabled),
|
||||
detected_location_json = case when ($4 is not null and $4 <> proxy_url) or ($5 is not null and $5 <> auto_location) then null else detected_location_json end,
|
||||
last_location_detection_json = case when ($4 is not null and $4 <> proxy_url) or ($5 is not null and $5 <> auto_location) then '{\"status\":\"notRequested\"}'::jsonb else last_location_detection_json end,
|
||||
last_test_at = case when $4 is not null and $4 <> proxy_url then null else last_test_at end,
|
||||
@@ -743,6 +755,7 @@ impl ProxyStore for PgProxyRepository {
|
||||
where id = $1 and revision = $2")
|
||||
.bind(&command.id).bind(i64::try_from(command.revision.get()).map_err(|_| store_error(invalid()))?)
|
||||
.bind(&command.name).bind(command.proxy.as_ref().map(OutboundProxy::expose_url)).bind(command.auto_location)
|
||||
.bind(command.enabled)
|
||||
.execute(&mut *transaction).await.map_err(|source| store_error(unavailable(source)))?;
|
||||
if changed.rows_affected() != 1 {
|
||||
return Err(store_error(conflict(&command.id)));
|
||||
@@ -767,7 +780,7 @@ impl ProxyStore for PgProxyRepository {
|
||||
context,
|
||||
MutationAuditOperation::OutboundProxyUpdate,
|
||||
&command.id,
|
||||
&["name", "proxy_url", "location", "auto_location"],
|
||||
&["name", "proxy_url", "location", "auto_location", "enabled"],
|
||||
revision,
|
||||
)
|
||||
.await
|
||||
@@ -799,6 +812,40 @@ impl ProxyStore for PgProxyRepository {
|
||||
let config_revision = bump_config_revision_in_transaction(&mut transaction)
|
||||
.await
|
||||
.map_err(store_error)?;
|
||||
// 先锁定代理并校验版本,阻止新绑定插入;后续解绑和删除共同提交
|
||||
let locked: Option<String> = sqlx::query_scalar(
|
||||
"select id from outbound_proxies where id = $1 and revision = $2 for update",
|
||||
)
|
||||
.bind(id)
|
||||
.bind(i64::try_from(revision.get()).map_err(|_| store_error(invalid()))?)
|
||||
.fetch_optional(&mut *transaction)
|
||||
.await
|
||||
.map_err(|source| store_error(unavailable(source)))?;
|
||||
if locked.is_none() {
|
||||
return Err(store_error(conflict(id)));
|
||||
}
|
||||
let account_ids: Vec<String> = sqlx::query_scalar(
|
||||
"update provider_accounts set outbound_proxy_id = null, outbound_proxy_url = null,
|
||||
updated_at = greatest(now(), updated_at) where outbound_proxy_id = $1 returning id",
|
||||
)
|
||||
.bind(id)
|
||||
.fetch_all(&mut *transaction)
|
||||
.await
|
||||
.map_err(|source| store_error(unavailable(source)))?;
|
||||
for account_id in account_ids {
|
||||
append_admin_audit_event_in_transaction(
|
||||
&mut transaction,
|
||||
mutation_audit(
|
||||
context,
|
||||
MutationAuditOperation::ProviderAccountUpdate,
|
||||
&account_id,
|
||||
vec!["outbound_proxy".to_owned()],
|
||||
),
|
||||
config_revision,
|
||||
)
|
||||
.await
|
||||
.map_err(store_error)?;
|
||||
}
|
||||
let result = sqlx::query("delete from outbound_proxies where id = $1 and revision = $2")
|
||||
.bind(id)
|
||||
.bind(i64::try_from(revision.get()).map_err(|_| store_error(invalid()))?)
|
||||
|
||||
@@ -0,0 +1,111 @@
|
||||
//! 客户端失败轮次的过期标记与原子条件清理
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use futures::future::BoxFuture;
|
||||
use gateway_core::{
|
||||
engine::{
|
||||
ModelRequestId,
|
||||
client_retry::{ClientRetryKey, ClientRetryStore},
|
||||
},
|
||||
error::{StoreError as CoreStoreError, StoreErrorKind},
|
||||
};
|
||||
use redis::aio::ConnectionManager;
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use super::namespace;
|
||||
use crate::StoreResult;
|
||||
|
||||
const CLEAR_SCRIPT: &str =
|
||||
"if redis.call('GET',KEYS[1])==ARGV[1] then return redis.call('DEL',KEYS[1]) else return 0 end";
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct RedisClientRetryRepository {
|
||||
connection: ConnectionManager,
|
||||
namespace: String,
|
||||
}
|
||||
|
||||
impl RedisClientRetryRepository {
|
||||
pub fn new(connection: ConnectionManager, key_namespace: &str) -> StoreResult<Self> {
|
||||
Ok(Self {
|
||||
connection,
|
||||
namespace: format!("{}:client-retry:v1", namespace(key_namespace)?),
|
||||
})
|
||||
}
|
||||
|
||||
fn key(&self, key: &ClientRetryKey) -> String {
|
||||
let mut digest = Sha256::new();
|
||||
for part in [
|
||||
key.client_key_id.as_str(),
|
||||
key.turn.thread_id(),
|
||||
key.turn.turn_id(),
|
||||
] {
|
||||
digest.update((part.len() as u64).to_be_bytes());
|
||||
digest.update(part.as_bytes());
|
||||
}
|
||||
format!("{}:{}", self.namespace, hex::encode(digest.finalize()))
|
||||
}
|
||||
}
|
||||
|
||||
impl ClientRetryStore for RedisClientRetryRepository {
|
||||
fn load<'a>(
|
||||
&'a self,
|
||||
key: &'a ClientRetryKey,
|
||||
) -> BoxFuture<'a, Result<Option<ModelRequestId>, CoreStoreError>> {
|
||||
Box::pin(async move {
|
||||
let mut connection = self.connection.clone();
|
||||
let value = redis::cmd("GET")
|
||||
.arg(self.key(key))
|
||||
.query_async::<Option<String>>(&mut connection)
|
||||
.await
|
||||
.map_err(unavailable)?;
|
||||
value
|
||||
.map(ModelRequestId::new)
|
||||
.transpose()
|
||||
.map_err(|error| CoreStoreError::caused_by(StoreErrorKind::InvalidData, error))
|
||||
})
|
||||
}
|
||||
|
||||
fn record_failure<'a>(
|
||||
&'a self,
|
||||
key: &'a ClientRetryKey,
|
||||
version: &'a ModelRequestId,
|
||||
ttl: Duration,
|
||||
) -> BoxFuture<'a, Result<(), CoreStoreError>> {
|
||||
Box::pin(async move {
|
||||
let mut connection = self.connection.clone();
|
||||
redis::cmd("SET")
|
||||
.arg(self.key(key))
|
||||
.arg(version.as_str())
|
||||
.arg("PX")
|
||||
.arg(u64::try_from(ttl.as_millis()).unwrap_or(u64::MAX))
|
||||
.query_async::<()>(&mut connection)
|
||||
.await
|
||||
.map_err(unavailable)
|
||||
})
|
||||
}
|
||||
|
||||
fn clear_failure<'a>(
|
||||
&'a self,
|
||||
key: &'a ClientRetryKey,
|
||||
expected: &'a ModelRequestId,
|
||||
) -> BoxFuture<'a, Result<(), CoreStoreError>> {
|
||||
Box::pin(async move {
|
||||
let mut connection = self.connection.clone();
|
||||
// 小脚本直接 EVAL,避免冷启动 NOSCRIPT 补发增加终态交付的网络往返
|
||||
redis::cmd("EVAL")
|
||||
.arg(CLEAR_SCRIPT)
|
||||
.arg(1)
|
||||
.arg(self.key(key))
|
||||
.arg(expected.as_str())
|
||||
.query_async::<u64>(&mut connection)
|
||||
.await
|
||||
.map(|_| ())
|
||||
.map_err(unavailable)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn unavailable(error: redis::RedisError) -> CoreStoreError {
|
||||
CoreStoreError::caused_by(StoreErrorKind::Unavailable, error)
|
||||
}
|
||||
@@ -6,6 +6,7 @@ mod admin_account_runtime;
|
||||
mod artifact_profile;
|
||||
mod auth;
|
||||
mod client_admission;
|
||||
mod client_retry;
|
||||
mod coordination_buffer;
|
||||
mod credential_cooldown;
|
||||
mod credential_leases;
|
||||
@@ -21,6 +22,7 @@ pub use admin_account_runtime::*;
|
||||
pub use artifact_profile::*;
|
||||
pub use auth::*;
|
||||
pub use client_admission::*;
|
||||
pub use client_retry::*;
|
||||
pub use coordination_buffer::*;
|
||||
pub use credential_cooldown::*;
|
||||
pub use credential_leases::*;
|
||||
|
||||
@@ -11,7 +11,9 @@ use gateway_admin::{
|
||||
ports::{plugins::PluginStore, proxy::ProxyStore, store::AdminStoreErrorKind},
|
||||
};
|
||||
use gateway_core::account::OutboundProxy;
|
||||
use gateway_store::postgres::{PgPluginStore, PgProxyRepository};
|
||||
use gateway_store::postgres::{
|
||||
PgPluginStore, PgProviderAccountRepository, PgProxyRepository, ProviderAccountRepository,
|
||||
};
|
||||
|
||||
use super::{
|
||||
super::TestDatabase,
|
||||
@@ -289,6 +291,7 @@ async fn source_proxy_is_revision_fenced_and_referenced_until_the_artifact_is_de
|
||||
let updated = proxies
|
||||
.update(
|
||||
UpdateProxy {
|
||||
enabled: None,
|
||||
auto_location: None,
|
||||
test: None,
|
||||
location: None,
|
||||
@@ -302,6 +305,25 @@ async fn source_proxy_is_revision_fenced_and_referenced_until_the_artifact_is_de
|
||||
.await
|
||||
.unwrap()
|
||||
.record;
|
||||
let accounts = PgProviderAccountRepository::new(database.pool.clone());
|
||||
let mut account = super::super::provider_accounts::account("acct_plugin_proxy", "Plugin proxy");
|
||||
account.outbound_proxy = Some(updated.proxy.clone());
|
||||
accounts.insert_provider_account(account).await.unwrap();
|
||||
let before: serde_json::Value = sqlx::query_scalar(
|
||||
"select to_jsonb(a) from provider_accounts a where id = 'acct_plugin_proxy'",
|
||||
)
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
let revision: i64 =
|
||||
sqlx::query_scalar("select config_revision from runtime_settings where id = 1")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
let audit_count: i64 = sqlx::query_scalar("select count(*) from admin_audit_events")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
let mut candidate = artifact('f', &["linux-x86_64"]);
|
||||
candidate.metadata.version = "2.0.0".into();
|
||||
let Err(error) = plugins
|
||||
@@ -319,6 +341,27 @@ async fn source_proxy_is_revision_fenced_and_referenced_until_the_artifact_is_de
|
||||
.kind(),
|
||||
AdminStoreErrorKind::Conflict
|
||||
);
|
||||
let after: serde_json::Value = sqlx::query_scalar(
|
||||
"select to_jsonb(a) from provider_accounts a where id = 'acct_plugin_proxy'",
|
||||
)
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(after, before);
|
||||
assert_eq!(
|
||||
sqlx::query_scalar::<_, i64>("select config_revision from runtime_settings where id = 1")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap(),
|
||||
revision
|
||||
);
|
||||
assert_eq!(
|
||||
sqlx::query_scalar::<_, i64>("select count(*) from admin_audit_events")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap(),
|
||||
audit_count
|
||||
);
|
||||
plugins
|
||||
.change_update_source(
|
||||
PluginSourceBinding {
|
||||
|
||||
@@ -2083,6 +2083,7 @@ async fn authorization_import_rejects_a_saved_proxy_changed_during_oauth() {
|
||||
let edited = proxies
|
||||
.update(
|
||||
UpdateProxy {
|
||||
enabled: None,
|
||||
auto_location: None,
|
||||
test: None,
|
||||
location: None,
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
//! 验证代理位置、账号绑定、分页查询与事务锁释放
|
||||
//! 验证代理启停、位置、账号绑定、自动解绑与事务锁释放
|
||||
|
||||
use gateway_admin::{
|
||||
model::{
|
||||
@@ -68,6 +68,7 @@ async fn proxy_location_is_shared_preserved_cleared_and_removed_with_binding() {
|
||||
saved = proxies
|
||||
.update(
|
||||
UpdateProxy {
|
||||
enabled: None,
|
||||
auto_location: None,
|
||||
test: None,
|
||||
id: saved.id.clone(),
|
||||
@@ -87,6 +88,7 @@ async fn proxy_location_is_shared_preserved_cleared_and_removed_with_binding() {
|
||||
saved = proxies
|
||||
.update(
|
||||
UpdateProxy {
|
||||
enabled: None,
|
||||
auto_location: None,
|
||||
test: None,
|
||||
id: saved.id.clone(),
|
||||
@@ -115,6 +117,7 @@ async fn proxy_location_is_shared_preserved_cleared_and_removed_with_binding() {
|
||||
proxies
|
||||
.update(
|
||||
UpdateProxy {
|
||||
enabled: None,
|
||||
auto_location: None,
|
||||
test: None,
|
||||
id: saved.id.clone(),
|
||||
@@ -132,6 +135,7 @@ async fn proxy_location_is_shared_preserved_cleared_and_removed_with_binding() {
|
||||
proxies
|
||||
.update(
|
||||
UpdateProxy {
|
||||
enabled: None,
|
||||
auto_location: None,
|
||||
test: None,
|
||||
id: saved.id.clone(),
|
||||
@@ -168,6 +172,201 @@ fn context() -> MutationContext {
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn disabled_proxy_projects_direct_while_preserving_bindings_and_credentials() {
|
||||
use gateway_core::{
|
||||
account::{ProviderAccountStore as _, ProviderRefreshQuery, RequestLocation},
|
||||
routing::ProviderKind,
|
||||
};
|
||||
use std::{
|
||||
num::NonZeroU32,
|
||||
time::{Duration, SystemTime},
|
||||
};
|
||||
|
||||
let Some(database) = TestDatabase::create("proxy_enabled").await else {
|
||||
return;
|
||||
};
|
||||
let proxies = PgProxyRepository::new(database.pool.clone());
|
||||
let accounts = PgProviderAccountRepository::new(database.pool.clone());
|
||||
let admin = admin_account_store(&database.pool);
|
||||
let location: RequestLocation = serde_json::from_value(serde_json::json!({
|
||||
"country":"JP", "region":"Tokyo", "city":"Tokyo", "timezone":"Asia/Tokyo"
|
||||
}))
|
||||
.unwrap();
|
||||
let mut saved = proxies
|
||||
.create(
|
||||
NewProxy {
|
||||
auto_location: false,
|
||||
test: None,
|
||||
location: Some(location.clone()),
|
||||
name: "启停测试".to_owned(),
|
||||
proxy: OutboundProxy::parse("http://user:secret@127.0.0.1:18080").unwrap(),
|
||||
},
|
||||
&context(),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.record;
|
||||
assert!(saved.enabled);
|
||||
proxies
|
||||
.record_test(&saved.id, saved.revision, success(), &context())
|
||||
.await
|
||||
.unwrap();
|
||||
for (id, provider) in [
|
||||
("acct_enabled_openai", "openai"),
|
||||
("acct_enabled_xai", "xai"),
|
||||
] {
|
||||
let mut seed = account(id, id);
|
||||
seed.provider_kind = provider.into();
|
||||
seed.outbound_proxy = Some(saved.proxy.clone());
|
||||
seed.has_refresh_token = true;
|
||||
seed.access_token_expires_at = Some(chrono::Utc::now());
|
||||
accounts.insert_provider_account(seed).await.unwrap();
|
||||
}
|
||||
let original: Vec<serde_json::Value> =
|
||||
sqlx::query_scalar("select to_jsonb(a) from provider_accounts a order by id")
|
||||
.fetch_all(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
let original_revision = accounts.list_accounts().await.unwrap()[0].revision();
|
||||
for enabled in [false, true] {
|
||||
let stale_revision = saved.revision;
|
||||
saved = proxies
|
||||
.update(
|
||||
UpdateProxy {
|
||||
enabled: Some(enabled),
|
||||
auto_location: None,
|
||||
test: None,
|
||||
location: None,
|
||||
id: saved.id.clone(),
|
||||
revision: saved.revision,
|
||||
name: saved.name.clone(),
|
||||
proxy: None,
|
||||
},
|
||||
&context(),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.record;
|
||||
assert_eq!(saved.enabled, enabled);
|
||||
assert_eq!(saved.account_count, 2);
|
||||
assert_eq!(saved.last_test, Some(success()));
|
||||
assert_eq!(saved.location.as_ref(), Some(&location));
|
||||
let current: Vec<serde_json::Value> =
|
||||
sqlx::query_scalar("select to_jsonb(a) from provider_accounts a order by id")
|
||||
.fetch_all(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(current, original);
|
||||
let mut projected = accounts.list_accounts().await.unwrap();
|
||||
for (id, provider) in [
|
||||
("acct_enabled_openai", "openai"),
|
||||
("acct_enabled_xai", "xai"),
|
||||
] {
|
||||
let id = ProviderAccountId::new(id).unwrap();
|
||||
projected.push(accounts.get_account(&id).await.unwrap().unwrap());
|
||||
projected.push(accounts.load_current_credential(&id).await.unwrap().account);
|
||||
let kind = ProviderKind::new(provider).unwrap();
|
||||
projected.extend(accounts.list_for_provider(&kind).await.unwrap());
|
||||
let now = SystemTime::now();
|
||||
projected.extend(
|
||||
accounts
|
||||
.list_refresh_candidates(ProviderRefreshQuery::new(
|
||||
kind,
|
||||
now + Duration::from_secs(300),
|
||||
now - Duration::from_secs(7200),
|
||||
now,
|
||||
vec![],
|
||||
NonZeroU32::new(10).unwrap(),
|
||||
))
|
||||
.await
|
||||
.unwrap()
|
||||
.into_iter()
|
||||
.map(|loaded| loaded.account),
|
||||
);
|
||||
let configured = admin
|
||||
.credential_details(&ProviderKind::new(provider).unwrap(), &id)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.credential;
|
||||
assert_eq!(configured.outbound_proxy.as_ref(), Some(&saved.proxy));
|
||||
assert_eq!(
|
||||
configured.effective_outbound_proxy(),
|
||||
enabled.then_some(&saved.proxy)
|
||||
);
|
||||
}
|
||||
assert!(projected.iter().all(|account| account.outbound_proxy()
|
||||
== enabled.then_some(&saved.proxy)
|
||||
&& account.request_location() == enabled.then_some(&location)
|
||||
&& account.revision() == original_revision
|
||||
&& account.enabled()));
|
||||
assert_eq!(
|
||||
proxies
|
||||
.update(
|
||||
UpdateProxy {
|
||||
enabled: Some(!enabled),
|
||||
auto_location: None,
|
||||
test: None,
|
||||
location: None,
|
||||
id: saved.id.clone(),
|
||||
revision: stale_revision,
|
||||
name: saved.name.clone(),
|
||||
proxy: None,
|
||||
},
|
||||
&context()
|
||||
)
|
||||
.await
|
||||
.unwrap_err()
|
||||
.kind(),
|
||||
AdminStoreErrorKind::Conflict
|
||||
);
|
||||
assert_eq!(proxies.get(&saved.id).await.unwrap().enabled, enabled);
|
||||
let reservation = proxies.reserve_import(&saved.id).await;
|
||||
assert_eq!(reservation.is_ok(), enabled);
|
||||
drop(reservation);
|
||||
if !enabled {
|
||||
let error = admin
|
||||
.update_account(
|
||||
update(
|
||||
"acct_enabled_openai",
|
||||
AccountProxySelection::Saved(saved.id.clone()),
|
||||
),
|
||||
&context(),
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert_eq!(error.kind(), AdminStoreErrorKind::Conflict);
|
||||
let error = admin
|
||||
.batch_update_accounts(
|
||||
BatchUpdateAccounts {
|
||||
account_ids: vec![
|
||||
"acct_enabled_openai".to_owned(),
|
||||
"acct_enabled_xai".to_owned(),
|
||||
],
|
||||
enabled: Some(false),
|
||||
concurrency_limit: None,
|
||||
weight: None,
|
||||
model_access: None,
|
||||
group_ids: None,
|
||||
outbound_proxy: Some(AccountProxySelection::Saved(saved.id.clone())),
|
||||
},
|
||||
&context(),
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert_eq!(error.kind(), AdminStoreErrorKind::Conflict);
|
||||
let unchanged: Vec<serde_json::Value> =
|
||||
sqlx::query_scalar("select to_jsonb(a) from provider_accounts a order by id")
|
||||
.fetch_all(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(unchanged, original);
|
||||
}
|
||||
}
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
fn success() -> ProxyTestResult {
|
||||
ProxyTestResult {
|
||||
location: Default::default(),
|
||||
@@ -575,6 +774,7 @@ async fn stale_proxy_test_releases_lock_before_next_update() {
|
||||
let update = store
|
||||
.update(
|
||||
UpdateProxy {
|
||||
enabled: None,
|
||||
id: saved.id.clone(),
|
||||
revision: saved.revision,
|
||||
name: saved.name.clone(),
|
||||
@@ -634,6 +834,7 @@ async fn import_reservation_blocks_proxy_mutations_until_rotated_credentials_are
|
||||
.unwrap();
|
||||
let reservation = store.reserve_import(&saved.id).await.unwrap();
|
||||
let replacement = UpdateProxy {
|
||||
enabled: None,
|
||||
auto_location: None,
|
||||
test: None,
|
||||
location: None,
|
||||
@@ -798,14 +999,6 @@ async fn managed_proxies_persist_bind_update_all_accounts_and_protect_stale_test
|
||||
.unwrap();
|
||||
}
|
||||
assert_eq!(store.get(&created.id).await.unwrap().account_count, 2);
|
||||
assert_eq!(
|
||||
store
|
||||
.delete(&created.id, created.revision, &context)
|
||||
.await
|
||||
.unwrap_err()
|
||||
.kind(),
|
||||
AdminStoreErrorKind::Conflict
|
||||
);
|
||||
assert!(
|
||||
read()
|
||||
.await
|
||||
@@ -822,6 +1015,7 @@ async fn managed_proxies_persist_bind_update_all_accounts_and_protect_stale_test
|
||||
let renamed = store
|
||||
.update(
|
||||
UpdateProxy {
|
||||
enabled: None,
|
||||
auto_location: None,
|
||||
test: None,
|
||||
id: created.id.clone(),
|
||||
@@ -842,6 +1036,7 @@ async fn managed_proxies_persist_bind_update_all_accounts_and_protect_stale_test
|
||||
let edited = store
|
||||
.update(
|
||||
UpdateProxy {
|
||||
enabled: None,
|
||||
auto_location: None,
|
||||
test: None,
|
||||
id: created.id.clone(),
|
||||
@@ -875,22 +1070,45 @@ async fn managed_proxies_persist_bind_update_all_accounts_and_protect_stale_test
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
for id in ["acct_one", "acct_two"] {
|
||||
admin
|
||||
.update_account(update(id, AccountProxySelection::Direct), &context)
|
||||
assert_eq!(
|
||||
store
|
||||
.delete(&created.id, created.revision, &context)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
.unwrap_err()
|
||||
.kind(),
|
||||
AdminStoreErrorKind::Conflict
|
||||
);
|
||||
assert!(
|
||||
read()
|
||||
.await
|
||||
.iter()
|
||||
.all(|row| row.1.as_deref() == Some(created.id.as_str()))
|
||||
);
|
||||
let preserved: Vec<serde_json::Value> = sqlx::query_scalar("select to_jsonb(a) - 'updated_at' - 'outbound_proxy_id' - 'outbound_proxy_url' from provider_accounts a order by id")
|
||||
.fetch_all(&database.pool).await.unwrap();
|
||||
let deleted_revision = store
|
||||
.delete(&created.id, edited.revision, &context)
|
||||
.await
|
||||
.unwrap();
|
||||
let after_delete: Vec<serde_json::Value> = sqlx::query_scalar("select to_jsonb(a) - 'updated_at' - 'outbound_proxy_id' - 'outbound_proxy_url' from provider_accounts a order by id")
|
||||
.fetch_all(&database.pool).await.unwrap();
|
||||
assert_eq!(preserved, after_delete);
|
||||
assert_eq!(
|
||||
sqlx::query_scalar::<_, i64>(
|
||||
"select count(*) from admin_audit_events where config_revision = $1"
|
||||
)
|
||||
.bind(i64::try_from(deleted_revision.get()).unwrap())
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap(),
|
||||
3
|
||||
);
|
||||
assert!(
|
||||
read()
|
||||
.await
|
||||
.iter()
|
||||
.all(|row| row.1.is_none() && row.2.is_none() && row.3 == 1)
|
||||
);
|
||||
store
|
||||
.delete(&created.id, edited.revision, &context)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(store.get(&created.id).await.is_err());
|
||||
let audits: Vec<serde_json::Value> =
|
||||
sqlx::query_scalar("select to_jsonb(a) from admin_audit_events a")
|
||||
@@ -984,6 +1202,12 @@ async fn migration_backfills_shared_proxies_without_changing_credentials() {
|
||||
.unwrap();
|
||||
}
|
||||
super::TEST_MIGRATOR.run(&database.pool).await.unwrap();
|
||||
assert!(
|
||||
sqlx::query_scalar::<_, bool>("select bool_and(enabled) from outbound_proxies")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap()
|
||||
);
|
||||
let accounts = PgProviderAccountRepository::new(database.pool.clone());
|
||||
let inherited = gateway_core::account::ProviderAccountStore::list_accounts(&accounts)
|
||||
.await
|
||||
@@ -1134,6 +1358,7 @@ async fn automatic_proxy_location_projects_to_accounts_and_guards_refresh_result
|
||||
saved = proxies
|
||||
.update(
|
||||
UpdateProxy {
|
||||
enabled: None,
|
||||
id: saved.id.clone(),
|
||||
revision: saved.revision,
|
||||
name: saved.name.clone(),
|
||||
@@ -1152,6 +1377,7 @@ async fn automatic_proxy_location_projects_to_accounts_and_guards_refresh_result
|
||||
saved = proxies
|
||||
.update(
|
||||
UpdateProxy {
|
||||
enabled: None,
|
||||
id: saved.id.clone(),
|
||||
revision: saved.revision,
|
||||
name: saved.name.clone(),
|
||||
@@ -1178,6 +1404,7 @@ async fn automatic_proxy_location_projects_to_accounts_and_guards_refresh_result
|
||||
saved = proxies
|
||||
.update(
|
||||
UpdateProxy {
|
||||
enabled: None,
|
||||
id: saved.id.clone(),
|
||||
revision: saved.revision,
|
||||
name: saved.name.clone(),
|
||||
|
||||
@@ -123,15 +123,17 @@ async fn privacy_policy_upgrade_defaults_off_and_round_trips_in_snapshot() {
|
||||
id: "remove-auth".into(),
|
||||
name: "配置者控制".into(),
|
||||
enabled: true,
|
||||
scope: PrivacyScope::RequestHeader,
|
||||
selector: "authorization".into(),
|
||||
action: PrivacyAction::RemoveField,
|
||||
pattern: None,
|
||||
replacement: String::new(),
|
||||
value: serde_json::Value::Null,
|
||||
replace_all: true,
|
||||
case_insensitive: false,
|
||||
multi_line: false,
|
||||
steps: vec![PrivacyStep {
|
||||
scope: PrivacyScope::RequestHeader,
|
||||
selector: "authorization".into(),
|
||||
action: PrivacyAction::RemoveField,
|
||||
pattern: None,
|
||||
replacement: String::new(),
|
||||
value: serde_json::Value::Null,
|
||||
replace_all: true,
|
||||
case_insensitive: false,
|
||||
multi_line: false,
|
||||
}],
|
||||
}],
|
||||
};
|
||||
let expected = update.values.codex_privacy_policy.clone();
|
||||
@@ -154,6 +156,137 @@ async fn privacy_policy_upgrade_defaults_off_and_round_trips_in_snapshot() {
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn privacy_steps_migration_preserves_rules_and_other_settings() {
|
||||
let Some(database) = TestDatabase::create_through("privacy_steps", 30).await else {
|
||||
return;
|
||||
};
|
||||
let legacy = serde_json::json!({
|
||||
"enabled":true, "onError":"skip_rule", "rules":[
|
||||
{"id":"first","name":"旧规则一","enabled":true,"scope":"request_header","selector":"x-private",
|
||||
"action":"remove_field","pattern":null,"replacement":"","value":null,
|
||||
"replaceAll":true,"caseInsensitive":false,"multiLine":false},
|
||||
{"id":"second","name":"旧规则二","enabled":false,"scope":"request_body","selector":"$.owner",
|
||||
"action":"regex_replace","pattern":"alice","replacement":"user","value":null,
|
||||
"replaceAll":false,"caseInsensitive":true,"multiLine":true}
|
||||
]
|
||||
});
|
||||
let mut expected = serde_json::json!({
|
||||
"enabled":true, "onError":"skip_rule", "rules":[
|
||||
{"id":"first","name":"旧规则一","enabled":true,"steps":[
|
||||
{"scope":"request_header","selector":"x-private","action":"remove_field",
|
||||
"pattern":null,"replacement":"","value":null,
|
||||
"replaceAll":true,"caseInsensitive":false,"multiLine":false}]},
|
||||
{"id":"second","name":"旧规则二","enabled":false,"steps":[
|
||||
{"scope":"request_body","selector":"$.owner","action":"regex_replace",
|
||||
"pattern":"alice","replacement":"user","value":null,
|
||||
"replaceAll":false,"caseInsensitive":true,"multiLine":true}]}
|
||||
]
|
||||
});
|
||||
// 同时包含聚合规则,迁移不能嵌套包装或拆散现有步骤
|
||||
let canonical = serde_json::json!({
|
||||
"id":"aggregate","name":"通用聚合规则","enabled":true,"steps":[
|
||||
{"scope":"request_body","selector":"$.metadata","action":"set_value",
|
||||
"pattern":null,"replacement":"","value":{"owner":"user","items":[1,2]},
|
||||
"replaceAll":true,"caseInsensitive":false,"multiLine":false},
|
||||
{"scope":"request_header","selector":"x-private","action":"remove_field",
|
||||
"pattern":null,"replacement":"","value":null,
|
||||
"replaceAll":true,"caseInsensitive":false,"multiLine":false}
|
||||
]
|
||||
});
|
||||
let mut mixed = legacy;
|
||||
mixed["rules"]
|
||||
.as_array_mut()
|
||||
.unwrap()
|
||||
.insert(1, canonical.clone());
|
||||
expected["rules"]
|
||||
.as_array_mut()
|
||||
.unwrap()
|
||||
.insert(1, canonical);
|
||||
sqlx::query("update runtime_settings set codex_privacy_policy_json = $1, config_revision = 70, refresh_margin_seconds = 1234 where id = 1")
|
||||
.bind(&mixed).execute(&database.pool).await.unwrap();
|
||||
let mut expected_row: serde_json::Value =
|
||||
sqlx::query_scalar("select to_jsonb(runtime_settings) from runtime_settings where id = 1")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
expected_row["codex_privacy_policy_json"] = expected.clone();
|
||||
|
||||
super::TEST_MIGRATOR.run(&database.pool).await.unwrap();
|
||||
let stored: serde_json::Value =
|
||||
sqlx::query_scalar("select to_jsonb(runtime_settings) from runtime_settings where id = 1")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(stored, expected_row);
|
||||
let repository = PgRuntimeSettingsRepository::new(database.pool.clone());
|
||||
let loaded = repository.load_runtime_settings().await.unwrap();
|
||||
assert_eq!(
|
||||
serde_json::to_value(&loaded.values.codex_privacy_policy).unwrap(),
|
||||
expected
|
||||
);
|
||||
// 服务重启后再次检查迁移,不产生写入或额外嵌套
|
||||
super::TEST_MIGRATOR.run(&database.pool).await.unwrap();
|
||||
let stored: serde_json::Value =
|
||||
sqlx::query_scalar("select to_jsonb(runtime_settings) from runtime_settings where id = 1")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(stored, expected_row);
|
||||
let snapshot = gateway_store::postgres::RuntimeSnapshotRepository::load_runtime_snapshot(
|
||||
&gateway_store::postgres::PgRuntimeSnapshotRepository::new(database.pool.clone()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
serde_json::to_value(snapshot.settings).unwrap()["codex_privacy_policy"],
|
||||
expected
|
||||
);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn privacy_steps_migration_rejects_non_object_without_losing_data() {
|
||||
let Some(database) = TestDatabase::create_through("privacy_invalid", 30).await else {
|
||||
return;
|
||||
};
|
||||
let invalid = serde_json::json!({
|
||||
"enabled":false,"onError":"reject_request","rules":[{},42]
|
||||
});
|
||||
sqlx::query("update runtime_settings set codex_privacy_policy_json = $1 where id = 1")
|
||||
.bind(&invalid)
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(super::TEST_MIGRATOR.run(&database.pool).await.is_err());
|
||||
let stored: serde_json::Value =
|
||||
sqlx::query_scalar("select codex_privacy_policy_json from runtime_settings where id = 1")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(stored, invalid);
|
||||
let applied: bool =
|
||||
sqlx::query_scalar("select exists(select 1 from _sqlx_migrations where version = 31)")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(!applied);
|
||||
sqlx::query("update runtime_settings set codex_privacy_policy_json = jsonb_set(codex_privacy_policy_json, '{rules}', '[]') where id = 1")
|
||||
.execute(&database.pool).await.unwrap();
|
||||
super::TEST_MIGRATOR.run(&database.pool).await.unwrap();
|
||||
let loaded = PgRuntimeSettingsRepository::new(database.pool.clone())
|
||||
.load_runtime_settings()
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(!loaded.values.codex_privacy_policy.enabled);
|
||||
assert!(loaded.values.codex_privacy_policy.rules.is_empty());
|
||||
assert_eq!(
|
||||
loaded.values.codex_privacy_policy.on_error,
|
||||
gateway_core::settings::privacy::PrivacyFailureMode::RejectRequest
|
||||
);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn smart_settings_upgrade_preserves_selection_and_publishes_custom_config() {
|
||||
use gateway_core::account::SmartSchedulingConfig;
|
||||
|
||||
@@ -0,0 +1,119 @@
|
||||
//! 验证失败轮次标记的隔离、过期与并发条件清理
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use gateway_core::{
|
||||
engine::{
|
||||
ModelRequestId,
|
||||
client_retry::{ClientRetryKey, ClientRetryStore},
|
||||
},
|
||||
operation::ClientTurn,
|
||||
policy::ClientApiKeyId,
|
||||
};
|
||||
use gateway_store::redis::RedisClientRetryRepository;
|
||||
use redis::aio::ConnectionManager;
|
||||
use uuid::Uuid;
|
||||
|
||||
async fn repository() -> Option<(RedisClientRetryRepository, ConnectionManager, String)> {
|
||||
let url = crate::support::test_env("CPR_TEST_REDIS_URL")?;
|
||||
let client = redis::Client::open(url).unwrap();
|
||||
let connection = client.get_connection_manager().await.unwrap();
|
||||
let namespace = format!("client-retry-test-{}", Uuid::new_v4());
|
||||
Some((
|
||||
RedisClientRetryRepository::new(connection.clone(), &namespace).unwrap(),
|
||||
connection,
|
||||
namespace,
|
||||
))
|
||||
}
|
||||
|
||||
fn key(client: &str, thread: &str, turn: &str) -> ClientRetryKey {
|
||||
ClientRetryKey {
|
||||
client_key_id: ClientApiKeyId::new(client).unwrap(),
|
||||
turn: ClientTurn::new(thread.to_owned(), turn.to_owned()).unwrap(),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn client_retry_versions_should_protect_newer_failures() {
|
||||
let Some((store, _, _)) = repository().await else {
|
||||
return;
|
||||
};
|
||||
let key = key("key_one", "thread", "turn");
|
||||
let old = ModelRequestId::new("req_old").unwrap();
|
||||
let newer = ModelRequestId::new("req_newer").unwrap();
|
||||
store
|
||||
.record_failure(&key, &old, Duration::from_secs(30))
|
||||
.await
|
||||
.unwrap();
|
||||
store
|
||||
.record_failure(&key, &newer, Duration::from_secs(30))
|
||||
.await
|
||||
.unwrap();
|
||||
store.clear_failure(&key, &old).await.unwrap();
|
||||
assert_eq!(store.load(&key).await.unwrap(), Some(newer.clone()));
|
||||
store.clear_failure(&key, &newer).await.unwrap();
|
||||
assert_eq!(store.load(&key).await.unwrap(), None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn client_retry_should_isolate_all_identity_parts() {
|
||||
let Some((store, _, _)) = repository().await else {
|
||||
return;
|
||||
};
|
||||
let first = key("key_one", "ab", "c");
|
||||
let version = ModelRequestId::new("req_first").unwrap();
|
||||
store
|
||||
.record_failure(&first, &version, Duration::from_secs(30))
|
||||
.await
|
||||
.unwrap();
|
||||
for other in [
|
||||
key("key_other", "ab", "c"),
|
||||
key("key_one", "a", "bc"),
|
||||
key("key_one", "ab", "other"),
|
||||
] {
|
||||
assert_eq!(store.load(&other).await.unwrap(), None);
|
||||
}
|
||||
assert_eq!(store.load(&first).await.unwrap(), Some(version.clone()));
|
||||
store.clear_failure(&first, &version).await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn client_retry_should_refresh_ttl_on_failure_and_expire_without_cleanup_worker() {
|
||||
let Some((store, mut connection, namespace)) = repository().await else {
|
||||
return;
|
||||
};
|
||||
let key = key("key_one", "thread", "turn");
|
||||
let old = ModelRequestId::new("req_old").unwrap();
|
||||
let newer = ModelRequestId::new("req_newer").unwrap();
|
||||
store
|
||||
.record_failure(&key, &old, Duration::from_secs(30))
|
||||
.await
|
||||
.unwrap();
|
||||
store
|
||||
.record_failure(&key, &newer, Duration::from_secs(300))
|
||||
.await
|
||||
.unwrap();
|
||||
let keys: Vec<String> = redis::cmd("KEYS")
|
||||
.arg(format!("{namespace}:*"))
|
||||
.query_async(&mut connection)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(keys.len(), 1);
|
||||
let ttl: i64 = redis::cmd("PTTL")
|
||||
.arg(&keys[0])
|
||||
.query_async(&mut connection)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(ttl > 290_000 && ttl <= 300_000);
|
||||
store
|
||||
.record_failure(&key, &newer, Duration::from_millis(100))
|
||||
.await
|
||||
.unwrap();
|
||||
tokio::time::timeout(Duration::from_secs(3), async {
|
||||
while store.load(&key).await.unwrap().is_some() {
|
||||
tokio::time::sleep(Duration::from_millis(20)).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("Redis expires marker");
|
||||
}
|
||||
@@ -3,6 +3,7 @@
|
||||
mod artifact_profile;
|
||||
mod auth;
|
||||
mod client_admission;
|
||||
mod client_retry;
|
||||
mod coordination_buffer;
|
||||
mod credential_cooldown;
|
||||
mod credential_leases;
|
||||
|
||||
@@ -23,7 +23,8 @@ gateway-core = { path = "../../gateway-core" }
|
||||
gateway-protocol = { path = "../../gateway-protocol" }
|
||||
getrandom = "0.4.3"
|
||||
hex = "0.4.3"
|
||||
# TLS 依赖按官方 Core 锁文件固定,包含用于约束传递依赖的直接声明。
|
||||
# TLS 依赖以官方 Core 锁文件为基线,安全补丁经握手特征核验后精确固定
|
||||
# 包含用于约束传递依赖的直接声明
|
||||
# hyper 受 workspace 的 SOCKS IPv6 修复补丁约束,保留补丁要求的兼容版本。
|
||||
# 升级前须重复采样 ClientHello,核对算法、扩展集合及排序约束、key_share 与 ALPN。
|
||||
aws-lc-rs = "=1.18.1"
|
||||
@@ -40,7 +41,7 @@ roxmltree = "0.21.1"
|
||||
semver = "1.0.27"
|
||||
# 与官方 rustls 路径一致:aws-lc 支持 P-521,优先提供 X25519MLKEM768。
|
||||
# reqwest 的 roots feature 仍会编译 ring;实际 provider 由 TLS 初始化显式选择。
|
||||
rustls = { version = "=0.23.45", default-features = false, features = ["aws_lc_rs", "prefer-post-quantum", "std", "tls12"] }
|
||||
rustls = { version = "=0.23.46", default-features = false, features = ["aws_lc_rs", "prefer-post-quantum", "std", "tls12"] }
|
||||
rustls-native-certs = "=0.8.3"
|
||||
rustls-pki-types = "=1.14.0"
|
||||
secrecy = "0.10.3"
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user