mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
Fix Codex image progress heartbeat merge regressions
This commit is contained in:
@@ -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,
|
||||
@@ -474,6 +477,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() {
|
||||
@@ -501,6 +531,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)
|
||||
|
||||
@@ -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))
|
||||
@@ -1669,7 +1688,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 {
|
||||
@@ -2268,10 +2289,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};
|
||||
|
||||
@@ -2359,6 +2380,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 {
|
||||
@@ -2538,6 +2586,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"));
|
||||
|
||||
Reference in New Issue
Block a user