mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-12 14:10:19 +08:00
Fix usage records filtering and pool trace display
This commit is contained in:
@@ -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()
|
||||
{
|
||||
|
||||
Reference in New Issue
Block a user