Merge branch 'fawney19:main' into main

This commit is contained in:
ZheFox
2026-06-16 12:17:28 +08:00
committed by GitHub
160 changed files with 24102 additions and 1311 deletions
@@ -719,6 +719,8 @@ mod tests {
provider_request_body: Some(json!({"model":"gpt-5","metadata":{}})),
provider_request_body_base64: None,
content_type: Some("application/json".to_string()),
content_encoding: None,
request_gzip: None,
proxy: None,
transport_profile: None,
timeouts: None,
@@ -20,6 +20,7 @@ mod pool_scheduler;
pub(crate) mod pool_scores;
mod redaction;
mod report_context;
mod request_gzip;
mod route;
mod runtime_miss;
mod spec_metadata;
@@ -46,6 +47,7 @@ pub(crate) use self::plan_builders::{
pub(crate) use self::pool_scores::{
build_provider_key_pool_score_upsert, provider_key_pool_score_id, provider_key_pool_score_scope,
};
pub(crate) use self::request_gzip::resolve_transport_request_gzip_policy;
pub(crate) use self::route::is_matching_stream_request as planner_is_matching_stream_request;
pub(crate) use self::runtime_miss::{
apply_local_runtime_candidate_terminal_reason, record_local_runtime_candidate_skip_reason,
@@ -17,7 +17,8 @@ use crate::ai_serving::planner::report_context::{
use crate::ai_serving::planner::spec_metadata::local_same_format_provider_spec_metadata;
use crate::ai_serving::planner::CandidateFailureDiagnostic;
use crate::ai_serving::planner::{
build_ai_execution_decision_response, AiExecutionDecisionResponseParts,
build_ai_execution_decision_response, resolve_transport_request_gzip_policy,
AiExecutionDecisionResponseParts,
};
use crate::ai_serving::transport::{
resolve_transport_execution_timeouts, resolve_transport_profile,
@@ -107,6 +108,11 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
json!(crate::ai_serving::transport::GEMINI_CLI_V1INTERNAL_ENVELOPE_NAME),
);
}
if !resolved.compatibility_edits.is_empty() {
if let Ok(value) = serde_json::to_value(&resolved.compatibility_edits) {
extra_fields.insert("request_body_compatibility_edits".to_string(), value);
}
}
let provider_api_format = resolved.provider_api_format.clone();
let effective_headers = input.effective_headers(&parts.headers);
let report_context = append_local_failover_policy_to_value(
@@ -175,8 +181,10 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
provider_request_headers,
provider_request_body,
transport_profile: _,
compatibility_edits: _,
request_redacted: _,
} = resolved;
let request_gzip = resolve_transport_request_gzip_policy(&transport);
let mut decision = build_ai_execution_decision_response(AiExecutionDecisionResponseParts {
decision_is_stream: spec_metadata.require_streaming,
@@ -203,6 +211,8 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
provider_request_body: Some(provider_request_body),
provider_request_body_base64: None,
content_type: Some("application/json".to_string()),
content_encoding: None,
request_gzip,
proxy,
transport_profile,
timeouts: resolve_transport_execution_timeouts(&transport),
@@ -19,7 +19,9 @@ use crate::ai_serving::transport::{
build_gemini_cli_v1internal_request, build_grok_browser_headers, build_grok_upstream_url,
build_same_format_provider_headers, resolve_local_gemini_cli_request_auth,
GeminiCliRequestAuth, GeminiCliRequestAuthSupport, GeminiCliRequestEnvelopeSupport,
GrokHeaderInput, SameFormatProviderHeadersInput, GEMINI_CLI_USER_AGENT, GROK_CHAT_PATH,
GrokHeaderInput, SameFormatProviderCompatibilityEdit,
SameFormatProviderCompatibilityEditAction, SameFormatProviderHeadersInput,
GEMINI_CLI_USER_AGENT, GROK_CHAT_PATH,
};
use crate::ai_serving::{CandidateFailureDiagnostic, GatewayProviderTransportSnapshot};
use crate::{AppState, GatewayError};
@@ -107,6 +109,7 @@ pub(crate) struct LocalSameFormatProviderCandidatePayloadParts {
pub(super) provider_request_headers: BTreeMap<String, String>,
pub(super) provider_request_body: Value,
pub(super) transport_profile: Option<ResolvedTransportProfile>,
pub(super) compatibility_edits: Vec<SameFormatProviderCompatibilityEdit>,
pub(super) request_redacted: bool,
}
@@ -153,8 +156,8 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
let body_json = redaction.body_json.as_ref();
let mut transport = Arc::clone(&prepared.transport);
let Some(mut base_provider_request_body) =
super::super::request::build_same_format_provider_request_body(
let Some(base_provider_request) =
super::super::request::build_same_format_provider_request_body_with_compatibility_report(
body_json,
prepared.provider_api_format.as_str(),
&prepared.mapped_model,
@@ -190,6 +193,8 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
.await;
return Ok(None);
};
let mut base_provider_request_body = base_provider_request.body;
let mut compatibility_edits = base_provider_request.compatibility_edits;
if let Some(mapping) =
crate::system_features::reasoning_model_directive_mapping_for_api_format_and_model(
state,
@@ -198,10 +203,18 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
)
.await
{
let before_mapping = base_provider_request_body.clone();
crate::ai_serving::apply_model_directive_mapping_patch(
&mut base_provider_request_body,
&mapping,
);
if before_mapping != base_provider_request_body {
compatibility_edits.push(SameFormatProviderCompatibilityEdit {
field: "model_directive_mapping".to_string(),
action: SameFormatProviderCompatibilityEditAction::RuntimeRewrite,
detail: "applied configured model directive mapping patch".to_string(),
});
}
// Directive mapping is a deep-merge patch and may overwrite/add `stream`;
// re-enforce stream-field policy afterward.
// Kiro behavior classification already hard-requires upstream streaming,
@@ -452,6 +465,7 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
provider_request_headers,
provider_request_body,
transport_profile,
compatibility_edits,
request_redacted: redaction.redacted,
}))
}
@@ -2,4 +2,5 @@ mod body;
mod url;
pub(super) use self::body::build_same_format_provider_request_body;
pub(super) use self::body::build_same_format_provider_request_body_with_compatibility_report;
pub(super) use self::url::build_same_format_upstream_url;
@@ -3,7 +3,9 @@ use serde_json::Value;
use super::super::LocalSameFormatProviderSpec;
use crate::ai_serving::transport::{
build_same_format_provider_request_body as build_same_format_provider_request_body_impl,
build_same_format_provider_request_body_with_compatibility_report as build_same_format_provider_request_body_with_compatibility_report_impl,
SameFormatProviderFamily, SameFormatProviderRequestBodyInput,
SameFormatProviderRequestBodyOutput,
};
pub(crate) fn build_same_format_provider_request_body(
@@ -36,6 +38,38 @@ pub(crate) fn build_same_format_provider_request_body(
})
}
pub(crate) fn build_same_format_provider_request_body_with_compatibility_report(
body_json: &Value,
provider_api_format: &str,
mapped_model: &str,
spec: LocalSameFormatProviderSpec,
body_rules: Option<&Value>,
request_headers: Option<&http::HeaderMap>,
upstream_is_stream: bool,
force_body_stream_field: bool,
kiro_auth: Option<&crate::ai_serving::transport::kiro::KiroRequestAuth>,
is_claude_code: bool,
enable_model_directives: bool,
) -> Option<SameFormatProviderRequestBodyOutput> {
build_same_format_provider_request_body_with_compatibility_report_impl(
SameFormatProviderRequestBodyInput {
body_json,
mapped_model,
client_api_format: spec.api_format,
provider_api_format,
source_model: body_json.get("model").and_then(Value::as_str),
family: same_format_provider_family(spec.family),
body_rules,
request_headers,
upstream_is_stream,
force_body_stream_field,
kiro_auth_config: kiro_auth.map(|auth| &auth.auth_config),
is_claude_code,
enable_model_directives,
},
)
}
fn same_format_provider_family(
family: super::super::LocalSameFormatProviderFamily,
) -> SameFormatProviderFamily {
@@ -29,7 +29,6 @@ impl<'a> ProviderRequestRedaction<'a> {
#[derive(Clone, Copy, Debug, Default)]
struct ChatPiiRedactionFeatureSettings {
enabled: Option<bool>,
inject_model_instruction: Option<bool>,
}
impl ChatPiiRedactionFeatureSettings {
@@ -44,21 +43,11 @@ impl ChatPiiRedactionFeatureSettings {
if let Some(enabled) = settings.get("enabled").and_then(Value::as_bool) {
self.enabled = Some(enabled);
}
if let Some(inject_model_instruction) = settings
.get("inject_model_instruction")
.and_then(Value::as_bool)
{
self.inject_model_instruction = Some(inject_model_instruction);
}
}
fn effective_enabled(self) -> bool {
self.enabled.unwrap_or(false)
}
fn effective_inject_model_instruction(self) -> bool {
self.inject_model_instruction.unwrap_or(true)
}
}
pub(crate) fn request_identity_response_encoding_when_redacted(
@@ -122,7 +111,7 @@ pub(crate) async fn resolve_provider_chat_pii_redaction<'a>(
&body_bytes,
format,
build_redaction_session_config(hmac_key, &runtime_config, now_unix_secs),
MaskChatRequestOptions::runtime(feature_settings.effective_inject_model_instruction()),
MaskChatRequestOptions::runtime(),
Some(&cache),
)
.await
@@ -190,3 +179,22 @@ fn redaction_mask_error_to_gateway_error(error: RedactionMaskError) -> GatewayEr
},
}
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::ChatPiiRedactionFeatureSettings;
#[test]
fn chat_pii_redaction_feature_settings_only_control_enablement() {
let mut settings = ChatPiiRedactionFeatureSettings::default();
settings.merge_from_value(Some(&json!({
"chat_pii_redaction": {
"enabled": true
}
})));
assert!(settings.effective_enabled());
}
}
@@ -0,0 +1,326 @@
use aether_ai_serving::AiRequestGzipPolicy;
use serde_json::Value;
use crate::ai_serving::is_openai_responses_family_format;
use super::state::GatewayProviderTransportSnapshot;
const DEFAULT_CODEX_REQUEST_GZIP_MIN_BYTES: usize = 64 * 1024;
pub(crate) fn resolve_transport_request_gzip_policy(
transport: &GatewayProviderTransportSnapshot,
) -> Option<AiRequestGzipPolicy> {
transport_request_gzip_policy_from_config(transport.endpoint.config.as_ref())
.or_else(|| transport_request_gzip_policy_from_config(transport.provider.config.as_ref()))
.or_else(|| default_transport_request_gzip_policy(transport))
}
fn default_transport_request_gzip_policy(
transport: &GatewayProviderTransportSnapshot,
) -> Option<AiRequestGzipPolicy> {
if !transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("codex")
{
return None;
}
if !is_codex_request_gzip_endpoint_api_format(transport.endpoint.api_format.as_str()) {
return None;
}
Some(AiRequestGzipPolicy {
enabled: Some(true),
min_bytes: Some(DEFAULT_CODEX_REQUEST_GZIP_MIN_BYTES),
})
}
fn is_codex_request_gzip_endpoint_api_format(api_format: &str) -> bool {
is_openai_responses_family_format(api_format)
|| api_format.trim().eq_ignore_ascii_case("openai:image")
}
fn transport_request_gzip_policy_from_config(
config: Option<&Value>,
) -> Option<AiRequestGzipPolicy> {
let object = config?.as_object()?;
for key in ["request_gzip", "request_body_gzip"] {
if let Some(policy) = object
.get(key)
.and_then(transport_request_gzip_policy_from_value)
{
return Some(policy);
}
}
let enabled = first_config_bool(
object,
&["request_gzip_enabled", "request_body_gzip_enabled"],
);
let min_bytes = first_config_usize(
object,
&["request_gzip_min_bytes", "request_body_gzip_min_bytes"],
);
match (enabled, min_bytes) {
(Some(false), _) => Some(AiRequestGzipPolicy {
enabled: Some(false),
min_bytes: None,
}),
(Some(true), min_bytes) => Some(AiRequestGzipPolicy {
enabled: Some(true),
min_bytes,
}),
(None, Some(min_bytes)) => Some(AiRequestGzipPolicy {
enabled: Some(true),
min_bytes: Some(min_bytes),
}),
(None, None) => None,
}
}
fn transport_request_gzip_policy_from_value(value: &Value) -> Option<AiRequestGzipPolicy> {
if let Some(enabled) = value.as_bool() {
return Some(AiRequestGzipPolicy {
enabled: Some(enabled),
min_bytes: None,
});
}
let object = value.as_object()?;
let enabled = first_config_bool(object, &["enabled"]);
let min_bytes = first_config_usize(object, &["min_bytes"]);
match (enabled, min_bytes) {
(Some(false), _) => Some(AiRequestGzipPolicy {
enabled: Some(false),
min_bytes: None,
}),
(Some(true), min_bytes) => Some(AiRequestGzipPolicy {
enabled: Some(true),
min_bytes,
}),
(None, Some(min_bytes)) => Some(AiRequestGzipPolicy {
enabled: Some(true),
min_bytes: Some(min_bytes),
}),
(None, None) => None,
}
}
fn first_config_bool(object: &serde_json::Map<String, Value>, keys: &[&str]) -> Option<bool> {
keys.iter()
.find_map(|key| object.get(*key).and_then(config_bool))
}
fn config_bool(value: &Value) -> Option<bool> {
value.as_bool().or_else(|| {
value.as_str().and_then(|text| {
let normalized = text.trim();
if normalized.eq_ignore_ascii_case("true") {
Some(true)
} else if normalized.eq_ignore_ascii_case("false") {
Some(false)
} else {
None
}
})
})
}
fn first_config_usize(object: &serde_json::Map<String, Value>, keys: &[&str]) -> Option<usize> {
keys.iter()
.find_map(|key| object.get(*key).and_then(config_usize))
}
fn config_usize(value: &Value) -> Option<usize> {
value
.as_u64()
.and_then(|number| usize::try_from(number).ok())
.or_else(|| {
value
.as_str()
.and_then(|text| text.trim().parse::<usize>().ok())
})
}
#[cfg(test)]
mod tests {
use super::*;
use aether_provider_transport::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
};
use serde_json::{json, Value};
fn sample_transport(
provider_type: &str,
endpoint_api_format: &str,
provider_config: Option<Value>,
endpoint_config: Option<Value>,
) -> GatewayProviderTransportSnapshot {
GatewayProviderTransportSnapshot {
provider: GatewayProviderTransportProvider {
id: "provider-1".to_string(),
name: "Provider".to_string(),
provider_type: provider_type.to_string(),
website: None,
is_active: true,
keep_priority_on_conversion: false,
enable_format_conversion: true,
concurrent_limit: None,
max_retries: None,
proxy: None,
request_timeout_secs: None,
stream_first_byte_timeout_secs: None,
config: provider_config,
},
endpoint: GatewayProviderTransportEndpoint {
id: "endpoint-1".to_string(),
provider_id: "provider-1".to_string(),
api_format: endpoint_api_format.to_string(),
api_family: None,
endpoint_kind: None,
is_active: true,
base_url: "https://api.example.test".to_string(),
header_rules: None,
body_rules: None,
max_retries: None,
custom_path: None,
config: endpoint_config,
format_acceptance_config: None,
proxy: None,
},
key: GatewayProviderTransportKey {
id: "key-1".to_string(),
provider_id: "provider-1".to_string(),
name: "key".to_string(),
auth_type: "api_key".to_string(),
is_active: true,
api_formats: None,
auth_type_by_format: None,
allow_auth_channel_mismatch_formats: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,
global_priority_by_format: None,
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
upstream_metadata: None,
decrypted_api_key: "secret".to_string(),
decrypted_auth_config: None,
},
}
}
#[test]
fn endpoint_request_gzip_policy_overrides_provider_policy() {
let transport = sample_transport(
"openai",
"openai:responses",
Some(json!({"request_gzip": false})),
Some(json!({"request_gzip": {"enabled": true, "min_bytes": 1024}})),
);
assert_eq!(
resolve_transport_request_gzip_policy(&transport),
Some(AiRequestGzipPolicy {
enabled: Some(true),
min_bytes: Some(1024),
})
);
}
#[test]
fn endpoint_request_gzip_false_disables_provider_and_codex_defaults() {
let transport = sample_transport(
"codex",
"openai:responses",
Some(json!({"request_gzip": {"enabled": true, "min_bytes": 1024}})),
Some(json!({"request_gzip": false})),
);
assert_eq!(
resolve_transport_request_gzip_policy(&transport),
Some(AiRequestGzipPolicy {
enabled: Some(false),
min_bytes: None,
})
);
}
#[test]
fn request_gzip_policy_supports_top_level_aliases() {
let transport = sample_transport(
"openai",
"openai:responses",
None,
Some(json!({
"request_body_gzip_enabled": true,
"request_body_gzip_min_bytes": "4096"
})),
);
assert_eq!(
resolve_transport_request_gzip_policy(&transport),
Some(AiRequestGzipPolicy {
enabled: Some(true),
min_bytes: Some(4096),
})
);
}
#[test]
fn request_gzip_policy_treats_min_bytes_only_as_enabled() {
let transport = sample_transport(
"openai",
"openai:responses",
None,
Some(json!({"request_gzip_min_bytes": 1})),
);
assert_eq!(
resolve_transport_request_gzip_policy(&transport),
Some(AiRequestGzipPolicy {
enabled: Some(true),
min_bytes: Some(1),
})
);
}
#[test]
fn codex_responses_endpoint_gets_default_request_gzip_policy() {
let transport = sample_transport("codex", "openai:responses", None, None);
assert_eq!(
resolve_transport_request_gzip_policy(&transport),
Some(AiRequestGzipPolicy {
enabled: Some(true),
min_bytes: Some(DEFAULT_CODEX_REQUEST_GZIP_MIN_BYTES),
})
);
}
#[test]
fn codex_image_endpoint_gets_default_request_gzip_policy() {
let transport = sample_transport("codex", "openai:image", None, None);
assert_eq!(
resolve_transport_request_gzip_policy(&transport),
Some(AiRequestGzipPolicy {
enabled: Some(true),
min_bytes: Some(DEFAULT_CODEX_REQUEST_GZIP_MIN_BYTES),
})
);
}
#[test]
fn non_codex_endpoint_does_not_get_default_request_gzip_policy() {
let transport = sample_transport("openai", "openai:responses", None, None);
assert_eq!(resolve_transport_request_gzip_policy(&transport), None);
}
}
@@ -7,7 +7,8 @@ use crate::ai_serving::planner::report_context::{
};
use crate::ai_serving::planner::spec_metadata::local_gemini_files_spec_metadata;
use crate::ai_serving::planner::{
build_ai_execution_decision_response, AiExecutionDecisionResponseParts,
build_ai_execution_decision_response, resolve_transport_request_gzip_policy,
AiExecutionDecisionResponseParts,
};
use crate::ai_serving::transport::{
resolve_transport_execution_timeouts, resolve_transport_profile,
@@ -123,6 +124,7 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
upstream_url,
file_name: _,
} = resolved;
let request_gzip = resolve_transport_request_gzip_policy(&transport);
let mut decision = build_ai_execution_decision_response(AiExecutionDecisionResponseParts {
decision_is_stream: spec_metadata.require_streaming,
@@ -154,6 +156,8 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
content_encoding: None,
request_gzip,
proxy,
transport_profile,
timeouts: resolve_transport_execution_timeouts(&transport),
@@ -5,7 +5,8 @@ use crate::ai_serving::planner::report_context::{
};
use crate::ai_serving::planner::spec_metadata::local_openai_image_spec_metadata;
use crate::ai_serving::planner::{
build_ai_execution_decision_response, AiExecutionDecisionResponseParts,
build_ai_execution_decision_response, resolve_transport_request_gzip_policy,
AiExecutionDecisionResponseParts,
};
use crate::ai_serving::transport::{
resolve_transport_execution_timeouts, resolve_transport_profile,
@@ -82,15 +83,6 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
"chatgpt_web_image".to_string(),
serde_json::Value::Bool(true),
);
extra_fields.insert(
"local_failover_policy".to_string(),
serde_json::json!({
"stop_status_codes": [400, 401, 403, 429, 500, 502, 503, 504],
"error_stop_patterns": [
{ "pattern": ".*" }
]
}),
);
}
let upstream_is_stream = resolved
.provider_request_body
@@ -143,6 +135,7 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
spec_metadata.api_format,
provider_api_format.as_str(),
);
let request_gzip = resolve_transport_request_gzip_policy(&transport);
let mut decision = build_ai_execution_decision_response(AiExecutionDecisionResponseParts {
decision_is_stream: spec_metadata.require_streaming,
@@ -169,6 +162,8 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
provider_request_body: Some(resolved.provider_request_body),
provider_request_body_base64: None,
content_type: Some("application/json".to_string()),
content_encoding: None,
request_gzip,
proxy,
transport_profile,
timeouts: resolve_transport_execution_timeouts(&transport),
@@ -5,7 +5,8 @@ use crate::ai_serving::planner::report_context::{
};
use crate::ai_serving::planner::spec_metadata::local_video_create_spec_metadata;
use crate::ai_serving::planner::{
build_ai_execution_decision_response, AiExecutionDecisionResponseParts,
build_ai_execution_decision_response, resolve_transport_request_gzip_policy,
AiExecutionDecisionResponseParts,
};
use crate::ai_serving::transport::{
resolve_transport_execution_timeouts, resolve_transport_profile,
@@ -103,6 +104,7 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat
provider_request_body,
upstream_url,
} = resolved;
let request_gzip = resolve_transport_request_gzip_policy(&transport);
let mut decision = build_ai_execution_decision_response(AiExecutionDecisionResponseParts {
decision_is_stream: false,
@@ -135,6 +137,8 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
content_encoding: None,
request_gzip,
proxy,
transport_profile,
timeouts: resolve_transport_execution_timeouts(&transport),
@@ -3,6 +3,7 @@ use std::collections::BTreeMap;
use super::{
apply_codex_openai_responses_special_body_edits, apply_codex_openai_responses_special_headers,
};
use crate::ai_serving::planner::standard::build_local_openai_responses_request_body;
use http::{HeaderMap, HeaderValue};
use serde_json::json;
@@ -36,6 +37,40 @@ fn applies_codex_defaults_when_body_rules_do_not_handle_fields() {
assert!(body.get("reasoning").is_none());
}
#[test]
fn local_openai_responses_codex_body_wraps_string_input_for_backend() {
let body = json!({
"model": "gpt-5",
"input": "hello"
});
let provider_request_body = build_local_openai_responses_request_body(
&body,
"gpt-5-upstream",
false,
false,
"codex",
"openai:responses",
None,
Some("key-123"),
&HeaderMap::new(),
false,
)
.expect("codex local openai responses body should build");
assert_eq!(
provider_request_body["input"],
json!([{
"type": "message",
"role": "user",
"content": [{
"type": "input_text",
"text": "hello"
}]
}])
);
}
#[test]
fn strips_store_for_compact_even_when_body_rules_handle_it() {
let body_rules = json!([
@@ -15,7 +15,8 @@ use crate::ai_serving::planner::report_context::{
use crate::ai_serving::planner::spec_metadata::local_standard_spec_metadata;
use crate::ai_serving::planner::CandidateFailureDiagnostic;
use crate::ai_serving::planner::{
build_ai_execution_decision_response, AiExecutionDecisionResponseParts,
build_ai_execution_decision_response, resolve_transport_request_gzip_policy,
AiExecutionDecisionResponseParts,
};
use crate::ai_serving::transport::{
resolve_transport_execution_timeouts, resolve_transport_profile,
@@ -175,6 +176,7 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
transport_profile: _,
request_redacted: _,
} = resolved;
let request_gzip = resolve_transport_request_gzip_policy(&transport);
let mut decision = build_ai_execution_decision_response(AiExecutionDecisionResponseParts {
decision_is_stream: spec_metadata.require_streaming,
@@ -201,6 +203,8 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
provider_request_body: Some(provider_request_body),
provider_request_body_base64: None,
content_type: Some("application/json".to_string()),
content_encoding: None,
request_gzip,
proxy,
transport_profile,
timeouts,
@@ -23,6 +23,7 @@ use crate::ai_serving::planner::spec_metadata::local_standard_spec_metadata;
use crate::ai_serving::planner::standard::{
apply_codex_openai_responses_special_headers, apply_deepseek_tool_call_thinking_compat,
is_deepseek_provider, request_body_build_failure_extra_data,
request_conversion_failure_extra_data,
};
use crate::ai_serving::transport::kiro::{
build_kiro_provider_headers, build_kiro_provider_request_body,
@@ -599,10 +600,14 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
attempt.candidate_index,
&attempt.candidate_id,
"provider_request_body_build_failed",
request_body_build_failure_extra_data(
request_conversion_failure_extra_data(
body_json,
spec_metadata.api_format,
provider_api_format,
Some(prepared_candidate.mapped_model.as_str()),
Some(parts.uri.path()),
upstream_is_stream,
"standard_family_request_conversion",
),
)
.await;
@@ -62,7 +62,8 @@ pub(crate) use crate::ai_serving::{
normalize_openai_responses_request_to_openai_chat_request, parse_openai_tool_result_content,
};
pub(crate) use aether_ai_serving::{
request_body_build_failure_extra_data, same_format_provider_request_body_failure_extra_data,
request_body_build_failure_extra_data, request_conversion_failure_extra_data,
same_format_provider_request_body_failure_extra_data,
};
pub(crate) fn build_standard_upstream_url(
@@ -6,7 +6,8 @@ use crate::ai_serving::planner::report_context::{
insert_provider_stream_event_api_format, LocalExecutionReportContextParts,
};
use crate::ai_serving::planner::{
build_ai_execution_decision_response, AiExecutionDecisionResponseParts,
build_ai_execution_decision_response, resolve_transport_request_gzip_policy,
AiExecutionDecisionResponseParts,
};
use crate::ai_serving::transport::{
resolve_transport_execution_timeouts, resolve_transport_profile,
@@ -104,15 +105,6 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
.eq_ignore_ascii_case("chatgpt_web")
{
extra_fields.insert("chatgpt_web_image".to_string(), serde_json::json!(true));
extra_fields.insert(
"local_failover_policy".to_string(),
serde_json::json!({
"stop_status_codes": [400, 401, 403, 429, 500, 502, 503, 504],
"error_stop_patterns": [
{ "pattern": ".*" }
]
}),
);
}
let super::request::LocalOpenAiChatCandidatePayloadParts {
client_api_format,
@@ -192,6 +184,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
),
&transport,
);
let request_gzip = resolve_transport_request_gzip_policy(&transport);
let mut decision = build_ai_execution_decision_response(AiExecutionDecisionResponseParts {
decision_is_stream,
@@ -218,6 +211,8 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
provider_request_body: Some(provider_request_body),
provider_request_body_base64: None,
content_type: Some("application/json".to_string()),
content_encoding: None,
request_gzip,
proxy,
transport_profile,
timeouts,
@@ -25,6 +25,7 @@ use crate::ai_serving::planner::standard::{
apply_deepseek_tool_call_thinking_compat, build_cross_format_openai_chat_request_body,
build_cross_format_openai_chat_upstream_url, build_local_openai_chat_request_body,
build_local_openai_chat_upstream_url, request_body_build_failure_extra_data,
request_conversion_failure_extra_data,
};
use crate::ai_serving::transport::auth::resolve_local_openai_bearer_auth;
use crate::ai_serving::transport::kiro::{
@@ -601,10 +602,14 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
candidate_index,
candidate_id,
"provider_request_body_build_failed",
request_body_build_failure_extra_data(
request_conversion_failure_extra_data(
body_json,
"openai:chat",
provider_api_format.as_str(),
Some(prepared_candidate.mapped_model.as_str()),
Some(parts.uri.path()),
upstream_is_stream,
"openai_chat_request_conversion",
),
)
.await;
@@ -326,6 +326,8 @@ mod tests {
})),
provider_request_body_base64: None,
content_type: Some("application/json".to_string()),
content_encoding: None,
request_gzip: None,
proxy: None,
transport_profile: None,
timeouts: None,
@@ -422,6 +424,8 @@ mod tests {
provider_request_body: Some(json!({"model":"gpt-5.4","messages":[],"stream":true})),
provider_request_body_base64: None,
content_type: Some("application/json".to_string()),
content_encoding: None,
request_gzip: None,
proxy: None,
transport_profile: None,
timeouts: None,
@@ -488,6 +492,8 @@ mod tests {
provider_request_body,
provider_request_body_base64: None,
content_type: Some("application/json".to_string()),
content_encoding: None,
request_gzip: None,
proxy: None,
transport_profile: None,
timeouts: None,
@@ -595,6 +601,8 @@ mod tests {
),
provider_request_body_base64: None,
content_type: Some("application/json".to_string()),
content_encoding: None,
request_gzip: None,
proxy: None,
transport_profile: None,
timeouts: None,
@@ -292,6 +292,8 @@ mod tests {
})),
provider_request_body_base64: None,
content_type: Some("application/json".to_string()),
content_encoding: None,
request_gzip: None,
proxy: None,
transport_profile: None,
timeouts: None,
@@ -387,6 +389,8 @@ mod tests {
provider_request_body: Some(json!({"model":"gpt-5.4","messages":[],"stream":false})),
provider_request_body_base64: None,
content_type: Some("application/json".to_string()),
content_encoding: None,
request_gzip: None,
proxy: None,
transport_profile: None,
timeouts: None,
@@ -458,6 +462,8 @@ mod tests {
),
provider_request_body_base64: None,
content_type: Some("application/json".to_string()),
content_encoding: None,
request_gzip: None,
proxy: None,
transport_profile: None,
timeouts: None,
@@ -9,7 +9,8 @@ use crate::ai_serving::planner::report_context::{
};
use crate::ai_serving::planner::spec_metadata::local_openai_responses_spec_metadata;
use crate::ai_serving::planner::{
build_ai_execution_decision_response, AiExecutionDecisionResponseParts,
build_ai_execution_decision_response, resolve_transport_request_gzip_policy,
AiExecutionDecisionResponseParts,
};
use crate::ai_serving::transport::{
resolve_transport_execution_timeouts, resolve_transport_profile,
@@ -101,15 +102,6 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
.eq_ignore_ascii_case("chatgpt_web")
{
extra_fields.insert("chatgpt_web_image".to_string(), json!(true));
extra_fields.insert(
"local_failover_policy".to_string(),
json!({
"stop_status_codes": [400, 401, 403, 429, 500, 502, 503, 504],
"error_stop_patterns": [
{ "pattern": ".*" }
]
}),
);
}
insert_provider_stream_event_api_format(
&mut extra_fields,
@@ -212,6 +204,7 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
image_request_summary: _,
request_redacted: _,
} = resolved;
let request_gzip = resolve_transport_request_gzip_policy(&transport);
let mut decision = build_ai_execution_decision_response(AiExecutionDecisionResponseParts {
decision_is_stream: spec_metadata.require_streaming,
@@ -238,6 +231,8 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
provider_request_body: Some(provider_request_body),
provider_request_body_base64: None,
content_type: Some("application/json".to_string()),
content_encoding: None,
request_gzip,
proxy,
transport_profile,
timeouts,
@@ -27,6 +27,7 @@ use crate::ai_serving::planner::standard::{
apply_deepseek_tool_call_thinking_compat, build_cross_format_openai_responses_request_body,
build_cross_format_openai_responses_upstream_url, build_local_openai_responses_request_body,
build_local_openai_responses_upstream_url, request_body_build_failure_extra_data,
request_conversion_failure_extra_data,
};
use crate::ai_serving::transport::antigravity::{
build_antigravity_safe_v1internal_request, build_antigravity_static_identity_headers,
@@ -371,10 +372,14 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
candidate_index,
candidate_id,
"provider_request_body_build_failed",
request_body_build_failure_extra_data(
request_conversion_failure_extra_data(
body_json,
spec_metadata.api_format,
provider_api_format,
Some(mapped_model.as_str()),
Some(parts.uri.path()),
upstream_is_stream,
"openai_responses_request_conversion",
),
)
.await;
@@ -72,8 +72,10 @@ pub(crate) use aether_provider_transport::{
build_local_openai_chat_upstream_url, build_local_openai_responses_upstream_url,
build_openai_image_headers, build_openai_image_upstream_url, build_passthrough_headers,
build_request_trace_proxy_value, build_same_format_provider_headers,
build_same_format_provider_request_body, build_same_format_provider_upstream_url,
build_standard_plan_fallback_headers, build_standard_plan_fallback_openai_chat_url,
build_same_format_provider_request_body,
build_same_format_provider_request_body_with_compatibility_report,
build_same_format_provider_upstream_url, build_standard_plan_fallback_headers,
build_standard_plan_fallback_openai_chat_url,
build_standard_plan_fallback_openai_responses_url, build_standard_provider_request_headers,
build_transport_request_url, build_transport_request_url_for_request_body,
build_video_create_headers, build_video_create_request_body, build_video_create_upstream_url,
@@ -106,12 +108,14 @@ pub(crate) use aether_provider_transport::{
GeminiCliRequestAuthUnsupportedReason, GeminiCliRequestEnvelopeSupport,
GeminiFilesHeadersInput, GeminiFilesRequestBodyError, GeminiFilesRequestBodyParts,
GrokHeaderInput, LocalResolvedOAuthRequestAuth, ProviderOpenAiImageHeadersInput,
ProviderVideoCreateFamily, ProviderVideoCreateHeadersInput, SameFormatProviderFamily,
SameFormatProviderHeadersInput, SameFormatProviderRequestBehavior,
ProviderVideoCreateFamily, ProviderVideoCreateHeadersInput,
SameFormatProviderCompatibilityEdit, SameFormatProviderCompatibilityEditAction,
SameFormatProviderFamily, SameFormatProviderHeadersInput, SameFormatProviderRequestBehavior,
SameFormatProviderRequestBehaviorParams, SameFormatProviderRequestBodyInput,
SameFormatProviderUpstreamUrlParams, StandardPlanFallbackAcceptPolicy,
StandardPlanFallbackHeadersInput, StandardProviderRequestHeaders,
StandardProviderRequestHeadersInput, TransportRequestBodySemanticsError,
TransportRequestUrlParams, GEMINI_CLI_USER_AGENT, GEMINI_CLI_V1INTERNAL_ENVELOPE_NAME,
GROK_CHAT_PATH, GROK_INTERNAL_HEADER, GROK_RATE_LIMITS_PATH, WINDSURF_ENVELOPE_NAME,
SameFormatProviderRequestBodyOutput, SameFormatProviderUpstreamUrlParams,
StandardPlanFallbackAcceptPolicy, StandardPlanFallbackHeadersInput,
StandardProviderRequestHeaders, StandardProviderRequestHeadersInput,
TransportRequestBodySemanticsError, TransportRequestUrlParams, GEMINI_CLI_USER_AGENT,
GEMINI_CLI_V1INTERNAL_ENVELOPE_NAME, GROK_CHAT_PATH, GROK_INTERNAL_HEADER,
GROK_RATE_LIMITS_PATH, WINDSURF_ENVELOPE_NAME,
};
+9 -3
View File
@@ -11,9 +11,15 @@ pub(crate) async fn read_request_candidate_trace(
request_id: &str,
attempted_only: bool,
) -> Result<Option<RequestCandidateTrace>, DataLayerError> {
let all_candidates = state
.list_request_candidates_by_request_id(request_id)
.await?;
let all_candidates = if attempted_only {
state
.list_attempted_request_candidates_by_request_id(request_id)
.await?
} else {
state
.list_request_candidates_by_request_id(request_id)
.await?
};
Ok(RequestCandidateTrace::from_candidates(
request_id,
all_candidates,
+18 -2
View File
@@ -19,6 +19,16 @@ impl GatewayDataState {
}
}
pub(crate) async fn list_attempted_request_candidates_by_request_id(
&self,
request_id: &str,
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
match &self.request_candidate_reader {
Some(repository) => repository.list_attempted_by_request_id(request_id).await,
None => Ok(Vec::new()),
}
}
pub(crate) async fn list_request_candidates_by_provider_id(
&self,
provider_id: &str,
@@ -416,18 +426,24 @@ impl GatewayDataState {
pub(crate) async fn cleanup_deleted_provider_catalog_refs(
&self,
provider_id: &str,
provider_deleted: bool,
endpoint_ids: &[String],
key_ids: &[String],
) -> Result<(), DataLayerError> {
let cleaned = match &self.provider_catalog_writer {
Some(repository) => {
repository
.cleanup_deleted_provider_refs(provider_id, endpoint_ids, key_ids)
.cleanup_deleted_provider_refs(
provider_id,
provider_deleted,
endpoint_ids,
key_ids,
)
.await
}
None => Ok(()),
};
if !endpoint_ids.is_empty() || !key_ids.is_empty() {
if provider_deleted || !endpoint_ids.is_empty() || !key_ids.is_empty() {
self.clear_provider_catalog_cache();
}
cleaned
@@ -1124,6 +1124,16 @@ impl GatewayDataState {
}
}
pub(crate) async fn find_request_usage_by_request_id_shallow(
&self,
request_id: &str,
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
match &self.usage_reader {
Some(repository) => repository.find_by_request_id_shallow(request_id).await,
None => Ok(None),
}
}
pub(crate) async fn find_request_usage_by_id(
&self,
usage_id: &str,
@@ -1958,6 +1968,14 @@ impl GatewayDataState {
self.find_request_usage_by_request_id(request_id).await
}
pub(crate) async fn read_request_usage_audit_shallow(
&self,
request_id: &str,
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
self.find_request_usage_by_request_id_shallow(request_id)
.await
}
pub(crate) async fn read_request_audit_bundle(
&self,
request_id: &str,
@@ -4,7 +4,9 @@ use std::sync::{
Arc, LazyLock,
};
use aether_admin::provider::pool as admin_provider_pool_pure;
use aether_admin::provider::{
pool as admin_provider_pool_pure, status as admin_provider_status_pure,
};
use aether_data_contracts::repository::candidate_selection::{
StoredMinimalCandidateSelectionRow, StoredPoolKeyCandidateOrder,
StoredPoolKeyCandidateRowsByKeyIdsQuery, StoredPoolKeyCandidateRowsQuery,
@@ -1217,9 +1219,15 @@ fn pool_key_requires_reauth_for_scheduling(
.map(str::trim)
.unwrap_or_default();
if !invalid_reason.is_empty() {
if pool_oauth_reason_has_tag(invalid_reason, "[OAUTH_EXPIRED]")
|| pool_oauth_reason_has_tag(invalid_reason, "[ACCOUNT_BLOCK]")
{
let account_state = admin_provider_status_pure::resolve_pool_account_state(
None,
key.upstream_metadata.as_ref(),
Some(invalid_reason),
);
if account_state.blocked && !account_state.recoverable {
return true;
}
if pool_oauth_reason_has_tag(invalid_reason, "[ACCOUNT_BLOCK]") {
return true;
}
if pool_oauth_reason_has_tag(invalid_reason, "[REQUEST_FAILED]") {
@@ -1230,6 +1238,9 @@ fn pool_key_requires_reauth_for_scheduling(
.expires_at_unix_secs
.is_none_or(|expires_at| expires_at == 0 || expires_at <= now_unix_secs);
}
if pool_oauth_reason_has_tag(invalid_reason, "[OAUTH_EXPIRED]") {
return false;
}
return true;
}
@@ -3265,8 +3276,7 @@ mod tests {
let mut key_a_invalid = sample_codex_pool_key("provider-a", "key-a-invalid");
key_a_invalid.oauth_invalid_at_unix_secs = Some(1_710_000_000);
key_a_invalid.oauth_invalid_reason =
Some("[OAUTH_EXPIRED] Codex Token 无效或已过期 (401)".to_string());
key_a_invalid.oauth_invalid_reason = Some("[OAUTH_EXPIRED] token invalidated".to_string());
let exhausted_status_snapshot = json!({
"quota": {
"provider_type": "codex",
@@ -3370,8 +3380,7 @@ mod tests {
let mut key_a_invalid = sample_codex_pool_key("provider-a", "key-a-invalid");
key_a_invalid.oauth_invalid_at_unix_secs = Some(1_710_000_000);
key_a_invalid.oauth_invalid_reason =
Some("[OAUTH_EXPIRED] Codex Token 无效或已过期 (401)".to_string());
key_a_invalid.oauth_invalid_reason = Some("[OAUTH_EXPIRED] token invalidated".to_string());
key_a_invalid.status_snapshot = Some(json!({
"quota": {
"provider_type": "codex",
@@ -3463,6 +3472,9 @@ mod tests {
key.oauth_invalid_reason = Some("[REQUEST_FAILED] 账号状态检查失败".to_string());
key.oauth_invalid_at_unix_secs = Some(100);
assert!(!pool_key_requires_reauth_for_scheduling(&key, 300));
key.oauth_invalid_reason = Some("[OAUTH_EXPIRED] session expired".to_string());
assert!(!pool_key_requires_reauth_for_scheduling(&key, 300));
}
#[test]
@@ -3471,6 +3483,10 @@ mod tests {
key.oauth_invalid_reason = Some("[ACCOUNT_BLOCK] account has been deactivated".to_string());
assert!(pool_key_requires_reauth_for_scheduling(&key, 100));
key.oauth_invalid_reason = Some("[OAUTH_EXPIRED] token invalidated".to_string());
key.oauth_invalid_at_unix_secs = None;
assert!(pool_key_requires_reauth_for_scheduling(&key, 100));
key.oauth_invalid_reason = Some("Kiro Token 无效或已过期".to_string());
key.oauth_invalid_at_unix_secs = None;
assert!(pool_key_requires_reauth_for_scheduling(&key, 100));
@@ -1007,6 +1007,100 @@ mod tests {
);
}
#[tokio::test]
async fn provider_failover_rules_can_stop_rate_limit_status() {
let result = ExecutionResult {
request_id: "req-1".to_string(),
candidate_id: None,
status_code: 429,
headers: Default::default(),
body: None,
telemetry: None,
error: None,
};
let local_report_context = serde_json::json!({
"candidate_index": 0,
"retry_index": 0,
});
let state = build_state_with_provider_config(Some(serde_json::json!({
"failover_rules": {
"stop_on_status_codes": [429]
}
})));
let plan = sample_plan();
assert!(
should_stop_local_candidate_failover_sync(
&state,
&plan,
"openai_chat_sync",
Some(&local_report_context),
&result,
Some("{\"error\":{\"message\":\"rate limited\"}}"),
)
.await
);
assert!(
!should_retry_next_local_candidate_sync(
&state,
&plan,
"openai_chat_sync",
Some(&local_report_context),
&result,
Some("{\"error\":{\"message\":\"rate limited\"}}"),
)
.await
);
}
#[tokio::test]
async fn status_only_error_stop_rule_can_stop_rate_limit_status() {
let result = ExecutionResult {
request_id: "req-1".to_string(),
candidate_id: None,
status_code: 429,
headers: Default::default(),
body: None,
telemetry: None,
error: None,
};
let local_report_context = serde_json::json!({
"candidate_index": 0,
"retry_index": 0,
});
let state = build_state_with_provider_config(Some(serde_json::json!({
"failover_rules": {
"error_stop_patterns": [
{"status_codes": [429]}
]
}
})));
let plan = sample_plan();
assert!(
should_stop_local_candidate_failover_sync(
&state,
&plan,
"openai_chat_sync",
Some(&local_report_context),
&result,
None,
)
.await
);
assert!(
!should_retry_next_local_candidate_sync(
&state,
&plan,
"openai_chat_sync",
Some(&local_report_context),
&result,
None,
)
.await
);
}
#[test]
fn resolve_local_failover_policy_reads_regex_rules() {
let state = build_state_with_provider_config(Some(serde_json::json!({
@@ -1039,6 +1133,32 @@ mod tests {
);
}
#[test]
fn resolve_local_failover_policy_reads_status_only_error_stop_rules() {
let state = build_state_with_provider_config(Some(serde_json::json!({
"failover_rules": {
"success_failover_patterns": [
{"status_codes": [200]}
],
"error_stop_patterns": [
{"status_codes": [429]}
]
}
})));
let plan = sample_plan();
let runtime = tokio::runtime::Runtime::new().expect("runtime should build");
let policy = runtime.block_on(resolve_local_failover_policy(&state, &plan, None));
assert!(policy.success_failover_patterns.is_empty());
assert_eq!(
policy.error_stop_patterns,
vec![LocalFailoverRegexRule {
pattern: String::new(),
status_codes: [429].into_iter().collect(),
}]
);
}
#[tokio::test]
async fn success_failover_pattern_can_retry_sync_candidate() {
let result = ExecutionResult {
@@ -1125,11 +1245,11 @@ mod tests {
}
#[tokio::test]
async fn chatgpt_web_report_context_stops_local_sync_failover_on_transport_errors() {
async fn report_context_failover_policy_does_not_override_provider_config() {
let result = ExecutionResult {
request_id: "req-1".to_string(),
candidate_id: None,
status_code: 503,
status_code: 429,
headers: Default::default(),
body: None,
telemetry: None,
@@ -1150,7 +1270,7 @@ mod tests {
let plan = sample_plan();
assert!(
should_stop_local_candidate_failover_sync(
should_retry_next_local_candidate_sync(
&state,
&plan,
"openai_image_sync",
@@ -1161,7 +1281,7 @@ mod tests {
.await
);
assert!(
!should_retry_next_local_candidate_sync(
!should_stop_local_candidate_failover_sync(
&state,
&plan,
"openai_image_sync",
@@ -3921,10 +3921,14 @@ mod tests {
StreamFrame, StreamFramePayload, StreamFrameType,
};
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data::repository::usage::InMemoryUsageReadRepository;
use aether_data_contracts::repository::candidates::{
RequestCandidateReadRepository, RequestCandidateStatus,
};
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_data_contracts::repository::usage::UsageReadRepository;
use aether_usage_runtime::UsageRuntimeConfig;
use async_stream::stream;
@@ -3954,6 +3958,77 @@ mod tests {
use crate::tunnel::{tunnel_protocol, TunnelProxyConn};
use crate::AppState;
fn provider_catalog_stop_429_for_plan(
plan: &ExecutionPlan,
) -> InMemoryProviderCatalogReadRepository {
let provider_type = plan.provider_name.as_deref().unwrap_or("custom");
let provider = StoredProviderCatalogProvider::new(
plan.provider_id.clone(),
plan.provider_id.clone(),
Some("https://provider.example".to_string()),
provider_type.to_string(),
)
.expect("provider should build")
.with_transport_fields(
true,
false,
false,
None,
Some(3),
None,
None,
None,
Some(json!({
"failover_rules": {
"stop_status_codes": [429]
}
})),
);
let endpoint = StoredProviderCatalogEndpoint::new(
plan.endpoint_id.clone(),
plan.provider_id.clone(),
plan.provider_api_format.clone(),
None,
None,
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://provider.example".to_string(),
None,
None,
Some(2),
None,
None,
None,
None,
)
.expect("endpoint transport should build");
let key = StoredProviderCatalogKey::new(
plan.key_id.clone(),
plan.provider_id.clone(),
plan.key_id.clone(),
"api_key".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(json!([plan.provider_api_format.clone()])),
"plain-upstream-key".to_string(),
None,
None,
Some(json!({ "openai:chat": 1 })),
None,
None,
None,
None,
)
.expect("key transport should build");
InMemoryProviderCatalogReadRepository::seed(vec![provider], vec![endpoint], vec![key])
}
fn test_decision() -> GatewayControlDecision {
GatewayControlDecision::synthetic(
"/v1/chat/completions",
@@ -5458,6 +5533,14 @@ mod tests {
transport_profile: None,
timeouts: None,
};
let state = state.with_data_state_for_tests(
crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests(
Arc::clone(&request_candidate_repository),
Arc::clone(&usage_repository),
)
.with_provider_catalog_reader(Arc::new(provider_catalog_stop_429_for_plan(&plan)))
.with_encryption_key_for_tests("development-key"),
);
let trailer_error = connect_json_frame(
2,
br#"{"error":{"code":"resource_exhausted","message":"an internal error occurred"}}"#,
@@ -5493,10 +5576,7 @@ mod tests {
"client_api_format": "claude:messages",
"needs_conversion": true,
"has_envelope": true,
"envelope_name": "windsurf:GetChatMessage",
"local_failover_policy": {
"stop_status_codes": [429]
}
"envelope_name": "windsurf:GetChatMessage"
})),
crate::clock::current_unix_ms(),
Instant::now(),
@@ -5590,6 +5670,14 @@ mod tests {
transport_profile: None,
timeouts: None,
};
let state = state.with_data_state_for_tests(
crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests(
Arc::clone(&request_candidate_repository),
Arc::clone(&usage_repository),
)
.with_provider_catalog_reader(Arc::new(provider_catalog_stop_429_for_plan(&plan)))
.with_encryption_key_for_tests("development-key"),
);
let connect_error = connect_json_frame(
2,
br#"{"error":{"code":"resource_exhausted","message":"quota exhausted"}}"#,
@@ -5624,10 +5712,7 @@ mod tests {
"client_api_format": "claude:messages",
"needs_conversion": true,
"has_envelope": true,
"envelope_name": "windsurf:GetChatMessage",
"local_failover_policy": {
"stop_status_codes": [429]
}
"envelope_name": "windsurf:GetChatMessage"
})),
crate::clock::current_unix_ms(),
Instant::now(),
@@ -16,6 +16,8 @@ pub(super) fn decode_execution_result_body(
};
if let Some(json_body) = body.json_body {
remove_header_case_insensitive(headers, "content-encoding");
remove_header_case_insensitive(headers, "content-length");
headers
.entry("content-type".to_string())
.or_insert_with(|| "application/json".to_string());
@@ -34,3 +36,49 @@ pub(super) fn decode_execution_result_body(
Ok((Vec::new(), None, None))
}
fn remove_header_case_insensitive(headers: &mut BTreeMap<String, String>, name: &str) {
if let Some(existing_key) = headers
.keys()
.find(|key| key.eq_ignore_ascii_case(name))
.cloned()
{
headers.remove(&existing_key);
}
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use aether_contracts::ResponseBody;
use serde_json::json;
use super::decode_execution_result_body;
#[test]
fn decoded_json_body_drops_stale_content_encoding_headers() {
let mut headers = BTreeMap::from([
("content-encoding".to_string(), "gzip".to_string()),
("content-length".to_string(), "999".to_string()),
]);
let (body_bytes, body_json, body_base64) = decode_execution_result_body(
Some(ResponseBody {
json_body: Some(json!({"ok": true})),
body_bytes_b64: None,
}),
&mut headers,
)
.expect("body should decode");
assert_eq!(body_json, Some(json!({"ok": true})));
assert_eq!(body_base64, None);
assert_eq!(body_bytes, br#"{"ok":true}"#);
assert_eq!(headers.get("content-encoding"), None);
assert_eq!(
headers.get("content-length").cloned(),
Some(body_bytes.len().to_string())
);
}
}
@@ -63,6 +63,8 @@ fn missing_exact_provider_request_payload(decision_kind: &str) -> AiExecutionDec
provider_request_body: None,
provider_request_body_base64: None,
content_type: Some("application/json".to_string()),
content_encoding: None,
request_gzip: None,
proxy: None,
transport_profile: None,
timeouts: None,
+133 -1
View File
@@ -111,6 +111,38 @@ impl LocalExecutionRuntimeMissContext {
}
Some(summaries.join(" | "))
}
pub(crate) fn all_provider_request_body_build_failures_detail(&self) -> Option<String> {
if self.candidate_contexts.is_empty()
|| !self.candidate_contexts.iter().all(|candidate| {
candidate.candidate.status == RequestCandidateStatus::Skipped
&& candidate
.candidate
.skip_reason
.as_deref()
.map(str::trim)
.is_some_and(|value| value == "provider_request_body_build_failed")
})
{
return None;
}
let diagnostic = self
.candidate_contexts
.iter()
.find_map(runtime_miss_candidate_failure_diagnostic)?;
let mut detail = format!("上游请求体转换失败:{}", diagnostic.message);
if diagnostic.path != "$" {
detail.push_str(&format!(";字段路径:{}", diagnostic.path));
}
detail.push_str("(原因代码: provider_request_body_build_failed)");
Some(detail)
}
}
struct RuntimeMissFailureDiagnostic {
path: String,
message: String,
}
pub(crate) async fn build_local_execution_exhaustion(
@@ -835,6 +867,41 @@ fn candidate_extra_data_string(candidate: &StoredRequestCandidate, key: &str) ->
.map(ToOwned::to_owned)
}
fn runtime_miss_candidate_failure_diagnostic(
candidate: &RuntimeMissCandidateContext,
) -> Option<RuntimeMissFailureDiagnostic> {
let extra_data = candidate.candidate.extra_data.as_ref()?.as_object()?;
let diagnostic = extra_data
.get("failure_diagnostic")
.and_then(Value::as_object)
.filter(|diagnostic| diagnostic.get("safe_to_show") != Some(&Value::Bool(false)))
.or_else(|| {
extra_data
.get("request_conversion_error")
.and_then(Value::as_object)
})
.or_else(|| {
extra_data
.get("request_body_build_error")
.and_then(Value::as_object)
})?;
let message = diagnostic
.get("message")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())?;
let path = diagnostic
.get("path")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("$");
Some(RuntimeMissFailureDiagnostic {
path: path.to_string(),
message: message.to_string(),
})
}
fn build_runtime_miss_candidate_endpoint_url(
candidate: &StoredRequestCandidate,
endpoint: &StoredProviderCatalogEndpoint,
@@ -1036,7 +1103,8 @@ mod tests {
use super::{
apply_runtime_miss_usage_routing, beautify_local_execution_client_error_message,
request_candidate_represents_provider_execution,
select_last_runtime_miss_executed_candidate, RuntimeMissCandidateContext,
select_last_runtime_miss_executed_candidate, LocalExecutionRuntimeMissContext,
RuntimeMissCandidateContext,
};
use crate::constants::EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS;
use crate::state::LocalExecutionRuntimeMissDiagnostic;
@@ -1161,4 +1229,68 @@ mod tests {
assert!(select_last_runtime_miss_executed_candidate(&contexts).is_none());
}
#[test]
fn runtime_miss_context_surfaces_request_conversion_field_diagnostic() {
let skipped_candidate = StoredRequestCandidate::new(
"cand-skipped".to_string(),
"req-1".to_string(),
Some("user-1".to_string()),
Some("api-key-1".to_string()),
Some("alice".to_string()),
Some("default".to_string()),
0,
0,
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("provider-key-1".to_string()),
RequestCandidateStatus::Skipped,
Some("provider_request_body_build_failed".to_string()),
false,
None,
None,
None,
None,
None,
Some(json!({
"failure_diagnostic": {
"kind": "request_conversion",
"path": "$.n",
"message": "openai:chat 字段 n 不能无损转换到 openai:responses:OpenAI Responses request has no canonical equivalent for this Chat field",
"safe_to_show": true
},
"request_conversion_error": {
"path": "$.n",
"message": "compat"
}
})),
None,
100,
None,
None,
)
.expect("candidate should build");
let context = LocalExecutionRuntimeMissContext {
candidate_contexts: vec![RuntimeMissCandidateContext {
candidate: skipped_candidate,
provider_name: Some("openai".to_string()),
key_name: Some("prod".to_string()),
client_api_format: Some("openai:chat".to_string()),
provider_api_format: Some("openai:responses".to_string()),
global_model_name: Some("gpt-5".to_string()),
selected_provider_model_name: Some("gpt-5-upstream".to_string()),
endpoint_url: Some("https://api.openai.example/v1/responses".to_string()),
}],
..LocalExecutionRuntimeMissContext::default()
};
let detail = context
.all_provider_request_body_build_failures_detail()
.expect("detail should include conversion diagnostic");
assert!(detail.contains("字段 n"));
assert!(detail.contains("字段路径:$.n"));
assert!(detail.contains("provider_request_body_build_failed"));
}
}
@@ -89,7 +89,7 @@ async fn resolve_admin_monitoring_trace(
{
let usage = app
.data
.read_request_usage_audit(request_id)
.read_request_usage_audit_shallow(request_id)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
return Ok(Some(ResolvedAdminMonitoringTrace { trace, usage }));
@@ -98,7 +98,7 @@ async fn resolve_admin_monitoring_trace(
let mut usage_candidates = Vec::new();
if let Some(usage) = app
.data
.read_request_usage_audit(request_id)
.read_request_usage_audit_shallow(request_id)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?
{
@@ -14,17 +14,76 @@ use aether_admin::observability::usage::{
admin_usage_bad_request_response, admin_usage_data_unavailable_response,
admin_usage_provider_key_name, ADMIN_USAGE_DATA_UNAVAILABLE_DETAIL,
};
use aether_data_contracts::repository::usage::UsageBodyField;
use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UsageBodyField};
use axum::{
body::Body,
http,
response::{IntoResponse, Response},
Json,
};
use serde_json::json;
use serde_json::{json, Value};
use std::collections::BTreeMap;
use tokio::try_join;
struct AdminUsageDetailBodyValue {
value: Option<Value>,
load_failed: bool,
}
async fn resolve_admin_usage_detail_request_body(
state: &AdminAppState<'_>,
item: &StoredRequestUsageAudit,
) -> AdminUsageDetailBodyValue {
match admin_usage_resolve_request_capture_body_for_item(state, item, None).await {
Ok(body) => AdminUsageDetailBodyValue {
value: body,
load_failed: false,
},
Err(err) => {
tracing::warn!(
error = ?err,
usage_id = %item.id,
request_id = %item.request_id,
field = UsageBodyField::RequestBody.as_storage_field(),
"failed to resolve admin usage detail body"
);
let value = admin_usage_resolve_request_capture_body(item, None);
AdminUsageDetailBodyValue {
load_failed: value.is_none(),
value,
}
}
}
}
async fn resolve_admin_usage_detail_body_value(
state: &AdminAppState<'_>,
item: &StoredRequestUsageAudit,
field: UsageBodyField,
) -> AdminUsageDetailBodyValue {
let inline_body = item.body_value(field);
match admin_usage_resolve_body_value(state, item, inline_body, field).await {
Ok(body) => AdminUsageDetailBodyValue {
value: body,
load_failed: false,
},
Err(err) => {
tracing::warn!(
error = ?err,
usage_id = %item.id,
request_id = %item.request_id,
field = field.as_storage_field(),
"failed to resolve admin usage detail body"
);
let value = inline_body.cloned();
AdminUsageDetailBodyValue {
load_failed: value.is_none(),
value,
}
}
}
}
pub(super) async fn maybe_build_local_admin_usage_detail_response(
state: &AdminAppState<'_>,
request_context: &AdminRequestContext<'_>,
@@ -174,32 +233,40 @@ pub(super) async fn maybe_build_local_admin_usage_detail_response(
let provider_key_name = admin_usage_provider_key_name(&item, &provider_key_names);
let mut detail_item = item.clone();
let mut body_load_errors = serde_json::Map::new();
let request_body = if include_bodies {
let (request_body, provider_request_body, response_body, client_response_body) = try_join!(
admin_usage_resolve_request_capture_body_for_item(state, &item, None),
admin_usage_resolve_body_value(
let (request_body, provider_request_body, response_body, client_response_body) = tokio::join!(
resolve_admin_usage_detail_request_body(state, &item),
resolve_admin_usage_detail_body_value(
state,
&item,
item.provider_request_body.as_ref(),
UsageBodyField::ProviderRequestBody,
),
admin_usage_resolve_body_value(
resolve_admin_usage_detail_body_value(
state,
&item,
item.response_body.as_ref(),
UsageBodyField::ResponseBody,
),
admin_usage_resolve_body_value(
resolve_admin_usage_detail_body_value(
state,
&item,
item.client_response_body.as_ref(),
UsageBodyField::ClientResponseBody,
),
)?;
detail_item.provider_request_body = provider_request_body;
detail_item.response_body = response_body;
detail_item.client_response_body = client_response_body;
request_body
);
for (field, resolved) in [
(UsageBodyField::RequestBody, &request_body),
(UsageBodyField::ProviderRequestBody, &provider_request_body),
(UsageBodyField::ResponseBody, &response_body),
(UsageBodyField::ClientResponseBody, &client_response_body),
] {
if resolved.load_failed {
body_load_errors.insert(field.as_storage_field().to_string(), json!(true));
}
}
detail_item.provider_request_body = provider_request_body.value;
detail_item.response_body = response_body.value;
detail_item.client_response_body = client_response_body.value;
request_body.value
} else {
None
};
@@ -207,7 +274,7 @@ pub(super) async fn maybe_build_local_admin_usage_detail_response(
// request_body 已通过 request capture 解析;其余 detached body 在上方并行加载。
}
let default_headers = admin_usage_curl_headers();
let payload = build_admin_usage_detail_payload(
let mut payload = build_admin_usage_detail_payload(
&detail_item,
&users_by_id,
&api_key_names,
@@ -218,6 +285,11 @@ pub(super) async fn maybe_build_local_admin_usage_detail_response(
request_body,
&default_headers,
);
payload["body_load_errors"] = if include_bodies && !body_load_errors.is_empty() {
Value::Object(body_load_errors)
} else {
Value::Null
};
return Ok(Some(attach_admin_audit_response(
Json(payload).into_response(),
@@ -5,13 +5,12 @@ use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::handlers::admin::shared::query_param_value;
use crate::GatewayError;
use aether_admin::observability::usage::{
admin_usage_bad_request_response, admin_usage_client_family,
admin_usage_data_unavailable_response, admin_usage_has_fallback, admin_usage_is_failed,
admin_usage_matches_search, admin_usage_matches_username, admin_usage_parse_ids,
admin_usage_parse_limit, admin_usage_parse_offset, admin_usage_provider_key_name,
admin_usage_record_json, build_admin_usage_active_requests_response,
build_admin_usage_records_response, build_admin_usage_summary_stats_response_from_summary,
ADMIN_USAGE_DATA_UNAVAILABLE_DETAIL,
admin_usage_bad_request_response, admin_usage_data_unavailable_response,
admin_usage_has_fallback, admin_usage_is_failed, admin_usage_matches_search,
admin_usage_matches_username, admin_usage_parse_ids, admin_usage_parse_limit,
admin_usage_parse_offset, admin_usage_provider_key_name, admin_usage_record_json,
build_admin_usage_active_requests_response, build_admin_usage_records_response,
build_admin_usage_summary_stats_response_from_summary, ADMIN_USAGE_DATA_UNAVAILABLE_DETAIL,
};
use aether_data::repository::users::StoredUserSummary;
use aether_data_contracts::repository::{
@@ -195,29 +194,51 @@ async fn resolve_admin_usage_attempt_flags_by_usage_id(
.collect())
}
async fn resolve_admin_usage_image_progress_by_request_id(
#[derive(Default)]
struct AdminUsageActiveCandidateState {
image_progress_by_request_id: BTreeMap<String, serde_json::Value>,
state_overrides_by_request_id: BTreeMap<String, serde_json::Value>,
}
async fn resolve_admin_usage_active_candidate_state(
state: &AdminAppState<'_>,
items: &[StoredRequestUsageAudit],
) -> Result<BTreeMap<String, serde_json::Value>, GatewayError> {
) -> Result<AdminUsageActiveCandidateState, GatewayError> {
if !state.has_request_candidate_data_reader() || items.is_empty() {
return Ok(BTreeMap::new());
return Ok(AdminUsageActiveCandidateState::default());
}
let request_ids = items
.iter()
.map(|item| item.request_id.clone())
.collect::<BTreeSet<_>>();
let mut progress_by_request_id = BTreeMap::new();
let active_usage_by_request_id = items
.iter()
.filter(|item| matches!(item.status.as_str(), "pending" | "streaming"))
.map(|item| (item.request_id.clone(), item))
.collect::<BTreeMap<_, _>>();
let mut candidate_state = AdminUsageActiveCandidateState::default();
for request_id in request_ids {
let candidates = state
.app()
.read_request_candidates_by_request_id(&request_id)
.await?;
if let Some(progress) = latest_admin_usage_image_progress(&candidates) {
progress_by_request_id.insert(request_id, progress);
candidate_state
.image_progress_by_request_id
.insert(request_id.clone(), progress);
}
if active_usage_by_request_id.contains_key(&request_id) {
if let Some(override_payload) =
admin_usage_terminal_candidate_state_override(&candidates)
{
candidate_state
.state_overrides_by_request_id
.insert(request_id, override_payload);
}
}
}
Ok(progress_by_request_id)
Ok(candidate_state)
}
fn latest_admin_usage_image_progress(
@@ -246,6 +267,82 @@ fn latest_admin_usage_image_progress(
.map(|(_, _, _, progress)| progress)
}
fn admin_usage_current_candidate(
candidates: &[StoredRequestCandidate],
) -> Option<&StoredRequestCandidate> {
candidates
.iter()
.filter(|candidate| {
!matches!(
candidate.status,
RequestCandidateStatus::Available
| RequestCandidateStatus::Unused
| RequestCandidateStatus::Skipped
)
})
.max_by_key(|candidate| {
(
candidate.candidate_index,
candidate.retry_index,
candidate
.started_at_unix_ms
.or(candidate.finished_at_unix_ms)
.unwrap_or(candidate.created_at_unix_ms),
)
})
}
fn admin_usage_unix_millis_to_rfc3339(unix_ms: u64) -> Option<String> {
let secs = i64::try_from(unix_ms / 1_000).ok()?;
let nanos = u32::try_from(unix_ms % 1_000)
.ok()?
.saturating_mul(1_000_000);
chrono::DateTime::<chrono::Utc>::from_timestamp(secs, nanos)
.map(|timestamp| timestamp.to_rfc3339())
}
fn admin_usage_terminal_candidate_state_override(
candidates: &[StoredRequestCandidate],
) -> Option<serde_json::Value> {
let candidate = admin_usage_current_candidate(candidates)?;
let status = match candidate.status {
RequestCandidateStatus::Success => "completed",
RequestCandidateStatus::Failed => "failed",
RequestCandidateStatus::Cancelled => "cancelled",
_ => return None,
};
let latency_ms = candidate.latency_ms.or_else(|| {
Some(
candidate
.finished_at_unix_ms?
.saturating_sub(candidate.started_at_unix_ms?),
)
});
let mut payload = json!({ "status": status });
if let Some(latency_ms) = latency_ms {
payload["response_time_ms"] = json!(latency_ms);
if let Some(response_time_updated_at) = candidate
.finished_at_unix_ms
.or_else(|| {
candidate
.started_at_unix_ms
.map(|started_at| started_at.saturating_add(latency_ms))
})
.and_then(admin_usage_unix_millis_to_rfc3339)
{
payload["response_time_updated_at"] = json!(response_time_updated_at);
}
}
if let Some(status_code) = candidate.status_code {
payload["status_code"] = json!(status_code);
}
if let Some(error_message) = candidate.error_message.as_ref() {
payload["error_message"] = json!(error_message);
}
Some(payload)
}
fn admin_usage_matches_attempt_status(
item: &StoredRequestUsageAudit,
status: &str,
@@ -264,19 +361,6 @@ fn admin_usage_matches_attempt_status(
}
}
fn admin_usage_matches_client_family(
item: &StoredRequestUsageAudit,
client_family: Option<&str>,
) -> bool {
let Some(client_family) = client_family
.map(str::trim)
.filter(|value| !value.is_empty())
else {
return true;
};
admin_usage_client_family(item).is_some_and(|value| value.eq_ignore_ascii_case(client_family))
}
fn admin_usage_bool_query_param(query: Option<&str>, name: &str) -> bool {
query_param_value(query, name)
.as_deref()
@@ -290,15 +374,24 @@ fn admin_usage_bool_query_param(query: Option<&str>, name: &str) -> bool {
.unwrap_or(false)
}
fn admin_usage_is_unknown_label(value: &str) -> bool {
matches!(
value.trim().to_ascii_lowercase().as_str(),
"unknown" | "unknow"
)
fn admin_usage_include_total_query_param(query: Option<&str>) -> bool {
query_param_value(query, "include_total")
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|value| {
!(value == "0"
|| value.eq_ignore_ascii_case("false")
|| value.eq_ignore_ascii_case("no")
|| value.eq_ignore_ascii_case("off"))
})
.unwrap_or(true)
}
fn admin_usage_has_unknown_model_or_provider(item: &StoredRequestUsageAudit) -> bool {
admin_usage_is_unknown_label(&item.model) || admin_usage_is_unknown_label(&item.provider_name)
fn admin_usage_fast_page_total(offset: usize, limit: usize, record_count: usize) -> usize {
offset
.saturating_add(record_count)
.saturating_add(usize::from(limit > 0 && record_count == limit))
}
#[allow(clippy::too_many_arguments)]
@@ -314,6 +407,7 @@ fn build_admin_usage_records_response_with_attempt_flags(
total: usize,
limit: usize,
offset: usize,
total_is_estimated: bool,
) -> Response<Body> {
let records: Vec<_> = items
.iter()
@@ -343,6 +437,7 @@ fn build_admin_usage_records_response_with_attempt_flags(
"total": total,
"limit": limit,
"offset": offset,
"total_is_estimated": total_is_estimated,
}))
.into_response()
}
@@ -485,6 +580,8 @@ fn build_admin_usage_keyword_search_query(
provider_name: base_query.provider_name.clone(),
model: base_query.model.clone(),
api_format: base_query.api_format.clone(),
client_family: base_query.client_family.clone(),
exclude_unknown_model_or_provider: base_query.exclude_unknown_model_or_provider,
statuses: base_query.statuses.clone(),
is_stream: base_query.is_stream,
error_only: base_query.error_only,
@@ -582,6 +679,7 @@ pub(super) async fn maybe_build_local_admin_usage_summary_response(
state.has_auth_api_key_data_reader(),
&BTreeMap::new(),
&BTreeMap::new(),
&BTreeMap::new(),
)));
};
state
@@ -605,15 +703,16 @@ pub(super) async fn maybe_build_local_admin_usage_summary_response(
};
let api_key_names = admin_usage_api_key_names(state, &items).await?;
let provider_key_names = admin_usage_provider_key_names(state, &items).await?;
let image_progress_by_request_id =
resolve_admin_usage_image_progress_by_request_id(state, &items).await?;
let active_candidate_state =
resolve_admin_usage_active_candidate_state(state, &items).await?;
return Ok(Some(build_admin_usage_active_requests_response(
&items,
&api_key_names,
state.has_auth_api_key_data_reader(),
&provider_key_names,
&image_progress_by_request_id,
&active_candidate_state.image_progress_by_request_id,
&active_candidate_state.state_overrides_by_request_id,
)));
}
Some("records")
@@ -641,6 +740,8 @@ pub(super) async fn maybe_build_local_admin_usage_summary_response(
let client_family_filter = query_param_value(query, "client_family");
let hide_unknown_records = admin_usage_bool_query_param(query, "hide_unknown")
|| admin_usage_bool_query_param(query, "hide_unknown_records");
let include_total = admin_usage_include_total_query_param(query);
let total_only = admin_usage_bool_query_param(query, "total_only");
let limit = match admin_usage_parse_limit(query) {
Ok(value) => value,
Err(detail) => return Ok(Some(admin_usage_bad_request_response(detail))),
@@ -664,13 +765,6 @@ pub(super) async fn maybe_build_local_admin_usage_summary_response(
offset,
)));
};
let base_query = build_admin_usage_records_query(
created_from_unix_secs,
created_until_unix_secs,
query,
None,
None,
);
let active_search = search.as_deref().filter(|value| !value.trim().is_empty());
let active_username_filter = username_filter
.as_deref()
@@ -678,10 +772,16 @@ pub(super) async fn maybe_build_local_admin_usage_summary_response(
let active_client_family_filter = client_family_filter
.as_deref()
.filter(|value| !value.trim().is_empty());
let (usage, total) = if hide_unknown_records
|| attempt_status_filter.is_some()
|| active_client_family_filter.is_some()
{
let mut base_query = build_admin_usage_records_query(
created_from_unix_secs,
created_until_unix_secs,
query,
None,
None,
);
base_query.client_family = active_client_family_filter.map(str::to_owned);
base_query.exclude_unknown_model_or_provider = hide_unknown_records;
let (usage, total, total_is_estimated) = if attempt_status_filter.is_some() {
let mut usage = state.list_usage_audits(&base_query).await?;
let user_ids: Vec<String> = usage
.iter()
@@ -718,18 +818,20 @@ pub(super) async fn maybe_build_local_admin_usage_summary_response(
&attempt_flags_by_usage_id,
request_candidate_reader_available,
)
}) && admin_usage_matches_client_family(item, active_client_family_filter)
&& (!hide_unknown_records
|| !admin_usage_has_unknown_model_or_provider(item))
})
});
sort_usage_newest_first(&mut usage);
let total = usage.len();
let records = usage
.into_iter()
.skip(offset)
.take(limit)
.collect::<Vec<_>>();
(records, total)
let records = if total_only {
Vec::new()
} else {
usage
.into_iter()
.skip(offset)
.take(limit)
.collect::<Vec<_>>()
};
(records, total, false)
} else if active_search.is_some() || active_username_filter.is_some() {
let keywords = active_search
.map(parse_admin_usage_search_keywords)
@@ -749,30 +851,53 @@ pub(super) async fn maybe_build_local_admin_usage_summary_response(
None,
None,
);
let total = usize::try_from(
state
.count_usage_audits_by_keyword_search(&keyword_query)
.await?,
)
.unwrap_or(usize::MAX);
let paged_query = UsageAuditKeywordSearchQuery {
limit: Some(limit),
offset: Some(offset),
..keyword_query
};
(
state
.list_usage_audits_by_keyword_search(&paged_query)
.await?,
total,
)
} else {
let total = usize::try_from(state.count_usage_audits(&base_query).await?)
if total_only {
let total = usize::try_from(
state
.count_usage_audits_by_keyword_search(&keyword_query)
.await?,
)
.unwrap_or(usize::MAX);
let mut paged_query = base_query.clone();
paged_query.limit = Some(limit);
paged_query.offset = Some(offset);
(state.list_usage_audits(&paged_query).await?, total)
(Vec::new(), total, false)
} else {
let paged_query = UsageAuditKeywordSearchQuery {
limit: Some(limit),
offset: Some(offset),
..keyword_query.clone()
};
let records = state
.list_usage_audits_by_keyword_search(&paged_query)
.await?;
let total = if include_total {
usize::try_from(
state
.count_usage_audits_by_keyword_search(&keyword_query)
.await?,
)
.unwrap_or(usize::MAX)
} else {
admin_usage_fast_page_total(offset, limit, records.len())
};
(records, total, !include_total)
}
} else {
if total_only {
let total = usize::try_from(state.count_usage_audits(&base_query).await?)
.unwrap_or(usize::MAX);
(Vec::new(), total, false)
} else {
let mut paged_query = base_query.clone();
paged_query.limit = Some(limit);
paged_query.offset = Some(offset);
let records = state.list_usage_audits(&paged_query).await?;
let total = if include_total {
usize::try_from(state.count_usage_audits(&base_query).await?)
.unwrap_or(usize::MAX)
} else {
admin_usage_fast_page_total(offset, limit, records.len())
};
(records, total, !include_total)
}
};
let user_ids: Vec<String> = usage
@@ -800,6 +925,7 @@ pub(super) async fn maybe_build_local_admin_usage_summary_response(
total,
limit,
offset,
total_is_estimated,
)));
}
_ => {}
@@ -807,3 +933,88 @@ pub(super) async fn maybe_build_local_admin_usage_summary_response(
Ok(None)
}
#[cfg(test)]
mod tests {
use aether_data_contracts::repository::candidates::{
RequestCandidateStatus, StoredRequestCandidate,
};
use super::admin_usage_terminal_candidate_state_override;
fn sample_candidate(
candidate_index: i32,
status: RequestCandidateStatus,
status_code: Option<i32>,
latency_ms: Option<i32>,
error_message: Option<&str>,
) -> StoredRequestCandidate {
StoredRequestCandidate::new(
format!("candidate-{candidate_index}"),
"req-1".to_string(),
Some("user-1".to_string()),
Some("api-key-1".to_string()),
Some("alice".to_string()),
Some("default".to_string()),
candidate_index,
0,
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("provider-key-1".to_string()),
status,
None,
false,
status_code,
None,
error_message.map(str::to_string),
latency_ms,
None,
None,
None,
1_000,
Some(1_000),
Some(10_210),
)
.expect("candidate should build")
}
#[test]
fn admin_usage_active_override_uses_current_terminal_candidate_latency() {
let candidate = sample_candidate(
0,
RequestCandidateStatus::Success,
Some(200),
Some(9_210),
None,
);
let payload =
admin_usage_terminal_candidate_state_override(&[candidate]).expect("override");
assert_eq!(payload["status"], "completed");
assert_eq!(payload["response_time_ms"], 9_210);
assert_eq!(
payload["response_time_updated_at"],
"1970-01-01T00:00:10.210+00:00"
);
}
#[test]
fn admin_usage_active_override_ignores_terminal_candidate_when_newer_attempt_is_live() {
let failed = sample_candidate(
0,
RequestCandidateStatus::Failed,
Some(503),
Some(1_000),
Some("first attempt failed"),
);
let mut streaming =
sample_candidate(1, RequestCandidateStatus::Streaming, None, None, None);
streaming.started_at_unix_ms = Some(10_500);
streaming.finished_at_unix_ms = None;
let payload = admin_usage_terminal_candidate_state_override(&[failed, streaming]);
assert!(payload.is_none());
}
}
@@ -94,7 +94,7 @@ pub(crate) async fn run_admin_provider_delete_task(
.map(|item| item.id.clone())
.collect::<Vec<_>>();
let key_ids = keys.iter().map(|item| item.id.clone()).collect::<Vec<_>>();
app.cleanup_deleted_provider_catalog_refs(&provider.id, &endpoint_ids, &key_ids)
app.cleanup_deleted_provider_catalog_refs(&provider.id, true, &endpoint_ids, &key_ids)
.await?;
task.stage = "deleting_models".to_string();
@@ -21,7 +21,10 @@ fn oauth_invalid_reason_is_account_level_block(reason: Option<&str>) -> bool {
snapshot.blocked
&& !matches!(
snapshot.code.trim().to_ascii_lowercase().as_str(),
"oauth_token_invalid" | "oauth_expired" | "oauth_refresh_failed"
"oauth_token_invalid"
| "oauth_token_expired"
| "oauth_expired"
| "oauth_refresh_failed"
)
}
@@ -13,6 +13,10 @@ pub(super) fn codex_looks_like_token_invalidated(message: Option<&str>) -> bool
admin_provider_quota_pure::codex_looks_like_token_invalidated(message)
}
pub(super) fn codex_looks_like_token_expired(message: Option<&str>) -> bool {
admin_provider_quota_pure::codex_looks_like_token_expired(message)
}
pub(super) fn codex_looks_like_workspace_deactivated(message: Option<&str>) -> bool {
admin_provider_quota_pure::codex_looks_like_workspace_deactivated(message)
}
@@ -3,7 +3,7 @@ mod parse;
mod plan;
use self::invalid::{
codex_build_invalid_state, codex_looks_like_token_invalidated,
codex_build_invalid_state, codex_looks_like_token_expired, codex_looks_like_token_invalidated,
codex_looks_like_workspace_deactivated, codex_soft_request_failure_reason,
codex_structured_invalid_reason,
};
@@ -275,6 +275,7 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
}
403 => {
let candidate_reason = if codex_looks_like_token_invalidated(err_msg.as_deref())
|| codex_looks_like_token_expired(err_msg.as_deref())
{
codex_structured_invalid_reason(403, err_msg.as_deref())
} else {
@@ -304,6 +304,7 @@ fn admin_pool_trimmed_string(value: Option<&Value>) -> Option<String> {
fn admin_pool_account_code_status_filter(code: &str) -> Option<&'static str> {
match code.trim().to_ascii_lowercase().as_str() {
"oauth_token_invalid" => Some("invalid"),
"oauth_token_expired" => Some("expired"),
"account_banned" | "account_suspended" => Some("account_banned"),
"account_disabled" => Some("account_disabled"),
"workspace_deactivated" => Some("workspace_deactivated"),
@@ -220,11 +220,17 @@ impl<'a> AdminAppState<'a> {
pub(crate) async fn cleanup_deleted_provider_catalog_refs(
&self,
provider_id: &str,
provider_deleted: bool,
endpoint_ids: &[String],
key_ids: &[String],
) -> Result<(), GatewayError> {
self.app
.cleanup_deleted_provider_catalog_refs(provider_id, endpoint_ids, key_ids)
.cleanup_deleted_provider_catalog_refs(
provider_id,
provider_deleted,
endpoint_ids,
key_ids,
)
.await
}
}
@@ -260,7 +260,7 @@ impl<'a> AdminAppState<'a> {
affected += 1;
}
}
self.cleanup_deleted_provider_catalog_refs(&provider.id, &[], &deleted_key_ids)
self.cleanup_deleted_provider_catalog_refs(&provider.id, false, &[], &deleted_key_ids)
.await?;
Ok(affected)
@@ -295,7 +295,7 @@ impl<'a> AdminAppState<'a> {
let deleted = self.delete_provider_catalog_key(&key.id).await?;
if deleted {
let deleted_key_ids = [key.id.clone()];
self.cleanup_deleted_provider_catalog_refs(&provider.id, &[], &deleted_key_ids)
self.cleanup_deleted_provider_catalog_refs(&provider.id, false, &[], &deleted_key_ids)
.await?;
}
Ok(deleted)
@@ -356,7 +356,7 @@ impl<'a> AdminAppState<'a> {
affected = affected.saturating_add(1);
}
}
self.cleanup_deleted_provider_catalog_refs(&provider.id, &[], &deleted_key_ids)
self.cleanup_deleted_provider_catalog_refs(&provider.id, false, &[], &deleted_key_ids)
.await?;
return Ok(Json(
+15 -7
View File
@@ -1797,13 +1797,21 @@ pub(crate) async fn proxy_request(
.all_candidates_skipped_for_reason(AUTH_API_KEY_CONCURRENCY_LIMIT_SKIP_REASON)
|| local_execution_runtime_miss_context
.all_candidates_skipped_for_reason(LEGACY_API_KEY_CONCURRENCY_LIMIT_SKIP_REASON);
let local_execution_runtime_miss_detail = local_execution_runtime_miss_detail(
control_decision,
local_execution_runtime_miss_diagnostic.as_ref(),
auth_api_key_concurrency_limited,
stream_request,
)
.unwrap_or_else(|| "当前 AI 请求无法在本地执行:没有匹配到可用的执行路径".to_string());
let local_execution_runtime_miss_detail = (!auth_api_key_concurrency_limited)
.then(|| {
local_execution_runtime_miss_context
.all_provider_request_body_build_failures_detail()
})
.flatten()
.or_else(|| {
local_execution_runtime_miss_detail(
control_decision,
local_execution_runtime_miss_diagnostic.as_ref(),
auth_api_key_concurrency_limited,
stream_request,
)
})
.unwrap_or_else(|| "当前 AI 请求无法在本地执行:没有匹配到可用的执行路径".to_string());
let local_execution_failure_path = if auth_api_key_concurrency_limited {
EXECUTION_PATH_LOCAL_API_KEY_CONCURRENCY_LIMITED
} else {
@@ -4,11 +4,14 @@ use aether_ai_serving::UPSTREAM_IS_STREAM_KEY;
use aether_billing::{
normalize_input_tokens_for_billing, normalize_total_input_context_for_cache_hit_rate,
};
use aether_data_contracts::repository::usage::{
StoredRequestUsageAudit, StoredUsageBreakdownSummaryRow, StoredUsageDailySummary,
UsageAuditKeywordSearchQuery, UsageAuditListQuery, UsageBreakdownGroupBy,
UsageBreakdownSummaryQuery, UsageCacheAffinityIntervalGroupBy, UsageCacheAffinityIntervalQuery,
UsageDashboardSummaryQuery,
use aether_data_contracts::repository::{
candidates::{RequestCandidateStatus, StoredRequestCandidate},
usage::{
StoredRequestUsageAudit, StoredUsageBreakdownSummaryRow, StoredUsageDailySummary,
UsageAuditKeywordSearchQuery, UsageAuditListQuery, UsageBreakdownGroupBy,
UsageBreakdownSummaryQuery, UsageCacheAffinityIntervalGroupBy,
UsageCacheAffinityIntervalQuery, UsageDashboardSummaryQuery,
},
};
use axum::{
body::Body,
@@ -17,7 +20,7 @@ use axum::{
Json,
};
use chrono::Utc;
use serde_json::json;
use serde_json::{json, Value};
use crate::GatewayError;
@@ -531,6 +534,8 @@ fn build_users_me_usage_active_payload(item: &StoredRequestUsageAudit) -> serde_
"rate_multiplier": item.settlement_rate_multiplier(),
"response_time_ms": item.response_time_ms,
"first_byte_time_ms": item.first_byte_time_ms,
"updated_at": unix_secs_to_rfc3339(item.updated_at_unix_secs),
"response_time_updated_at": users_me_usage_response_time_updated_at(item),
"status_code": item.status_code,
"error_message": item.error_message,
"api_format": item.api_format,
@@ -573,6 +578,118 @@ fn build_users_me_usage_active_payload(item: &StoredRequestUsageAudit) -> serde_
payload
}
fn users_me_usage_response_time_updated_at(item: &StoredRequestUsageAudit) -> Option<String> {
item.response_time_ms?;
if matches!(item.status.as_str(), "pending" | "streaming")
&& item.updated_at_unix_secs <= item.created_at_unix_ms
{
return None;
}
unix_secs_to_rfc3339(item.updated_at_unix_secs)
}
fn unix_millis_to_rfc3339(unix_ms: u64) -> Option<String> {
let secs = i64::try_from(unix_ms / 1_000).ok()?;
let nanos = u32::try_from(unix_ms % 1_000)
.ok()?
.saturating_mul(1_000_000);
chrono::DateTime::<Utc>::from_timestamp(secs, nanos).map(|timestamp| timestamp.to_rfc3339())
}
fn users_me_usage_current_candidate(
candidates: &[StoredRequestCandidate],
) -> Option<&StoredRequestCandidate> {
candidates
.iter()
.filter(|candidate| {
!matches!(
candidate.status,
RequestCandidateStatus::Available
| RequestCandidateStatus::Unused
| RequestCandidateStatus::Skipped
)
})
.max_by_key(|candidate| {
(
candidate.candidate_index,
candidate.retry_index,
candidate
.started_at_unix_ms
.or(candidate.finished_at_unix_ms)
.unwrap_or(candidate.created_at_unix_ms),
)
})
}
fn users_me_usage_terminal_candidate_state_override(
candidates: &[StoredRequestCandidate],
) -> Option<Value> {
let candidate = users_me_usage_current_candidate(candidates)?;
let status = match candidate.status {
RequestCandidateStatus::Success => "completed",
RequestCandidateStatus::Failed => "failed",
RequestCandidateStatus::Cancelled => "cancelled",
_ => return None,
};
let latency_ms = candidate.latency_ms.or_else(|| {
Some(
candidate
.finished_at_unix_ms?
.saturating_sub(candidate.started_at_unix_ms?),
)
});
let mut payload = json!({ "status": status });
if let Some(latency_ms) = latency_ms {
payload["response_time_ms"] = json!(latency_ms);
if let Some(response_time_updated_at) = candidate
.finished_at_unix_ms
.or_else(|| {
candidate
.started_at_unix_ms
.map(|started_at| started_at.saturating_add(latency_ms))
})
.and_then(unix_millis_to_rfc3339)
{
payload["response_time_updated_at"] = json!(response_time_updated_at);
}
}
if let Some(status_code) = candidate.status_code {
payload["status_code"] = json!(status_code);
}
if let Some(error_message) = candidate.error_message.as_ref() {
payload["error_message"] = json!(error_message);
}
Some(payload)
}
async fn resolve_users_me_usage_active_state_overrides_by_request_id(
state: &AppState,
items: &[StoredRequestUsageAudit],
) -> Result<BTreeMap<String, Value>, GatewayError> {
if !state.has_request_candidate_data_reader() || items.is_empty() {
return Ok(BTreeMap::new());
}
let active_request_ids = items
.iter()
.filter(|item| matches!(item.status.as_str(), "pending" | "streaming"))
.map(|item| item.request_id.clone())
.collect::<BTreeSet<_>>();
let mut overrides = BTreeMap::new();
for request_id in active_request_ids {
let candidates = state
.read_request_candidates_by_request_id(&request_id)
.await?;
if let Some(override_payload) =
users_me_usage_terminal_candidate_state_override(&candidates)
{
overrides.insert(request_id, override_payload);
}
}
Ok(overrides)
}
fn users_me_usage_is_failed(item: &StoredRequestUsageAudit) -> bool {
let has_failure_signal = item.status_code.is_some_and(|value| value >= 400)
|| item
@@ -925,6 +1042,8 @@ pub(super) async fn handle_users_me_usage_get(
provider_name: None,
model: None,
api_format: None,
client_family: None,
exclude_unknown_model_or_provider: false,
statuses: None,
is_stream: None,
error_only: false,
@@ -978,6 +1097,8 @@ pub(super) async fn handle_users_me_usage_get(
provider_name: None,
model: None,
api_format: None,
client_family: None,
exclude_unknown_model_or_provider: false,
statuses: None,
is_stream: None,
error_only: false,
@@ -1004,6 +1125,8 @@ pub(super) async fn handle_users_me_usage_get(
provider_name: None,
model: None,
api_format: None,
client_family: None,
exclude_unknown_model_or_provider: false,
statuses: None,
is_stream: None,
error_only: false,
@@ -1140,6 +1263,8 @@ pub(super) async fn handle_users_me_usage_active_get(
provider_name: None,
model: None,
api_format: None,
client_family: None,
exclude_unknown_model_or_provider: false,
statuses: Some(vec!["pending".to_string(), "streaming".to_string()]),
is_stream: None,
error_only: false,
@@ -1168,11 +1293,35 @@ pub(super) async fn handle_users_me_usage_active_get(
.filter(|item| !users_me_usage_is_failed(item))
.collect::<Vec<_>>()
};
let active_state_overrides =
match resolve_users_me_usage_active_state_overrides_by_request_id(state, &items).await {
Ok(value) => value,
Err(err) => {
return build_auth_error_response(
http::StatusCode::INTERNAL_SERVER_ERROR,
format!("user active usage candidate lookup failed: {err:?}"),
false,
);
}
};
Json(json!({
"requests": items
.iter()
.map(build_users_me_usage_active_payload)
.map(|item| {
let mut payload = build_users_me_usage_active_payload(item);
if let (Some(payload), Some(overrides)) = (
payload.as_object_mut(),
active_state_overrides
.get(&item.request_id)
.and_then(Value::as_object),
) {
for (key, value) in overrides {
payload.insert(key.clone(), value.clone());
}
}
payload
})
.collect::<Vec<_>>(),
}))
.into_response()
@@ -1364,13 +1513,16 @@ async fn build_usage_heatmap_summaries(
mod tests {
use std::collections::BTreeMap;
use aether_data_contracts::repository::usage::StoredRequestUsageAudit;
use aether_data_contracts::repository::{
candidates::{RequestCandidateStatus, StoredRequestCandidate},
usage::StoredRequestUsageAudit,
};
use serde_json::json;
use super::{
build_users_me_usage_active_payload, build_users_me_usage_record_payload,
users_me_usage_client_is_stream, users_me_usage_is_failed,
users_me_usage_upstream_is_stream,
users_me_usage_terminal_candidate_state_override, users_me_usage_upstream_is_stream,
};
fn sample_usage(status: &str) -> StoredRequestUsageAudit {
@@ -1415,6 +1567,41 @@ mod tests {
.expect("usage should build")
}
fn sample_candidate(
status: RequestCandidateStatus,
status_code: Option<i32>,
latency_ms: Option<i32>,
error_message: Option<&str>,
) -> StoredRequestCandidate {
StoredRequestCandidate::new(
"candidate-1".to_string(),
"req-1".to_string(),
Some("user-1".to_string()),
Some("api-key-1".to_string()),
Some("alice".to_string()),
Some("default".to_string()),
0,
0,
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("provider-key-1".to_string()),
status,
None,
false,
status_code,
None,
error_message.map(str::to_string),
latency_ms,
None,
None,
None,
1_000,
Some(1_000),
Some(10_210),
)
.expect("candidate should build")
}
#[test]
fn user_usage_record_payload_rehydrates_cache_creation_total_from_classified_fields() {
let item = StoredRequestUsageAudit {
@@ -1447,6 +1634,45 @@ mod tests {
assert_eq!(payload["cache_creation_ephemeral_1h_input_tokens"], 6);
}
#[test]
fn user_usage_active_override_uses_terminal_candidate_latency() {
let candidate = sample_candidate(
RequestCandidateStatus::Success,
Some(200),
Some(9_210),
None,
);
let payload =
users_me_usage_terminal_candidate_state_override(&[candidate]).expect("override");
assert_eq!(payload["status"], "completed");
assert_eq!(payload["response_time_ms"], 9_210);
assert_eq!(payload["status_code"], 200);
assert_eq!(
payload["response_time_updated_at"],
"1970-01-01T00:00:10.210+00:00"
);
}
#[test]
fn user_usage_active_override_ignores_terminal_candidate_when_newer_attempt_is_live() {
let failed = sample_candidate(
RequestCandidateStatus::Failed,
Some(503),
Some(1_000),
Some("first attempt failed"),
);
let mut streaming = sample_candidate(RequestCandidateStatus::Streaming, None, None, None);
streaming.candidate_index = 1;
streaming.started_at_unix_ms = Some(10_500);
streaming.finished_at_unix_ms = None;
let payload = users_me_usage_terminal_candidate_state_override(&[failed, streaming]);
assert!(payload.is_none());
}
#[test]
fn user_usage_payload_keeps_claude_effective_input_when_cache_read_is_large() {
let item = StoredRequestUsageAudit {
@@ -265,14 +265,16 @@ fn build_provider_key_oauth_status_snapshot(key: &StoredProviderCatalogKey) -> V
if let Some(reason) =
tagged_oauth_invalid_reason(invalid_reason.as_deref(), OAUTH_EXPIRED_PREFIX)
{
let (code, label) =
admin_provider_status_pure::oauth_token_snapshot_status_parts(reason.as_str());
return json!({
"code": "invalid",
"label": "已失效",
"code": code,
"label": label,
"reason": reason,
"expires_at": expires_at_unix_secs,
"invalid_at": invalid_at_unix_secs,
"source": "oauth_invalid",
"requires_reauth": true,
"requires_reauth": code == "invalid",
"expiring_soon": false,
});
}
@@ -352,7 +352,8 @@ fn normalize_chat_pii_redaction_feature_settings(
fn normalize_chat_pii_redaction_feature_object(
feature: &mut Map<String, Value>,
) -> Result<(), String> {
for key in ["enabled", "inject_model_instruction"] {
feature.remove("inject_model_instruction");
for key in ["enabled"] {
if let Some(value) = feature.get(key) {
if !value.is_boolean() {
return Err(format!("chat_pii_redaction.{key} 必须是布尔值"));
@@ -450,7 +451,7 @@ mod tests {
fn user_self_feature_update_preserves_notification_push_permission() {
let normalized = normalize_user_self_feature_settings_update(
Some(json!({
"chat_pii_redaction": {"enabled": true, "inject_model_instruction": false},
"chat_pii_redaction": {"enabled": true},
"notification_push_service": {"enabled": false}
})),
Some(json!({
@@ -61,11 +61,8 @@ pub(crate) fn classify_local_failover(
}
if input.status_code >= 400
&& input.response_text.is_some_and(|text| {
policy
.error_stop_patterns
.iter()
.any(|rule| local_failover_regex_rule_matches(rule, text, input.status_code))
&& policy.error_stop_patterns.iter().any(|rule| {
local_failover_regex_rule_matches(rule, input.response_text, input.status_code)
})
{
return LocalFailoverClassification::StopErrorPattern;
@@ -76,7 +73,7 @@ pub(crate) fn classify_local_failover(
policy
.success_failover_patterns
.iter()
.any(|rule| local_failover_regex_rule_matches(rule, text, input.status_code))
.any(|rule| local_failover_regex_rule_matches(rule, Some(text), input.status_code))
})
{
return LocalFailoverClassification::RetrySuccessPattern;
@@ -190,14 +187,23 @@ fn first_non_empty_json_text(
fn local_failover_regex_rule_matches(
rule: &LocalFailoverRegexRule,
response_text: &str,
response_text: Option<&str>,
status_code: u16,
) -> bool {
if !rule.status_codes.is_empty() && !rule.status_codes.contains(&status_code) {
return false;
}
Regex::new(&rule.pattern)
let pattern = rule.pattern.trim();
if pattern.is_empty() {
return !rule.status_codes.is_empty();
}
let Some(response_text) = response_text else {
return false;
};
Regex::new(pattern)
.ok()
.is_some_and(|regex| regex.is_match(response_text))
}
@@ -260,6 +266,50 @@ mod tests {
);
}
#[test]
fn classifier_detects_error_stop_pattern_without_status_codes_on_any_error_status() {
let policy = LocalFailoverPolicy {
error_stop_patterns: vec![LocalFailoverRegexRule {
pattern: "content_policy_violation".to_string(),
status_codes: BTreeSet::new(),
}],
..LocalFailoverPolicy::default()
};
for status_code in [400, 429, 503] {
assert_eq!(
classify_local_failover(
&policy,
LocalFailoverInput::new(
status_code,
Some("{\"error\":\"content_policy_violation\"}")
)
),
LocalFailoverClassification::StopErrorPattern
);
}
}
#[test]
fn classifier_detects_status_only_error_stop_rule_without_response_text() {
let policy = LocalFailoverPolicy {
error_stop_patterns: vec![LocalFailoverRegexRule {
pattern: String::new(),
status_codes: [429].into_iter().collect(),
}],
..LocalFailoverPolicy::default()
};
assert_eq!(
classify_local_failover(&policy, LocalFailoverInput::new(429, None)),
LocalFailoverClassification::StopErrorPattern
);
assert_eq!(
classify_local_failover(&policy, LocalFailoverInput::new(503, None)),
LocalFailoverClassification::RetryUpstreamFailure
);
}
#[test]
fn classifier_detects_success_continue_status_code() {
let policy = LocalFailoverPolicy {
@@ -2031,7 +2031,7 @@ mod tests {
}
#[tokio::test]
async fn oauth_invalidation_marks_codex_key_invalid() {
async fn oauth_invalidation_marks_codex_key_expired() {
let state = codex_state();
let plan = sample_codex_plan();
@@ -2069,7 +2069,7 @@ mod tests {
.and_then(|value| value.get("oauth"))
.and_then(|value| value.get("code"))
.and_then(Value::as_str),
Some("invalid")
Some("expired")
);
}
+19 -30
View File
@@ -25,27 +25,8 @@ pub(crate) struct LocalFailoverRegexRule {
pub(crate) async fn resolve_local_failover_policy(
state: &AppState,
plan: &ExecutionPlan,
report_context: Option<&serde_json::Value>,
_report_context: Option<&serde_json::Value>,
) -> LocalFailoverPolicy {
if let Some(policy) = local_failover_policy_from_report_context(report_context) {
debug!(
event_name = "local_failover_policy_loaded",
log_type = "debug",
request_id = %plan.request_id,
provider_id = %plan.provider_id,
endpoint_id = %plan.endpoint_id,
key_id = %plan.key_id,
source = "report_context",
max_retries = ?policy.max_retries,
stop_status_code_count = policy.stop_status_codes.len(),
continue_status_code_count = policy.continue_status_codes.len(),
success_failover_pattern_count = policy.success_failover_patterns.len(),
error_stop_pattern_count = policy.error_stop_patterns.len(),
"gateway loaded local failover policy from report context"
);
return policy;
}
let transport = match state
.read_provider_transport_snapshot(&plan.provider_id, &plan.endpoint_id, &plan.key_id)
.await
@@ -201,31 +182,39 @@ fn parse_regex_rules(
rules: &serde_json::Map<String, serde_json::Value>,
key: &str,
) -> Vec<LocalFailoverRegexRule> {
let allow_status_only = key == "error_stop_patterns";
rules
.get(key)
.and_then(Value::as_array)
.into_iter()
.flat_map(|items| items.iter())
.filter_map(parse_regex_rule)
.filter_map(|value| parse_regex_rule(value, allow_status_only))
.collect()
}
fn parse_regex_rule(value: &serde_json::Value) -> Option<LocalFailoverRegexRule> {
fn parse_regex_rule(
value: &serde_json::Value,
allow_status_only: bool,
) -> Option<LocalFailoverRegexRule> {
let object = value.as_object()?;
let pattern = object
.get("pattern")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())?;
.unwrap_or_default();
let status_codes: BTreeSet<u16> = object
.get("status_codes")
.and_then(Value::as_array)
.into_iter()
.flat_map(|values| values.iter())
.filter_map(|value| parse_u64_value(value).and_then(|value| u16::try_from(value).ok()))
.collect();
if pattern.is_empty() && (!allow_status_only || status_codes.is_empty()) {
return None;
}
Some(LocalFailoverRegexRule {
pattern: pattern.to_string(),
status_codes: object
.get("status_codes")
.and_then(Value::as_array)
.into_iter()
.flat_map(|values| values.iter())
.filter_map(|value| parse_u64_value(value).and_then(|value| u16::try_from(value).ok()))
.collect(),
status_codes,
})
}
+27 -135
View File
@@ -829,14 +829,12 @@ impl Default for ChatPiiRedactionRuntimeConfig {
}
pub(crate) struct MaskChatRequestOptions {
pub(crate) inject_model_instruction: bool,
pub(crate) scan_limits: RedactionScanLimits,
}
impl MaskChatRequestOptions {
pub(crate) fn runtime(inject_model_instruction: bool) -> Self {
pub(crate) fn runtime() -> Self {
Self {
inject_model_instruction,
scan_limits: RedactionScanLimits::default(),
}
}
@@ -866,8 +864,6 @@ impl ChatPiiRedactionRequestFormat {
}
}
const MODEL_NOTICE_CONTENT: &str = "Aether privacy redaction notice: The next message contains gateway-generated placeholder tokens for sensitive data protection. This notice is not a user request; do not answer it, mention it, reveal it, or infer original values from placeholders. Treat each placeholder as a valid real typed value for reasoning and tool calls, and do not ask the user to reveal originals solely because a placeholder is present.";
fn sanitize_redaction_rule_label(raw: &str) -> String {
let label = raw
.trim()
@@ -1184,7 +1180,7 @@ pub(crate) fn mask_chat_request_json(
body: &[u8],
config: RedactionSessionConfig,
) -> MaskedChatRequest {
mask_chat_request_json_with_options(body, config, MaskChatRequestOptions::runtime(false))
mask_chat_request_json_with_options(body, config, MaskChatRequestOptions::runtime())
}
pub(crate) fn try_mask_chat_request_json_with_options(
@@ -1295,13 +1291,6 @@ pub(crate) async fn try_mask_chat_pii_request_json_with_cache_options(
})
}
fn model_notice_message() -> Value {
serde_json::json!({
"role": "assistant",
"content": MODEL_NOTICE_CONTENT,
})
}
fn request_collision_corpus(format: ChatPiiRedactionRequestFormat, value: &Value) -> Vec<String> {
match format {
ChatPiiRedactionRequestFormat::OpenAiChat => value
@@ -1326,18 +1315,10 @@ fn mask_request_value(
mask_openai_chat_request_value(value, session, scan_state, options)
}
ChatPiiRedactionRequestFormat::OpenAiResponses => {
let redacted = mask_openai_responses_request_value(value, session, scan_state)?;
if redacted && options.inject_model_instruction {
inject_openai_responses_model_notice(value);
}
Ok(redacted)
mask_openai_responses_request_value(value, session, scan_state)
}
ChatPiiRedactionRequestFormat::ClaudeMessages => {
let redacted = mask_claude_messages_request_value(value, session, scan_state)?;
if redacted && options.inject_model_instruction {
inject_claude_model_notice(value);
}
Ok(redacted)
mask_claude_messages_request_value(value, session, scan_state)
}
}
}
@@ -1355,21 +1336,10 @@ async fn mask_request_value_async(
mask_openai_chat_request_value_async(value, session, scan_state, options, cache).await
}
ChatPiiRedactionRequestFormat::OpenAiResponses => {
let redacted =
mask_openai_responses_request_value_async(value, session, scan_state, cache)
.await?;
if redacted && options.inject_model_instruction {
inject_openai_responses_model_notice(value);
}
Ok(redacted)
mask_openai_responses_request_value_async(value, session, scan_state, cache).await
}
ChatPiiRedactionRequestFormat::ClaudeMessages => {
let redacted =
mask_claude_messages_request_value_async(value, session, scan_state, cache).await?;
if redacted && options.inject_model_instruction {
inject_claude_model_notice(value);
}
Ok(redacted)
mask_claude_messages_request_value_async(value, session, scan_state, cache).await
}
}
}
@@ -1385,16 +1355,10 @@ fn mask_openai_chat_request_value(
};
let mut redacted = false;
let mut notice_inserted = false;
let mut index = 0;
while index < messages.len() {
let message_redacted = mask_chat_message_value(&mut messages[index], session, scan_state)?;
redacted |= message_redacted;
if options.inject_model_instruction && message_redacted && !notice_inserted {
messages.insert(index, model_notice_message());
notice_inserted = true;
index += 1;
}
index += 1;
}
Ok(redacted)
@@ -1412,17 +1376,11 @@ async fn mask_openai_chat_request_value_async(
};
let mut redacted = false;
let mut notice_inserted = false;
let mut index = 0;
while index < messages.len() {
let message_redacted =
mask_chat_message_value_async(&mut messages[index], session, scan_state, cache).await?;
redacted |= message_redacted;
if options.inject_model_instruction && message_redacted && !notice_inserted {
messages.insert(index, model_notice_message());
notice_inserted = true;
index += 1;
}
index += 1;
}
Ok(redacted)
@@ -2150,56 +2108,6 @@ async fn mask_json_string_async(
Ok(true)
}
fn inject_openai_responses_model_notice(value: &mut Value) {
let Some(request) = value.as_object_mut() else {
return;
};
match request.get_mut("instructions") {
Some(Value::String(instructions)) => prepend_model_notice(instructions),
Some(_) => {}
None => {
request.insert(
"instructions".to_string(),
Value::String(MODEL_NOTICE_CONTENT.to_string()),
);
}
}
}
fn inject_claude_model_notice(value: &mut Value) {
let Some(request) = value.as_object_mut() else {
return;
};
match request.get_mut("system") {
Some(Value::String(system)) => prepend_model_notice(system),
Some(Value::Array(parts)) => parts.insert(
0,
serde_json::json!({
"type": "text",
"text": MODEL_NOTICE_CONTENT,
}),
),
Some(_) => {}
None => {
request.insert(
"system".to_string(),
Value::String(MODEL_NOTICE_CONTENT.to_string()),
);
}
}
}
fn prepend_model_notice(text: &mut String) {
if text.contains(MODEL_NOTICE_CONTENT) {
return;
}
if text.trim().is_empty() {
*text = MODEL_NOTICE_CONTENT.to_string();
} else {
*text = format!("{MODEL_NOTICE_CONTENT}\n\n{text}");
}
}
pub(crate) struct RestoredSyncResponseBody {
pub(crate) body: Vec<u8>,
pub(crate) restored: bool,
@@ -4408,7 +4316,7 @@ mod tests {
&raw,
ChatPiiRedactionRequestFormat::ClaudeMessages,
test_config(),
MaskChatRequestOptions::runtime(true),
MaskChatRequestOptions::runtime(),
)
.expect("claude messages request should mask");
@@ -4417,11 +4325,7 @@ mod tests {
let masked_json: serde_json::Value =
serde_json::from_slice(&masked.body).expect("masked request should stay valid JSON");
assert_eq!(masked_json["metadata"]["owner"], "[email protected]");
assert!(masked_json["system"][0]["text"]
.as_str()
.expect("notice should remain a string")
.contains("Aether privacy redaction notice"));
assert!(!masked_json["system"][1]["text"]
assert!(!masked_json["system"][0]["text"]
.as_str()
.expect("system text should remain a string")
.contains("[email protected]"));
@@ -4454,7 +4358,7 @@ mod tests {
&raw,
ChatPiiRedactionRequestFormat::OpenAiChat,
test_config(),
MaskChatRequestOptions::runtime(false),
MaskChatRequestOptions::runtime(),
)
.expect("chat request should mask");
@@ -4497,7 +4401,7 @@ mod tests {
&raw,
ChatPiiRedactionRequestFormat::OpenAiResponses,
test_config(),
MaskChatRequestOptions::runtime(true),
MaskChatRequestOptions::runtime(),
)
.expect("responses request should mask");
@@ -4509,7 +4413,6 @@ mod tests {
let instructions = masked_json["instructions"]
.as_str()
.expect("instructions should remain a string");
assert!(instructions.contains("Aether privacy redaction notice"));
assert!(!instructions.contains("[email protected]"));
assert!(!masked_json["input"][0]["content"][0]["text"]
.as_str()
@@ -4995,7 +4898,7 @@ mod tests {
let masked = mask_chat_request_json_with_options(
&serde_json::to_vec(&request).expect("request should serialize"),
build_redaction_session_config(b"redaction-test-key".to_vec(), &config, 600),
MaskChatRequestOptions::runtime(false),
MaskChatRequestOptions::runtime(),
);
let masked_json: serde_json::Value =
@@ -5032,7 +4935,7 @@ mod tests {
let masked = mask_chat_request_json_with_options(
&serde_json::to_vec(&request).expect("request should serialize"),
build_redaction_session_config(b"redaction-test-key".to_vec(), &config, 600),
MaskChatRequestOptions::runtime(true),
MaskChatRequestOptions::runtime(),
);
assert!(!masked.redacted);
@@ -5041,9 +4944,6 @@ mod tests {
assert_eq!(masked_json, request);
assert!(masked_json.to_string().contains("[email protected]"));
assert!(!masked_json.to_string().contains("<AETHER:"));
assert!(!masked_json
.to_string()
.contains("Aether privacy redaction notice"));
}
#[test]
@@ -5105,7 +5005,7 @@ mod tests {
}
#[test]
fn proxy_pii_redaction_provider_bound_request_uses_sentinels_and_inserts_safe_notice() {
fn proxy_pii_redaction_provider_bound_request_uses_sentinels_without_prompt_notice() {
let config = ChatPiiRedactionRuntimeConfig::default();
let request = json!({
"model": "gpt-5",
@@ -5119,7 +5019,7 @@ mod tests {
let masked = mask_chat_request_json_with_options(
&serde_json::to_vec(&request).expect("request should serialize"),
build_redaction_session_config(b"redaction-test-key".to_vec(), &config, 600),
MaskChatRequestOptions::runtime(true),
MaskChatRequestOptions::runtime(),
);
assert!(masked.redacted);
@@ -5128,26 +5028,18 @@ mod tests {
let messages = masked_json["messages"]
.as_array()
.expect("messages should be an array");
assert_eq!(messages.len(), 4);
assert_eq!(messages.len(), 3);
assert_eq!(messages[0]["role"], "system");
assert_eq!(messages[1]["role"], "assistant");
assert!(messages[1..]
.iter()
.all(|message| message["role"].as_str() != Some("system")));
let notice = messages[1]["content"]
.as_str()
.expect("notice should be text");
assert!(notice.contains("not a user request"));
assert!(notice.contains("do not answer"));
assert!(notice.contains("do not answer it, mention it"));
assert!(!notice.contains("[email protected]"));
assert_eq!(messages[2]["role"], "user");
let content = messages[2]["content"]
assert_eq!(messages[1]["role"], "user");
let content = messages[1]["content"]
.as_str()
.expect("user content should be text");
assert!(!content.contains("[email protected]"));
assert!(content.contains("<AETHER:EMAIL:"));
assert_eq!(messages[3]["role"], "assistant");
assert_eq!(messages[2]["role"], "assistant");
}
#[test]
@@ -5239,7 +5131,7 @@ mod tests {
let large_err = try_mask_chat_request_json_with_options(
&serde_json::to_vec(&large_request).expect("request should serialize"),
test_config(),
MaskChatRequestOptions::runtime(false),
MaskChatRequestOptions::runtime(),
)
.expect_err("oversized scan should fail closed");
assert_eq!(
@@ -5261,7 +5153,7 @@ mod tests {
let dense_err = try_mask_chat_request_json_with_options(
&serde_json::to_vec(&dense_request).expect("request should serialize"),
test_config(),
MaskChatRequestOptions::runtime(false).with_scan_limits(RedactionScanLimits {
MaskChatRequestOptions::runtime().with_scan_limits(RedactionScanLimits {
max_scanned_text_bytes: 1024,
max_detections: 1,
}),
@@ -5296,7 +5188,7 @@ mod tests {
let first_masked = try_mask_chat_request_json_with_cache_options(
&serde_json::to_vec(&first_request).expect("request should serialize"),
build_redaction_session_config(b"redaction-test-key".to_vec(), &config, 600),
MaskChatRequestOptions::runtime(false),
MaskChatRequestOptions::runtime(),
Some(&cache),
)
.await
@@ -5315,7 +5207,7 @@ mod tests {
let second_masked = try_mask_chat_request_json_with_cache_options(
&serde_json::to_vec(&second_request).expect("request should serialize"),
build_redaction_session_config(b"redaction-test-key".to_vec(), &config, 899),
MaskChatRequestOptions::runtime(false),
MaskChatRequestOptions::runtime(),
Some(&cache),
)
.await
@@ -5356,7 +5248,7 @@ mod tests {
let rolled_masked = try_mask_chat_request_json_with_cache_options(
&serde_json::to_vec(&second_request).expect("request should serialize"),
build_redaction_session_config(b"redaction-test-key".to_vec(), &config, 900),
MaskChatRequestOptions::runtime(false),
MaskChatRequestOptions::runtime(),
Some(&cache),
)
.await
@@ -5420,7 +5312,7 @@ mod tests {
let first_masked = try_mask_chat_request_json_with_cache_options(
&serde_json::to_vec(&first_request).expect("request should serialize"),
build_redaction_session_config(b"redaction-test-key".to_vec(), &config, 600),
MaskChatRequestOptions::runtime(false),
MaskChatRequestOptions::runtime(),
Some(&cache),
)
.await
@@ -5454,7 +5346,7 @@ mod tests {
let second_masked = try_mask_chat_request_json_with_cache_options(
&serde_json::to_vec(&colliding_request).expect("request should serialize"),
build_redaction_session_config(b"redaction-test-key".to_vec(), &config, 899),
MaskChatRequestOptions::runtime(false),
MaskChatRequestOptions::runtime(),
Some(&cache),
)
.await
@@ -5525,7 +5417,7 @@ mod tests {
let masked = try_mask_chat_request_json_with_cache_options(
&serde_json::to_vec(&request).expect("request should serialize"),
build_redaction_session_config(b"redaction-test-key".to_vec(), &config, 600),
MaskChatRequestOptions::runtime(false),
MaskChatRequestOptions::runtime(),
Some(&cache),
)
.await
@@ -5570,7 +5462,7 @@ mod tests {
let masked = try_mask_chat_request_json_with_cache_options(
&serde_json::to_vec(&request).expect("request should serialize"),
build_redaction_session_config(b"redaction-test-key".to_vec(), &config, 600),
MaskChatRequestOptions::runtime(false),
MaskChatRequestOptions::runtime(),
Some(&cache),
)
.await
@@ -376,9 +376,6 @@ fn oauth_invalid_reason_blocks_scheduling(
now_unix_secs: u64,
) -> bool {
let trimmed_reason = invalid_reason.trim();
if oauth_invalid_reason_has_tag(trimmed_reason, "[OAUTH_EXPIRED]") {
return true;
}
let account_state = admin_provider_status_pure::resolve_pool_account_state(
Some(provider_type),
@@ -433,6 +430,7 @@ fn oauth_account_state_code_is_hard_block(code: &str) -> bool {
| "account_forbidden"
| "account_blocked"
| "account_verification"
| "oauth_token_invalid"
)
}
@@ -2064,7 +2064,7 @@ async fn keeps_codex_candidate_selectable_when_oauth_token_is_expired() {
let mut key = sample_key("key-codex", "provider-codex", Some(10));
key.auth_type = "oauth".to_string();
key.oauth_invalid_at_unix_secs = Some(1_710_000_000);
key.oauth_invalid_reason = Some("Codex Token 无效或已过期".to_string());
key.oauth_invalid_reason = Some("[OAUTH_EXPIRED] session expired".to_string());
key
}],
));
+26 -1
View File
@@ -564,11 +564,17 @@ impl AppState {
pub(crate) async fn cleanup_deleted_provider_catalog_refs(
&self,
provider_id: &str,
provider_deleted: bool,
endpoint_ids: &[String],
key_ids: &[String],
) -> Result<(), GatewayError> {
self.data
.cleanup_deleted_provider_catalog_refs(provider_id, endpoint_ids, key_ids)
.cleanup_deleted_provider_catalog_refs(
provider_id,
provider_deleted,
endpoint_ids,
key_ids,
)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
for key_id in key_ids {
@@ -770,6 +776,25 @@ impl AppState {
tasks.insert(task.task_id.clone(), task);
}
pub(crate) fn reserve_provider_delete_task(
&self,
task: LocalProviderDeleteTaskState,
) -> LocalProviderDeleteTaskState {
let mut tasks = self
.provider_delete_tasks
.lock()
.expect("provider delete tasks cache should lock");
if let Some(existing) = tasks
.values()
.find(|existing| existing.provider_id == task.provider_id && existing.is_active())
.cloned()
{
return existing;
}
tasks.insert(task.task_id.clone(), task.clone());
task
}
pub(crate) fn get_provider_delete_task(
&self,
task_id: &str,
+10 -4
View File
@@ -171,7 +171,10 @@ fn oauth_invalid_reason_is_account_block(reason: Option<&str>) -> bool {
snapshot.blocked
&& !matches!(
snapshot.code.trim().to_ascii_lowercase().as_str(),
"oauth_token_invalid" | "oauth_expired" | "oauth_refresh_failed"
"oauth_token_invalid"
| "oauth_token_expired"
| "oauth_expired"
| "oauth_refresh_failed"
)
}
@@ -356,14 +359,16 @@ fn build_oauth_status_snapshot_value(key: &StoredProviderCatalogKey) -> Value {
let invalid_reason = trimmed_reason(key.oauth_invalid_reason.as_deref());
if let Some(reason) = tagged_reason(invalid_reason.as_deref(), OAUTH_EXPIRED_PREFIX) {
let (code, label) =
aether_admin::provider::status::oauth_token_snapshot_status_parts(reason.as_str());
return json!({
"code": "invalid",
"label": "已失效",
"code": code,
"label": label,
"reason": reason,
"expires_at": expires_at_unix_secs,
"invalid_at": invalid_at_unix_secs,
"source": "oauth_invalid",
"requires_reauth": true,
"requires_reauth": code == "invalid",
"expiring_soon": false,
});
}
@@ -1291,6 +1296,7 @@ impl AppState {
let deleted_key_ids = [key_id.to_string()];
self.cleanup_deleted_provider_catalog_refs(
&transport.provider.id,
false,
&[],
&deleted_key_ids,
)
+6
View File
@@ -11,6 +11,12 @@ pub(crate) struct LocalProviderDeleteTaskState {
pub message: String,
}
impl LocalProviderDeleteTaskState {
pub(crate) fn is_active(&self) -> bool {
matches!(self.status.as_str(), "pending" | "running")
}
}
#[derive(Debug, Clone, PartialEq)]
pub(crate) enum LocalMutationOutcome<T> {
Applied(T),
+30 -12
View File
@@ -46,6 +46,7 @@ pub(crate) const TASK_KEY_STATS_HOURLY_AGG: &str = "maintenance.stats.hourly.agg
pub(crate) const TASK_KEY_USAGE_SYNC_REPORT: &str = "usage.sync.report";
pub(crate) const TASK_KEY_PROVIDER_OAUTH_ACCOUNT_REFRESH: &str = "provider.oauth.account.refresh";
pub(crate) const TASK_KEY_PROVIDER_BALANCE_REFRESH: &str = "provider.ops.balance.refresh";
const PROVIDER_DELETE_LOCK_TTL_SECS: u64 = 60 * 60 * 6;
const RETRY_ONCE: RetryPolicy = RetryPolicy { max_attempts: 1 };
@@ -501,17 +502,23 @@ pub(crate) async fn submit_provider_delete_task(
};
let task_id = Uuid::new_v4().simple().to_string()[..16].to_string();
state.put_provider_delete_task(crate::LocalProviderDeleteTaskState {
task_id: task_id.clone(),
provider_id: provider.id.clone(),
status: "pending".to_string(),
stage: "queued".to_string(),
total_keys: 0,
deleted_keys: 0,
total_endpoints: 0,
deleted_endpoints: 0,
message: "delete task submitted".to_string(),
});
let reserved =
state
.as_ref()
.reserve_provider_delete_task(crate::LocalProviderDeleteTaskState {
task_id: task_id.clone(),
provider_id: provider.id.clone(),
status: "pending".to_string(),
stage: "queued".to_string(),
total_keys: 0,
deleted_keys: 0,
total_endpoints: 0,
deleted_endpoints: 0,
message: "delete task submitted".to_string(),
});
if reserved.task_id != task_id {
return Ok(Some(reserved.task_id));
}
let app = state.cloned_app();
let provider_id = provider.id.clone();
@@ -555,7 +562,7 @@ pub(crate) async fn submit_provider_delete_task(
spawn_named("task-runtime-provider-delete", async move {
let lock_key = format!("task_runtime:lock:{TASK_KEY_PROVIDER_DELETE}:{provider_id}");
let lock_ttl = std::time::Duration::from_secs(60 * 15);
let lock_ttl = std::time::Duration::from_secs(PROVIDER_DELETE_LOCK_TTL_SECS);
let lock = app
.runtime_state
.lock_try_acquire(&lock_key, app.tunnel.local_instance_id(), lock_ttl)
@@ -563,6 +570,17 @@ pub(crate) async fn submit_provider_delete_task(
.ok()
.flatten();
if lock.is_none() {
app.put_provider_delete_task(crate::LocalProviderDeleteTaskState {
task_id: run_id.clone(),
provider_id: provider_id.clone(),
status: "failed".to_string(),
stage: "skipped".to_string(),
total_keys: 0,
deleted_keys: 0,
total_endpoints: 0,
deleted_endpoints: 0,
message: "provider delete skipped: another node is running this task".to_string(),
});
let _ = update_run_status(
&app,
&run_id,
@@ -112,7 +112,6 @@ fn auth_repository_with_redaction_feature_settings() -> Arc<InMemoryAuthApiKeySn
Some(json!({
"chat_pii_redaction": {
"enabled": true,
"inject_model_instruction": true,
}
})),
)]),
@@ -112,7 +112,6 @@ async fn proxy_pii_redaction_local_openai_chat_runtime_masks_headers_and_restore
Some(json!({
"chat_pii_redaction": {
"enabled": true,
"inject_model_instruction": true,
}
})),
)]),
@@ -363,13 +362,7 @@ async fn proxy_pii_redaction_local_openai_chat_runtime_masks_headers_and_restore
let provider_body_text = serde_json::to_string(&seen.body).expect("body should serialize");
assert!(!provider_body_text.contains("[email protected]"));
assert!(provider_body_text.contains("<AETHER:EMAIL:"));
assert_eq!(seen.body["messages"][0]["role"], "assistant");
let notice = seen.body["messages"][0]["content"]
.as_str()
.expect("notice should be text");
assert!(notice.contains("not a user request"));
assert!(notice.contains("do not answer"));
assert_eq!(seen.body["messages"][1]["role"], "user");
assert_eq!(seen.body["messages"][0]["role"], "user");
let stored_candidates = request_candidate_repository
.list_by_request_id("trace-proxy-pii-redaction-sync")
@@ -230,14 +230,10 @@ fn redaction_test_rules() -> serde_json::Value {
])
}
fn chat_pii_redaction_feature_settings(
enabled: bool,
inject_model_instruction: bool,
) -> serde_json::Value {
fn chat_pii_redaction_feature_settings(enabled: bool) -> serde_json::Value {
json!({
"chat_pii_redaction": {
"enabled": enabled,
"inject_model_instruction": inject_model_instruction,
}
})
}
@@ -245,7 +241,6 @@ fn chat_pii_redaction_feature_settings(
fn auth_repository_with_redaction_feature_settings(
test_id: &str,
feature_enabled: bool,
inject_model_instruction: bool,
) -> Arc<InMemoryAuthApiKeySnapshotRepository> {
let snapshot = auth_snapshot(&format!("api-key-{test_id}"), &format!("user-{test_id}"));
let key_hash = hash_api_key(&format!("sk-client-{test_id}"));
@@ -257,10 +252,7 @@ fn auth_repository_with_redaction_feature_settings(
.with_export_records(vec![auth_export_record(
&snapshot,
key_hash,
Some(chat_pii_redaction_feature_settings(
feature_enabled,
inject_model_instruction,
)),
Some(chat_pii_redaction_feature_settings(feature_enabled)),
)]),
)
}
@@ -372,8 +364,7 @@ async fn run_sync_redaction_case_with_system_config(
}),
);
let (provider_url, provider_handle) = start_server(provider_app).await;
let auth_repository =
auth_repository_with_redaction_feature_settings(test_id, feature_enabled, true);
let auth_repository = auth_repository_with_redaction_feature_settings(test_id, feature_enabled);
let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
candidate_row(test_id),
@@ -522,14 +513,9 @@ async fn ai_execute_sync_pii_redaction_round_trip_impl() {
assert!(provider_body_text.contains("<AETHER:ACCESS_TOKEN:"));
assert!(provider_body_text.contains("<AETHER:SECRET_KEY:"));
assert_eq!(seen.body["messages"][0]["role"], "system");
assert_eq!(seen.body["messages"][1]["role"], "assistant");
let notice = seen.body["messages"][1]["content"]
.as_str()
.expect("notice should be text");
assert!(notice.contains("not a user request"));
assert_eq!(seen.body["messages"][2]["role"], "user");
assert_eq!(seen.body["messages"][3]["role"], "assistant");
assert_eq!(seen.body["messages"][4]["role"], "tool");
assert_eq!(seen.body["messages"][1]["role"], "user");
assert_eq!(seen.body["messages"][2]["role"], "assistant");
assert_eq!(seen.body["messages"][3]["role"], "tool");
let response_content = response_json["choices"][0]["message"]["content"]
.as_str()
@@ -702,7 +688,7 @@ async fn ai_execute_pii_redaction_restores_executed_candidate_session_after_late
);
let (provider_url, provider_handle) = start_server(provider_app).await;
let auth_repository =
auth_repository_with_redaction_feature_settings("redaction-candidate-session", true, true);
auth_repository_with_redaction_feature_settings("redaction-candidate-session", true);
let mut later_candidate = candidate_row("redaction-candidate-session");
later_candidate.provider_id = "provider-redaction-candidate-session-later".to_string();
later_candidate.endpoint_id = "endpoint-redaction-candidate-session-later".to_string();
@@ -817,7 +803,7 @@ async fn pii_redaction_performance_limits_do_not_forward_unredacted_body_upstrea
);
let (provider_url, provider_handle) = start_server(provider_app).await;
let auth_repository =
auth_repository_with_redaction_feature_settings("pii-redaction-limit", true, true);
auth_repository_with_redaction_feature_settings("pii-redaction-limit", true);
let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
candidate_row("pii-redaction-limit"),
@@ -893,7 +879,7 @@ async fn ai_execute_pii_redaction_missing_encryption_key_fails_closed_before_pro
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let test_id = "ai-execute-pii-redaction-missing-encryption-key";
let auth_repository = auth_repository_with_redaction_feature_settings(test_id, true, true);
let auth_repository = auth_repository_with_redaction_feature_settings(test_id, true);
let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
candidate_row(test_id),
@@ -498,8 +498,7 @@ fn auth_repository(case: &RedactionFormatCase) -> Arc<InMemoryAuthApiKeySnapshot
key_hash,
Some(json!({
"chat_pii_redaction": {
"enabled": true,
"inject_model_instruction": true
"enabled": true
}
})),
)]),
@@ -117,6 +117,48 @@ fn admin_provider_oauth_complete_dispatch_remains_thin() {
}
}
#[test]
fn postgres_provider_cleanup_preserves_usage_history() {
let postgres_provider_catalog =
read_workspace_file("crates/aether-data/src/repository/provider_catalog/postgres.rs");
for forbidden in [
"UPDATE usage SET provider_id = NULL",
"UPDATE usage SET provider_endpoint_id = NULL",
"UPDATE usage SET provider_api_key_id = NULL",
] {
assert!(
!postgres_provider_catalog.contains(forbidden),
"provider cleanup must not rewrite usage history with {forbidden}"
);
}
}
#[test]
fn provider_cleanup_keeps_common_backends_in_sync() {
for path in [
"crates/aether-data/src/repository/provider_catalog/postgres.rs",
"crates/aether-data/src/repository/provider_catalog/mysql.rs",
"crates/aether-data/src/repository/provider_catalog/sqlite.rs",
] {
let source = read_workspace_file(path);
for required in [
"UPDATE user_preferences SET default_provider_id = NULL WHERE default_provider_id =",
"UPDATE video_tasks SET provider_id = NULL WHERE provider_id =",
"DELETE FROM request_candidates WHERE provider_id =",
"UPDATE video_tasks SET endpoint_id = NULL WHERE endpoint_id =",
"DELETE FROM request_candidates WHERE endpoint_id =",
"DELETE FROM gemini_file_mappings WHERE key_id =",
"UPDATE video_tasks SET key_id = NULL WHERE key_id =",
] {
assert!(
source.contains(required),
"{path} should keep provider cleanup behavior in sync with {required}"
);
}
}
}
#[test]
fn admin_provider_oauth_complete_helpers_are_split() {
let complete_mod = read_workspace_file(
@@ -2506,7 +2506,7 @@ fn ai_serving_same_format_provider_root_request_separates_body_and_url_policy()
"apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request.rs",
);
for pattern in [
"super::super::request::build_same_format_provider_request_body(",
"super::super::request::build_same_format_provider_request_body_with_compatibility_report(",
"super::super::request::build_same_format_upstream_url(",
] {
assert!(
@@ -5403,7 +5403,7 @@ async fn gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_p
assert!(stored_key.oauth_invalid_at_unix_secs.is_some());
assert_eq!(
stored_key.oauth_invalid_reason.as_deref(),
Some("[OAUTH_EXPIRED] Codex Token 无效或已过期 (401)")
Some("[OAUTH_EXPIRED] Codex Token 已过期 (401)")
);
} else if account_state_recheck_attempted
&& payload["account_state_recheck_error"] == "wham/usage API 返回状态码 403"
@@ -5449,7 +5449,40 @@ async fn gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_p
oauth_snapshot.get("expires_at"),
auth_config.get("expires_at")
);
if stored_key.oauth_invalid_reason.is_some() {
if stored_key
.oauth_invalid_reason
.as_deref()
.is_some_and(|reason| reason.starts_with("[OAUTH_EXPIRED]"))
{
assert_eq!(
oauth_snapshot
.get("code")
.and_then(serde_json::Value::as_str),
Some("expired")
);
assert_eq!(
oauth_snapshot
.get("label")
.and_then(serde_json::Value::as_str),
Some("已过期")
);
assert_eq!(
oauth_snapshot.get("reason"),
Some(&json!("Codex Token 已过期 (401)"))
);
assert_eq!(
oauth_snapshot
.get("requires_reauth")
.and_then(serde_json::Value::as_bool),
Some(false)
);
assert_eq!(
oauth_snapshot
.get("expiring_soon")
.and_then(serde_json::Value::as_bool),
Some(false)
);
} else if stored_key.oauth_invalid_reason.is_some() {
assert_eq!(
oauth_snapshot
.get("code")
@@ -5463,8 +5496,11 @@ async fn gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_p
Some("已失效")
);
assert_eq!(
oauth_snapshot.get("reason"),
Some(&json!("Codex Token 无效或已过期 (401)"))
oauth_snapshot
.get("reason")
.and_then(serde_json::Value::as_str)
.is_some_and(|reason| !reason.trim().is_empty()),
true
);
assert_eq!(
oauth_snapshot
@@ -3526,6 +3526,15 @@ async fn gateway_cleans_up_admin_pool_banned_keys_locally_with_trusted_admin_pri
);
banned_key.name = "banned".to_string();
banned_key.oauth_invalid_reason = Some("account_banned".to_string());
let mut oauth_invalidated_key = sample_key(
"key-openai-oauth-invalidated",
"provider-openai",
"openai:chat",
"sk-oauth-invalidated",
);
oauth_invalidated_key.name = "oauth-invalidated".to_string();
oauth_invalidated_key.oauth_invalid_reason =
Some("[OAUTH_EXPIRED] token invalidated".to_string());
let mut oauth_expired_key = sample_key(
"key-openai-oauth-expired",
"provider-openai",
@@ -3533,7 +3542,7 @@ async fn gateway_cleans_up_admin_pool_banned_keys_locally_with_trusted_admin_pri
"sk-oauth-expired",
);
oauth_expired_key.name = "oauth-expired".to_string();
oauth_expired_key.oauth_invalid_reason = Some("[OAUTH_EXPIRED] token invalidated".to_string());
oauth_expired_key.oauth_invalid_reason = Some("[OAUTH_EXPIRED] session expired".to_string());
let mut healthy_key = sample_key(
"key-openai-healthy",
"provider-openai",
@@ -3545,7 +3554,12 @@ async fn gateway_cleans_up_admin_pool_banned_keys_locally_with_trusted_admin_pri
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
Vec::new(),
vec![banned_key, oauth_expired_key, healthy_key],
vec![
banned_key,
oauth_invalidated_key,
oauth_expired_key,
healthy_key,
],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
@@ -3575,8 +3589,8 @@ async fn gateway_cleans_up_admin_pool_banned_keys_locally_with_trusted_admin_pri
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["affected"], 1);
assert_eq!(payload["message"], "已清理 1 个异常账号");
assert_eq!(payload["affected"], 2);
assert_eq!(payload["message"], "已清理 2 个异常账号");
let remaining_keys = provider_catalog_repository
.list_keys_by_provider_ids(&["provider-openai".to_string()])
@@ -3586,6 +3600,9 @@ async fn gateway_cleans_up_admin_pool_banned_keys_locally_with_trusted_admin_pri
assert!(remaining_keys
.iter()
.any(|key| key.id == "key-openai-oauth-expired"));
assert!(!remaining_keys
.iter()
.any(|key| key.id == "key-openai-oauth-invalidated"));
assert!(remaining_keys
.iter()
.any(|key| key.id == "key-openai-healthy"));
@@ -1703,6 +1703,52 @@ async fn gateway_submits_admin_provider_delete_task_locally_with_trusted_admin_p
upstream_handle.abort();
}
#[test]
fn provider_delete_task_reservation_reuses_active_provider_task() {
let state = AppState::new().expect("gateway should build");
let first = crate::LocalProviderDeleteTaskState {
task_id: "task-first".to_string(),
provider_id: "provider-openai".to_string(),
status: "pending".to_string(),
stage: "queued".to_string(),
total_keys: 0,
deleted_keys: 0,
total_endpoints: 0,
deleted_endpoints: 0,
message: "delete task submitted".to_string(),
};
let second = crate::LocalProviderDeleteTaskState {
task_id: "task-second".to_string(),
provider_id: "provider-openai".to_string(),
status: "pending".to_string(),
stage: "queued".to_string(),
total_keys: 0,
deleted_keys: 0,
total_endpoints: 0,
deleted_endpoints: 0,
message: "delete task submitted".to_string(),
};
assert_eq!(
state.reserve_provider_delete_task(first.clone()).task_id,
"task-first"
);
assert_eq!(
state.reserve_provider_delete_task(second.clone()).task_id,
"task-first"
);
state.put_provider_delete_task(crate::LocalProviderDeleteTaskState {
status: "completed".to_string(),
stage: "completed".to_string(),
..first
});
assert_eq!(
state.reserve_provider_delete_task(second).task_id,
"task-second"
);
}
#[tokio::test]
async fn local_admin_provider_delete_task_status_attaches_audit_only_for_terminal_states() {
let mut completed_state = AppState::new().expect("gateway should build");
@@ -1172,7 +1172,7 @@ async fn gateway_handles_admin_usage_active_ids_for_terminal_updates() {
#[tokio::test]
async fn gateway_handles_admin_usage_records_locally_with_trusted_admin_principal() {
let (upstream_url, upstream_hits, upstream_handle) =
let (_upstream_url, upstream_hits, upstream_handle) =
start_usage_upstream("/api/admin/usage/records").await;
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
@@ -1336,6 +1336,100 @@ async fn gateway_filters_admin_usage_records_with_unknown_model_or_provider() {
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_supports_fast_admin_usage_record_totals() {
let (upstream_url, upstream_hits, upstream_handle) =
start_usage_upstream("/api/admin/usage/records").await;
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
sample_usage_row(
"usage-a",
"req-a",
Some("user-1"),
Some("key-1"),
Some("primary"),
"OpenAI",
"gpt-5",
"completed",
120,
30,
0.3,
0.36,
DAY_2_UNIX_SECS,
),
sample_usage_row(
"usage-b",
"req-b",
Some("user-1"),
Some("key-1"),
Some("primary"),
"OpenAI",
"gpt-5-mini",
"completed",
80,
20,
0.2,
0.24,
DAY_2_UNIX_SECS - 1,
),
sample_usage_row(
"usage-c",
"req-c",
Some("user-1"),
Some("key-1"),
Some("primary"),
"Anthropic",
"claude-sonnet",
"completed",
60,
10,
0.1,
0.12,
DAY_1_UNIX_SECS,
),
]));
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_usage_reader_for_tests(
usage_repository,
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let fast_response = admin_request(reqwest::Client::new().get(format!(
"{gateway_url}/api/admin/usage/records?start_date=2024-03-21&end_date=2024-03-22&tz_offset_minutes=0&include_total=false&limit=2&offset=0"
)))
.send()
.await
.expect("request should succeed");
assert_eq!(fast_response.status(), StatusCode::OK);
let fast_payload: serde_json::Value =
fast_response.json().await.expect("json body should parse");
assert_eq!(fast_payload["records"].as_array().unwrap().len(), 2);
assert_eq!(fast_payload["total"], 3);
assert_eq!(fast_payload["total_is_estimated"], true);
let total_response = admin_request(reqwest::Client::new().get(format!(
"{gateway_url}/api/admin/usage/records?start_date=2024-03-21&end_date=2024-03-22&tz_offset_minutes=0&total_only=true&limit=2&offset=0"
)))
.send()
.await
.expect("request should succeed");
assert_eq!(total_response.status(), StatusCode::OK);
let total_payload: serde_json::Value =
total_response.json().await.expect("json body should parse");
assert_eq!(total_payload["records"].as_array().unwrap().len(), 0);
assert_eq!(total_payload["total"], 3);
assert_eq!(total_payload["total_is_estimated"], false);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_usage_records_with_provider_key_name_fallback_from_request_metadata()
{
@@ -4999,8 +4999,7 @@ async fn gateway_updates_users_me_detail_locally_without_proxying_upstream() {
"username": "alice-updated",
"feature_settings": {
"chat_pii_redaction": {
"enabled": true,
"inject_model_instruction": false
"enabled": true
}
}
}))
@@ -5034,10 +5033,6 @@ async fn gateway_updates_users_me_detail_locally_without_proxying_upstream() {
get_payload["feature_settings"]["chat_pii_redaction"]["enabled"],
true
);
assert_eq!(
get_payload["feature_settings"]["chat_pii_redaction"]["inject_model_instruction"],
false
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
@@ -7167,8 +7162,7 @@ async fn gateway_handles_users_me_api_key_writes_locally_without_proxying_upstre
"concurrent_limit": 4,
"feature_settings": {
"chat_pii_redaction": {
"enabled": true,
"inject_model_instruction": false
"enabled": true
}
}
}))
@@ -7187,10 +7181,6 @@ async fn gateway_handles_users_me_api_key_writes_locally_without_proxying_upstre
update_payload["feature_settings"]["chat_pii_redaction"]["enabled"],
true
);
assert_eq!(
update_payload["feature_settings"]["chat_pii_redaction"]["inject_model_instruction"],
false
);
assert_eq!(update_payload["message"], "API密钥已更新");
let toggle_response = client