mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
Merge pull request #391 from RWDai/fix/affinity-list-display
Fix affinity list display
This commit is contained in:
@@ -4,59 +4,81 @@ use crate::handlers::admin::request::AdminAppState;
|
||||
use crate::GatewayError;
|
||||
use std::time::Duration;
|
||||
|
||||
fn parse_admin_monitoring_cache_affinity_key(raw_key: &str) -> Option<(String, String, String)> {
|
||||
let parts = raw_key.split(':').collect::<Vec<_>>();
|
||||
let start = parts
|
||||
.iter()
|
||||
.position(|segment| *segment == "cache_affinity")?;
|
||||
let affinity_key = parts.get(start + 1)?.trim();
|
||||
if affinity_key.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let api_format = parts
|
||||
.get(start + 2)
|
||||
.map(|value| value.trim())
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or("unknown")
|
||||
.to_string();
|
||||
let model_name = parts
|
||||
.get(start + 3..)
|
||||
.filter(|segments| !segments.is_empty())
|
||||
.map(|segments| segments.join(":"))
|
||||
.unwrap_or_else(|| "unknown".to_string());
|
||||
Some((affinity_key.to_string(), api_format, model_name))
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
enum AdminMonitoringAffinityKeyKind {
|
||||
Cache,
|
||||
Scheduler,
|
||||
}
|
||||
|
||||
fn parse_admin_monitoring_scheduler_affinity_key(
|
||||
raw_key: &str,
|
||||
) -> Option<(String, String, String)> {
|
||||
let parts = raw_key.split(':').collect::<Vec<_>>();
|
||||
let start = parts
|
||||
.iter()
|
||||
.position(|segment| *segment == "scheduler_affinity")?;
|
||||
let affinity_key = parts.get(start + 1)?.trim();
|
||||
if affinity_key.is_empty() {
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
struct ParsedAdminMonitoringAffinityKey {
|
||||
affinity_key: String,
|
||||
api_format: String,
|
||||
model_name: String,
|
||||
client_family: Option<String>,
|
||||
session_hash: Option<String>,
|
||||
}
|
||||
|
||||
fn split_admin_monitoring_api_format_and_model(
|
||||
segments: &[&str],
|
||||
key_kind: AdminMonitoringAffinityKeyKind,
|
||||
) -> Option<(String, String)> {
|
||||
if segments.len() < 2 {
|
||||
return None;
|
||||
}
|
||||
|
||||
let remaining = parts.get(start + 2..)?;
|
||||
if remaining.len() < 2 {
|
||||
return None;
|
||||
let max_api_format_segments = segments.len().saturating_sub(1).min(3);
|
||||
for segment_count in (2..=max_api_format_segments).rev() {
|
||||
let candidate = segments[..segment_count]
|
||||
.iter()
|
||||
.map(|segment| segment.trim())
|
||||
.collect::<Vec<_>>()
|
||||
.join(":");
|
||||
if is_known_admin_monitoring_api_format(&candidate) {
|
||||
let model_name = segments[segment_count..]
|
||||
.iter()
|
||||
.map(|segment| segment.trim())
|
||||
.filter(|segment| !segment.is_empty())
|
||||
.collect::<Vec<_>>()
|
||||
.join(":");
|
||||
if model_name.is_empty() {
|
||||
return None;
|
||||
}
|
||||
return Some((candidate, model_name));
|
||||
}
|
||||
}
|
||||
|
||||
let (api_format, model_name_parts) = if remaining.len() == 2 {
|
||||
(remaining[0].trim().to_string(), &remaining[1..])
|
||||
} else {
|
||||
(
|
||||
format!("{}:{}", remaining[0].trim(), remaining[1].trim()),
|
||||
&remaining[2..],
|
||||
)
|
||||
};
|
||||
if api_format.trim().is_empty() {
|
||||
return None;
|
||||
if matches!(key_kind, AdminMonitoringAffinityKeyKind::Scheduler)
|
||||
&& segments.len() >= 3
|
||||
&& is_known_admin_monitoring_api_format_family(segments[0])
|
||||
{
|
||||
let api_format_kind = segments.get(1)?.trim();
|
||||
if api_format_kind.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let model_name = segments[2..]
|
||||
.iter()
|
||||
.map(|segment| segment.trim())
|
||||
.filter(|segment| !segment.is_empty())
|
||||
.collect::<Vec<_>>()
|
||||
.join(":");
|
||||
if model_name.is_empty() {
|
||||
return None;
|
||||
}
|
||||
return Some((
|
||||
format!(
|
||||
"{}:{api_format_kind}",
|
||||
segments[0].trim().to_ascii_lowercase()
|
||||
),
|
||||
model_name,
|
||||
));
|
||||
}
|
||||
|
||||
let model_name = model_name_parts
|
||||
let api_format = segments.first()?.trim();
|
||||
if api_format.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let model_name = segments[1..]
|
||||
.iter()
|
||||
.map(|segment| segment.trim())
|
||||
.filter(|segment| !segment.is_empty())
|
||||
@@ -66,7 +88,126 @@ fn parse_admin_monitoring_scheduler_affinity_key(
|
||||
return None;
|
||||
}
|
||||
|
||||
Some((affinity_key.to_string(), api_format, model_name))
|
||||
Some((api_format.to_string(), model_name))
|
||||
}
|
||||
|
||||
fn is_known_admin_monitoring_api_format_family(value: &str) -> bool {
|
||||
matches!(
|
||||
value.trim().to_ascii_lowercase().as_str(),
|
||||
"openai" | "claude" | "gemini" | "jina" | "doubao"
|
||||
)
|
||||
}
|
||||
|
||||
fn is_known_admin_monitoring_api_format(value: &str) -> bool {
|
||||
matches!(
|
||||
value.trim().to_ascii_lowercase().as_str(),
|
||||
"openai:chat"
|
||||
| "openai:responses"
|
||||
| "openai:responses:compact"
|
||||
| "openai:image"
|
||||
| "openai:video"
|
||||
| "openai:embedding"
|
||||
| "openai:rerank"
|
||||
| "claude:messages"
|
||||
| "gemini:generate_content"
|
||||
| "gemini:video"
|
||||
| "gemini:files"
|
||||
| "gemini:embedding"
|
||||
| "jina:embedding"
|
||||
| "jina:rerank"
|
||||
| "doubao:embedding"
|
||||
)
|
||||
}
|
||||
|
||||
fn parse_admin_monitoring_cache_affinity_key(
|
||||
raw_key: &str,
|
||||
) -> Option<ParsedAdminMonitoringAffinityKey> {
|
||||
let parts = raw_key.split(':').collect::<Vec<_>>();
|
||||
let start = parts
|
||||
.iter()
|
||||
.position(|segment| *segment == "cache_affinity")?;
|
||||
let affinity_key = parts.get(start + 1)?.trim();
|
||||
if affinity_key.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let remaining = parts.get(start + 2..)?;
|
||||
let (api_format, model_name) = split_admin_monitoring_api_format_and_model(
|
||||
remaining,
|
||||
AdminMonitoringAffinityKeyKind::Cache,
|
||||
)
|
||||
.unwrap_or_else(|| ("unknown".to_string(), "unknown".to_string()));
|
||||
Some(ParsedAdminMonitoringAffinityKey {
|
||||
affinity_key: affinity_key.to_string(),
|
||||
api_format,
|
||||
model_name,
|
||||
client_family: None,
|
||||
session_hash: None,
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_admin_monitoring_scheduler_affinity_key(
|
||||
raw_key: &str,
|
||||
) -> Option<ParsedAdminMonitoringAffinityKey> {
|
||||
let parts = raw_key.split(':').collect::<Vec<_>>();
|
||||
let start = parts
|
||||
.iter()
|
||||
.position(|segment| *segment == "scheduler_affinity")?;
|
||||
if parts.get(start + 1).is_some_and(|segment| *segment == "v2") {
|
||||
return parse_admin_monitoring_scheduler_affinity_v2_key(&parts, start);
|
||||
}
|
||||
|
||||
let affinity_key = parts.get(start + 1)?.trim();
|
||||
if affinity_key.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let remaining = parts.get(start + 2..)?;
|
||||
let (api_format, model_name) = split_admin_monitoring_api_format_and_model(
|
||||
remaining,
|
||||
AdminMonitoringAffinityKeyKind::Scheduler,
|
||||
)?;
|
||||
|
||||
Some(ParsedAdminMonitoringAffinityKey {
|
||||
affinity_key: affinity_key.to_string(),
|
||||
api_format,
|
||||
model_name,
|
||||
client_family: None,
|
||||
session_hash: None,
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_admin_monitoring_scheduler_affinity_v2_key(
|
||||
parts: &[&str],
|
||||
start: usize,
|
||||
) -> Option<ParsedAdminMonitoringAffinityKey> {
|
||||
let affinity_key = parts.get(start + 2)?.trim();
|
||||
if affinity_key.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let remaining = parts.get(start + 3..)?;
|
||||
if remaining.len() < 4 {
|
||||
return None;
|
||||
}
|
||||
let session_hash = remaining.last()?.trim();
|
||||
let client_family = remaining.get(remaining.len().saturating_sub(2))?.trim();
|
||||
if session_hash.is_empty() || client_family.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let api_model_segments = &remaining[..remaining.len() - 2];
|
||||
let (api_format, model_name) = split_admin_monitoring_api_format_and_model(
|
||||
api_model_segments,
|
||||
AdminMonitoringAffinityKeyKind::Scheduler,
|
||||
)?;
|
||||
|
||||
Some(ParsedAdminMonitoringAffinityKey {
|
||||
affinity_key: affinity_key.to_string(),
|
||||
api_format,
|
||||
model_name,
|
||||
client_family: Some(client_family.to_ascii_lowercase()),
|
||||
session_hash: Some(session_hash.to_string()),
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) fn admin_monitoring_scheduler_affinity_cache_key(
|
||||
@@ -78,9 +219,47 @@ pub(super) fn admin_monitoring_scheduler_affinity_cache_key(
|
||||
if affinity_key.is_empty() || api_format.is_empty() || model_name.is_empty() {
|
||||
return None;
|
||||
}
|
||||
Some(format!(
|
||||
"scheduler_affinity:{affinity_key}:{api_format}:{model_name}"
|
||||
))
|
||||
match (
|
||||
record
|
||||
.client_family
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty()),
|
||||
record
|
||||
.session_hash
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty()),
|
||||
) {
|
||||
(Some(client_family), Some(session_hash)) => Some(format!(
|
||||
"scheduler_affinity:v2:{affinity_key}:{api_format}:{model_name}:{client_family}:{session_hash}"
|
||||
)),
|
||||
_ => Some(format!(
|
||||
"scheduler_affinity:{affinity_key}:{api_format}:{model_name}"
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn admin_monitoring_json_string_field(
|
||||
object: &serde_json::Map<String, serde_json::Value>,
|
||||
field: &str,
|
||||
) -> Option<String> {
|
||||
object
|
||||
.get(field)
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
fn admin_monitoring_json_request_count(
|
||||
object: &serde_json::Map<String, serde_json::Value>,
|
||||
) -> Option<u64> {
|
||||
object.get("request_count").and_then(|value| {
|
||||
value
|
||||
.as_u64()
|
||||
.or_else(|| value.as_i64().and_then(|number| u64::try_from(number).ok()))
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) fn admin_monitoring_cache_affinity_record(
|
||||
@@ -89,35 +268,22 @@ pub(super) fn admin_monitoring_cache_affinity_record(
|
||||
) -> Option<AdminMonitoringCacheAffinityRecord> {
|
||||
let payload = serde_json::from_str::<serde_json::Value>(raw_value).ok()?;
|
||||
let object = payload.as_object()?;
|
||||
let (affinity_key, parsed_api_format, parsed_model_name) =
|
||||
parse_admin_monitoring_cache_affinity_key(raw_key)?;
|
||||
let api_format = object
|
||||
.get("api_format")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or(parsed_api_format.as_str())
|
||||
.to_string();
|
||||
let model_name = object
|
||||
.get("model_name")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or(parsed_model_name.as_str())
|
||||
.to_string();
|
||||
let request_count = object
|
||||
.get("request_count")
|
||||
.and_then(|value| {
|
||||
value
|
||||
.as_u64()
|
||||
.or_else(|| value.as_i64().and_then(|number| u64::try_from(number).ok()))
|
||||
})
|
||||
.unwrap_or(0);
|
||||
let parsed = parse_admin_monitoring_cache_affinity_key(raw_key)?;
|
||||
let api_format = admin_monitoring_json_string_field(object, "api_format")
|
||||
.unwrap_or_else(|| parsed.api_format.clone());
|
||||
let model_name = admin_monitoring_json_string_field(object, "model_name")
|
||||
.unwrap_or_else(|| parsed.model_name.clone());
|
||||
let request_count = admin_monitoring_json_request_count(object);
|
||||
Some(AdminMonitoringCacheAffinityRecord {
|
||||
raw_key: raw_key.to_string(),
|
||||
affinity_key,
|
||||
affinity_key: parsed.affinity_key,
|
||||
api_format,
|
||||
model_name,
|
||||
client_family: admin_monitoring_json_string_field(object, "client_family")
|
||||
.map(|value| value.to_ascii_lowercase())
|
||||
.or(parsed.client_family),
|
||||
session_hash: admin_monitoring_json_string_field(object, "session_hash")
|
||||
.or(parsed.session_hash),
|
||||
provider_id: object
|
||||
.get("provider_id")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
@@ -132,7 +298,8 @@ pub(super) fn admin_monitoring_cache_affinity_record(
|
||||
.map(ToOwned::to_owned),
|
||||
created_at: object.get("created_at").cloned(),
|
||||
expire_at: object.get("expire_at").cloned(),
|
||||
request_count,
|
||||
request_count: request_count.unwrap_or(0),
|
||||
request_count_known: request_count.is_some(),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -143,23 +310,25 @@ pub(super) fn admin_monitoring_scheduler_affinity_record(
|
||||
ttl: Duration,
|
||||
now_unix_secs: u64,
|
||||
) -> Option<AdminMonitoringCacheAffinityRecord> {
|
||||
let (affinity_key, api_format, model_name) =
|
||||
parse_admin_monitoring_scheduler_affinity_key(cache_key)?;
|
||||
let parsed = parse_admin_monitoring_scheduler_affinity_key(cache_key)?;
|
||||
let age_secs = age.as_secs();
|
||||
let created_at = now_unix_secs.saturating_sub(age_secs);
|
||||
let expire_at = created_at.saturating_add(ttl.as_secs());
|
||||
|
||||
Some(AdminMonitoringCacheAffinityRecord {
|
||||
raw_key: cache_key.to_string(),
|
||||
affinity_key,
|
||||
api_format,
|
||||
model_name,
|
||||
affinity_key: parsed.affinity_key,
|
||||
api_format: parsed.api_format,
|
||||
model_name: parsed.model_name,
|
||||
client_family: parsed.client_family,
|
||||
session_hash: parsed.session_hash,
|
||||
provider_id: Some(target.provider_id.clone()),
|
||||
endpoint_id: Some(target.endpoint_id.clone()),
|
||||
key_id: Some(target.key_id.clone()),
|
||||
created_at: Some(serde_json::json!(created_at)),
|
||||
expire_at: Some(serde_json::json!(expire_at)),
|
||||
request_count: 0,
|
||||
request_count_known: false,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -169,36 +338,23 @@ pub(super) fn admin_monitoring_scheduler_affinity_record_from_raw(
|
||||
) -> Option<AdminMonitoringCacheAffinityRecord> {
|
||||
let payload = serde_json::from_str::<serde_json::Value>(raw_value).ok()?;
|
||||
let object = payload.as_object()?;
|
||||
let (affinity_key, parsed_api_format, parsed_model_name) =
|
||||
parse_admin_monitoring_scheduler_affinity_key(raw_key)?;
|
||||
let api_format = object
|
||||
.get("api_format")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or(parsed_api_format.as_str())
|
||||
.to_string();
|
||||
let model_name = object
|
||||
.get("model_name")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or(parsed_model_name.as_str())
|
||||
.to_string();
|
||||
let request_count = object
|
||||
.get("request_count")
|
||||
.and_then(|value| {
|
||||
value
|
||||
.as_u64()
|
||||
.or_else(|| value.as_i64().and_then(|number| u64::try_from(number).ok()))
|
||||
})
|
||||
.unwrap_or(0);
|
||||
let parsed = parse_admin_monitoring_scheduler_affinity_key(raw_key)?;
|
||||
let api_format = admin_monitoring_json_string_field(object, "api_format")
|
||||
.unwrap_or_else(|| parsed.api_format.clone());
|
||||
let model_name = admin_monitoring_json_string_field(object, "model_name")
|
||||
.unwrap_or_else(|| parsed.model_name.clone());
|
||||
let request_count = admin_monitoring_json_request_count(object);
|
||||
|
||||
Some(AdminMonitoringCacheAffinityRecord {
|
||||
raw_key: raw_key.to_string(),
|
||||
affinity_key,
|
||||
affinity_key: parsed.affinity_key,
|
||||
api_format,
|
||||
model_name,
|
||||
client_family: admin_monitoring_json_string_field(object, "client_family")
|
||||
.map(|value| value.to_ascii_lowercase())
|
||||
.or(parsed.client_family),
|
||||
session_hash: admin_monitoring_json_string_field(object, "session_hash")
|
||||
.or(parsed.session_hash),
|
||||
provider_id: object
|
||||
.get("provider_id")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
@@ -213,7 +369,8 @@ pub(super) fn admin_monitoring_scheduler_affinity_record_from_raw(
|
||||
.map(ToOwned::to_owned),
|
||||
created_at: object.get("created_at").cloned(),
|
||||
expire_at: object.get("expire_at").cloned(),
|
||||
request_count,
|
||||
request_count: request_count.unwrap_or(0),
|
||||
request_count_known: request_count.is_some(),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -227,10 +384,15 @@ pub(super) fn clear_admin_monitoring_scheduler_affinity_entries(
|
||||
state: &AdminAppState<'_>,
|
||||
records: &[AdminMonitoringCacheAffinityRecord],
|
||||
) {
|
||||
let scheduler_keys = records
|
||||
.iter()
|
||||
.filter_map(admin_monitoring_scheduler_affinity_cache_key)
|
||||
.collect::<std::collections::BTreeSet<_>>();
|
||||
let mut scheduler_keys = std::collections::BTreeSet::new();
|
||||
for record in records {
|
||||
if record.raw_key.contains("scheduler_affinity:") {
|
||||
scheduler_keys.insert(record.raw_key.clone());
|
||||
}
|
||||
if let Some(cache_key) = admin_monitoring_scheduler_affinity_cache_key(record) {
|
||||
scheduler_keys.insert(cache_key);
|
||||
}
|
||||
}
|
||||
for scheduler_key in scheduler_keys {
|
||||
let _ = state
|
||||
.as_ref()
|
||||
@@ -286,3 +448,132 @@ pub(super) async fn delete_admin_monitoring_cache_affinity_raw_keys(
|
||||
state, raw_keys,
|
||||
))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
admin_monitoring_scheduler_affinity_record_from_raw,
|
||||
parse_admin_monitoring_cache_affinity_key, parse_admin_monitoring_scheduler_affinity_key,
|
||||
};
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn parses_legacy_scheduler_affinity_key_with_multi_segment_api_format() {
|
||||
let parsed = parse_admin_monitoring_scheduler_affinity_key(
|
||||
"scheduler_affinity:user-key-1:openai:chat:gpt-4.1",
|
||||
)
|
||||
.expect("legacy scheduler key should parse");
|
||||
|
||||
assert_eq!(parsed.affinity_key, "user-key-1");
|
||||
assert_eq!(parsed.api_format, "openai:chat");
|
||||
assert_eq!(parsed.model_name, "gpt-4.1");
|
||||
assert_eq!(parsed.client_family, None);
|
||||
assert_eq!(parsed.session_hash, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_v2_scheduler_affinity_key_with_client_family() {
|
||||
let parsed = parse_admin_monitoring_scheduler_affinity_key(
|
||||
"scheduler_affinity:v2:user-key-1:openai:responses:gpt-5.5:codex:d35efdccd572e9c5430e17d663dfde4cce83ea7e6665ac332f565780c98b1dff",
|
||||
)
|
||||
.expect("v2 scheduler key should parse");
|
||||
|
||||
assert_eq!(parsed.affinity_key, "user-key-1");
|
||||
assert_eq!(parsed.api_format, "openai:responses");
|
||||
assert_eq!(parsed.model_name, "gpt-5.5");
|
||||
assert_eq!(parsed.client_family.as_deref(), Some("codex"));
|
||||
assert_eq!(
|
||||
parsed.session_hash.as_deref(),
|
||||
Some("d35efdccd572e9c5430e17d663dfde4cce83ea7e6665ac332f565780c98b1dff")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_v2_scheduler_affinity_key_with_three_segment_api_format() {
|
||||
let parsed = parse_admin_monitoring_scheduler_affinity_key(
|
||||
"scheduler_affinity:v2:user-key-1:openai:responses:compact:gpt-5.5:opencode:sessionhash",
|
||||
)
|
||||
.expect("compact v2 scheduler key should parse");
|
||||
|
||||
assert_eq!(parsed.affinity_key, "user-key-1");
|
||||
assert_eq!(parsed.api_format, "openai:responses:compact");
|
||||
assert_eq!(parsed.model_name, "gpt-5.5");
|
||||
assert_eq!(parsed.client_family.as_deref(), Some("opencode"));
|
||||
assert_eq!(parsed.session_hash.as_deref(), Some("sessionhash"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_scheduler_affinity_key_with_unregistered_two_segment_api_format() {
|
||||
let parsed = parse_admin_monitoring_scheduler_affinity_key(
|
||||
"scheduler_affinity:v2:user-key-1:openai:image:gpt-image-1:opencode:sessionhash",
|
||||
)
|
||||
.expect("image v2 scheduler key should parse");
|
||||
|
||||
assert_eq!(parsed.affinity_key, "user-key-1");
|
||||
assert_eq!(parsed.api_format, "openai:image");
|
||||
assert_eq!(parsed.model_name, "gpt-image-1");
|
||||
assert_eq!(parsed.client_family.as_deref(), Some("opencode"));
|
||||
assert_eq!(parsed.session_hash.as_deref(), Some("sessionhash"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keeps_legacy_cache_affinity_single_segment_api_format() {
|
||||
let parsed = parse_admin_monitoring_cache_affinity_key(
|
||||
"cache_affinity:user-key-1:openai:model-alpha",
|
||||
)
|
||||
.expect("cache affinity key should parse");
|
||||
|
||||
assert_eq!(parsed.affinity_key, "user-key-1");
|
||||
assert_eq!(parsed.api_format, "openai");
|
||||
assert_eq!(parsed.model_name, "model-alpha");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn scheduler_raw_record_exposes_client_family_and_known_count() {
|
||||
let raw_value = json!({
|
||||
"provider_id": "provider-1",
|
||||
"endpoint_id": "endpoint-1",
|
||||
"key_id": "provider-key-1",
|
||||
"created_at": 1710000000u64,
|
||||
"expire_at": 1710000300u64,
|
||||
"request_count": 3u64,
|
||||
})
|
||||
.to_string();
|
||||
|
||||
let record = admin_monitoring_scheduler_affinity_record_from_raw(
|
||||
"scheduler_affinity:v2:user-key-1:openai:responses:gpt-5.5:codex:sessionhash",
|
||||
&raw_value,
|
||||
)
|
||||
.expect("raw scheduler affinity should parse");
|
||||
|
||||
assert_eq!(record.affinity_key, "user-key-1");
|
||||
assert_eq!(record.api_format, "openai:responses");
|
||||
assert_eq!(record.model_name, "gpt-5.5");
|
||||
assert_eq!(record.client_family.as_deref(), Some("codex"));
|
||||
assert_eq!(record.session_hash.as_deref(), Some("sessionhash"));
|
||||
assert_eq!(record.request_count, 3);
|
||||
assert!(record.request_count_known);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn scheduler_raw_record_treats_zero_request_count_as_known() {
|
||||
let raw_value = json!({
|
||||
"provider_id": "provider-1",
|
||||
"endpoint_id": "endpoint-1",
|
||||
"key_id": "provider-key-1",
|
||||
"created_at": 1710000000u64,
|
||||
"expire_at": 1710000300u64,
|
||||
"request_count": 0u64,
|
||||
})
|
||||
.to_string();
|
||||
|
||||
let record = admin_monitoring_scheduler_affinity_record_from_raw(
|
||||
"scheduler_affinity:v2:user-key-1:openai:responses:gpt-5.5:codex:sessionhash",
|
||||
&raw_value,
|
||||
)
|
||||
.expect("raw scheduler affinity should parse");
|
||||
|
||||
assert_eq!(record.request_count, 0);
|
||||
assert!(record.request_count_known);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -197,10 +197,13 @@ pub(super) async fn build_admin_monitoring_cache_affinities_response(
|
||||
"global_model_id": affinity.model_name,
|
||||
"model_name": affinity.model_name,
|
||||
"model_display_name": serde_json::Value::Null,
|
||||
"client_family": affinity.client_family,
|
||||
"session_hash": affinity.session_hash,
|
||||
"api_format": affinity.api_format,
|
||||
"created_at": affinity.created_at,
|
||||
"expire_at": affinity.expire_at,
|
||||
"request_count": affinity.request_count,
|
||||
"request_count_known": affinity.request_count_known,
|
||||
}));
|
||||
}
|
||||
|
||||
@@ -318,9 +321,12 @@ pub(super) async fn build_admin_monitoring_cache_affinity_response(
|
||||
"key_id": item.key_id,
|
||||
"api_format": item.api_format,
|
||||
"model_name": item.model_name,
|
||||
"client_family": item.client_family,
|
||||
"session_hash": item.session_hash,
|
||||
"created_at": item.created_at,
|
||||
"expire_at": item.expire_at,
|
||||
"request_count": item.request_count,
|
||||
"request_count_known": item.request_count_known,
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
@@ -20,6 +20,50 @@ use aether_admin::observability::monitoring::{
|
||||
};
|
||||
use axum::{body::Body, response::Response};
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct AdminMonitoringCacheAffinityDeleteFilter {
|
||||
client_family: Option<String>,
|
||||
session_hash: Option<String>,
|
||||
}
|
||||
|
||||
fn admin_monitoring_cache_affinity_delete_filter_from_query(
|
||||
query: Option<&str>,
|
||||
) -> AdminMonitoringCacheAffinityDeleteFilter {
|
||||
let Some(query) = query else {
|
||||
return AdminMonitoringCacheAffinityDeleteFilter::default();
|
||||
};
|
||||
let mut filter = AdminMonitoringCacheAffinityDeleteFilter::default();
|
||||
for (key, value) in url::form_urlencoded::parse(query.as_bytes()) {
|
||||
let value = value.trim();
|
||||
if value.is_empty() {
|
||||
continue;
|
||||
}
|
||||
match key.as_ref() {
|
||||
"client_family" => filter.client_family = Some(value.to_ascii_lowercase()),
|
||||
"session_hash" => filter.session_hash = Some(value.to_string()),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
filter
|
||||
}
|
||||
|
||||
fn admin_monitoring_cache_affinity_matches_delete_filter(
|
||||
item: &super::super::cache_types::AdminMonitoringCacheAffinityRecord,
|
||||
filter: &AdminMonitoringCacheAffinityDeleteFilter,
|
||||
) -> bool {
|
||||
if let Some(client_family) = filter.client_family.as_deref() {
|
||||
if item.client_family.as_deref() != Some(client_family) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
if let Some(session_hash) = filter.session_hash.as_deref() {
|
||||
if item.session_hash.as_deref() != Some(session_hash) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
pub(in super::super) async fn build_admin_monitoring_cache_affinity_delete_response(
|
||||
state: &AdminAppState<'_>,
|
||||
request_context: &AdminRequestContext<'_>,
|
||||
@@ -41,31 +85,33 @@ pub(in super::super) async fn build_admin_monitoring_cache_affinity_delete_respo
|
||||
|
||||
let target_affinity_keys =
|
||||
std::iter::once(affinity_key.clone()).collect::<std::collections::BTreeSet<_>>();
|
||||
let target_affinity =
|
||||
let delete_filter = admin_monitoring_cache_affinity_delete_filter_from_query(
|
||||
request_context.request_query_string.as_deref(),
|
||||
);
|
||||
let target_affinities =
|
||||
list_admin_monitoring_cache_affinity_records_by_affinity_keys(state, &target_affinity_keys)
|
||||
.await?
|
||||
.into_iter()
|
||||
.find(|item| {
|
||||
.filter(|item| {
|
||||
item.affinity_key == affinity_key
|
||||
&& item.endpoint_id.as_deref() == Some(endpoint_id.as_str())
|
||||
&& item.model_name == model_id
|
||||
&& item.api_format.eq_ignore_ascii_case(&api_format)
|
||||
});
|
||||
let Some(target_affinity) = target_affinity else {
|
||||
&& admin_monitoring_cache_affinity_matches_delete_filter(item, &delete_filter)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
if target_affinities.is_empty() {
|
||||
return Ok(admin_monitoring_not_found_response(
|
||||
"未找到指定的缓存亲和性记录",
|
||||
));
|
||||
};
|
||||
}
|
||||
|
||||
let _ = delete_admin_monitoring_cache_affinity_raw_keys(
|
||||
state,
|
||||
std::slice::from_ref(&target_affinity.raw_key),
|
||||
)
|
||||
.await?;
|
||||
clear_admin_monitoring_scheduler_affinity_entries(
|
||||
state,
|
||||
std::slice::from_ref(&target_affinity),
|
||||
);
|
||||
let raw_keys = target_affinities
|
||||
.iter()
|
||||
.map(|item| item.raw_key.clone())
|
||||
.collect::<Vec<_>>();
|
||||
let _ = delete_admin_monitoring_cache_affinity_raw_keys(state, &raw_keys).await?;
|
||||
clear_admin_monitoring_scheduler_affinity_entries(state, &target_affinities);
|
||||
|
||||
let mut api_key_by_id = admin_monitoring_list_export_api_key_records_by_ids(
|
||||
state,
|
||||
|
||||
@@ -216,6 +216,9 @@ async fn list_admin_monitoring_cache_affinity_records_matching(
|
||||
runner
|
||||
.keyspace()
|
||||
.key(&format!("scheduler_affinity:{affinity_key}:*")),
|
||||
runner
|
||||
.keyspace()
|
||||
.key(&format!("scheduler_affinity:v2:{affinity_key}:*")),
|
||||
]
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
|
||||
@@ -4,12 +4,15 @@ pub(super) struct AdminMonitoringCacheAffinityRecord {
|
||||
pub(super) affinity_key: String,
|
||||
pub(super) api_format: String,
|
||||
pub(super) model_name: String,
|
||||
pub(super) client_family: Option<String>,
|
||||
pub(super) session_hash: Option<String>,
|
||||
pub(super) provider_id: Option<String>,
|
||||
pub(super) endpoint_id: Option<String>,
|
||||
pub(super) key_id: Option<String>,
|
||||
pub(super) created_at: Option<serde_json::Value>,
|
||||
pub(super) expire_at: Option<serde_json::Value>,
|
||||
pub(super) request_count: u64,
|
||||
pub(super) request_count_known: bool,
|
||||
}
|
||||
|
||||
pub(super) struct AdminMonitoringCacheSnapshot {
|
||||
|
||||
@@ -339,6 +339,146 @@ async fn admin_monitoring_cache_affinities_and_delete_use_runtime_scheduler_affi
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_monitoring_cache_affinities_parse_session_scoped_scheduler_affinity_cache() {
|
||||
let user_repository = Arc::new(
|
||||
InMemoryUserReadRepository::seed_auth_users(vec![sample_monitoring_auth_user("user-1")])
|
||||
.with_export_users(vec![sample_monitoring_export_user("user-1")]),
|
||||
);
|
||||
let auth_repository = Arc::new(
|
||||
InMemoryAuthApiKeySnapshotRepository::default().with_export_records(vec![
|
||||
sample_monitoring_export_api_key("user-1", "user-key-1"),
|
||||
]),
|
||||
);
|
||||
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider()],
|
||||
vec![sample_monitoring_catalog_endpoint()],
|
||||
vec![sample_monitoring_catalog_key()],
|
||||
));
|
||||
let state = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_provider_catalog_reader_for_tests(provider_catalog)
|
||||
.with_user_reader(user_repository)
|
||||
.with_auth_api_key_reader(auth_repository),
|
||||
);
|
||||
let client_session = aether_scheduler_core::ClientSessionAffinity::new(
|
||||
Some("Codex".to_string()),
|
||||
Some("account=acct-1;session=session-1".to_string()),
|
||||
);
|
||||
let other_client_session = aether_scheduler_core::ClientSessionAffinity::new(
|
||||
Some("Codex".to_string()),
|
||||
Some("account=acct-1;session=session-2".to_string()),
|
||||
);
|
||||
let affinity_cache_key =
|
||||
aether_scheduler_core::build_scheduler_affinity_cache_key_for_api_key_id_with_client_session(
|
||||
"user-key-1",
|
||||
"openai:responses",
|
||||
"gpt-5.5",
|
||||
Some(&client_session),
|
||||
)
|
||||
.expect("session scheduler affinity cache key should build");
|
||||
let other_affinity_cache_key =
|
||||
aether_scheduler_core::build_scheduler_affinity_cache_key_for_api_key_id_with_client_session(
|
||||
"user-key-1",
|
||||
"openai:responses",
|
||||
"gpt-5.5",
|
||||
Some(&other_client_session),
|
||||
)
|
||||
.expect("other session scheduler affinity cache key should build");
|
||||
assert!(affinity_cache_key
|
||||
.starts_with("scheduler_affinity:v2:user-key-1:openai:responses:gpt-5.5:codex:"));
|
||||
let session_hash = affinity_cache_key
|
||||
.rsplit(':')
|
||||
.next()
|
||||
.expect("session hash should exist")
|
||||
.to_string();
|
||||
state.scheduler_affinity_cache.insert(
|
||||
affinity_cache_key.clone(),
|
||||
crate::cache::SchedulerAffinityTarget {
|
||||
provider_id: "provider-1".to_string(),
|
||||
endpoint_id: "endpoint-1".to_string(),
|
||||
key_id: "provider-key-1".to_string(),
|
||||
},
|
||||
crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL,
|
||||
128,
|
||||
);
|
||||
state.scheduler_affinity_cache.insert(
|
||||
other_affinity_cache_key.clone(),
|
||||
crate::cache::SchedulerAffinityTarget {
|
||||
provider_id: "provider-1".to_string(),
|
||||
endpoint_id: "endpoint-1".to_string(),
|
||||
key_id: "provider-key-1".to_string(),
|
||||
},
|
||||
crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL,
|
||||
128,
|
||||
);
|
||||
|
||||
let list_response = local_monitoring_response(
|
||||
&state,
|
||||
&request_context(
|
||||
http::Method::GET,
|
||||
"/api/admin/monitoring/cache/affinities?keyword=alice&limit=20&offset=0",
|
||||
),
|
||||
)
|
||||
.await
|
||||
.expect("handler should not error")
|
||||
.expect("route should be handled locally");
|
||||
assert_eq!(list_response.status(), http::StatusCode::OK);
|
||||
let list_body = to_bytes(list_response.into_body(), usize::MAX)
|
||||
.await
|
||||
.expect("body should read");
|
||||
let list_payload: serde_json::Value =
|
||||
serde_json::from_slice(&list_body).expect("json body should parse");
|
||||
let items = list_payload["data"]["items"]
|
||||
.as_array()
|
||||
.expect("items should be an array");
|
||||
let item = items
|
||||
.iter()
|
||||
.find(|item| item["session_hash"] == json!(session_hash))
|
||||
.expect("session-scoped item should be listed");
|
||||
assert_eq!(list_payload["data"]["meta"]["total"], json!(2));
|
||||
assert_eq!(item["affinity_key"], json!("user-key-1"));
|
||||
assert_eq!(item["username"], json!("alice"));
|
||||
assert_eq!(item["api_format"], json!("openai:responses"));
|
||||
assert_eq!(item["model_name"], json!("gpt-5.5"));
|
||||
assert_eq!(item["client_family"], json!("codex"));
|
||||
assert_eq!(item["provider_name"], json!("OpenAI"));
|
||||
assert_eq!(item["key_name"], json!("prod-key"));
|
||||
assert_eq!(item["request_count"], json!(0));
|
||||
assert_eq!(item["request_count_known"], json!(false));
|
||||
|
||||
let delete_response = local_monitoring_response(
|
||||
&state,
|
||||
&request_context(
|
||||
http::Method::DELETE,
|
||||
&format!(
|
||||
"/api/admin/monitoring/cache/affinity/user-key-1/endpoint-1/gpt-5.5/openai:responses?client_family=codex&session_hash={session_hash}"
|
||||
),
|
||||
),
|
||||
)
|
||||
.await
|
||||
.expect("handler should not error")
|
||||
.expect("route should be handled locally");
|
||||
assert_eq!(delete_response.status(), http::StatusCode::OK);
|
||||
assert_eq!(
|
||||
state.read_scheduler_affinity_target(
|
||||
&affinity_cache_key,
|
||||
crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL,
|
||||
),
|
||||
None
|
||||
);
|
||||
assert!(
|
||||
state
|
||||
.read_scheduler_affinity_target(
|
||||
&other_affinity_cache_key,
|
||||
crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL,
|
||||
)
|
||||
.is_some(),
|
||||
"deleting one session-scoped row should keep sibling sessions"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_monitoring_cache_users_delete_returns_local_payload_from_test_store() {
|
||||
let user_repository = Arc::new(
|
||||
|
||||
@@ -67,6 +67,7 @@ impl AppState {
|
||||
};
|
||||
|
||||
let cache_key = cache_key.to_string();
|
||||
let namespaced_cache_key = runner.keyspace().key(&cache_key);
|
||||
let provider_id = target.provider_id.clone();
|
||||
let endpoint_id = target.endpoint_id.clone();
|
||||
let key_id = target.key_id.clone();
|
||||
@@ -75,19 +76,48 @@ impl AppState {
|
||||
let expire_at = now_unix_secs.saturating_add(ttl_seconds);
|
||||
|
||||
handle.spawn(async move {
|
||||
let payload = serde_json::json!({
|
||||
"provider_id": provider_id,
|
||||
"endpoint_id": endpoint_id,
|
||||
"key_id": key_id,
|
||||
"created_at": now_unix_secs,
|
||||
"expire_at": expire_at,
|
||||
"request_count": 0,
|
||||
});
|
||||
let Ok(serialized) = serde_json::to_string(&payload) else {
|
||||
let Ok(mut connection) = runner.client().get_multiplexed_async_connection().await
|
||||
else {
|
||||
return;
|
||||
};
|
||||
let _ = runner
|
||||
.setex(&cache_key, &serialized, Some(ttl_seconds))
|
||||
let script = r#"
|
||||
local existing = redis.call('GET', KEYS[1])
|
||||
local request_count = 0
|
||||
local created_at = tonumber(ARGV[4])
|
||||
if existing then
|
||||
local ok, payload = pcall(cjson.decode, existing)
|
||||
if ok and type(payload) == 'table' then
|
||||
if type(payload['request_count']) == 'number' then
|
||||
request_count = payload['request_count']
|
||||
end
|
||||
if type(payload['created_at']) == 'number' then
|
||||
created_at = payload['created_at']
|
||||
end
|
||||
end
|
||||
end
|
||||
request_count = request_count + 1
|
||||
local payload = {
|
||||
provider_id = ARGV[1],
|
||||
endpoint_id = ARGV[2],
|
||||
key_id = ARGV[3],
|
||||
created_at = created_at,
|
||||
expire_at = tonumber(ARGV[5]),
|
||||
request_count = request_count
|
||||
}
|
||||
redis.call('SETEX', KEYS[1], tonumber(ARGV[6]), cjson.encode(payload))
|
||||
return request_count
|
||||
"#;
|
||||
let _ = redis::cmd("EVAL")
|
||||
.arg(script)
|
||||
.arg(1)
|
||||
.arg(&namespaced_cache_key)
|
||||
.arg(&provider_id)
|
||||
.arg(&endpoint_id)
|
||||
.arg(&key_id)
|
||||
.arg(now_unix_secs)
|
||||
.arg(expire_at)
|
||||
.arg(ttl_seconds)
|
||||
.query_async::<i64>(&mut connection)
|
||||
.await;
|
||||
});
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user