Fix Codex image progress heartbeat merge regressions

This commit is contained in:
fawney19
2026-05-10 02:10:23 +08:00
156 changed files with 11741 additions and 943 deletions

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,
@@ -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)

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))
@@ -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"));