i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-01 17:00:21 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor: 优化调度候选排序与用量写入链路并改进 Fernet 缓存与前端批量列表

Browse Source
This commit is contained in:
fawney19
2026-04-21 16:19:07 +08:00
parent c5c56ff92f
commit 25a2b417be
77 changed files with 4801 additions and 2881 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
crates/aether-ai-pipeline/src/api.rs
View File

@@ -80,7 +80,7 @@ pub use crate::conversion::{
pub use crate::finalize::common::{
build_generated_tool_call_id, build_local_success_background_report,
build_local_success_conversion_background_report, canonicalize_tool_arguments,
prepare_local_success_response_parts,
prepare_local_success_response_parts, prepare_local_success_response_parts_owned,
};
pub use crate::finalize::sse::{encode_done_sse, encode_json_sse, map_claude_stop_reason};
pub use crate::finalize::standard::claude::stream::{ClaudeClientEmitter, ClaudeProviderState};

41
crates/aether-ai-pipeline/src/finalize/common.rs
View File

@@ -21,7 +21,13 @@ pub fn prepare_local_success_response_parts(
headers: &BTreeMap<String, String>,
body_json: &Value,
) -> serde_json::Result<(Vec<u8>, BTreeMap<String, String>)> {
let mut headers = headers.clone();
prepare_local_success_response_parts_owned(headers.clone(), body_json)
}
pub fn prepare_local_success_response_parts_owned(
mut headers: BTreeMap<String, String>,
body_json: &Value,
) -> serde_json::Result<(Vec<u8>, BTreeMap<String, String>)> {
headers.remove("content-encoding");
headers.remove("content-length");
headers.insert("content-type".to_string(), "application/json".to_string());
@@ -77,7 +83,7 @@ mod tests {
use super::{
build_generated_tool_call_id, build_local_success_background_report,
build_local_success_conversion_background_report, canonicalize_tool_arguments,
prepare_local_success_response_parts,
prepare_local_success_response_parts, prepare_local_success_response_parts_owned,
};
use aether_usage_runtime::GatewaySyncReportRequest;
use std::collections::BTreeMap;
@@ -127,6 +133,37 @@ mod tests {
);
}
#[test]
fn prepare_local_success_response_parts_owned_normalizes_headers() {
let headers = BTreeMap::from([
("content-encoding".to_string(), "gzip".to_string()),
("content-length".to_string(), "999".to_string()),
("x-test".to_string(), "1".to_string()),
]);
let (body_bytes, normalized_headers) =
prepare_local_success_response_parts_owned(headers, &serde_json::json!({"ok": true}))
.expect("response parts should serialize");
assert_eq!(
serde_json::from_slice::<Value>(&body_bytes).expect("json body"),
serde_json::json!({"ok": true})
);
assert_eq!(
normalized_headers.get("content-type").map(String::as_str),
Some("application/json")
);
assert!(!normalized_headers.contains_key("content-encoding"));
let expected_length = body_bytes.len().to_string();
assert_eq!(
normalized_headers.get("content-length").map(String::as_str),
Some(expected_length.as_str())
);
assert_eq!(
normalized_headers.get("x-test").map(String::as_str),
Some("1")
);
}
#[test]
fn build_local_success_background_report_maps_finalize_kind() {
let payload = GatewaySyncReportRequest {

29
crates/aether-ai-pipeline/src/planner/standard/matrix.rs
View File

@@ -1,3 +1,5 @@
use std::borrow::Cow;
use aether_provider_transport::url::{
build_claude_messages_url, build_gemini_content_url, build_openai_chat_url,
build_openai_cli_url, build_passthrough_path_url,
@@ -30,13 +32,13 @@ pub fn build_standard_request_body(
body_rules: Option<&Value>,
user_api_key_id: Option<&str>,
) -> Option<Value> {
let canonical_request = normalize_standard_request_to_openai_chat_request(
let canonical_request = normalize_standard_request_to_openai_chat_request_cow(
body_json,
client_api_format,
request_path,
)?;
let mut provider_request_body = build_standard_request_body_from_canonical(
&canonical_request,
canonical_request.as_ref(),
mapped_model,
provider_api_format,
upstream_is_stream,
@@ -99,14 +101,29 @@ pub fn normalize_standard_request_to_openai_chat_request(
client_api_format: &str,
request_path: &str,
) -> Option<Value> {
normalize_standard_request_to_openai_chat_request_cow(
body_json,
client_api_format,
request_path,
)
.map(Cow::into_owned)
}
fn normalize_standard_request_to_openai_chat_request_cow<'a>(
body_json: &'a Value,
client_api_format: &str,
request_path: &str,
) -> Option<Cow<'a, Value>> {
match client_api_format.trim().to_ascii_lowercase().as_str() {
"openai:chat" => Some(body_json.clone()),
"openai:chat" => Some(Cow::Borrowed(body_json)),
"openai:cli" | "openai:compact" => {
normalize_openai_cli_request_to_openai_chat_request(body_json)
normalize_openai_cli_request_to_openai_chat_request(body_json).map(Cow::Owned)
}
"claude:chat" | "claude:cli" => {
normalize_claude_request_to_openai_chat_request(body_json).map(Cow::Owned)
}
"claude:chat" | "claude:cli" => normalize_claude_request_to_openai_chat_request(body_json),
"gemini:chat" | "gemini:cli" => {
normalize_gemini_request_to_openai_chat_request(body_json, request_path)
normalize_gemini_request_to_openai_chat_request(body_json, request_path).map(Cow::Owned)
}
_ => None,
}

173
crates/aether-crypto/src/python_fernet.rs
View File

@@ -1,8 +1,9 @@
use std::collections::HashMap;
use std::sync::{LazyLock, Mutex};
use std::collections::{HashMap, VecDeque};
use std::sync::{Arc, LazyLock, RwLock};
use std::time::{SystemTime, UNIX_EPOCH};
use aes::cipher::{block_padding::Pkcs7, BlockDecryptMut, BlockEncryptMut, KeyIvInit};
use base64::decoded_len_estimate;
use base64::engine::general_purpose::{STANDARD, URL_SAFE, URL_SAFE_NO_PAD};
use base64::Engine as _;
use cbc::{Decryptor, Encryptor};
@@ -16,7 +17,8 @@ const HMAC_SIZE: usize = 32;
const IV_SIZE: usize = 16;
const SIGNING_KEY_SIZE: usize = 16;
const ENCRYPTION_KEY_SIZE: usize = 16;
const MIN_TOKEN_SIZE: usize = 1 + 8 + IV_SIZE + HMAC_SIZE;
const MIN_CIPHERTEXT_SIZE: usize = 16;
const MIN_TOKEN_SIZE: usize = 1 + 8 + IV_SIZE + MIN_CIPHERTEXT_SIZE + HMAC_SIZE;
const PBKDF2_ITERATIONS: u32 = 100_000;
const MAX_CACHED_DERIVED_KEYS: usize = 16;
@@ -24,13 +26,63 @@ pub const APP_SALT_SEED: &[u8] = b"aether-v1";
pub const APP_SALT_HEX: &str = "8797080a7a4b45b4810e934d1af36261";
pub const DEVELOPMENT_ENCRYPTION_KEY: &str = "dev-encryption-key-do-not-use-in-production";
static RAW_FERNET_KEY_CACHE: LazyLock<Mutex<HashMap<Box<str>, [u8; 32]>>> =
LazyLock::new(|| Mutex::new(HashMap::new()));
static RAW_FERNET_KEY_CACHE: LazyLock<RwLock<RawFernetKeyCache>> =
LazyLock::new(|| RwLock::new(RawFernetKeyCache::default()));
static APP_SALT: LazyLock<[u8; 16]> = LazyLock::new(|| {
let mut salt = [0u8; 16];
salt.copy_from_slice(&Sha256::digest(APP_SALT_SEED)[..16]);
salt
});
type Aes128CbcDec = Decryptor<aes::Aes128>;
type Aes128CbcEnc = Encryptor<aes::Aes128>;
type HmacSha256 = Hmac<Sha256>;
fn base64_encoded_len(input_len: usize) -> usize {
input_len.div_ceil(3) * 4
}
fn base64_unpadded_len(input_len: usize) -> usize {
let full_chunks = (input_len / 3) * 4;
match input_len % 3 {
0 => full_chunks,
1 => full_chunks + 2,
_ => full_chunks + 3,
}
}
fn minimum_wrapped_token_len() -> usize {
base64_unpadded_len(base64_unpadded_len(MIN_TOKEN_SIZE))
}
#[derive(Debug, Default)]
struct RawFernetKeyCache {
entries: HashMap<Arc<str>, [u8; 32]>,
insertion_order: VecDeque<Arc<str>>,
}
impl RawFernetKeyCache {
fn get(&self, secret: &str) -> Option<[u8; 32]> {
self.entries.get(secret).copied()
}
fn insert(&mut self, secret: &str, raw_key: [u8; 32]) {
if self.entries.contains_key(secret) {
return;
}
if self.entries.len() >= MAX_CACHED_DERIVED_KEYS {
if let Some(oldest) = self.insertion_order.pop_front() {
self.entries.remove(oldest.as_ref());
}
}
let secret: Arc<str> = Arc::from(secret);
self.insertion_order.push_back(secret.clone());
self.entries.insert(secret, raw_key);
}
}
#[derive(Debug, thiserror::Error)]
pub enum PythonFernetError {
#[error("invalid Python Fernet outer base64 payload")]
@@ -68,9 +120,9 @@ impl PythonFernetCompat {
let outer =
decode_urlsafe(ciphertext).map_err(|_| PythonFernetError::InvalidOuterBase64)?;
let inner =
let token =
decode_urlsafe_bytes(&outer).map_err(|_| PythonFernetError::InvalidInnerBase64)?;
let plaintext = self.decrypt_token_bytes(&inner)?;
let plaintext = self.decrypt_token_bytes(token)?;
String::from_utf8(plaintext).map_err(PythonFernetError::InvalidUtf8)
}
@@ -93,7 +145,7 @@ impl PythonFernetCompat {
}
}
fn decrypt_token_bytes(&self, token: &[u8]) -> Result<Vec<u8>, PythonFernetError> {
fn decrypt_token_bytes(&self, mut token: Vec<u8>) -> Result<Vec<u8>, PythonFernetError> {
if token.len() < MIN_TOKEN_SIZE {
return Err(PythonFernetError::InvalidTokenStructure);
}
@@ -102,22 +154,30 @@ impl PythonFernetCompat {
}
let signed_len = token.len() - HMAC_SIZE;
let (signed, signature) = token.split_at(signed_len);
let mut mac = HmacSha256::new_from_slice(&self.signing_key)
.map_err(|_| PythonFernetError::InvalidTokenSignature)?;
mac.update(signed);
mac.verify_slice(signature)
.map_err(|_| PythonFernetError::InvalidTokenSignature)?;
{
let (signed, signature) = token.split_at(signed_len);
let mut mac = HmacSha256::new_from_slice(&self.signing_key)
.map_err(|_| PythonFernetError::InvalidTokenSignature)?;
mac.update(signed);
mac.verify_slice(signature)
.map_err(|_| PythonFernetError::InvalidTokenSignature)?;
}
let iv_offset = 1 + 8;
let ciphertext_offset = iv_offset + IV_SIZE;
let iv = &token[iv_offset..ciphertext_offset];
let mut ciphertext = token[ciphertext_offset..signed_len].to_vec();
let plaintext = Aes128CbcDec::new((&self.encryption_key).into(), iv.into())
.decrypt_padded_mut::<Pkcs7>(&mut ciphertext)
.map_err(|_| PythonFernetError::InvalidPadding)?;
Ok(plaintext.to_vec())
let plaintext_len = {
let (_, payload) = token.split_at_mut(iv_offset);
let (iv, ciphertext_and_signature) = payload.split_at_mut(IV_SIZE);
let ciphertext = &mut ciphertext_and_signature[..signed_len - ciphertext_offset];
Aes128CbcDec::new((&self.encryption_key).into(), (&iv[..]).into())
.decrypt_padded_mut::<Pkcs7>(ciphertext)
.map_err(|_| PythonFernetError::InvalidPadding)?
.len()
};
token.copy_within(ciphertext_offset..ciphertext_offset + plaintext_len, 0);
token.truncate(plaintext_len);
Ok(token)
}
fn encrypt_token(
@@ -131,14 +191,13 @@ impl PythonFernetCompat {
padded[..plaintext.len()].copy_from_slice(plaintext);
let ciphertext = Aes128CbcEnc::new((&self.encryption_key).into(), (&iv).into())
.encrypt_padded_mut::<Pkcs7>(&mut padded, plaintext.len())
.map_err(|_| PythonFernetError::InvalidPadding)?
.to_vec();
.map_err(|_| PythonFernetError::InvalidPadding)?;
let mut signed = Vec::with_capacity(1 + 8 + IV_SIZE + ciphertext.len() + HMAC_SIZE);
signed.push(FERNET_VERSION);
signed.extend_from_slice(&timestamp.to_be_bytes());
signed.extend_from_slice(&iv);
signed.extend_from_slice(&ciphertext);
signed.extend_from_slice(ciphertext);
let mut mac = HmacSha256::new_from_slice(&self.signing_key)
.map_err(|_| PythonFernetError::InvalidTokenSignature)?;
@@ -146,8 +205,12 @@ impl PythonFernetCompat {
let signature = mac.finalize().into_bytes();
signed.extend_from_slice(&signature);
let inner = URL_SAFE.encode(signed);
Ok(URL_SAFE.encode(inner.as_bytes()))
let mut inner = String::with_capacity(base64_encoded_len(signed.len()));
URL_SAFE.encode_string(&signed, &mut inner);
let mut outer = String::with_capacity(base64_encoded_len(inner.len()));
URL_SAFE.encode_string(inner.as_bytes(), &mut outer);
Ok(outer)
}
}
@@ -164,7 +227,7 @@ pub fn decrypt_python_fernet_ciphertext(
pub fn looks_like_python_fernet_ciphertext(ciphertext: &str) -> bool {
let ciphertext = ciphertext.trim();
if ciphertext.is_empty() {
if ciphertext.is_empty() || ciphertext.len() < minimum_wrapped_token_len() {
return false;
}
@@ -190,32 +253,36 @@ pub fn warm_python_fernet_secret(secret: &str) {
}
fn raw_fernet_key(secret: &str) -> [u8; 32] {
if let Ok(raw_key) = decode_direct_fernet_key(secret) {
return raw_key;
}
if let Some(raw_key) = RAW_FERNET_KEY_CACHE
.lock()
.read()
.expect("raw fernet key cache should lock")
.get(secret)
.copied()
{
return raw_key;
}
let mut salt = [0u8; 16];
salt.copy_from_slice(&Sha256::digest(APP_SALT_SEED)[..16]);
let mut raw_key = [0u8; 32];
pbkdf2_hmac::<Sha256>(secret.as_bytes(), &salt, PBKDF2_ITERATIONS, &mut raw_key);
let mut cache = RAW_FERNET_KEY_CACHE
.lock()
.write()
.expect("raw fernet key cache should lock");
if cache.len() >= MAX_CACHED_DERIVED_KEYS && !cache.contains_key(secret) {
cache.clear();
if let Some(raw_key) = cache.get(secret) {
return raw_key;
}
cache.insert(secret.into(), raw_key);
let raw_key = match decode_direct_fernet_key(secret) {
Ok(raw_key) => raw_key,
Err(_) => {
let mut raw_key = [0u8; 32];
pbkdf2_hmac::<Sha256>(
secret.as_bytes(),
&*APP_SALT,
PBKDF2_ITERATIONS,
&mut raw_key,
);
raw_key
}
};
cache.insert(secret, raw_key);
raw_key
}
@@ -232,15 +299,23 @@ fn decode_direct_fernet_key(secret: &str) -> Result<[u8; 32], PythonFernetError>
}
fn decode_urlsafe(value: &str) -> Result<Vec<u8>, base64::DecodeError> {
URL_SAFE
.decode(value)
.or_else(|_| URL_SAFE_NO_PAD.decode(value))
decode_with_engine_fallback(value.as_bytes())
}
fn decode_urlsafe_bytes(value: &[u8]) -> Result<Vec<u8>, base64::DecodeError> {
URL_SAFE
.decode(value)
.or_else(|_| URL_SAFE_NO_PAD.decode(value))
decode_with_engine_fallback(value)
}
fn decode_with_engine_fallback(value: &[u8]) -> Result<Vec<u8>, base64::DecodeError> {
let mut decoded = Vec::with_capacity(decoded_len_estimate(value.len()));
match URL_SAFE.decode_vec(value, &mut decoded) {
Ok(()) => Ok(decoded),
Err(_) => {
decoded.clear();
URL_SAFE_NO_PAD.decode_vec(value, &mut decoded)?;
Ok(decoded)
}
}
}
#[cfg(test)]

22
crates/aether-provider-transport/src/snapshot_mapping.rs
View File

@@ -113,11 +113,12 @@ fn decrypt_secret(
ciphertext: &str,
field_name: &str,
) -> Result<String, DataLayerError> {
if should_use_plaintext_secret(ciphertext, field_name) {
return Ok(ciphertext.trim().to_string());
}
match decrypt_python_fernet_ciphertext(encryption_key, ciphertext) {
Ok(value) => Ok(value),
Err(_error) if should_use_plaintext_secret(ciphertext, field_name) => {
Ok(ciphertext.trim().to_string())
}
Err(error) => {
for fallback_encryption_key in fallback_encryption_keys {
if let Ok(value) =
@@ -156,14 +157,19 @@ fn should_use_plaintext_secret(ciphertext: &str, field_name: &str) -> bool {
if ciphertext.is_empty() {
return false;
}
if looks_like_python_fernet_ciphertext(ciphertext) {
return false;
}
match field_name {
"provider_api_keys.api_key" => !ciphertext.starts_with('{') && !ciphertext.starts_with('['),
"provider_api_keys.api_key" => {
if ciphertext.starts_with('{') || ciphertext.starts_with('[') {
return false;
}
!looks_like_python_fernet_ciphertext(ciphertext)
}
"provider_api_keys.auth_config" => {
ciphertext.starts_with('{') || ciphertext.starts_with('[')
if ciphertext.starts_with('{') || ciphertext.starts_with('[') {
return true;
}
false
}
_ => false,
}

333
crates/aether-scheduler-core/src/candidate.rs
View File

@@ -45,6 +45,20 @@ pub struct BuildMinimalCandidateSelectionInput<'a> {
pub priority_mode: SchedulerPriorityMode,
}
#[derive(Debug, Clone, Copy)]
struct RequiredCapabilityDescriptor<'a> {
name: &'a str,
compatible: bool,
}
#[derive(Debug, Clone, Copy)]
struct CandidateOrderingState {
capability_priority: (u32, u32),
affinity_hash: Option<u64>,
health_bucket: Option<crate::ProviderKeyHealthBucket>,
health_score: f64,
}
pub fn candidate_supports_required_capability(
candidate: &SchedulerMinimalCandidateSelectionCandidate,
required_capability: &str,
@@ -91,23 +105,17 @@ pub fn requested_capability_priority_for_candidate(
return (0, 0);
};
let mut exclusive_misses = 0u32;
let mut compatible_misses = 0u32;
for (capability, value) in required_capabilities {
if !requested_capability_is_enabled(value) {
continue;
}
if candidate_supports_required_capability(candidate, capability) {
continue;
}
if requested_capability_is_compatible(capability) {
compatible_misses += 1;
} else {
exclusive_misses += 1;
}
}
(exclusive_misses, compatible_misses)
requested_capability_priority_for_candidate_descriptors(
required_capabilities
.iter()
.filter_map(|(capability, value)| {
requested_capability_is_enabled(value).then_some(RequiredCapabilityDescriptor {
name: capability.as_str(),
compatible: requested_capability_is_compatible(capability),
})
}),
candidate,
)
}
pub fn auth_api_key_concurrency_limit_reached(
@@ -153,7 +161,8 @@ pub fn build_minimal_candidate_selection(
return Ok(Vec::new());
}
let mut candidates = Vec::new();
let required_capabilities = enabled_required_capabilities(required_capabilities);
let mut candidates = Vec::with_capacity(rows.len());
for row in rows {
if !crate::auth_constraints_allow_provider(
auth_constraints,
@@ -196,16 +205,9 @@ pub fn build_minimal_candidate_selection(
});
}
candidates.sort_by(|left, right| {
requested_capability_priority_for_candidate(required_capabilities, left)
.cmp(&requested_capability_priority_for_candidate(
required_capabilities,
right,
))
.then_with(|| {
compare_candidates_by_priority_mode(left, right, priority_mode, affinity_key)
})
});
let ordering_states =
build_candidate_ordering_states(&candidates, &required_capabilities, affinity_key, None);
sort_candidates_by_ordering_state(&mut candidates, &ordering_states, priority_mode, false);
Ok(candidates)
}
@@ -298,28 +300,27 @@ pub fn collect_selectable_candidates_from_keys(
selectable_keys: &BTreeSet<(String, String, String)>,
cached_affinity_target: Option<&crate::SchedulerAffinityTarget>,
) -> Vec<SchedulerMinimalCandidateSelectionCandidate> {
let mut selected = Vec::new();
let mut promoted = None;
let mut selected = Vec::with_capacity(candidates.len());
let mut emitted_keys = BTreeSet::new();
if let Some(target) = cached_affinity_target {
if let Some(candidate) = candidates
.iter()
.find(|candidate| crate::matches_affinity_target(candidate, target))
.cloned()
{
let key = crate::candidate_key(&candidate);
if selectable_keys.contains(&key) && emitted_keys.insert(key) {
selected.push(candidate);
}
}
}
for candidate in candidates {
let key = crate::candidate_key(&candidate);
if !selectable_keys.contains(&key) || !emitted_keys.insert(key) {
continue;
}
selected.push(candidate);
if promoted.is_none()
&& cached_affinity_target
.is_some_and(|target| crate::matches_affinity_target(&candidate, target))
{
promoted = Some(candidate);
} else {
selected.push(candidate);
}
}
if let Some(candidate) = promoted {
selected.insert(0, candidate);
}
selected
@@ -332,35 +333,14 @@ pub fn reorder_candidates_by_scheduler_health(
affinity_key: Option<&str>,
priority_mode: SchedulerPriorityMode,
) {
candidates.sort_by(|left, right| {
requested_capability_priority_for_candidate(required_capabilities, left)
.cmp(&requested_capability_priority_for_candidate(
required_capabilities,
right,
))
.then_with(|| match priority_mode {
SchedulerPriorityMode::Provider => left
.provider_priority
.cmp(&right.provider_priority)
.then(left.key_internal_priority.cmp(&right.key_internal_priority))
.then_with(|| {
compare_provider_key_health_order(left, right, provider_key_rpm_states)
})
.then_with(|| crate::compare_affinity_order(left, right, affinity_key))
.then_with(|| compare_candidate_identity(left, right)),
SchedulerPriorityMode::GlobalKey => left
.key_global_priority_for_format
.unwrap_or(i32::MAX)
.cmp(&right.key_global_priority_for_format.unwrap_or(i32::MAX))
.then_with(|| {
compare_provider_key_health_order(left, right, provider_key_rpm_states)
})
.then_with(|| crate::compare_affinity_order(left, right, affinity_key))
.then(left.provider_priority.cmp(&right.provider_priority))
.then(left.key_internal_priority.cmp(&right.key_internal_priority))
.then_with(|| compare_candidate_identity(left, right)),
})
});
let required_capabilities = enabled_required_capabilities(required_capabilities);
let ordering_states = build_candidate_ordering_states(
candidates,
&required_capabilities,
affinity_key,
Some(provider_key_rpm_states),
);
sort_candidates_by_ordering_state(candidates, &ordering_states, priority_mode, true);
}
#[derive(Clone, Copy, Debug)]
@@ -456,33 +436,6 @@ pub fn candidate_runtime_skip_reason_with_state(
None
}
fn compare_provider_key_health_order(
left: &SchedulerMinimalCandidateSelectionCandidate,
right: &SchedulerMinimalCandidateSelectionCandidate,
provider_key_rpm_states: &BTreeMap<String, StoredProviderCatalogKey>,
) -> std::cmp::Ordering {
let left_bucket = candidate_provider_key_health_bucket(left, provider_key_rpm_states);
let right_bucket = candidate_provider_key_health_bucket(right, provider_key_rpm_states);
right_bucket.cmp(&left_bucket).then_with(|| {
let left_score = candidate_provider_key_health_score(left, provider_key_rpm_states);
let right_score = candidate_provider_key_health_score(right, provider_key_rpm_states);
right_score
.partial_cmp(&left_score)
.unwrap_or(std::cmp::Ordering::Equal)
})
}
fn candidate_provider_key_health_bucket(
candidate: &SchedulerMinimalCandidateSelectionCandidate,
provider_key_rpm_states: &BTreeMap<String, StoredProviderCatalogKey>,
) -> Option<crate::ProviderKeyHealthBucket> {
provider_key_rpm_states
.get(&candidate.key_id)
.and_then(|key| {
crate::provider_key_health_bucket(key, candidate.endpoint_api_format.as_str())
})
}
fn compare_candidate_identity(
left: &SchedulerMinimalCandidateSelectionCandidate,
right: &SchedulerMinimalCandidateSelectionCandidate,
@@ -497,12 +450,190 @@ fn compare_candidate_identity(
)
}
fn enabled_required_capabilities(
required_capabilities: Option<&serde_json::Value>,
) -> Vec<RequiredCapabilityDescriptor<'_>> {
let Some(required_capabilities) = required_capabilities.and_then(serde_json::Value::as_object)
else {
return Vec::new();
};
required_capabilities
.iter()
.filter_map(|(capability, value)| {
requested_capability_is_enabled(value).then_some(RequiredCapabilityDescriptor {
name: capability.as_str(),
compatible: requested_capability_is_compatible(capability),
})
})
.collect()
}
fn requested_capability_priority_for_candidate_descriptors<'a, I>(
required_capabilities: I,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) -> (u32, u32)
where
I: IntoIterator<Item = RequiredCapabilityDescriptor<'a>>,
{
let mut exclusive_misses = 0u32;
let mut compatible_misses = 0u32;
for capability in required_capabilities {
if candidate_supports_required_capability(candidate, capability.name) {
continue;
}
if capability.compatible {
compatible_misses += 1;
} else {
exclusive_misses += 1;
}
}
(exclusive_misses, compatible_misses)
}
fn build_candidate_ordering_states(
candidates: &[SchedulerMinimalCandidateSelectionCandidate],
required_capabilities: &[RequiredCapabilityDescriptor<'_>],
affinity_key: Option<&str>,
provider_key_rpm_states: Option<&BTreeMap<String, StoredProviderCatalogKey>>,
) -> Vec<CandidateOrderingState> {
candidates
.iter()
.map(|candidate| CandidateOrderingState {
capability_priority: requested_capability_priority_for_candidate_descriptors(
required_capabilities.iter().copied(),
candidate,
),
affinity_hash: affinity_key.map(|key| crate::candidate_affinity_hash(key, candidate)),
health_bucket: provider_key_rpm_states.and_then(|states| {
states.get(&candidate.key_id).and_then(|key| {
crate::provider_key_health_bucket(key, candidate.endpoint_api_format.as_str())
})
}),
health_score: candidate_provider_key_health_score(candidate, provider_key_rpm_states),
})
.collect()
}
fn sort_candidates_by_ordering_state(
candidates: &mut [SchedulerMinimalCandidateSelectionCandidate],
ordering_states: &[CandidateOrderingState],
priority_mode: SchedulerPriorityMode,
include_health: bool,
) {
if candidates.len() < 2 {
return;
}
let mut order = (0..candidates.len()).collect::<Vec<_>>();
order.sort_by(|left, right| {
compare_candidates_with_ordering_state(
&ordering_states[*left],
&candidates[*left],
&ordering_states[*right],
&candidates[*right],
priority_mode,
include_health,
)
});
apply_candidate_order(candidates, order);
}
fn compare_candidates_with_ordering_state(
left_state: &CandidateOrderingState,
left_candidate: &SchedulerMinimalCandidateSelectionCandidate,
right_state: &CandidateOrderingState,
right_candidate: &SchedulerMinimalCandidateSelectionCandidate,
priority_mode: SchedulerPriorityMode,
include_health: bool,
) -> std::cmp::Ordering {
left_state
.capability_priority
.cmp(&right_state.capability_priority)
.then_with(|| {
compare_priority_before_health(left_candidate, right_candidate, priority_mode)
})
.then_with(|| {
if include_health {
compare_provider_key_health_state(left_state, right_state)
} else {
std::cmp::Ordering::Equal
}
})
.then_with(|| left_state.affinity_hash.cmp(&right_state.affinity_hash))
.then_with(|| {
compare_priority_after_affinity(left_candidate, right_candidate, priority_mode)
})
.then_with(|| compare_candidate_identity(left_candidate, right_candidate))
}
fn compare_priority_before_health(
left: &SchedulerMinimalCandidateSelectionCandidate,
right: &SchedulerMinimalCandidateSelectionCandidate,
priority_mode: SchedulerPriorityMode,
) -> std::cmp::Ordering {
match priority_mode {
SchedulerPriorityMode::Provider => left
.provider_priority
.cmp(&right.provider_priority)
.then(left.key_internal_priority.cmp(&right.key_internal_priority)),
SchedulerPriorityMode::GlobalKey => left
.key_global_priority_for_format
.unwrap_or(i32::MAX)
.cmp(&right.key_global_priority_for_format.unwrap_or(i32::MAX)),
}
}
fn compare_priority_after_affinity(
left: &SchedulerMinimalCandidateSelectionCandidate,
right: &SchedulerMinimalCandidateSelectionCandidate,
priority_mode: SchedulerPriorityMode,
) -> std::cmp::Ordering {
match priority_mode {
SchedulerPriorityMode::Provider => std::cmp::Ordering::Equal,
SchedulerPriorityMode::GlobalKey => left
.provider_priority
.cmp(&right.provider_priority)
.then(left.key_internal_priority.cmp(&right.key_internal_priority)),
}
}
fn compare_provider_key_health_state(
left: &CandidateOrderingState,
right: &CandidateOrderingState,
) -> std::cmp::Ordering {
right
.health_bucket
.cmp(&left.health_bucket)
.then_with(|| right.health_score.total_cmp(&left.health_score))
}
fn apply_candidate_order(
candidates: &mut [SchedulerMinimalCandidateSelectionCandidate],
sorted_old_indices: Vec<usize>,
) {
let mut target_positions = vec![0usize; sorted_old_indices.len()];
for (new_position, old_position) in sorted_old_indices.into_iter().enumerate() {
target_positions[old_position] = new_position;
}
for index in 0..candidates.len() {
let current = index;
while target_positions[current] != current {
let target = target_positions[current];
candidates.swap(current, target);
target_positions.swap(current, target);
}
}
}
fn candidate_provider_key_health_score(
candidate: &SchedulerMinimalCandidateSelectionCandidate,
provider_key_rpm_states: &BTreeMap<String, StoredProviderCatalogKey>,
provider_key_rpm_states: Option<&BTreeMap<String, StoredProviderCatalogKey>>,
) -> f64 {
provider_key_rpm_states
.get(&candidate.key_id)
.and_then(|states| states.get(&candidate.key_id))
.and_then(|key| {
crate::effective_provider_key_health_score(key, candidate.endpoint_api_format.as_str())
})

84
crates/aether-scheduler-core/src/model.rs
View File

@@ -44,15 +44,17 @@ fn resolve_global_model_name_by<F>(
where
F: Fn(&StoredMinimalCandidateSelectionRow) -> bool,
{
let mut matches = rows
.iter()
.filter(|row| matches(row))
.map(|row| row.global_model_name.trim())
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.collect::<BTreeSet<_>>()
.into_iter();
matches.next()
let mut best_match = None::<&str>;
for row in rows.iter().filter(|row| matches(row)) {
let candidate = row.global_model_name.trim();
if candidate.is_empty() {
continue;
}
if best_match.is_none_or(|current| candidate < current) {
best_match = Some(candidate);
}
}
best_match.map(ToOwned::to_owned)
}
pub fn resolve_provider_model_name(
@@ -75,25 +77,26 @@ pub fn resolve_provider_model_name(
return Some((selected_provider_model_name, None));
}
let candidate_models = candidate_model_names(row, api_format);
let mut sorted_allowed_models = key_allowed_models
.iter()
.map(|value| value.trim())
.map(String::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.collect::<Vec<_>>();
sorted_allowed_models.sort();
sorted_allowed_models.sort_unstable();
for allowed_model in &sorted_allowed_models {
if candidate_models.contains(allowed_model.as_str()) {
return Some((allowed_model.clone(), Some(allowed_model.clone())));
for &allowed_model in &sorted_allowed_models {
if row_has_candidate_model_name(row, api_format, allowed_model) {
let allowed_model = allowed_model.to_owned();
return Some((allowed_model.clone(), Some(allowed_model)));
}
}
let global_model_mappings = row.global_model_mappings.as_ref()?;
for allowed_model in sorted_allowed_models {
for &allowed_model in &sorted_allowed_models {
for pattern in global_model_mappings {
if matches_model_mapping(pattern, &allowed_model) {
if matches_model_mapping(pattern, allowed_model) {
let allowed_model = allowed_model.to_owned();
return Some((allowed_model.clone(), Some(allowed_model)));
}
}
@@ -110,23 +113,14 @@ pub fn select_provider_model_name(
return row.model_provider_model_name.clone();
};
let mut scoped = mappings
mappings
.iter()
.filter(|mapping| mapping_scope_matches(mapping, api_format))
.collect::<Vec<_>>();
if scoped.is_empty() {
return row.model_provider_model_name.clone();
}
scoped.sort_by(|left, right| {
left.priority
.cmp(&right.priority)
.then(left.name.cmp(&right.name))
});
let top_priority = scoped[0].priority;
scoped
.into_iter()
.find(|mapping| mapping.priority == top_priority)
.min_by(|left, right| {
left.priority
.cmp(&right.priority)
.then(left.name.cmp(&right.name))
})
.map(|mapping| mapping.name.clone())
.unwrap_or_else(|| row.model_provider_model_name.clone())
}
@@ -153,7 +147,7 @@ fn mapping_scope_matches(mapping: &StoredProviderModelMapping, api_format: &str)
api_formats
.iter()
.any(|value| normalize_api_format(value) == api_format)
.any(|value| api_format_matches(value, api_format))
}
pub fn row_supports_required_capability(
@@ -230,7 +224,7 @@ pub fn extract_global_priority_for_format(
let Some(value) = object
.iter()
.find(|(key, _)| normalize_api_format(key) == api_format)
.find(|(key, _)| api_format_matches(key, api_format))
.map(|(_, value)| value)
else {
return Ok(None);
@@ -262,6 +256,26 @@ pub fn normalize_api_format(value: &str) -> String {
value.trim().to_ascii_lowercase()
}
fn row_has_candidate_model_name(
row: &StoredMinimalCandidateSelectionRow,
api_format: &str,
model_name: &str,
) -> bool {
row.model_provider_model_name == model_name
|| row
.model_provider_model_mappings
.as_ref()
.is_some_and(|mappings| {
mappings.iter().any(|mapping| {
mapping_scope_matches(mapping, api_format) && mapping.name == model_name
})
})
}
fn api_format_matches(left: &str, right: &str) -> bool {
left.trim().eq_ignore_ascii_case(right.trim())
}
#[cfg(test)]
mod tests {
use super::matches_model_mapping;

108
crates/aether-scheduler-core/src/request_candidate.rs
View File

@@ -2,7 +2,7 @@ use aether_contracts::{ExecutionError, ExecutionPlan};
use aether_data_contracts::repository::candidates::{
RequestCandidateStatus, StoredRequestCandidate, UpsertRequestCandidateRecord,
};
use serde_json::Value;
use serde_json::{Map, Value};
#[derive(Debug, Clone, PartialEq)]
pub struct SchedulerRequestCandidateReportContext {
@@ -91,20 +91,13 @@ pub fn parse_request_candidate_report_context(
report_context: Option<&Value>,
) -> Option<SchedulerRequestCandidateReportContext> {
let report_context = report_context?;
let retry_index = report_context
.get("retry_index")
.and_then(Value::as_u64)
.unwrap_or_default();
Some(SchedulerRequestCandidateReportContext {
request_id: string_field(report_context, "request_id"),
candidate_id: string_field(report_context, "candidate_id"),
user_id: string_field(report_context, "user_id"),
api_key_id: string_field(report_context, "api_key_id"),
candidate_index: report_context
.get("candidate_index")
.and_then(Value::as_u64)
.and_then(|value| u32::try_from(value).ok()),
retry_index: u32::try_from(retry_index).unwrap_or(u32::MAX),
candidate_index: u32_field(report_context, "candidate_index"),
retry_index: u32_field(report_context, "retry_index").unwrap_or_default(),
provider_id: string_field(report_context, "provider_id"),
endpoint_id: string_field(report_context, "endpoint_id"),
key_id: string_field(report_context, "key_id"),
@@ -123,9 +116,24 @@ pub fn resolve_report_request_candidate_slot(
now_unix_ms: u64,
generated_candidate_id: String,
) -> Option<SchedulerResolvedReportRequestCandidateSlot> {
let request_id = metadata.request_id.clone()?;
let matched_candidate = match_existing_report_candidate(existing_candidates, &metadata);
let synthesized_extra_data = build_report_candidate_extra_data(&metadata);
let SchedulerRequestCandidateReportContext {
request_id,
candidate_id,
user_id,
api_key_id,
candidate_index: metadata_candidate_index,
retry_index,
provider_id,
endpoint_id,
key_id,
client_api_format,
provider_api_format,
proxy,
} = metadata;
let request_id = request_id?;
let synthesized_extra_data =
build_report_candidate_extra_data(client_api_format, provider_api_format, proxy);
let created_at_unix_ms = matched_candidate
.as_ref()
.map(|candidate| candidate.created_at_unix_ms)
@@ -133,42 +141,42 @@ pub fn resolve_report_request_candidate_slot(
let candidate_index = matched_candidate
.as_ref()
.map(|candidate| candidate.candidate_index)
.or(metadata.candidate_index)
.or(metadata_candidate_index)
.unwrap_or_else(|| next_candidate_index(existing_candidates));
let retry_index = matched_candidate
.as_ref()
.map(|candidate| candidate.retry_index)
.unwrap_or(metadata.retry_index);
.unwrap_or(retry_index);
Some(SchedulerResolvedReportRequestCandidateSlot {
id: matched_candidate
.as_ref()
.map(|candidate| candidate.id.clone())
.or(metadata.candidate_id)
.or(candidate_id)
.unwrap_or(generated_candidate_id),
request_id,
user_id: matched_candidate
.as_ref()
.and_then(|candidate| candidate.user_id.clone())
.or(metadata.user_id),
.or(user_id),
api_key_id: matched_candidate
.as_ref()
.and_then(|candidate| candidate.api_key_id.clone())
.or(metadata.api_key_id),
.or(api_key_id),
candidate_index,
retry_index,
provider_id: matched_candidate
.as_ref()
.and_then(|candidate| candidate.provider_id.clone())
.or(metadata.provider_id),
.or(provider_id),
endpoint_id: matched_candidate
.as_ref()
.and_then(|candidate| candidate.endpoint_id.clone())
.or(metadata.endpoint_id),
.or(endpoint_id),
key_id: matched_candidate
.as_ref()
.and_then(|candidate| candidate.key_id.clone())
.or(metadata.key_id),
.or(key_id),
extra_data: merge_request_candidate_extra_data(
matched_candidate
.as_ref()
@@ -195,22 +203,14 @@ pub fn build_execution_request_candidate_seed(
.and_then(Value::as_object)
.cloned()
.unwrap_or_default();
let request_id = string_field(&Value::Object(context.clone()), "request_id")
.unwrap_or_else(|| plan.request_id.clone());
let candidate_index = context
.get("candidate_index")
.and_then(Value::as_u64)
.and_then(|value| u32::try_from(value).ok())
.unwrap_or(0);
let retry_index = context
.get("retry_index")
.and_then(Value::as_u64)
.and_then(|value| u32::try_from(value).ok())
.unwrap_or(0);
let candidate_id = string_field(&Value::Object(context.clone()), "candidate_id")
.unwrap_or(generated_candidate_id);
let user_id = string_field(&Value::Object(context.clone()), "user_id");
let api_key_id = string_field(&Value::Object(context.clone()), "api_key_id");
let request_id =
string_field_from_object(&context, "request_id").unwrap_or_else(|| plan.request_id.clone());
let candidate_index = u32_field_from_object(&context, "candidate_index").unwrap_or(0);
let retry_index = u32_field_from_object(&context, "retry_index").unwrap_or(0);
let candidate_id =
string_field_from_object(&context, "candidate_id").unwrap_or(generated_candidate_id);
let user_id = string_field_from_object(&context, "user_id");
let api_key_id = string_field_from_object(&context, "api_key_id");
context.insert("request_id".to_string(), Value::String(request_id.clone()));
context.insert(
@@ -374,7 +374,10 @@ pub fn finalize_execution_request_candidate_report_context(
report_context: Value,
candidate_id: &str,
) -> Value {
let mut context = report_context.as_object().cloned().unwrap_or_default();
let mut context = match report_context {
Value::Object(context) => context,
_ => Map::new(),
};
let candidate_id = candidate_id.trim();
if !candidate_id.is_empty() {
context.insert(
@@ -410,6 +413,12 @@ fn extract_error_message(body_json: &Value) -> Option<&str> {
fn string_field(value: &Value, key: &str) -> Option<String> {
value
.as_object()
.and_then(|object| string_field_from_object(object, key))
}
fn string_field_from_object(object: &Map<String, Value>, key: &str) -> Option<String> {
object
.get(key)
.and_then(Value::as_str)
.map(str::trim)
@@ -417,6 +426,19 @@ fn string_field(value: &Value, key: &str) -> Option<String> {
.map(ToOwned::to_owned)
}
fn u32_field(value: &Value, key: &str) -> Option<u32> {
value
.as_object()
.and_then(|object| u32_field_from_object(object, key))
}
fn u32_field_from_object(object: &Map<String, Value>, key: &str) -> Option<u32> {
object
.get(key)
.and_then(Value::as_u64)
.and_then(|value| u32::try_from(value).ok())
}
fn match_existing_report_candidate<'a>(
candidates: &'a [StoredRequestCandidate],
metadata: &SchedulerRequestCandidateReportContext,
@@ -465,24 +487,26 @@ fn next_candidate_index(candidates: &[StoredRequestCandidate]) -> u32 {
}
fn build_report_candidate_extra_data(
metadata: &SchedulerRequestCandidateReportContext,
client_api_format: Option<String>,
provider_api_format: Option<String>,
proxy: Option<Value>,
) -> Option<Value> {
let mut extra_data = serde_json::Map::new();
let mut extra_data = Map::with_capacity(5);
extra_data.insert("gateway_execution_runtime".to_string(), Value::Bool(true));
extra_data.insert("phase".to_string(), Value::String("3c_trial".to_string()));
if let Some(client_api_format) = metadata.client_api_format.clone() {
if let Some(client_api_format) = client_api_format {
extra_data.insert(
"client_api_format".to_string(),
Value::String(client_api_format),
);
}
if let Some(provider_api_format) = metadata.provider_api_format.clone() {
if let Some(provider_api_format) = provider_api_format {
extra_data.insert(
"provider_api_format".to_string(),
Value::String(provider_api_format),
);
}
if let Some(proxy) = metadata.proxy.clone() {
if let Some(proxy) = proxy {
extra_data.insert("proxy".to_string(), proxy);
}
(!extra_data.is_empty()).then_some(Value::Object(extra_data))

416
crates/aether-usage-runtime/src/body_capture.rs
View File

@@ -1,6 +1,9 @@
use std::io::{self, Write};
use aether_data_contracts::repository::usage::{
UpsertUsageRecord, UsageBodyCaptureState, UsageBodyField,
};
use serde::Serialize;
use serde_json::{json, Map, Value};
use crate::event::UsageEvent;
@@ -76,6 +79,22 @@ pub struct UsageBodyCaptureEngine {
policy: UsageBodyCapturePolicy,
}
#[derive(Default)]
struct CountingWriter {
bytes: u64,
}
impl Write for CountingWriter {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
self.bytes = self.bytes.saturating_add(buf.len() as u64);
Ok(buf.len())
}
fn flush(&mut self) -> io::Result<()> {
Ok(())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct RuntimeBodyCaptureStates {
pub request: UsageBodyCaptureState,
@@ -286,9 +305,7 @@ fn limit_usage_body_capture_value(
value: Value,
max_bytes: Option<usize>,
) -> LimitedUsageBodyCapture {
let source_bytes = serde_json::to_vec(&value)
.ok()
.map(|bytes| bytes.len() as u64);
let source_bytes = json_serialized_len(&value);
let Some(limit) = max_bytes.filter(|value| *value > 0) else {
return LimitedUsageBodyCapture {
stored_bytes: source_bytes,
@@ -327,9 +344,7 @@ fn limit_usage_body_capture_value(
"value_kind": usage_value_kind(&other),
}),
};
let stored_bytes = serde_json::to_vec(&truncated_value)
.ok()
.map(|bytes| bytes.len() as u64);
let stored_bytes = json_serialized_len(&truncated_value);
LimitedUsageBodyCapture {
value: truncated_value,
source_bytes: Some(source_len),
@@ -340,13 +355,6 @@ fn limit_usage_body_capture_value(
}
fn truncate_usage_body_string(value: &str, max_bytes: usize) -> String {
if serde_json::to_vec(&Value::String(value.to_string()))
.ok()
.is_some_and(|bytes| bytes.len() <= max_bytes)
{
return value.to_string();
}
let mut end = value.len();
while end > 0 {
while end > 0 && !value.is_char_boundary(end) {
@@ -354,10 +362,7 @@ fn truncate_usage_body_string(value: &str, max_bytes: usize) -> String {
}
let mut candidate = value[..end].to_string();
candidate.push_str(TRUNCATED_BODY_STRING_SUFFIX);
if serde_json::to_vec(&Value::String(candidate.clone()))
.ok()
.is_some_and(|bytes| bytes.len() <= max_bytes)
{
if json_serialized_len(&candidate).is_some_and(|bytes| bytes <= max_bytes as u64) {
return candidate;
}
end = value[..end]
@@ -379,27 +384,45 @@ fn truncate_usage_body_string(value: &str, max_bytes: usize) -> String {
.to_string()
}
fn json_serialized_len<T: Serialize>(value: &T) -> Option<u64> {
let mut writer = CountingWriter::default();
serde_json::to_writer(&mut writer, value).ok()?;
Some(writer.bytes)
}
pub(crate) fn sync_usage_body_ref_metadata(
metadata: &mut Option<Value>,
field: UsageBodyField,
body_ref: Option<&str>,
) {
let key = field.as_ref_key();
let Some(body_ref) = body_ref.map(str::trim).filter(|value| !value.is_empty()) else {
if let Some(object) = metadata.as_mut().and_then(Value::as_object_mut) {
object.remove(field.as_ref_key());
let clear_metadata = match metadata.as_mut() {
Some(Value::Object(object)) => {
object.remove(key);
object.is_empty()
}
_ => false,
};
if clear_metadata {
*metadata = None;
}
return;
};
if let Some(Value::Object(object)) = metadata.as_mut() {
if object.get(key).and_then(Value::as_str) == Some(body_ref) {
return;
}
object.insert(key.to_owned(), Value::String(body_ref.to_owned()));
return;
}
let object = metadata
.get_or_insert_with(|| Value::Object(Map::new()))
.as_object_mut();
let Some(object) = object else {
return;
};
object.insert(
field.as_ref_key().to_string(),
Value::String(body_ref.to_string()),
);
object.insert(key.to_owned(), Value::String(body_ref.to_owned()));
}
pub(crate) fn build_payload_body_capture_metadata(
@@ -408,36 +431,44 @@ pub(crate) fn build_payload_body_capture_metadata(
provider_body_state: Option<UsageBodyCaptureState>,
client_body_state: Option<UsageBodyCaptureState>,
) -> Option<Value> {
let mut metadata = Map::new();
if let Some(decoded_len) = provider_body_base64.and_then(decoded_base64_len_hint) {
let provider_decoded_len = provider_body_base64.and_then(decoded_base64_len_hint);
let client_decoded_len = client_body_base64.and_then(decoded_base64_len_hint);
let body_capture_capacity =
usize::from(provider_body_state.is_some()) + usize::from(client_body_state.is_some());
let mut metadata = Map::with_capacity(
usize::from(provider_decoded_len.is_some())
+ usize::from(client_decoded_len.is_some())
+ usize::from(body_capture_capacity > 0),
);
if let Some(decoded_len) = provider_decoded_len {
metadata.insert(
"provider_response_body_base64_bytes".to_string(),
Value::Number(decoded_len.into()),
);
}
if let Some(decoded_len) = client_body_base64.and_then(decoded_base64_len_hint) {
if let Some(decoded_len) = client_decoded_len {
metadata.insert(
"client_response_body_base64_bytes".to_string(),
Value::Number(decoded_len.into()),
);
}
let mut body_capture = Map::new();
append_body_capture_metadata_entry(
&mut body_capture,
"response",
provider_body_state,
provider_body_base64.and_then(decoded_base64_len_hint),
provider_body_base64.and_then(decoded_base64_len_hint),
);
append_body_capture_metadata_entry(
&mut body_capture,
"client_response",
client_body_state,
client_body_base64.and_then(decoded_base64_len_hint),
client_body_base64.and_then(decoded_base64_len_hint),
);
if !body_capture.is_empty() {
if body_capture_capacity > 0 {
let mut body_capture = Map::with_capacity(body_capture_capacity);
append_body_capture_metadata_entry(
&mut body_capture,
"response",
provider_body_state,
provider_decoded_len,
provider_decoded_len,
);
append_body_capture_metadata_entry(
&mut body_capture,
"client_response",
client_body_state,
client_decoded_len,
client_decoded_len,
);
metadata.insert("body_capture".to_string(), Value::Object(body_capture));
}
@@ -476,21 +507,29 @@ pub(crate) fn append_runtime_body_capture_metadata(
input.provider_request_body_ref,
input.provider_request_unavailable,
);
upsert_body_capture_metadata_entry(metadata, "request", Some(states.request), None, None, None);
upsert_body_capture_metadata_entry(
metadata,
"provider_request",
Some(states.provider_request),
input.provider_request_source_bytes,
input.provider_request_source_bytes,
input.provider_request_unavailable_reason,
let Some(body_capture_object) = body_capture_object_mut(metadata, 2) else {
return;
};
body_capture_object.insert(
"request".to_string(),
build_body_capture_metadata_entry(states.request, None, None, None),
);
body_capture_object.insert(
"provider_request".to_string(),
build_body_capture_metadata_entry(
states.provider_request,
input.provider_request_source_bytes,
input.provider_request_source_bytes,
input.provider_request_unavailable_reason,
),
);
}
pub(crate) fn build_plan_body_capture_metadata(
provider_request_body_base64: Option<&str>,
) -> Option<Value> {
let mut metadata = Map::new();
provider_request_body_base64?;
let mut metadata = Map::with_capacity(2);
append_plan_body_capture_metadata(&mut metadata, provider_request_body_base64);
(!metadata.is_empty()).then_some(Value::Object(metadata))
}
@@ -507,13 +546,17 @@ pub(crate) fn append_plan_body_capture_metadata(
Value::Number(decoded_len.into()),
);
}
upsert_body_capture_metadata_entry(
metadata,
"provider_request",
Some(UsageBodyCaptureState::Unavailable),
decoded_len,
decoded_len,
Some("body_bytes_base64_only"),
let Some(body_capture_object) = body_capture_object_mut(metadata, 1) else {
return;
};
body_capture_object.insert(
"provider_request".to_string(),
build_body_capture_metadata_entry(
UsageBodyCaptureState::Unavailable,
decoded_len,
decoded_len,
Some("body_bytes_base64_only"),
),
);
}
}
@@ -528,58 +571,16 @@ fn append_body_capture_metadata_entry(
let Some(state) = state else {
return;
};
let mut entry = Map::new();
entry.insert(
"state".to_string(),
Value::String(state.as_str().to_string()),
target.insert(
key.to_string(),
build_body_capture_metadata_entry(
state,
stored_bytes,
source_bytes,
matches!(state, UsageBodyCaptureState::Truncated)
.then_some("body_capture_limit_exceeded"),
),
);
if let Some(stored_bytes) = stored_bytes {
entry.insert("stored_bytes".to_string(), json!(stored_bytes));
}
if let Some(source_bytes) = source_bytes {
entry.insert("source_bytes".to_string(), json!(source_bytes));
}
if matches!(state, UsageBodyCaptureState::Truncated) {
entry.insert(
"reason".to_string(),
Value::String("body_capture_limit_exceeded".to_string()),
);
}
target.insert(key.to_string(), Value::Object(entry));
}
pub(crate) fn upsert_body_capture_metadata_entry(
metadata: &mut Map<String, Value>,
key: &str,
state: Option<UsageBodyCaptureState>,
stored_bytes: Option<u64>,
source_bytes: Option<u64>,
reason: Option<&str>,
) {
let body_capture = metadata
.entry("body_capture".to_string())
.or_insert_with(|| Value::Object(Map::new()));
let Some(body_capture_object) = body_capture.as_object_mut() else {
return;
};
let Some(state) = state else {
return;
};
let mut entry = Map::new();
entry.insert(
"state".to_string(),
Value::String(state.as_str().to_string()),
);
if let Some(bytes) = stored_bytes {
entry.insert("stored_bytes".to_string(), json!(bytes));
}
if let Some(bytes) = source_bytes {
entry.insert("source_bytes".to_string(), json!(bytes));
}
if let Some(reason) = reason {
entry.insert("reason".to_string(), Value::String(reason.to_string()));
}
body_capture_object.insert(key.to_string(), Value::Object(entry));
}
fn upsert_body_capture_metadata_value_entry(
@@ -593,22 +594,63 @@ fn upsert_body_capture_metadata_value_entry(
let Some(state) = state else {
return;
};
let metadata_object = metadata
.get_or_insert_with(|| Value::Object(Map::new()))
.as_object_mut();
let Some(metadata_object) = metadata_object else {
let Some(body_capture_object) = body_capture_value_object_mut(metadata, 1) else {
return;
};
upsert_body_capture_metadata_entry(
metadata_object,
key,
Some(state),
stored_bytes,
source_bytes,
reason,
body_capture_object.insert(
key.to_string(),
build_body_capture_metadata_entry(state, stored_bytes, source_bytes, reason),
);
}
fn body_capture_object_mut(
metadata: &mut Map<String, Value>,
capacity: usize,
) -> Option<&mut Map<String, Value>> {
let body_capture = metadata
.entry("body_capture".to_string())
.or_insert_with(|| Value::Object(Map::with_capacity(capacity)));
body_capture.as_object_mut()
}
fn body_capture_value_object_mut(
metadata: &mut Option<Value>,
capacity: usize,
) -> Option<&mut Map<String, Value>> {
let metadata_object = metadata
.get_or_insert_with(|| Value::Object(Map::with_capacity(1)))
.as_object_mut();
let metadata_object = metadata_object?;
body_capture_object_mut(metadata_object, capacity)
}
fn build_body_capture_metadata_entry(
state: UsageBodyCaptureState,
stored_bytes: Option<u64>,
source_bytes: Option<u64>,
reason: Option<&str>,
) -> Value {
let mut entry = Map::with_capacity(
1 + usize::from(stored_bytes.is_some())
+ usize::from(source_bytes.is_some())
+ usize::from(reason.is_some()),
);
entry.insert(
"state".to_string(),
Value::String(state.as_str().to_owned()),
);
if let Some(bytes) = stored_bytes {
entry.insert("stored_bytes".to_string(), json!(bytes));
}
if let Some(bytes) = source_bytes {
entry.insert("source_bytes".to_string(), json!(bytes));
}
if let Some(reason) = reason {
entry.insert("reason".to_string(), Value::String(reason.to_owned()));
}
Value::Object(entry)
}
pub(crate) fn decoded_base64_len_hint(body_base64: &str) -> Option<u64> {
let body_base64 = body_base64.trim();
if body_base64.is_empty() {
@@ -642,9 +684,18 @@ pub(crate) fn decoded_base64_len_hint(body_base64: &str) -> Option<u64> {
}
fn sanitize_usage_body_ref(value: Option<String>) -> Option<String> {
value
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
value.and_then(trim_owned_non_empty_string)
}
fn trim_owned_non_empty_string(value: String) -> Option<String> {
let trimmed = value.trim();
if trimmed.is_empty() {
return None;
}
if trimmed.len() == value.len() {
return Some(value);
}
Some(trimmed.to_string())
}
fn usage_value_kind(value: &Value) -> &'static str {
@@ -657,3 +708,122 @@ fn usage_value_kind(value: &Value) -> &'static str {
Value::Object(_) => "object",
}
}
#[cfg(test)]
mod tests {
use super::{
build_plan_body_capture_metadata, sync_usage_body_ref_metadata,
trim_owned_non_empty_string, truncate_usage_body_string,
upsert_body_capture_metadata_value_entry,
};
use aether_data_contracts::repository::usage::UsageBodyCaptureState;
use aether_data_contracts::repository::usage::UsageBodyField;
use serde_json::{Map, Value};
#[test]
fn build_plan_body_capture_metadata_returns_none_without_base64_body() {
assert!(build_plan_body_capture_metadata(None).is_none());
}
#[test]
fn trim_owned_non_empty_string_preserves_clean_values_and_drops_blank_ones() {
assert_eq!(
trim_owned_non_empty_string("blob://body-ref-1".to_string()),
Some("blob://body-ref-1".to_string()),
);
assert_eq!(
trim_owned_non_empty_string(" blob://body-ref-1 ".to_string()),
Some("blob://body-ref-1".to_string()),
);
assert_eq!(trim_owned_non_empty_string(" ".to_string()), None);
}
#[test]
fn upsert_body_capture_metadata_value_entry_ignores_none_state() {
let mut metadata = Some(Value::Object(Map::<String, Value>::new()));
upsert_body_capture_metadata_value_entry(&mut metadata, "response", None, None, None, None);
assert_eq!(metadata, Some(Value::Object(Map::new())));
}
#[test]
fn upsert_body_capture_metadata_value_entry_preserves_existing_metadata_fields() {
let mut metadata = Some(Value::Object(Map::from_iter([(
"request_body_ref".to_string(),
Value::String("blob://body-ref-1".to_string()),
)])));
upsert_body_capture_metadata_value_entry(
&mut metadata,
"response",
Some(UsageBodyCaptureState::Reference),
None,
None,
None,
);
assert_eq!(
metadata,
Some(Value::Object(Map::from_iter([
(
"request_body_ref".to_string(),
Value::String("blob://body-ref-1".to_string()),
),
(
"body_capture".to_string(),
Value::Object(Map::from_iter([(
"response".to_string(),
Value::Object(Map::from_iter([(
"state".to_string(),
Value::String("reference".to_string()),
)])),
)])),
),
]))),
);
}
#[test]
fn sync_usage_body_ref_metadata_clears_empty_metadata_object() {
let mut metadata = Some(Value::Object(Map::from_iter([(
"request_body_ref".to_string(),
Value::String("blob://body-ref-1".to_string()),
)])));
sync_usage_body_ref_metadata(&mut metadata, UsageBodyField::RequestBody, None);
assert!(metadata.is_none());
}
#[test]
fn sync_usage_body_ref_metadata_preserves_existing_ref_value() {
let mut metadata = Some(Value::Object(Map::from_iter([(
"request_body_ref".to_string(),
Value::String("blob://body-ref-1".to_string()),
)])));
sync_usage_body_ref_metadata(
&mut metadata,
UsageBodyField::RequestBody,
Some("blob://body-ref-1"),
);
assert_eq!(
metadata,
Some(Value::Object(Map::from_iter([(
"request_body_ref".to_string(),
Value::String("blob://body-ref-1".to_string()),
)]))),
);
}
#[test]
fn truncate_usage_body_string_respects_json_byte_limit() {
let limit = 32usize;
let truncated = truncate_usage_body_string("x".repeat(256).as_str(), limit);
assert!(truncated.ends_with("...[truncated]"));
assert!(serde_json::to_vec(&truncated)
.ok()
.is_some_and(|bytes| bytes.len() <= limit));
}
}

147
crates/aether-usage-runtime/src/request_metadata.rs
View File

@@ -31,13 +31,36 @@ pub(crate) fn merge_usage_request_metadata(
(!metadata.is_empty()).then_some(Value::Object(metadata))
}
pub(crate) fn merge_usage_request_metadata_owned(
base: Option<Value>,
override_value: Option<Value>,
) -> Option<Value> {
let mut metadata = match base {
Some(Value::Object(base)) => base,
_ => Map::new(),
};
if let Some(Value::Object(override_object)) = override_value {
move_allowed_metadata_fields(override_object, &mut metadata);
}
(!metadata.is_empty()).then_some(Value::Object(metadata))
}
pub(crate) fn sanitize_usage_request_metadata(value: Option<Value>) -> Option<Value> {
let Value::Object(object) = value? else {
return None;
};
let mut filtered = Map::new();
copy_allowed_metadata_fields(&object, &mut filtered);
move_allowed_metadata_fields(object, &mut filtered);
(!filtered.is_empty()).then_some(Value::Object(filtered))
}
pub(crate) fn sanitize_usage_request_metadata_ref(value: Option<&Value>) -> Option<Value> {
let object = value.and_then(Value::as_object)?;
let mut filtered = Map::new();
copy_allowed_metadata_fields(object, &mut filtered);
(!filtered.is_empty()).then_some(Value::Object(filtered))
}
@@ -62,6 +85,26 @@ fn copy_allowed_metadata_fields(source: &Map<String, Value>, target: &mut Map<St
copy_number(source, target, "price_per_request");
}
fn move_allowed_metadata_fields(mut source: Map<String, Value>, target: &mut Map<String, Value>) {
remove_non_empty_string(&mut source, target, "trace_id");
remove_number(&mut source, target, "provider_request_body_base64_bytes");
remove_number(&mut source, target, "provider_response_body_base64_bytes");
remove_number(&mut source, target, "client_response_body_base64_bytes");
remove_non_null_value(&mut source, target, "billing_snapshot");
remove_non_empty_string(&mut source, target, "billing_snapshot_schema_version");
remove_non_empty_string(&mut source, target, "billing_snapshot_status");
remove_non_null_value(&mut source, target, "dimensions");
remove_non_null_value(&mut source, target, "billing_rule_snapshot");
remove_non_null_value(&mut source, target, "scheduling_audit");
remove_number(&mut source, target, "rate_multiplier");
remove_bool(&mut source, target, "is_free_tier");
remove_number(&mut source, target, "input_price_per_1m");
remove_number(&mut source, target, "output_price_per_1m");
remove_number(&mut source, target, "cache_creation_price_per_1m");
remove_number(&mut source, target, "cache_read_price_per_1m");
remove_number(&mut source, target, "price_per_request");
}
fn copy_non_empty_string(source: &Map<String, Value>, target: &mut Map<String, Value>, key: &str) {
let Some(value) = source
.get(key)
@@ -77,6 +120,20 @@ fn copy_non_empty_string(source: &Map<String, Value>, target: &mut Map<String, V
);
}
fn remove_non_empty_string(
source: &mut Map<String, Value>,
target: &mut Map<String, Value>,
key: &str,
) {
let Some(Value::String(value)) = source.remove(key) else {
return;
};
let Some(value) = trim_and_truncate_usage_request_metadata_string_owned(value) else {
return;
};
target.insert(key.to_string(), Value::String(value));
}
fn copy_number(source: &Map<String, Value>, target: &mut Map<String, Value>, key: &str) {
let Some(value) = source.get(key).filter(|value| value.is_number()) else {
return;
@@ -84,6 +141,13 @@ fn copy_number(source: &Map<String, Value>, target: &mut Map<String, Value>, key
target.insert(key.to_string(), value.clone());
}
fn remove_number(source: &mut Map<String, Value>, target: &mut Map<String, Value>, key: &str) {
let Some(value) = source.remove(key).filter(|value| value.is_number()) else {
return;
};
target.insert(key.to_string(), value);
}
fn copy_bool(source: &Map<String, Value>, target: &mut Map<String, Value>, key: &str) {
let Some(value) = source.get(key).filter(|value| value.is_boolean()) else {
return;
@@ -91,6 +155,13 @@ fn copy_bool(source: &Map<String, Value>, target: &mut Map<String, Value>, key:
target.insert(key.to_string(), value.clone());
}
fn remove_bool(source: &mut Map<String, Value>, target: &mut Map<String, Value>, key: &str) {
let Some(value) = source.remove(key).filter(|value| value.is_boolean()) else {
return;
};
target.insert(key.to_string(), value);
}
fn copy_non_null_value(source: &Map<String, Value>, target: &mut Map<String, Value>, key: &str) {
let Some(value) = source.get(key).filter(|value| !value.is_null()) else {
return;
@@ -101,6 +172,20 @@ fn copy_non_null_value(source: &Map<String, Value>, target: &mut Map<String, Val
);
}
fn remove_non_null_value(
source: &mut Map<String, Value>,
target: &mut Map<String, Value>,
key: &str,
) {
let Some(value) = source.remove(key).filter(|value| !value.is_null()) else {
return;
};
target.insert(
key.to_string(),
sanitize_usage_request_metadata_value_owned(value),
);
}
fn sanitize_usage_request_metadata_value(value: &Value) -> Value {
match value {
Value::String(text) => Value::String(truncate_usage_request_metadata_string(text)),
@@ -109,6 +194,14 @@ fn sanitize_usage_request_metadata_value(value: &Value) -> Value {
}
}
fn sanitize_usage_request_metadata_value_owned(value: Value) -> Value {
match value {
Value::String(text) => Value::String(truncate_usage_request_metadata_string_owned(text)),
_ if usage_request_metadata_within_limits(&value) => value,
_ => truncated_usage_request_metadata_value(&value),
}
}
fn truncate_usage_request_metadata_string(value: &str) -> String {
const TRUNCATED_SUFFIX: &str = "...[truncated]";
@@ -134,6 +227,24 @@ fn truncate_usage_request_metadata_string(value: &str) -> String {
format!("{}{TRUNCATED_SUFFIX}", &value[..end])
}
fn trim_and_truncate_usage_request_metadata_string_owned(value: String) -> Option<String> {
let trimmed = value.trim();
if trimmed.is_empty() {
return None;
}
if trimmed.len() == value.len() {
return Some(truncate_usage_request_metadata_string_owned(value));
}
Some(truncate_usage_request_metadata_string(trimmed))
}
fn truncate_usage_request_metadata_string_owned(value: String) -> String {
if value.len() <= MAX_USAGE_REQUEST_METADATA_STRING_BYTES {
return value;
}
truncate_usage_request_metadata_string(value.as_str())
}
fn truncated_usage_request_metadata_value(value: &Value) -> Value {
json!({
"truncated": true,
@@ -217,7 +328,8 @@ mod tests {
use super::{
build_usage_request_metadata_seed, merge_usage_request_metadata,
sanitize_usage_request_metadata, MAX_USAGE_REQUEST_METADATA_BYTES,
merge_usage_request_metadata_owned, sanitize_usage_request_metadata,
sanitize_usage_request_metadata_ref, MAX_USAGE_REQUEST_METADATA_BYTES,
MAX_USAGE_REQUEST_METADATA_DEPTH, MAX_USAGE_REQUEST_METADATA_NODES,
};
@@ -363,4 +475,35 @@ mod tests {
assert_eq!(metadata, None);
}
#[test]
fn owned_merge_matches_filtered_merge_for_trusted_objects() {
let base = Some(json!({
"trace_id": "trace-1",
"provider_request_body_base64_bytes": 128
}));
let override_value = Some(json!({
"billing_snapshot_status": "complete",
"trace_id": "trace-2"
}));
assert_eq!(
merge_usage_request_metadata_owned(base.clone(), override_value.clone()),
merge_usage_request_metadata(base, override_value)
);
}
#[test]
fn borrowed_sanitize_matches_owned_sanitize() {
let value = json!({
"trace_id": "trace-1",
"billing_snapshot": {"status": "complete"},
"provider_name": "OpenAI"
});
assert_eq!(
sanitize_usage_request_metadata_ref(Some(&value)),
sanitize_usage_request_metadata(Some(value))
);
}
}

86
crates/aether-usage-runtime/src/runtime.rs
View File

@@ -12,8 +12,7 @@ use tracing::warn;
use crate::executor::spawn_on_usage_background_runtime;
use crate::{
apply_usage_body_capture_policy_to_event, apply_usage_body_capture_policy_to_record,
build_pending_usage_record_from_seed, build_stream_terminal_usage_seed,
build_streaming_usage_record_from_seed, build_sync_terminal_usage_seed,
build_stream_terminal_usage_seed, build_sync_terminal_usage_seed,
build_terminal_usage_event_from_seed, build_upsert_usage_record_from_event,
build_usage_queue_worker, settle_usage_if_needed, LifecycleUsageSeed,
StreamTerminalUsagePayloadSeed, SyncTerminalUsagePayloadSeed, TerminalUsageContextSeed,
@@ -74,17 +73,6 @@ pub struct UsageRuntime {
config: UsageRuntimeConfig,
}
struct SyncTerminalUsageTaskInput {
context_seed: TerminalUsageContextSeed,
payload_seed: SyncTerminalUsagePayloadSeed,
}
struct StreamTerminalUsageTaskInput {
context_seed: TerminalUsageContextSeed,
payload_seed: StreamTerminalUsagePayloadSeed,
cancelled: bool,
}
impl Default for UsageRuntime {
fn default() -> Self {
Self::disabled()
@@ -126,7 +114,7 @@ impl UsageRuntime {
Some(worker.spawn())
}
pub fn record_pending<T>(&self, data: &T, seed: &LifecycleUsageSeed)
pub fn record_pending<T>(&self, data: &T, seed: LifecycleUsageSeed)
where
T: UsageRuntimeAccess + Clone + 'static,
{
@@ -134,11 +122,10 @@ impl UsageRuntime {
return;
}
let data = T::clone(data);
let seed = seed.clone();
let request_id = seed.request_id.clone();
spawn_on_usage_background_runtime(boxed_usage_task(async move {
let now_unix_secs = now_unix_secs();
match build_pending_usage_record_offthread(&seed, now_unix_secs).await {
match build_pending_usage_record_offthread(seed, now_unix_secs).await {
Ok(mut record) => {
apply_body_capture_policy_to_record_from_data(&data, &mut record).await;
if let Err(err) = data.upsert_usage_record(record).await {
@@ -183,9 +170,9 @@ impl UsageRuntime {
spawn_on_usage_background_runtime(boxed_usage_task(async move {
let now_unix_secs = now_unix_secs();
match build_streaming_usage_record_offthread(
&seed,
seed,
status_code,
telemetry.as_ref(),
telemetry,
now_unix_secs,
)
.await
@@ -218,8 +205,8 @@ impl UsageRuntime {
pub fn record_sync_terminal<T>(
&self,
data: &T,
context_seed: &TerminalUsageContextSeed,
payload_seed: &SyncTerminalUsagePayloadSeed,
context_seed: TerminalUsageContextSeed,
payload_seed: SyncTerminalUsagePayloadSeed,
) where
T: UsageRuntimeAccess + Clone + 'static,
{
@@ -229,12 +216,8 @@ impl UsageRuntime {
let runtime = self.clone();
let data = T::clone(data);
let request_id = context_seed.request_id.clone();
let input = Box::new(SyncTerminalUsageTaskInput {
context_seed: context_seed.clone(),
payload_seed: payload_seed.clone(),
});
spawn_on_usage_background_runtime(boxed_usage_task(async move {
match build_sync_terminal_usage_event_offthread(input).await {
match build_sync_terminal_usage_event_offthread(context_seed, payload_seed).await {
Ok(mut event) => {
apply_body_capture_policy_from_data(&data, &mut event).await;
if let Err(err) = data.enrich_usage_event(&mut event).await {
@@ -264,8 +247,8 @@ impl UsageRuntime {
pub fn record_stream_terminal<T>(
&self,
data: &T,
context_seed: &TerminalUsageContextSeed,
payload_seed: &StreamTerminalUsagePayloadSeed,
context_seed: TerminalUsageContextSeed,
payload_seed: StreamTerminalUsagePayloadSeed,
cancelled: bool,
) where
T: UsageRuntimeAccess + Clone + 'static,
@@ -276,13 +259,10 @@ impl UsageRuntime {
let runtime = self.clone();
let data = T::clone(data);
let request_id = context_seed.request_id.clone();
let input = Box::new(StreamTerminalUsageTaskInput {
context_seed: context_seed.clone(),
payload_seed: payload_seed.clone(),
cancelled,
});
spawn_on_usage_background_runtime(boxed_usage_task(async move {
match build_stream_terminal_usage_event_offthread(input).await {
match build_stream_terminal_usage_event_offthread(context_seed, payload_seed, cancelled)
.await
{
Ok(mut event) => {
apply_body_capture_policy_from_data(&data, &mut event).await;
if let Err(err) = data.enrich_usage_event(&mut event).await {
@@ -413,28 +393,27 @@ impl UsageRuntime {
}
async fn build_pending_usage_record_offthread(
seed: &LifecycleUsageSeed,
seed: LifecycleUsageSeed,
now_unix_secs: u64,
) -> Result<UpsertUsageRecord, DataLayerError> {
let seed = seed.clone();
tokio::task::spawn_blocking(move || build_pending_usage_record_from_seed(&seed, now_unix_secs))
.await
.map_err(join_error_to_data_layer)?
tokio::task::spawn_blocking(move || {
crate::write::build_pending_usage_record_from_owned_seed(seed, now_unix_secs)
})
.await
.map_err(join_error_to_data_layer)?
}
async fn build_streaming_usage_record_offthread(
seed: &LifecycleUsageSeed,
seed: LifecycleUsageSeed,
status_code: u16,
telemetry: Option<&ExecutionTelemetry>,
telemetry: Option<ExecutionTelemetry>,
now_unix_secs: u64,
) -> Result<UpsertUsageRecord, DataLayerError> {
let seed = seed.clone();
let telemetry = telemetry.cloned();
tokio::task::spawn_blocking(move || {
build_streaming_usage_record_from_seed(
&seed,
crate::write::build_streaming_usage_record_from_owned_seed(
seed,
status_code,
telemetry.as_ref(),
telemetry,
now_unix_secs,
)
})
@@ -443,12 +422,13 @@ async fn build_streaming_usage_record_offthread(
}
async fn build_sync_terminal_usage_event_offthread(
input: Box<SyncTerminalUsageTaskInput>,
context_seed: TerminalUsageContextSeed,
payload_seed: SyncTerminalUsagePayloadSeed,
) -> Result<UsageEvent, DataLayerError> {
tokio::task::spawn_blocking(move || {
build_terminal_usage_event_from_seed(build_sync_terminal_usage_seed(
input.context_seed,
input.payload_seed,
context_seed,
payload_seed,
))
})
.await
@@ -456,13 +436,15 @@ async fn build_sync_terminal_usage_event_offthread(
}
async fn build_stream_terminal_usage_event_offthread(
input: Box<StreamTerminalUsageTaskInput>,
context_seed: TerminalUsageContextSeed,
payload_seed: StreamTerminalUsagePayloadSeed,
cancelled: bool,
) -> Result<UsageEvent, DataLayerError> {
tokio::task::spawn_blocking(move || {
build_terminal_usage_event_from_seed(build_stream_terminal_usage_seed(
input.context_seed,
input.payload_seed,
input.cancelled,
context_seed,
payload_seed,
cancelled,
))
})
.await

1530
crates/aether-usage-runtime/src/write.rs
View File

File diff suppressed because it is too large Load Diff