i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 01:10:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

Improve request trace upstream diagnostics

Browse Source
This commit is contained in:
fawney19
2026-05-09 10:46:04 +08:00
parent 4cf0de681a
commit 79d07ac79a
33 changed files with 2315 additions and 427 deletions
Download Patch File Download Diff File Expand all files Collapse all files

120
crates/aether-admin/src/observability/monitoring.rs
View File

@@ -301,6 +301,9 @@ pub fn build_admin_monitoring_trace_request_payload_response_with_key_accounts(
.collect::<Vec<_>>();
Json(json!({
"request_id": trace.request_id,
"request_path": admin_monitoring_trace_request_path(usage),
"request_query_string": admin_monitoring_trace_request_query_string(usage),
"request_path_and_query": admin_monitoring_trace_request_path_and_query(usage),
"total_candidates": trace.total_candidates,
"final_status": trace.final_status,
"total_latency_ms": trace.total_latency_ms,
@@ -470,6 +473,33 @@ fn build_admin_monitoring_trace_candidate_extra_data(
.entry("first_byte_time_ms".to_string())
.or_insert_with(|| json!(first_byte_time_ms));
}
if let Some(request_path) = admin_monitoring_usage_request_path(usage) {
extra_object
.entry("request_path".to_string())
.or_insert_with(|| json!(request_path));
}
if let Some(request_query_string) = admin_monitoring_usage_request_query_string(usage) {
extra_object
.entry("request_query_string".to_string())
.or_insert_with(|| json!(request_query_string));
}
if let Some(request_path_and_query) = admin_monitoring_usage_request_path_and_query(usage) {
extra_object
.entry("request_path_and_query".to_string())
.or_insert_with(|| json!(request_path_and_query));
}
if admin_monitoring_usage_is_error_node(usage) {
if let Some(response) = admin_monitoring_trace_response_data(
"upstream_response",
usage.status_code,
usage.response_headers.as_ref(),
usage.response_body.as_ref(),
usage.response_body_ref.as_deref(),
usage.response_body_state,
) {
extra_object.insert("upstream_response".to_string(), response);
}
}
if let Some(proxy_value) = extra_object.get_mut("proxy") {
if let Some(proxy_object) = proxy_value.as_object_mut() {
@@ -497,6 +527,96 @@ fn build_admin_monitoring_trace_candidate_extra_data(
}
}
fn admin_monitoring_trace_response_data(
source: &str,
status_code: Option<u16>,
headers: Option<&Value>,
body: Option<&Value>,
body_ref: Option<&str>,
body_state: Option<aether_data_contracts::repository::usage::UsageBodyCaptureState>,
) -> Option<Value> {
if status_code.is_none()
&& headers.is_none()
&& body.is_none()
&& body_ref.is_none()
&& body_state.is_none()
{
return None;
}
Some(json!({
"source": source,
"status_code": status_code,
"headers": headers.cloned().unwrap_or(Value::Null),
"body": body.cloned().unwrap_or(Value::Null),
"body_ref": body_ref,
"body_state": body_state.map(|state| state.as_str()),
}))
}
fn admin_monitoring_usage_is_error_node(usage: &StoredRequestUsageAudit) -> bool {
!usage.status.eq_ignore_ascii_case("completed")
|| usage
.status_code
.is_some_and(|status| !(200..300).contains(&status))
}
fn admin_monitoring_trace_request_path(usage: Option<&StoredRequestUsageAudit>) -> Option<String> {
usage.and_then(admin_monitoring_usage_request_path)
}
fn admin_monitoring_trace_request_query_string(
usage: Option<&StoredRequestUsageAudit>,
) -> Option<String> {
usage.and_then(admin_monitoring_usage_request_query_string)
}
fn admin_monitoring_trace_request_path_and_query(
usage: Option<&StoredRequestUsageAudit>,
) -> Option<String> {
usage.and_then(admin_monitoring_usage_request_path_and_query)
}
fn admin_monitoring_usage_request_path(usage: &StoredRequestUsageAudit) -> Option<String> {
admin_monitoring_usage_metadata_string(usage, "request_path")
}
fn admin_monitoring_usage_request_query_string(usage: &StoredRequestUsageAudit) -> Option<String> {
admin_monitoring_usage_metadata_string(usage, "request_query_string")
.map(|value| value.trim_start_matches('?').to_string())
.filter(|value| !value.is_empty())
}
fn admin_monitoring_usage_request_path_and_query(
usage: &StoredRequestUsageAudit,
) -> Option<String> {
admin_monitoring_usage_metadata_string(usage, "request_path_and_query").or_else(|| {
let path = admin_monitoring_usage_metadata_string(usage, "request_path")?;
let query = admin_monitoring_usage_metadata_string(usage, "request_query_string")
.map(|value| value.trim_start_matches('?').to_string())
.filter(|value| !value.is_empty());
Some(match query {
Some(query) if !path.contains('?') => format!("{path}?{query}"),
_ => path,
})
})
}
fn admin_monitoring_usage_metadata_string(
usage: &StoredRequestUsageAudit,
key: &str,
) -> Option<String> {
usage
.request_metadata
.as_ref()
.and_then(Value::as_object)
.and_then(|metadata| metadata.get(key))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
fn json_string_field(object: &serde_json::Map<String, Value>, key: &str) -> Option<String> {
object
.get(key)

96
crates/aether-admin/src/observability/usage.rs
View File

@@ -1,4 +1,5 @@
use crate::observability::stats::{aggregate_usage_stats, parse_bounded_u32, round_to};
use aether_ai_formats::api::request_path_implies_stream_request;
use aether_billing::{
normalize_input_tokens_for_billing, normalize_total_input_context_for_cache_hit_rate,
};
@@ -227,7 +228,9 @@ pub fn admin_usage_matches_api_format(
}
pub fn admin_usage_is_failed(item: &StoredRequestUsageAudit) -> bool {
let has_failure_signal = item.status_code.is_some_and(|value| value >= 400)
let has_failure_signal = item
.status_code
.is_some_and(|value| !(200..300).contains(&value))
|| item
.error_message
.as_deref()
@@ -256,7 +259,9 @@ pub fn admin_usage_matches_status(item: &StoredRequestUsageAudit, status: Option
"stream" => item.is_stream,
"standard" => !item.is_stream,
"error" => {
item.status_code.is_some_and(|value| value >= 400) || item.error_message.is_some()
item.status_code
.is_some_and(|value| !(200..300).contains(&value))
|| item.error_message.is_some()
}
"pending" | "streaming" | "completed" | "cancelled" => item.status == status,
"failed" => admin_usage_is_failed(item),
@@ -955,12 +960,26 @@ fn admin_usage_infer_upstream_stream_from_captured_bodies(
}
}
fn admin_usage_request_path_implies_client_stream(item: &StoredRequestUsageAudit) -> bool {
let Some(metadata) = item.request_metadata.as_ref().and_then(Value::as_object) else {
return false;
};
["request_path", "request_path_and_query"]
.into_iter()
.filter_map(|field| metadata.get(field).and_then(Value::as_str))
.any(request_path_implies_stream_request)
}
pub fn admin_usage_client_is_stream(item: &StoredRequestUsageAudit) -> bool {
item.request_metadata
.as_ref()
.and_then(Value::as_object)
.and_then(|metadata| metadata.get("client_requested_stream"))
.and_then(Value::as_bool)
admin_usage_request_path_implies_client_stream(item)
.then_some(true)
.or_else(|| {
item.request_metadata
.as_ref()
.and_then(Value::as_object)
.and_then(|metadata| metadata.get("client_requested_stream"))
.and_then(Value::as_bool)
})
.or_else(|| admin_usage_request_body_stream_flag(item))
.or_else(|| admin_usage_headers_stream_flag(item.client_response_headers.as_ref()))
.or_else(|| admin_usage_request_body_implies_default_non_stream(item).then_some(false))
@@ -1665,7 +1684,9 @@ pub fn admin_usage_is_success(item: &StoredRequestUsageAudit) -> bool {
matches!(
item.status.as_str(),
"completed" | "success" | "ok" | "billed" | "settled"
) && item.status_code.is_none_or(|code| code < 400)
) && item
.status_code
.is_none_or(|code| (200..300).contains(&code))
}
pub fn admin_usage_matches_optional_id(value: Option<&str>, expected: Option<&str>) -> bool {
@@ -2258,10 +2279,10 @@ mod tests {
use super::{
admin_usage_active_request_json, admin_usage_client_is_stream, admin_usage_has_body_value,
admin_usage_has_fallback, admin_usage_is_failed, admin_usage_matches_search,
admin_usage_matches_status, admin_usage_matches_username, admin_usage_record_json,
admin_usage_resolve_request_capture_body, admin_usage_upstream_is_stream,
build_admin_usage_detail_payload,
admin_usage_has_fallback, admin_usage_is_failed, admin_usage_is_success,
admin_usage_matches_search, admin_usage_matches_status, admin_usage_matches_username,
admin_usage_record_json, admin_usage_resolve_request_capture_body,
admin_usage_upstream_is_stream, build_admin_usage_detail_payload,
};
use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UsageBodyField};
@@ -2349,6 +2370,33 @@ mod tests {
assert_eq!(record["client_is_stream"], false);
}
#[test]
fn client_requested_stream_uses_stream_generate_content_path_over_stale_metadata_flag() {
let item = StoredRequestUsageAudit {
is_stream: true,
request_metadata: Some(json!({
"client_requested_stream": false,
"request_path": "/v1beta/models/gemini-3.1-flash-image-preview:streamGenerateContent",
"request_path_and_query": "/v1beta/models/gemini-3.1-flash-image-preview:streamGenerateContent?alt=sse"
})),
..sample_usage("completed", Some(200), None)
};
assert!(admin_usage_client_is_stream(&item));
let record = admin_usage_record_json(
&item,
&BTreeMap::new(),
&BTreeMap::new(),
false,
false,
None,
);
assert_eq!(record["upstream_is_stream"], true);
assert_eq!(record["client_requested_stream"], true);
assert_eq!(record["client_is_stream"], true);
}
#[test]
fn client_requested_stream_falls_back_to_request_body_stream_flag() {
let item = StoredRequestUsageAudit {
@@ -2528,6 +2576,30 @@ mod tests {
assert!(admin_usage_matches_status(&item, Some("failed")));
}
#[test]
fn redirect_status_is_not_admin_usage_success() {
let item = sample_usage("completed", Some(302), None);
assert!(!admin_usage_is_success(&item));
assert!(!admin_usage_is_failed(&item));
assert!(admin_usage_matches_status(&item, Some("error")));
assert!(admin_usage_matches_status(&item, Some("completed")));
}
#[test]
fn failed_redirect_status_counts_as_admin_usage_failed() {
let item = sample_usage(
"failed",
Some(302),
Some("execution runtime stream returned non-success status 302"),
);
assert!(admin_usage_is_failed(&item));
assert!(admin_usage_matches_status(&item, Some("failed")));
assert!(admin_usage_matches_status(&item, Some("error")));
assert!(!admin_usage_is_success(&item));
}
#[test]
fn active_status_with_failure_signal_counts_as_failed() {
let item = sample_usage("pending", Some(503), Some("upstream failed"));

1
crates/aether-ai-formats/Cargo.toml
View File

@@ -15,4 +15,5 @@ serde.workspace = true
serde_json.workspace = true
sha1 = "0.10"
sha2.workspace = true
url.workspace = true
uuid.workspace = true

4
crates/aether-ai-formats/src/api.rs
View File

@@ -96,7 +96,9 @@ pub use crate::formats::shared::response::{
};
pub use crate::formats::shared::routing::{
is_matching_stream_http_request, is_matching_stream_request,
resolve_execution_runtime_stream_plan_kind, resolve_execution_runtime_sync_plan_kind,
request_path_implies_stream_request, resolve_execution_runtime_stream_plan_kind,
resolve_execution_runtime_sync_plan_kind, sanitize_request_path,
sanitize_request_path_and_query, sanitize_request_query_string,
supports_stream_execution_decision_kind, supports_sync_execution_decision_kind,
};
pub use crate::formats::shared::sse::{encode_done_sse, encode_json_sse, map_claude_stop_reason};

103
crates/aether-ai-formats/src/formats/shared/routing.rs
View File

@@ -1,4 +1,5 @@
use http::Method;
use url::form_urlencoded;
use crate::contracts::{
CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND,
@@ -304,6 +305,67 @@ fn resolve_gemini_generate_content_plan_kind(
}
}
pub fn request_path_implies_stream_request(path: &str) -> bool {
let trimmed = path.trim();
let path = trimmed
.split_once('?')
.map(|(path, _)| path)
.unwrap_or(trimmed);
path.ends_with(":streamGenerateContent")
}
pub fn sanitize_request_path(path: &str) -> Option<String> {
let path = path
.trim()
.split_once('?')
.map(|(path, _)| path)
.unwrap_or_else(|| path.trim())
.trim();
(!path.is_empty()).then(|| path.to_string())
}
pub fn sanitize_request_query_string(query: &str) -> Option<String> {
let query = query.trim().trim_start_matches('?').trim();
if query.is_empty() {
return None;
}
let mut serializer = form_urlencoded::Serializer::new(String::new());
for (key, value) in form_urlencoded::parse(query.as_bytes()) {
if request_query_key_is_safe_to_trace(key.as_ref()) {
serializer.append_pair(key.as_ref(), value.as_ref());
}
}
let sanitized = serializer.finish();
(!sanitized.is_empty()).then_some(sanitized)
}
pub fn sanitize_request_path_and_query(path: &str, query: Option<&str>) -> Option<String> {
let trimmed = path.trim();
let (path, embedded_query) = trimmed
.split_once('?')
.map(|(path, query)| (path.trim(), Some(query)))
.unwrap_or((trimmed, None));
if path.is_empty() {
return None;
}
let sanitized_query = query
.and_then(sanitize_request_query_string)
.or_else(|| embedded_query.and_then(sanitize_request_query_string));
Some(match sanitized_query {
Some(query) => format!("{path}?{query}"),
None => path.to_string(),
})
}
fn request_query_key_is_safe_to_trace(key: &str) -> bool {
matches!(
key.to_ascii_lowercase().as_str(),
"alt" | "view" | "pagesize" | "page_size" | "limit" | "offset"
)
}
pub fn is_matching_stream_request(
plan_kind: &str,
path: &str,
@@ -320,7 +382,7 @@ pub fn is_matching_stream_request(
.and_then(|value| value.as_bool())
.unwrap_or(false),
GEMINI_CHAT_STREAM_PLAN_KIND | GEMINI_CLI_STREAM_PLAN_KIND => {
path.ends_with(":streamGenerateContent")
request_path_implies_stream_request(path)
}
_ => true,
}
@@ -388,7 +450,9 @@ mod tests {
use super::{
is_matching_stream_http_request, is_matching_stream_request,
resolve_execution_runtime_stream_plan_kind, resolve_execution_runtime_sync_plan_kind,
request_path_implies_stream_request, resolve_execution_runtime_stream_plan_kind,
resolve_execution_runtime_sync_plan_kind, sanitize_request_path,
sanitize_request_path_and_query, sanitize_request_query_string,
supports_stream_execution_decision_kind, supports_sync_execution_decision_kind,
};
use crate::contracts::{
@@ -609,6 +673,41 @@ mod tests {
);
}
#[test]
fn stream_path_detection_handles_gemini_method_paths_with_query() {
assert!(request_path_implies_stream_request(
"/v1beta/models/gemini-2.5-pro:streamGenerateContent?alt=sse"
));
assert!(request_path_implies_stream_request(
" /v1internal:streamGenerateContent?alt=sse "
));
assert!(!request_path_implies_stream_request(
"/v1beta/models/gemini-2.5-pro:generateContent?alt=sse"
));
}
#[test]
fn request_path_metadata_sanitizer_drops_sensitive_query_parameters() {
assert_eq!(
sanitize_request_path("/v1beta/models/gemini-2.5-pro:generateContent?key=secret")
.as_deref(),
Some("/v1beta/models/gemini-2.5-pro:generateContent")
);
assert_eq!(
sanitize_request_query_string("?key=secret&alt=sse&pageSize=10&token=hidden")
.as_deref(),
Some("alt=sse&pageSize=10")
);
assert_eq!(
sanitize_request_path_and_query(
"/v1beta/models/gemini-2.5-pro:streamGenerateContent?key=secret&alt=sse",
None
)
.as_deref(),
Some("/v1beta/models/gemini-2.5-pro:streamGenerateContent?alt=sse")
);
}
#[test]
fn stream_matching_requires_openai_stream_flag() {
assert!(!is_matching_stream_request(

90
crates/aether-provider-transport/src/request_url/mod.rs
View File

@@ -12,7 +12,7 @@ use crate::claude_code::build_claude_code_messages_url;
use crate::snapshot::GatewayProviderTransportSnapshot;
use crate::url::{
build_claude_messages_url, build_gemini_content_url, build_openai_chat_url,
build_openai_responses_url, build_passthrough_path_url,
build_openai_responses_url, build_passthrough_path_url, normalize_gemini_content_action_path,
};
use crate::vertex::{
build_vertex_api_key_gemini_content_url, resolve_local_vertex_api_key_query_auth,
@@ -36,6 +36,8 @@ pub fn build_transport_request_url(
}
let provider_api_format = params.provider_api_format.trim().to_ascii_lowercase();
let normalized_provider_api_format =
aether_ai_formats::normalize_api_format_alias(&provider_api_format);
let custom_path = transport
.endpoint
.custom_path
@@ -45,14 +47,19 @@ pub fn build_transport_request_url(
.map(|path| expand_custom_path_template(path, build_path_params(params)));
if let Some(path) = custom_path.as_deref() {
let blocked_keys = if provider_api_format.starts_with("gemini:") {
let blocked_keys = if normalized_provider_api_format.starts_with("gemini:") {
&["key"][..]
} else {
&[][..]
};
let normalized_path = if normalized_provider_api_format == "gemini:generate_content" {
normalize_gemini_content_action_path(path, params.upstream_is_stream)
} else {
path.to_string()
};
let url = build_passthrough_path_url(
&transport.endpoint.base_url,
path,
normalized_path.as_str(),
params.request_query,
blocked_keys,
)?;
@@ -63,7 +70,7 @@ pub fn build_transport_request_url(
));
}
let url = match aether_ai_formats::normalize_api_format_alias(&provider_api_format).as_str() {
let url = match normalized_provider_api_format.as_str() {
"openai:chat" => Some(build_openai_chat_url(
&transport.endpoint.base_url,
params.request_query,
@@ -570,6 +577,81 @@ mod tests {
);
}
#[test]
fn rewrites_hardcoded_gemini_custom_path_action_to_match_stream_mode() {
let stream_transport = sample_transport(
"custom",
"gemini:generate_content",
"https://generativelanguage.googleapis.com",
Some("/v1beta/models/{model}:generateContent"),
);
let stream_url = build_transport_request_url(
&stream_transport,
TransportRequestUrlParams {
provider_api_format: "gemini:generate_content",
mapped_model: Some("gemini-2.5-pro"),
upstream_is_stream: true,
request_query: Some("key=client-key&foo=bar"),
kiro_api_region: None,
},
)
.expect("stream custom path url");
assert_eq!(
stream_url,
"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-pro:streamGenerateContent?foo=bar&alt=sse"
);
let sync_transport = sample_transport(
"custom",
"gemini:generate_content",
"https://generativelanguage.googleapis.com",
Some("/v1beta/models/{model}:streamGenerateContent"),
);
let sync_url = build_transport_request_url(
&sync_transport,
TransportRequestUrlParams {
provider_api_format: "gemini:generate_content",
mapped_model: Some("gemini-2.5-pro"),
upstream_is_stream: false,
request_query: Some("foo=bar"),
kiro_api_region: None,
},
)
.expect("sync custom path url");
assert_eq!(
sync_url,
"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-pro:generateContent?foo=bar"
);
let v1_transport = sample_transport(
"custom",
"gemini:generate_content",
"https://generativelanguage.googleapis.com",
Some("/v1/models/{model}:generateContent"),
);
let v1_stream_url = build_transport_request_url(
&v1_transport,
TransportRequestUrlParams {
provider_api_format: "gemini:generate_content",
mapped_model: Some("gemini-2.5-pro"),
upstream_is_stream: true,
request_query: None,
kiro_api_region: None,
},
)
.expect("v1 stream custom path url");
assert_eq!(
v1_stream_url,
"https://generativelanguage.googleapis.com/v1/models/gemini-2.5-pro:streamGenerateContent?alt=sse"
);
}
#[test]
fn keeps_original_custom_path_when_template_params_are_missing() {
let transport = sample_transport(

99
crates/aether-provider-transport/src/url.rs
View File

@@ -69,9 +69,10 @@ pub fn build_gemini_content_url(
} else {
"generateContent"
};
let mut url = if trimmed_base_url.ends_with("/v1beta") {
let mut url = if trimmed_base_url.ends_with("/v1") || trimmed_base_url.ends_with("/v1beta") {
format!("{trimmed_base_url}/models/{trimmed_model}:{operation}")
} else if trimmed_base_url.contains("/v1beta/models/") {
} else if gemini_content_base_url_contains_model_path(trimmed_base_url) {
let trimmed_base_url = strip_gemini_content_action(trimmed_base_url);
format!("{trimmed_base_url}:{operation}")
} else {
format!("{trimmed_base_url}/v1beta/models/{trimmed_model}:{operation}")
@@ -80,6 +81,37 @@ pub fn build_gemini_content_url(
Some(url)
}
pub fn normalize_gemini_content_action_path(path: &str, stream: bool) -> String {
let trimmed = path.trim();
let (path, query) = split_path_query(trimmed);
let action = if stream {
"streamGenerateContent"
} else {
"generateContent"
};
let normalized = strip_gemini_content_action(path);
let normalized = if normalized.len() == path.len() {
path.to_string()
} else {
format!("{normalized}:{action}")
};
match query {
Some(query) => format!("{normalized}?{query}"),
None => normalized,
}
}
fn strip_gemini_content_action(value: &str) -> &str {
value
.strip_suffix(":streamGenerateContent")
.or_else(|| value.strip_suffix(":generateContent"))
.unwrap_or(value)
}
fn gemini_content_base_url_contains_model_path(value: &str) -> bool {
value.contains("/v1/models/") || value.contains("/v1beta/models/")
}
pub fn build_gemini_video_predict_long_running_url(
upstream_base_url: &str,
model: &str,
@@ -92,9 +124,9 @@ pub fn build_gemini_video_predict_long_running_url(
return None;
}
let mut url = if trimmed_base_url.ends_with("/v1beta") {
let mut url = if trimmed_base_url.ends_with("/v1") || trimmed_base_url.ends_with("/v1beta") {
format!("{trimmed_base_url}/models/{trimmed_model}:predictLongRunning")
} else if trimmed_base_url.contains("/v1beta/models/") {
} else if gemini_content_base_url_contains_model_path(trimmed_base_url) {
format!("{trimmed_base_url}:predictLongRunning")
} else {
format!("{trimmed_base_url}/v1beta/models/{trimmed_model}:predictLongRunning")
@@ -248,6 +280,7 @@ mod tests {
build_gemini_content_url, build_gemini_files_passthrough_url,
build_gemini_video_predict_long_running_url, build_openai_chat_url,
build_openai_responses_url, build_passthrough_path_url,
normalize_gemini_content_action_path,
};
#[test]
@@ -289,6 +322,64 @@ mod tests {
);
}
#[test]
fn gemini_content_urls_rewrite_existing_base_action_for_stream_mode() {
assert_eq!(
build_gemini_content_url(
"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-pro:generateContent",
"ignored-model",
true,
Some("foo=bar")
)
.as_deref(),
Some(
"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-pro:streamGenerateContent?foo=bar"
)
);
assert_eq!(
build_gemini_content_url(
"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-pro:streamGenerateContent",
"ignored-model",
false,
Some("foo=bar")
)
.as_deref(),
Some(
"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-pro:generateContent?foo=bar"
)
);
assert_eq!(
build_gemini_content_url(
"https://generativelanguage.googleapis.com/v1/models/gemini-2.5-pro:generateContent",
"ignored-model",
true,
Some("foo=bar")
)
.as_deref(),
Some(
"https://generativelanguage.googleapis.com/v1/models/gemini-2.5-pro:streamGenerateContent?foo=bar"
)
);
}
#[test]
fn normalizes_gemini_content_action_in_custom_paths() {
assert_eq!(
normalize_gemini_content_action_path(
"/v1beta/models/gemini-2.5-pro:generateContent?alt=sse",
true
),
"/v1beta/models/gemini-2.5-pro:streamGenerateContent?alt=sse"
);
assert_eq!(
normalize_gemini_content_action_path(
"/v1beta/models/gemini-2.5-pro:streamGenerateContent",
false
),
"/v1beta/models/gemini-2.5-pro:generateContent"
);
}
#[test]
fn merges_base_path_and_request_query_for_passthrough_paths() {
assert_eq!(

81
crates/aether-scheduler-core/src/request_candidate.rs
View File

@@ -17,11 +17,15 @@ pub struct SchedulerRequestCandidateReportContext {
pub key_id: Option<String>,
pub client_api_format: Option<String>,
pub provider_api_format: Option<String>,
pub request_path: Option<String>,
pub request_query_string: Option<String>,
pub request_path_and_query: Option<String>,
pub upstream_url: Option<String>,
pub mapped_model: Option<String>,
pub key_name: Option<String>,
pub header_rules: Option<Value>,
pub body_rules: Option<Value>,
pub upstream_response: Option<Value>,
pub proxy: Option<Value>,
pub error_flow: Option<Value>,
pub candidate_group_id: Option<String>,
@@ -60,11 +64,15 @@ pub struct SchedulerExecutionRequestCandidateSeed {
struct ReportCandidateExtraDataInput {
client_api_format: Option<String>,
provider_api_format: Option<String>,
request_path: Option<String>,
request_query_string: Option<String>,
request_path_and_query: Option<String>,
upstream_url: Option<String>,
mapped_model: Option<String>,
key_name: Option<String>,
header_rules: Option<Value>,
body_rules: Option<Value>,
upstream_response: Option<Value>,
proxy: Option<Value>,
error_flow: Option<Value>,
candidate_group_id: Option<String>,
@@ -138,6 +146,9 @@ pub fn parse_request_candidate_report_context(
key_id: string_field(report_context, "key_id"),
client_api_format: string_field(report_context, "client_api_format"),
provider_api_format: string_field(report_context, "provider_api_format"),
request_path: string_field(report_context, "request_path"),
request_query_string: string_field(report_context, "request_query_string"),
request_path_and_query: string_field(report_context, "request_path_and_query"),
upstream_url: string_field(report_context, "upstream_url"),
mapped_model: string_field(report_context, "mapped_model"),
key_name: string_field(report_context, "key_name"),
@@ -149,6 +160,10 @@ pub fn parse_request_candidate_report_context(
.get("body_rules")
.cloned()
.filter(|value| !value.is_null()),
upstream_response: report_context
.get("upstream_response")
.cloned()
.filter(|value| !value.is_null()),
proxy: report_context
.get("proxy")
.cloned()
@@ -187,11 +202,15 @@ pub fn resolve_report_request_candidate_slot(
key_id,
client_api_format,
provider_api_format,
request_path,
request_query_string,
request_path_and_query,
upstream_url,
mapped_model,
key_name,
header_rules,
body_rules,
upstream_response,
proxy,
error_flow,
candidate_group_id,
@@ -207,11 +226,15 @@ pub fn resolve_report_request_candidate_slot(
let synthesized_extra_data = build_report_candidate_extra_data(ReportCandidateExtraDataInput {
client_api_format,
provider_api_format,
request_path,
request_query_string,
request_path_and_query,
upstream_url,
mapped_model,
key_name,
header_rules,
body_rules,
upstream_response,
proxy,
error_flow,
candidate_group_id,
@@ -326,11 +349,15 @@ pub fn build_execution_request_candidate_seed(
build_report_candidate_extra_data(ReportCandidateExtraDataInput {
client_api_format: metadata.client_api_format,
provider_api_format: metadata.provider_api_format,
request_path: metadata.request_path,
request_query_string: metadata.request_query_string,
request_path_and_query: metadata.request_path_and_query,
upstream_url: metadata.upstream_url,
mapped_model: metadata.mapped_model,
key_name: metadata.key_name,
header_rules: metadata.header_rules,
body_rules: metadata.body_rules,
upstream_response: metadata.upstream_response,
proxy: metadata.proxy,
error_flow: metadata.error_flow,
candidate_group_id: metadata.candidate_group_id,
@@ -431,11 +458,15 @@ pub fn build_local_request_candidate_status_record(
let extra_data = build_report_candidate_extra_data(ReportCandidateExtraDataInput {
client_api_format: metadata.client_api_format.clone(),
provider_api_format: metadata.provider_api_format.clone(),
request_path: metadata.request_path.clone(),
request_query_string: metadata.request_query_string.clone(),
request_path_and_query: metadata.request_path_and_query.clone(),
upstream_url: metadata.upstream_url.clone(),
mapped_model: metadata.mapped_model.clone(),
key_name: metadata.key_name.clone(),
header_rules: metadata.header_rules.clone(),
body_rules: metadata.body_rules.clone(),
upstream_response: metadata.upstream_response.clone(),
proxy: metadata.proxy.clone(),
error_flow: metadata.error_flow.clone(),
candidate_group_id: metadata.candidate_group_id.clone(),
@@ -667,11 +698,15 @@ fn build_report_candidate_extra_data(input: ReportCandidateExtraDataInput) -> Op
let ReportCandidateExtraDataInput {
client_api_format,
provider_api_format,
request_path,
request_query_string,
request_path_and_query,
upstream_url,
mapped_model,
key_name,
header_rules,
body_rules,
upstream_response,
proxy,
error_flow,
candidate_group_id,
@@ -698,6 +733,21 @@ fn build_report_candidate_extra_data(input: ReportCandidateExtraDataInput) -> Op
Value::String(provider_api_format),
);
}
if let Some(request_path) = request_path {
extra_data.insert("request_path".to_string(), Value::String(request_path));
}
if let Some(request_query_string) = request_query_string {
extra_data.insert(
"request_query_string".to_string(),
Value::String(request_query_string),
);
}
if let Some(request_path_and_query) = request_path_and_query {
extra_data.insert(
"request_path_and_query".to_string(),
Value::String(request_path_and_query),
);
}
if let Some(upstream_url) = upstream_url {
extra_data.insert("upstream_url".to_string(), Value::String(upstream_url));
}
@@ -713,6 +763,9 @@ fn build_report_candidate_extra_data(input: ReportCandidateExtraDataInput) -> Op
if let Some(body_rules) = body_rules {
extra_data.insert("body_rules".to_string(), body_rules);
}
if let Some(upstream_response) = upstream_response {
extra_data.insert("upstream_response".to_string(), upstream_response);
}
if let Some(proxy) = proxy {
extra_data.insert("proxy".to_string(), proxy);
}
@@ -903,6 +956,11 @@ mod tests {
"body_rules": [
{"op": "remove", "path": "/store"}
],
"upstream_response": {
"status_code": 503,
"headers": {"retry-after": "2"},
"body": {"error": {"message": "overloaded"}}
},
"proxy": {
"node_id": "proxy-node-1",
"node_name": "edge-1",
@@ -962,6 +1020,13 @@ mod tests {
.map(Vec::len),
Some(1)
);
assert_eq!(
slot.extra_data
.as_ref()
.and_then(|value| value.get("upstream_response"))
.and_then(|value| value.get("status_code")),
Some(&json!(503))
);
assert_eq!(
slot.extra_data
.as_ref()
@@ -1063,6 +1128,8 @@ mod tests {
"api_key_id": "api-key-1",
"client_api_format": "openai:chat",
"provider_api_format": "openai:responses",
"request_path": "/v1/responses",
"request_query_string": "debug=true",
"upstream_url": "https://example.com/v1/responses",
"mapped_model": "gpt-5-upstream",
"key_name": "primary",
@@ -1098,6 +1165,20 @@ mod tests {
.and_then(|value| value.get("provider_api_format")),
Some(&json!("openai:responses"))
);
assert_eq!(
record
.extra_data
.as_ref()
.and_then(|value| value.get("request_path")),
Some(&json!("/v1/responses"))
);
assert_eq!(
record
.extra_data
.as_ref()
.and_then(|value| value.get("request_query_string")),
Some(&json!("debug=true"))
);
assert_eq!(
record
.extra_data

63
crates/aether-usage-runtime/src/request_metadata.rs
View File

@@ -1,3 +1,6 @@
use aether_ai_formats::api::{
sanitize_request_path, sanitize_request_path_and_query, sanitize_request_query_string,
};
use aether_contracts::ExecutionPlan;
use serde_json::{json, Map, Value};
@@ -72,9 +75,13 @@ fn copy_allowed_metadata_fields(source: &Map<String, Value>, target: &mut Map<St
copy_bool(source, target, "client_requested_stream");
copy_bool(source, target, "upstream_is_stream");
copy_bool(source, target, "api_key_is_standalone");
copy_non_empty_string(source, target, "request_path");
copy_non_empty_string(source, target, "request_query_string");
copy_non_empty_string(source, target, "request_path_and_query");
copy_number(source, target, "provider_request_body_base64_bytes");
copy_number(source, target, "provider_response_body_base64_bytes");
copy_number(source, target, "client_response_body_base64_bytes");
copy_number(source, target, "client_response_status_code");
copy_non_null_value(source, target, "billing_snapshot");
copy_non_empty_string(source, target, "billing_snapshot_schema_version");
copy_non_empty_string(source, target, "billing_snapshot_status");
@@ -96,6 +103,7 @@ fn copy_allowed_metadata_fields(source: &Map<String, Value>, target: &mut Map<St
copy_number(source, target, "cache_read_price_per_1m");
copy_number(source, target, "price_per_request");
copy_non_null_value(source, target, "proxy");
sanitize_request_path_metadata_fields(target);
}
fn move_allowed_metadata_fields(mut source: Map<String, Value>, target: &mut Map<String, Value>) {
@@ -105,9 +113,13 @@ fn move_allowed_metadata_fields(mut source: Map<String, Value>, target: &mut Map
remove_bool(&mut source, target, "client_requested_stream");
remove_bool(&mut source, target, "upstream_is_stream");
remove_bool(&mut source, target, "api_key_is_standalone");
remove_non_empty_string(&mut source, target, "request_path");
remove_non_empty_string(&mut source, target, "request_query_string");
remove_non_empty_string(&mut source, target, "request_path_and_query");
remove_number(&mut source, target, "provider_request_body_base64_bytes");
remove_number(&mut source, target, "provider_response_body_base64_bytes");
remove_number(&mut source, target, "client_response_body_base64_bytes");
remove_number(&mut source, target, "client_response_status_code");
remove_non_null_value(&mut source, target, "billing_snapshot");
remove_non_empty_string(&mut source, target, "billing_snapshot_schema_version");
remove_non_empty_string(&mut source, target, "billing_snapshot_status");
@@ -129,6 +141,38 @@ fn move_allowed_metadata_fields(mut source: Map<String, Value>, target: &mut Map
remove_number(&mut source, target, "cache_read_price_per_1m");
remove_number(&mut source, target, "price_per_request");
remove_non_null_value(&mut source, target, "proxy");
sanitize_request_path_metadata_fields(target);
}
fn sanitize_request_path_metadata_fields(target: &mut Map<String, Value>) {
let path = target
.get("request_path")
.and_then(Value::as_str)
.and_then(sanitize_request_path);
let query = target
.get("request_query_string")
.and_then(Value::as_str)
.and_then(sanitize_request_query_string);
let path_and_query = target
.get("request_path_and_query")
.and_then(Value::as_str)
.and_then(|value| sanitize_request_path_and_query(value, None))
.or_else(|| {
path.as_deref()
.and_then(|path| sanitize_request_path_and_query(path, query.as_deref()))
});
apply_optional_string_field(target, "request_path", path.as_deref());
apply_optional_string_field(target, "request_query_string", query.as_deref());
apply_optional_string_field(target, "request_path_and_query", path_and_query.as_deref());
}
fn apply_optional_string_field(target: &mut Map<String, Value>, key: &str, value: Option<&str>) {
if let Some(value) = value {
target.insert(key.to_string(), Value::String(value.to_string()));
} else {
target.remove(key);
}
}
fn copy_non_empty_string(source: &Map<String, Value>, target: &mut Map<String, Value>, key: &str) {
@@ -451,6 +495,25 @@ mod tests {
);
}
#[test]
fn sanitizes_request_path_query_metadata() {
let metadata = sanitize_usage_request_metadata(Some(json!({
"request_path": "/v1beta/models/gemini-2.5-pro:streamGenerateContent?key=secret",
"request_query_string": "key=secret&alt=sse&pageSize=10&token=hidden",
"request_path_and_query": "/v1beta/models/gemini-2.5-pro:streamGenerateContent?key=secret&alt=sse&pageSize=10&token=hidden",
})))
.expect("metadata should remain");
assert_eq!(
metadata,
json!({
"request_path": "/v1beta/models/gemini-2.5-pro:streamGenerateContent",
"request_query_string": "alt=sse&pageSize=10",
"request_path_and_query": "/v1beta/models/gemini-2.5-pro:streamGenerateContent?alt=sse&pageSize=10",
})
);
}
#[test]
fn sanitizes_large_allowed_metadata_values_to_bounded_representations() {
let metadata = sanitize_usage_request_metadata(Some(json!({

89
crates/aether-usage-runtime/src/write.rs
View File

@@ -529,7 +529,8 @@ fn build_terminal_usage_event_from_seed_impl(
};
let routing = merge_routing_seed_with_metadata_owned(routing, request_metadata.as_ref());
let body_refs = merge_body_refs_seed_with_metadata_owned(body_refs, request_metadata.as_ref());
let error_message = resolve_error_message(status_code, provider_response.as_ref(), None);
let error_message = resolve_error_message(status_code, provider_response.as_ref(), None)
.or_else(|| resolve_error_message(status_code, client_response.as_ref(), None));
let api_family = infer_api_family(&client_contract).map(ToOwned::to_owned);
let endpoint_kind = infer_endpoint_kind(&client_contract).map(ToOwned::to_owned);
let provider_api_family = infer_api_family(&provider_contract).map(ToOwned::to_owned);
@@ -734,7 +735,8 @@ pub fn build_sync_terminal_usage_payload_seed(
let context = payload.report_context.as_ref().and_then(Value::as_object);
let provider_response_headers = context_usage_value(context, "provider_response_headers")
.or_else(|| headers_to_json(&payload.headers));
let client_response_headers = headers_to_json(&payload.headers);
let client_response_headers = context_usage_value(context, "client_response_headers")
.or_else(|| headers_to_json(&payload.headers));
SyncTerminalUsagePayloadSeed {
report_kind: payload.report_kind.clone(),
status_code: payload.status_code,
@@ -941,7 +943,7 @@ fn infer_sync_terminal_state(
) -> UsageTerminalState {
if status_code == 499 || report_kind.contains("cancel") {
UsageTerminalState::Cancelled
} else if status_code >= 400
} else if !(200..300).contains(&status_code)
|| provider_response
.and_then(|value| value.get("error"))
.is_some_and(|value| !value.is_null())
@@ -959,7 +961,7 @@ fn infer_stream_terminal_state(
) -> UsageTerminalState {
if cancelled || status_code == 499 || report_kind.contains("cancel") {
UsageTerminalState::Cancelled
} else if status_code >= 400 {
} else if !(200..300).contains(&status_code) {
UsageTerminalState::Failed
} else {
UsageTerminalState::Completed
@@ -1982,6 +1984,8 @@ fn resolve_error_category(status_code: u16, event_type: UsageEventType) -> Optio
UsageEventType::Cancelled => Some("cancelled".to_string()),
UsageEventType::Failed if status_code >= 500 => Some("server_error".to_string()),
UsageEventType::Failed if status_code >= 400 => Some("client_error".to_string()),
UsageEventType::Failed if status_code >= 300 => Some("redirect".to_string()),
UsageEventType::Failed => Some("non_success_status".to_string()),
_ => None,
}
}
@@ -2000,7 +2004,7 @@ fn resolve_error_message(
if explicit_error_message.is_some() {
return explicit_error_message;
}
if status_code < 400 {
if (200..300).contains(&status_code) {
return None;
}
@@ -2798,6 +2802,81 @@ mod tests {
);
}
#[test]
fn stream_terminal_usage_marks_redirect_status_as_failed() {
let plan = ExecutionPlan {
request_id: "req-stream-redirect-usage".to_string(),
candidate_id: Some("cand-stream-redirect-usage".to_string()),
provider_name: Some("ChatGPTWeb".to_string()),
provider_id: "provider-redirect".to_string(),
endpoint_id: "endpoint-redirect".to_string(),
key_id: "key-redirect".to_string(),
method: "POST".to_string(),
url: "https://example.com/v1beta/models/gemini:streamGenerateContent".to_string(),
headers: BTreeMap::new(),
content_type: None,
content_encoding: None,
body: RequestBody {
json_body: None,
body_bytes_b64: None,
body_ref: None,
},
stream: true,
client_api_format: "gemini:generate_content".to_string(),
provider_api_format: "gemini:generate_content".to_string(),
model_name: Some("gemini".to_string()),
proxy: None,
transport_profile: None,
timeouts: None,
};
let client_body = json!({
"error": {
"type": "execution_runtime_non_success_status",
"message": "execution runtime stream returned non-success status 302",
"code": 302,
"upstream_status": 302,
"location": "/"
}
});
let payload = GatewayStreamReportRequest {
trace_id: "trace-stream-redirect-usage".to_string(),
report_kind: "gemini_chat_stream_success".to_string(),
report_context: Some(json!({
"client_api_format": "gemini:generate_content",
"provider_api_format": "gemini:generate_content"
})),
status_code: 302,
headers: BTreeMap::from([
("content-type".to_string(), "application/json".to_string()),
("x-aether-upstream-status".to_string(), "302".to_string()),
]),
provider_body_base64: Some(
base64::engine::general_purpose::STANDARD
.encode(br#"{"error":{"message":"raw redirect body"}}"#),
),
provider_body_state: Some(UsageBodyCaptureState::Inline),
client_body_base64: Some(
base64::engine::general_purpose::STANDARD
.encode(serde_json::to_vec(&client_body).expect("body should encode")),
),
client_body_state: Some(UsageBodyCaptureState::Inline),
terminal_summary: None,
telemetry: None,
};
let event =
build_stream_terminal_usage_event(&plan, payload.report_context.as_ref(), &payload)
.expect("usage event should build");
assert_eq!(event.event_type, UsageEventType::Failed);
assert_eq!(event.data.status_code, Some(302));
assert_eq!(event.data.error_category.as_deref(), Some("redirect"));
assert_eq!(
event.data.error_message.as_deref(),
Some("raw redirect body")
);
}
#[test]
fn builds_stream_terminal_usage_from_terminal_summary_usage_without_decoding_bodies() {
let plan = ExecutionPlan {