Fix usage records filtering and pool trace display

This commit is contained in:
fawney19
2026-05-28 20:24:48 +08:00
parent df518ad668
commit ef2953038e
11 changed files with 309 additions and 15 deletions
@@ -795,6 +795,10 @@ impl<'a> PoolKeyCursor<'a> {
if self.skip_candidate_if_runtime_cooldown(&candidate).await {
continue;
}
if candidate.orchestration.candidate_group_id.is_none() {
candidate.orchestration.candidate_group_id =
Some(pool_cursor_candidate_group_id(&self.group));
}
candidate.orchestration.pool_key_index = Some(self.next_pool_key_index);
self.next_pool_key_index = self.next_pool_key_index.saturating_add(1);
return Some(candidate);
@@ -1539,6 +1543,17 @@ fn pool_candidate_facts(candidate: &EligibleLocalExecutionCandidate) -> PoolCand
}
}
fn pool_cursor_candidate_group_id(group: &EligibleLocalExecutionCandidate) -> String {
format!(
"provider={}|endpoint={}|model={}|selected_model={}|api_format={}|singleton_key=*",
group.candidate.provider_id,
group.candidate.endpoint_id,
group.candidate.model_id,
group.candidate.selected_provider_model_name,
group.provider_api_format,
)
}
fn pool_scheduling_config(
config: AdminProviderPoolConfig,
provider_type: &str,
@@ -2813,6 +2828,12 @@ mod tests {
.expect("cursor should skip disallowed pool key and return allowed key");
assert_eq!(candidate.candidate.key_id, "key-b");
assert_eq!(candidate.orchestration.pool_key_index, Some(0));
assert_eq!(
candidate.orchestration.candidate_group_id.as_deref(),
Some(
"provider=provider-pool|endpoint=endpoint-1|model=model-1|selected_model=gpt-5|api_format=openai:chat|singleton_key=*"
)
);
assert_eq!(
cursor
.skip_reason_counts
@@ -277,6 +277,30 @@ fn admin_usage_matches_client_family(
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()
.map(str::trim)
.map(|value| {
value == "1"
|| value.eq_ignore_ascii_case("true")
|| value.eq_ignore_ascii_case("yes")
|| value.eq_ignore_ascii_case("on")
})
.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_has_unknown_model_or_provider(item: &StoredRequestUsageAudit) -> bool {
admin_usage_is_unknown_label(&item.model) || admin_usage_is_unknown_label(&item.provider_name)
}
#[allow(clippy::too_many_arguments)]
fn build_admin_usage_records_response_with_attempt_flags(
items: &[StoredRequestUsageAudit],
@@ -615,6 +639,8 @@ pub(super) async fn maybe_build_local_admin_usage_summary_response(
let search = query_param_value(query, "search");
let username_filter = query_param_value(query, "username");
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 limit = match admin_usage_parse_limit(query) {
Ok(value) => value,
Err(detail) => return Ok(Some(admin_usage_bad_request_response(detail))),
@@ -652,7 +678,8 @@ 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 attempt_status_filter.is_some()
let (usage, total) = if hide_unknown_records
|| attempt_status_filter.is_some()
|| active_client_family_filter.is_some()
{
let mut usage = state.list_usage_audits(&base_query).await?;
@@ -692,6 +719,8 @@ pub(super) async fn maybe_build_local_admin_usage_summary_response(
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();
@@ -1258,6 +1258,84 @@ async fn gateway_handles_admin_usage_records_locally_with_trusted_admin_principa
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_filters_admin_usage_records_with_unknown_model_or_provider() {
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-visible",
"req-visible",
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-unknown-provider",
"req-unknown-provider",
Some("user-1"),
Some("key-1"),
Some("primary"),
"unknow",
"gpt-5",
"completed",
120,
30,
0.3,
0.36,
DAY_2_UNIX_SECS,
),
sample_usage_row(
"usage-unknown-model",
"req-unknown-model",
Some("user-1"),
Some("key-1"),
Some("primary"),
"OpenAI",
"unknown",
"completed",
120,
30,
0.3,
0.36,
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 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&hide_unknown=true&limit=10&offset=0"
)))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["total"], 1);
assert_eq!(payload["records"][0]["id"], "usage-visible");
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()
{