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

Merge upstream main into feat/500-api-key-ip-whitelist

Browse Source
This commit is contained in:
RWDai
2026-05-20 10:26:56 +08:00
parent d19bf71343 ed75ae6d56
commit bdc848b19e
501 changed files with 47013 additions and 3667 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-data/src/repository/announcements/memory.rs
+51
View File
@@ -135,6 +135,48 @@ impl AnnouncementReadRepository for InMemoryAnnouncementReadRepository {
Ok(total)
}
async fn list_required_unread_active_announcements(
&self,
user_id: &str,
now_unix_secs: u64,
limit: usize,
) -> Result<Vec<StoredAnnouncement>, DataLayerError> {
let announcements = self
.announcements
.read()
.expect("announcement repository lock");
let reads = self
.announcement_reads
.read()
.expect("announcement reads repository lock");
let mut items = announcements
.iter()
.filter(|announcement| {
announcement.requires_ack
&& announcement.is_active
&& announcement
.start_time_unix_secs
.is_none_or(|value| value <= now_unix_secs)
&& announcement
.end_time_unix_secs
.is_none_or(|value| value >= now_unix_secs)
&& !reads.contains(&(user_id.to_string(), announcement.id.clone()))
})
.cloned()
.collect::<Vec<_>>();
items.sort_by(|left, right| {
right
.is_pinned
.cmp(&left.is_pinned)
.then_with(|| right.priority.cmp(&left.priority))
.then_with(|| right.created_at_unix_ms.cmp(&left.created_at_unix_ms))
.then_with(|| left.id.cmp(&right.id))
});
items.truncate(limit);
Ok(items)
}
}
#[async_trait]
@@ -153,6 +195,7 @@ impl AnnouncementWriteRepository for InMemoryAnnouncementReadRepository {
record.priority,
true,
record.is_pinned,
record.requires_ack,
Some(record.author_id),
None,
record.start_time_unix_secs.map(|value| value as i64),
@@ -201,6 +244,9 @@ impl AnnouncementWriteRepository for InMemoryAnnouncementReadRepository {
if let Some(is_pinned) = record.is_pinned {
announcement.is_pinned = is_pinned;
}
if let Some(requires_ack) = record.requires_ack {
announcement.requires_ack = requires_ack;
}
if let Some(start_time_unix_secs) = record.start_time_unix_secs {
announcement.start_time_unix_secs = Some(start_time_unix_secs);
}
@@ -261,6 +307,7 @@ mod tests {
10,
true,
true,
false,
Some("admin-1".to_string()),
Some("admin".to_string()),
None,
@@ -291,6 +338,7 @@ mod tests {
kind: "maintenance".to_string(),
priority: 10,
is_pinned: true,
requires_ack: false,
author_id: "admin-1".to_string(),
start_time_unix_secs: None,
end_time_unix_secs: None,
@@ -308,6 +356,7 @@ mod tests {
priority: Some(99),
is_active: Some(false),
is_pinned: Some(false),
requires_ack: Some(true),
start_time_unix_secs: None,
end_time_unix_secs: None,
})
@@ -337,6 +386,7 @@ mod tests {
10,
true,
true,
false,
Some("admin-1".to_string()),
Some("admin".to_string()),
None,
@@ -353,6 +403,7 @@ mod tests {
5,
true,
false,
false,
Some("admin-1".to_string()),
Some("admin".to_string()),
None,
crates/aether-data/src/repository/announcements/mysql.rs
+40 -2
View File
@@ -18,6 +18,7 @@ SELECT
a.priority,
a.is_active,
a.is_pinned,
a.requires_ack,
a.author_id,
u.username AS author_username,
a.start_time AS start_time_unix_secs,
@@ -144,6 +145,39 @@ WHERE a.is_active = 1
.map_sql_err()?;
Ok(row.try_get::<i64, _>("total").map_sql_err()?.max(0) as u64)
}
async fn list_required_unread_active_announcements(
&self,
user_id: &str,
now_unix_secs: u64,
limit: usize,
) -> Result<Vec<StoredAnnouncement>, DataLayerError> {
let rows = sqlx::query(&format!(
r#"
{ANNOUNCEMENT_SELECT}
WHERE a.is_active = 1
AND a.requires_ack = 1
AND (a.start_time IS NULL OR a.start_time <= ?)
AND (a.end_time IS NULL OR a.end_time >= ?)
AND NOT EXISTS (
SELECT 1
FROM announcement_reads r
WHERE r.user_id = ?
AND r.announcement_id = a.id
)
ORDER BY a.is_pinned DESC, a.priority DESC, a.created_at DESC, a.id ASC
LIMIT ?
"#
))
.bind(now_unix_secs as i64)
.bind(now_unix_secs as i64)
.bind(user_id)
.bind(limit as i64)
.fetch_all(&self.pool)
.await
.map_sql_err()?;
rows.iter().map(map_announcement_row).collect()
}
}
#[async_trait]
@@ -159,9 +193,9 @@ impl AnnouncementWriteRepository for MysqlAnnouncementRepository {
r#"
INSERT INTO announcements (
id, title, content, `type`, priority, author_id, is_active, is_pinned,
start_time, end_time, created_at, updated_at
requires_ack, start_time, end_time, created_at, updated_at
)
VALUES (?, ?, ?, ?, ?, ?, 1, ?, ?, ?, ?, ?)
VALUES (?, ?, ?, ?, ?, ?, 1, ?, ?, ?, ?, ?, ?)
"#,
)
.bind(&id)
@@ -171,6 +205,7 @@ VALUES (?, ?, ?, ?, ?, ?, 1, ?, ?, ?, ?, ?)
.bind(record.priority)
.bind(record.author_id)
.bind(record.is_pinned)
.bind(record.requires_ack)
.bind(optional_i64_from_u64(
record.start_time_unix_secs,
"announcements.start_time",
@@ -204,6 +239,7 @@ SET title = COALESCE(?, title),
priority = COALESCE(?, priority),
is_active = COALESCE(?, is_active),
is_pinned = COALESCE(?, is_pinned),
requires_ack = COALESCE(?, requires_ack),
start_time = COALESCE(?, start_time),
end_time = COALESCE(?, end_time),
updated_at = ?
@@ -216,6 +252,7 @@ WHERE id = ?
.bind(record.priority)
.bind(record.is_active)
.bind(record.is_pinned)
.bind(record.requires_ack)
.bind(optional_i64_from_u64(
record.start_time_unix_secs,
"announcements.start_time",
@@ -303,6 +340,7 @@ fn map_announcement_row(row: &MySqlRow) -> Result<StoredAnnouncement, DataLayerE
row.try_get("priority").map_sql_err()?,
row.try_get("is_active").map_sql_err()?,
row.try_get("is_pinned").map_sql_err()?,
row.try_get("requires_ack").map_sql_err()?,
row.try_get("author_id").map_sql_err()?,
row.try_get("author_username").map_sql_err()?,
row.try_get("start_time_unix_secs").map_sql_err()?,
crates/aether-data/src/repository/announcements/postgres.rs
+59 -2
View File
@@ -18,6 +18,7 @@ SELECT
a.priority,
a.is_active,
a.is_pinned,
a.requires_ack,
a.author_id,
u.username AS author_username,
EXTRACT(EPOCH FROM a.start_time)::bigint AS start_time_unix_secs,
@@ -28,6 +29,38 @@ FROM announcements a
LEFT JOIN users u ON u.id = a.author_id
"#;
const LIST_REQUIRED_UNREAD_ACTIVE_ANNOUNCEMENTS_SQL: &str = r#"
SELECT
a.id,
a.title,
a.content,
a.type,
a.priority,
a.is_active,
a.is_pinned,
a.requires_ack,
a.author_id,
u.username AS author_username,
EXTRACT(EPOCH FROM a.start_time)::bigint AS start_time_unix_secs,
EXTRACT(EPOCH FROM a.end_time)::bigint AS end_time_unix_secs,
EXTRACT(EPOCH FROM a.created_at)::bigint AS created_at_unix_ms,
EXTRACT(EPOCH FROM a.updated_at)::bigint AS updated_at_unix_secs
FROM announcements a
LEFT JOIN users u ON u.id = a.author_id
WHERE a.is_active = TRUE
AND a.requires_ack = TRUE
AND (a.start_time IS NULL OR a.start_time <= TO_TIMESTAMP($2::double precision))
AND (a.end_time IS NULL OR a.end_time >= TO_TIMESTAMP($2::double precision))
AND NOT EXISTS (
SELECT 1
FROM announcement_reads r
WHERE r.user_id = $1
AND r.announcement_id = a.id
)
ORDER BY a.is_pinned DESC, a.priority DESC, a.created_at DESC, a.id ASC
LIMIT $3
"#;
const CREATE_ANNOUNCEMENT_SQL: &str = r#"
INSERT INTO announcements (
id,
@@ -38,6 +71,7 @@ INSERT INTO announcements (
author_id,
is_active,
is_pinned,
requires_ack,
start_time,
end_time,
created_at,
@@ -54,6 +88,7 @@ VALUES (
$7,
$8,
$9,
$10,
NOW(),
NOW()
)
@@ -65,6 +100,7 @@ RETURNING
priority,
is_active,
is_pinned,
requires_ack,
author_id,
(SELECT username FROM users WHERE id = announcements.author_id) AS author_username,
EXTRACT(EPOCH FROM start_time)::bigint AS start_time_unix_secs,
@@ -82,8 +118,9 @@ SET
priority = COALESCE($5, priority),
is_active = COALESCE($6, is_active),
is_pinned = COALESCE($7, is_pinned),
start_time = COALESCE($8, start_time),
end_time = COALESCE($9, end_time),
requires_ack = COALESCE($8, requires_ack),
start_time = COALESCE($9, start_time),
end_time = COALESCE($10, end_time),
updated_at = NOW()
WHERE id = $1
RETURNING
@@ -94,6 +131,7 @@ RETURNING
priority,
is_active,
is_pinned,
requires_ack,
author_id,
(SELECT username FROM users WHERE id = announcements.author_id) AS author_username,
EXTRACT(EPOCH FROM start_time)::bigint AS start_time_unix_secs,
@@ -247,6 +285,22 @@ impl AnnouncementReadRepository for SqlxAnnouncementReadRepository {
.max(0) as u64;
Ok(total)
}
async fn list_required_unread_active_announcements(
&self,
user_id: &str,
now_unix_secs: u64,
limit: usize,
) -> Result<Vec<StoredAnnouncement>, DataLayerError> {
let rows = sqlx::query(LIST_REQUIRED_UNREAD_ACTIVE_ANNOUNCEMENTS_SQL)
.bind(user_id)
.bind(now_unix_secs as f64)
.bind(limit as i64)
.fetch_all(&self.pool)
.await
.map_postgres_err()?;
rows.iter().map(map_announcement_row).collect()
}
}
#[async_trait]
@@ -264,6 +318,7 @@ impl AnnouncementWriteRepository for SqlxAnnouncementReadRepository {
.bind(record.priority)
.bind(record.author_id)
.bind(record.is_pinned)
.bind(record.requires_ack)
.bind(optional_datetime(record.start_time_unix_secs))
.bind(optional_datetime(record.end_time_unix_secs))
.fetch_one(&self.pool)
@@ -285,6 +340,7 @@ impl AnnouncementWriteRepository for SqlxAnnouncementReadRepository {
.bind(record.priority)
.bind(record.is_active)
.bind(record.is_pinned)
.bind(record.requires_ack)
.bind(optional_datetime(record.start_time_unix_secs))
.bind(optional_datetime(record.end_time_unix_secs))
.fetch_optional(&self.pool)
@@ -351,6 +407,7 @@ fn map_announcement_row(row: &PgRow) -> Result<StoredAnnouncement, DataLayerErro
row.try_get("priority").map_postgres_err()?,
row.try_get("is_active").map_postgres_err()?,
row.try_get("is_pinned").map_postgres_err()?,
row.try_get("requires_ack").map_postgres_err()?,
row.try_get("author_id").map_postgres_err()?,
row.try_get("author_username").map_postgres_err()?,
row.try_get("start_time_unix_secs").map_postgres_err()?,
crates/aether-data/src/repository/announcements/sqlite.rs
+42 -2
View File
@@ -19,6 +19,7 @@ SELECT
a.priority,
a.is_active,
a.is_pinned,
a.requires_ack,
a.author_id,
u.username AS author_username,
a.start_time AS start_time_unix_secs,
@@ -158,6 +159,39 @@ impl AnnouncementReadRepository for SqliteAnnouncementRepository {
.max(0) as u64;
Ok(total)
}
async fn list_required_unread_active_announcements(
&self,
user_id: &str,
now_unix_secs: u64,
limit: usize,
) -> Result<Vec<StoredAnnouncement>, DataLayerError> {
let rows = sqlx::query(&format!(
r#"
{ANNOUNCEMENT_SELECT}
WHERE a.is_active = 1
AND a.requires_ack = 1
AND (a.start_time IS NULL OR a.start_time <= ?)
AND (a.end_time IS NULL OR a.end_time >= ?)
AND NOT EXISTS (
SELECT 1
FROM announcement_reads r
WHERE r.user_id = ?
AND r.announcement_id = a.id
)
ORDER BY a.is_pinned DESC, a.priority DESC, a.created_at DESC, a.id ASC
LIMIT ?
"#
))
.bind(now_unix_secs as i64)
.bind(now_unix_secs as i64)
.bind(user_id)
.bind(limit as i64)
.fetch_all(&self.pool)
.await
.map_sql_err()?;
rows.iter().map(map_announcement_row).collect()
}
}
#[async_trait]
@@ -173,9 +207,9 @@ impl AnnouncementWriteRepository for SqliteAnnouncementRepository {
r#"
INSERT INTO announcements (
id, title, content, type, priority, author_id, is_active, is_pinned,
start_time, end_time, created_at, updated_at
requires_ack, start_time, end_time, created_at, updated_at
)
VALUES (?, ?, ?, ?, ?, ?, 1, ?, ?, ?, ?, ?)
VALUES (?, ?, ?, ?, ?, ?, 1, ?, ?, ?, ?, ?, ?)
"#,
)
.bind(&id)
@@ -185,6 +219,7 @@ VALUES (?, ?, ?, ?, ?, ?, 1, ?, ?, ?, ?, ?)
.bind(record.priority)
.bind(record.author_id)
.bind(record.is_pinned)
.bind(record.requires_ack)
.bind(optional_i64_from_u64(
record.start_time_unix_secs,
"announcements.start_time",
@@ -218,6 +253,7 @@ SET title = COALESCE(?, title),
priority = COALESCE(?, priority),
is_active = COALESCE(?, is_active),
is_pinned = COALESCE(?, is_pinned),
requires_ack = COALESCE(?, requires_ack),
start_time = COALESCE(?, start_time),
end_time = COALESCE(?, end_time),
updated_at = ?
@@ -230,6 +266,7 @@ WHERE id = ?
.bind(record.priority)
.bind(record.is_active)
.bind(record.is_pinned)
.bind(record.requires_ack)
.bind(optional_i64_from_u64(
record.start_time_unix_secs,
"announcements.start_time",
@@ -317,6 +354,7 @@ fn map_announcement_row(row: &SqliteRow) -> Result<StoredAnnouncement, DataLayer
row.try_get("priority").map_sql_err()?,
row.try_get("is_active").map_sql_err()?,
row.try_get("is_pinned").map_sql_err()?,
row.try_get("requires_ack").map_sql_err()?,
row.try_get("author_id").map_sql_err()?,
row.try_get("author_username").map_sql_err()?,
row.try_get("start_time_unix_secs").map_sql_err()?,
@@ -355,6 +393,7 @@ mod tests {
kind: "info".to_string(),
priority: 10,
is_pinned: true,
requires_ack: false,
author_id: "user-1".to_string(),
start_time_unix_secs: Some(100),
end_time_unix_secs: Some(300),
@@ -406,6 +445,7 @@ mod tests {
priority: Some(20),
is_active: Some(false),
is_pinned: Some(false),
requires_ack: Some(true),
start_time_unix_secs: None,
end_time_unix_secs: None,
})
crates/aether-data/src/repository/announcements/types.rs
+12
View File
@@ -9,6 +9,7 @@ pub struct StoredAnnouncement {
pub priority: i32,
pub is_active: bool,
pub is_pinned: bool,
pub requires_ack: bool,
pub author_id: Option<String>,
pub author_username: Option<String>,
pub start_time_unix_secs: Option<u64>,
@@ -27,6 +28,7 @@ impl StoredAnnouncement {
priority: i32,
is_active: bool,
is_pinned: bool,
requires_ack: bool,
author_id: Option<String>,
author_username: Option<String>,
start_time_unix_secs: Option<i64>,
@@ -63,6 +65,7 @@ impl StoredAnnouncement {
priority,
is_active,
is_pinned,
requires_ack,
author_id,
author_username,
start_time_unix_secs: start_time_unix_secs
@@ -117,6 +120,13 @@ pub trait AnnouncementReadRepository: Send + Sync {
user_id: &str,
now_unix_secs: u64,
) -> Result<u64, crate::DataLayerError>;
async fn list_required_unread_active_announcements(
&self,
user_id: &str,
now_unix_secs: u64,
limit: usize,
) -> Result<Vec<StoredAnnouncement>, crate::DataLayerError>;
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
@@ -126,6 +136,7 @@ pub struct CreateAnnouncementRecord {
pub kind: String,
pub priority: i32,
pub is_pinned: bool,
pub requires_ack: bool,
pub author_id: String,
pub start_time_unix_secs: Option<u64>,
pub end_time_unix_secs: Option<u64>,
@@ -166,6 +177,7 @@ pub struct UpdateAnnouncementRecord {
pub priority: Option<i32>,
pub is_active: Option<bool>,
pub is_pinned: Option<bool>,
pub requires_ack: Option<bool>,
pub start_time_unix_secs: Option<u64>,
pub end_time_unix_secs: Option<u64>,
}
crates/aether-data/src/repository/auth/postgres.rs
+11 -4
View File
@@ -529,7 +529,7 @@ SET
name = COALESCE($3, name),
rate_limit = COALESCE($4, rate_limit),
concurrent_limit = COALESCE($5, concurrent_limit),
allowed_ips = CASE WHEN $6 THEN $7::json ELSE allowed_ips END,
allowed_ips = CASE WHEN $6 THEN $7::jsonb ELSE allowed_ips END,
updated_at = NOW()
WHERE user_id = $1
AND id = $2
@@ -569,7 +569,7 @@ SET
allowed_providers = CASE WHEN $7 THEN $8::json ELSE allowed_providers END,
allowed_api_formats = CASE WHEN $9 THEN $10::json ELSE allowed_api_formats END,
allowed_models = CASE WHEN $11 THEN $12::json ELSE allowed_models END,
allowed_ips = CASE WHEN $13 THEN $14::json ELSE allowed_ips END,
allowed_ips = CASE WHEN $13 THEN $14::jsonb ELSE allowed_ips END,
expires_at = CASE WHEN $15 THEN $16::timestamptz ELSE expires_at END,
auto_delete_on_expiry = CASE WHEN $17 THEN $18 ELSE auto_delete_on_expiry END,
updated_at = NOW()
@@ -1572,7 +1572,7 @@ fn map_auth_api_key_export_row(
row_get(row, "is_active")?,
row_get(row, "expires_at_unix_secs")?,
row_get(row, "auto_delete_on_expiry")?,
row_get::<i32>(row, "total_requests")?.into(),
row_get::<i64>(row, "total_requests")?,
row_get::<i64>(row, "total_tokens")?,
row_get(row, "total_cost_usd")?,
row_get(row, "is_standalone")?,
@@ -1593,6 +1593,7 @@ mod tests {
use super::{
SqlxAuthApiKeySnapshotReadRepository, CREATE_STANDALONE_API_KEY_SQL,
CREATE_USER_API_KEY_SQL, UPDATE_STANDALONE_API_KEY_BASIC_SQL,
UPDATE_USER_API_KEY_BASIC_SQL,
};
use crate::driver::postgres::{PostgresPoolConfig, PostgresPoolFactory};
@@ -1621,7 +1622,7 @@ mod tests {
assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL
.contains("allowed_models = CASE WHEN $11 THEN $12::json ELSE allowed_models END"));
assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL
.contains("allowed_ips = CASE WHEN $13 THEN $14::json ELSE allowed_ips END"));
.contains("allowed_ips = CASE WHEN $13 THEN $14::jsonb ELSE allowed_ips END"));
assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL
.contains("rate_limit = CASE WHEN $3 THEN $4 ELSE rate_limit END"));
assert!(UPDATE_STANDALONE_API_KEY_BASIC_SQL
@@ -1631,6 +1632,12 @@ mod tests {
));
}
#[test]
fn update_user_api_key_basic_sql_casts_allowed_ips_as_jsonb() {
assert!(UPDATE_USER_API_KEY_BASIC_SQL
.contains("allowed_ips = CASE WHEN $6 THEN $7::jsonb ELSE allowed_ips END"));
}
#[tokio::test]
async fn repository_constructs_from_lazy_pool() {
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
crates/aether-data/src/repository/candidate_selection/memory.rs
+42 -2
View File
@@ -316,12 +316,23 @@ fn key_auth_channel_matches(row: &StoredMinimalCandidateSelectionRow, api_format
"gemini_cli" | "antigravity" => {
auth_type == "oauth" && api_format == "gemini:generate_content"
}
"grok" => {
auth_type == "oauth"
&& matches!(
api_format.as_str(),
"openai:chat" | "openai:responses" | "claude:messages" | "openai:image"
)
}
"vertex_ai" => {
(auth_type == "api_key" && api_format == "gemini:generate_content")
(auth_type == "api_key"
&& matches!(
api_format.as_str(),
"gemini:generate_content" | "gemini:embedding"
))
|| (matches!(auth_type.as_str(), "service_account" | "vertex_ai")
&& matches!(
api_format.as_str(),
"claude:messages" | "gemini:generate_content"
"claude:messages" | "gemini:generate_content" | "gemini:embedding"
))
}
_ => auth_type != "oauth",
@@ -419,6 +430,35 @@ mod tests {
assert_eq!(rows[0].provider_id, "provider-1");
}
#[tokio::test]
async fn includes_grok_oauth_rows_for_chat_models() {
let mut row = sample_row(
"provider-grok",
"openai:chat",
"grok-4.20-0309-non-reasoning",
10,
);
row.provider_type = "grok".to_string();
row.provider_name = "grok".to_string();
row.key_auth_type = "oauth".to_string();
row.key_api_formats = Some(vec![
"openai:chat".to_string(),
"openai:responses".to_string(),
"claude:messages".to_string(),
"openai:image".to_string(),
]);
let repository = InMemoryMinimalCandidateSelectionReadRepository::seed(vec![row]);
let rows = repository
.list_for_exact_api_format("openai:chat")
.await
.expect("list should succeed");
assert_eq!(rows.len(), 1);
assert_eq!(rows[0].provider_type, "grok");
assert_eq!(rows[0].global_model_name, "grok-4.20-0309-non-reasoning");
}
#[tokio::test]
async fn requested_model_filter_respects_endpoint_scoped_default_mapping() {
let mut selected = sample_row("provider-1", "openai:chat", "deepseek-v4-pro", 10);
crates/aether-data/src/repository/candidate_selection/mysql.rs
+13 -2
View File
@@ -445,12 +445,23 @@ fn key_auth_channel_matches(row: &CandidateSelectionRow, api_format: &str) -> bo
"gemini_cli" | "antigravity" => {
auth_type == "oauth" && api_format == "gemini:generate_content"
}
"grok" => {
auth_type == "oauth"
&& matches!(
api_format.as_str(),
"openai:chat" | "openai:responses" | "claude:messages" | "openai:image"
)
}
"vertex_ai" => {
(auth_type == "api_key" && api_format == "gemini:generate_content")
(auth_type == "api_key"
&& matches!(
api_format.as_str(),
"gemini:generate_content" | "gemini:embedding"
))
|| (matches!(auth_type.as_str(), "service_account" | "vertex_ai")
&& matches!(
api_format.as_str(),
"claude:messages" | "gemini:generate_content"
"claude:messages" | "gemini:generate_content" | "gemini:embedding"
))
}
_ => auth_type != "oauth",
crates/aether-data/src/repository/candidate_selection/postgres.rs
+57 -6
View File
@@ -103,6 +103,11 @@ WHERE p.is_active = TRUE
)
)
)
OR (
LOWER(BTRIM(p.provider_type)) = 'grok'
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($3) IN ('openai:chat', 'openai:responses', 'claude:messages', 'openai:image')
)
OR (
LOWER(BTRIM(p.provider_type)) IN ('gemini_cli', 'antigravity')
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
@@ -113,11 +118,11 @@ WHERE p.is_active = TRUE
AND (
(
LOWER(BTRIM(pak.auth_type)) = 'api_key'
AND LOWER($3) = 'gemini:generate_content'
AND LOWER($3) IN ('gemini:generate_content', 'gemini:embedding')
)
OR (
LOWER(BTRIM(pak.auth_type)) IN ('service_account', 'vertex_ai')
AND LOWER($3) IN ('claude:messages', 'gemini:generate_content')
AND LOWER($3) IN ('claude:messages', 'gemini:generate_content', 'gemini:embedding')
)
)
)
@@ -127,6 +132,7 @@ WHERE p.is_active = TRUE
'claude_code',
'codex',
'gemini_cli',
'grok',
'vertex_ai',
'antigravity',
'kiro'
@@ -286,6 +292,11 @@ WHERE p.is_active = TRUE
)
)
)
OR (
LOWER(BTRIM(p.provider_type)) = 'grok'
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($4) IN ('openai:chat', 'openai:responses', 'claude:messages', 'openai:image')
)
OR (
LOWER(BTRIM(p.provider_type)) IN ('gemini_cli', 'antigravity')
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
@@ -296,11 +307,11 @@ WHERE p.is_active = TRUE
AND (
(
LOWER(BTRIM(pak.auth_type)) = 'api_key'
AND LOWER($4) = 'gemini:generate_content'
AND LOWER($4) IN ('gemini:generate_content', 'gemini:embedding')
)
OR (
LOWER(BTRIM(pak.auth_type)) IN ('service_account', 'vertex_ai')
AND LOWER($4) IN ('claude:messages', 'gemini:generate_content')
AND LOWER($4) IN ('claude:messages', 'gemini:generate_content', 'gemini:embedding')
)
)
)
@@ -310,6 +321,7 @@ WHERE p.is_active = TRUE
'claude_code',
'codex',
'gemini_cli',
'grok',
'vertex_ai',
'antigravity',
'kiro'
@@ -468,6 +480,11 @@ WHERE p.is_active = TRUE
)
)
)
OR (
LOWER(BTRIM(p.provider_type)) = 'grok'
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($6) IN ('openai:chat', 'openai:responses', 'claude:messages', 'openai:image')
)
OR (
LOWER(BTRIM(p.provider_type)) IN ('gemini_cli', 'antigravity')
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
@@ -478,11 +495,11 @@ WHERE p.is_active = TRUE
AND (
(
LOWER(BTRIM(pak.auth_type)) = 'api_key'
AND LOWER($6) = 'gemini:generate_content'
AND LOWER($6) IN ('gemini:generate_content', 'gemini:embedding')
)
OR (
LOWER(BTRIM(pak.auth_type)) IN ('service_account', 'vertex_ai')
AND LOWER($6) IN ('claude:messages', 'gemini:generate_content')
AND LOWER($6) IN ('claude:messages', 'gemini:generate_content', 'gemini:embedding')
)
)
)
@@ -492,6 +509,7 @@ WHERE p.is_active = TRUE
'claude_code',
'codex',
'gemini_cli',
'grok',
'vertex_ai',
'antigravity',
'kiro'
@@ -1287,6 +1305,39 @@ mod tests {
}
}
#[test]
fn candidate_selection_sql_allows_grok_oauth_chat_auth() {
let requested_model_sql = requested_model_selection_sql();
for sql in [
LIST_FOR_EXACT_API_FORMAT_SQL,
LIST_FOR_EXACT_API_FORMAT_AND_GLOBAL_MODEL_SQL,
LIST_POOL_KEYS_FOR_GROUP_SQL,
requested_model_sql.as_str(),
] {
assert!(sql.contains("LOWER(BTRIM(p.provider_type)) = 'grok'"));
assert!(sql.contains("LOWER(BTRIM(pak.auth_type)) = 'oauth'"));
assert!(sql
.contains("'openai:chat', 'openai:responses', 'claude:messages', 'openai:image'"));
assert!(sql.contains("'grok',"));
}
}
#[test]
fn candidate_selection_sql_allows_vertex_embedding_auth() {
let requested_model_sql = requested_model_selection_sql();
for sql in [
LIST_FOR_EXACT_API_FORMAT_SQL,
LIST_FOR_EXACT_API_FORMAT_AND_GLOBAL_MODEL_SQL,
LIST_POOL_KEYS_FOR_GROUP_SQL,
requested_model_sql.as_str(),
] {
assert!(sql.contains("LOWER(BTRIM(p.provider_type)) = 'vertex_ai'"));
assert!(sql.contains("gemini:embedding"));
assert!(sql.contains("gemini:generate_content"));
assert!(sql.contains("claude:messages"));
}
}
#[test]
fn requested_model_selection_page_sql_adds_limit_and_offset() {
let sql = requested_model_selection_page_sql();
crates/aether-data/src/repository/candidate_selection/sqlite.rs
+13 -2
View File
@@ -824,12 +824,23 @@ fn key_auth_channel_matches(row: &CandidateSelectionRow, api_format: &str) -> bo
"gemini_cli" | "antigravity" => {
auth_type == "oauth" && api_format == "gemini:generate_content"
}
"grok" => {
auth_type == "oauth"
&& matches!(
api_format.as_str(),
"openai:chat" | "openai:responses" | "claude:messages" | "openai:image"
)
}
"vertex_ai" => {
(auth_type == "api_key" && api_format == "gemini:generate_content")
(auth_type == "api_key"
&& matches!(
api_format.as_str(),
"gemini:generate_content" | "gemini:embedding"
))
|| (matches!(auth_type.as_str(), "service_account" | "vertex_ai")
&& matches!(
api_format.as_str(),
"claude:messages" | "gemini:generate_content"
"claude:messages" | "gemini:generate_content" | "gemini:embedding"
))
}
_ => auth_type != "oauth",
crates/aether-data/src/repository/management_tokens/postgres.rs
+1 -1
View File
@@ -589,7 +589,7 @@ fn map_token_row(row: &PgRow) -> Result<StoredManagementToken, DataLayerError> {
optional_unix_secs(row.try_get("expires_at_unix_secs").map_postgres_err()?),
optional_unix_secs(row.try_get("last_used_at_unix_secs").map_postgres_err()?),
row.try_get("last_used_ip").map_postgres_err()?,
u64::try_from(row.try_get::<i32, _>("usage_count").map_postgres_err()?).unwrap_or(0),
u64::try_from(row.try_get::<i64, _>("usage_count").map_postgres_err()?).unwrap_or(0),
row.try_get("is_active").map_postgres_err()?,
)
.with_timestamps(
crates/aether-data/src/repository/oauth_providers/memory.rs
+2
View File
@@ -100,6 +100,7 @@ impl OAuthProviderWriteRepository for InMemoryOAuthProviderRepository {
record.scopes.clone(),
record.attribute_mapping.clone(),
record.extra_config.clone(),
record.icon_url.clone(),
record.is_enabled,
)
.with_timestamps(created_at, now);
@@ -150,6 +151,7 @@ mod tests {
frontend_callback_url: "https://frontend.example.com/auth/callback".to_string(),
attribute_mapping: Some(serde_json::json!({"email": "email"})),
extra_config: Some(serde_json::json!({"team": true})),
icon_url: None,
is_enabled: true,
}
}
crates/aether-data/src/repository/oauth_providers/mysql.rs
+7 -1
View File
@@ -46,6 +46,7 @@ SELECT
frontend_callback_url,
attribute_mapping,
extra_config,
icon_url,
is_enabled,
created_at AS created_at_unix_ms,
updated_at AS updated_at_unix_secs
@@ -67,6 +68,7 @@ SELECT
frontend_callback_url,
attribute_mapping,
extra_config,
icon_url,
is_enabled,
created_at AS created_at_unix_ms,
updated_at AS updated_at_unix_secs
@@ -176,13 +178,14 @@ INSERT INTO oauth_providers (
frontend_callback_url,
attribute_mapping,
extra_config,
icon_url,
is_enabled,
created_at,
updated_at
) VALUES (
?, ?, ?,
CASE ? WHEN 'set' THEN ? WHEN 'clear' THEN NULL ELSE NULL END,
?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?
?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?
)
ON DUPLICATE KEY UPDATE
display_name = VALUES(display_name),
@@ -200,6 +203,7 @@ ON DUPLICATE KEY UPDATE
frontend_callback_url = VALUES(frontend_callback_url),
attribute_mapping = VALUES(attribute_mapping),
extra_config = VALUES(extra_config),
icon_url = VALUES(icon_url),
is_enabled = VALUES(is_enabled),
updated_at = VALUES(updated_at)
"#,
@@ -217,6 +221,7 @@ ON DUPLICATE KEY UPDATE
.bind(&record.frontend_callback_url)
.bind(json_to_string(record.attribute_mapping.as_ref())?)
.bind(json_to_string(record.extra_config.as_ref())?)
.bind(record.icon_url.as_deref())
.bind(record.is_enabled)
.bind(now as i64)
.bind(now as i64)
@@ -361,6 +366,7 @@ fn map_oauth_provider_row(row: &MySqlRow) -> Result<StoredOAuthProviderConfig, D
row.try_get("extra_config").map_sql_err()?,
"oauth_providers.extra_config",
)?,
row.try_get("icon_url").map_sql_err()?,
row.try_get("is_enabled").map_sql_err()?,
)
.with_timestamps(
crates/aether-data/src/repository/oauth_providers/postgres.rs
+7
View File
@@ -22,6 +22,7 @@ SELECT
frontend_callback_url,
attribute_mapping,
extra_config,
icon_url,
is_enabled,
EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_ms,
EXTRACT(EPOCH FROM updated_at)::bigint AS updated_at_unix_secs
@@ -77,6 +78,7 @@ INSERT INTO oauth_providers (
frontend_callback_url,
attribute_mapping,
extra_config,
icon_url,
is_enabled,
created_at,
updated_at
@@ -99,6 +101,7 @@ VALUES (
$12,
$13,
$14,
$15,
NOW(),
NOW()
)
@@ -118,6 +121,7 @@ SET display_name = EXCLUDED.display_name,
frontend_callback_url = EXCLUDED.frontend_callback_url,
attribute_mapping = EXCLUDED.attribute_mapping,
extra_config = EXCLUDED.extra_config,
icon_url = EXCLUDED.icon_url,
is_enabled = EXCLUDED.is_enabled,
updated_at = NOW()
RETURNING
@@ -133,6 +137,7 @@ RETURNING
frontend_callback_url,
attribute_mapping,
extra_config,
icon_url,
is_enabled,
EXTRACT(EPOCH FROM created_at)::bigint AS created_at_unix_ms,
EXTRACT(EPOCH FROM updated_at)::bigint AS updated_at_unix_secs
@@ -230,6 +235,7 @@ impl OAuthProviderWriteRepository for SqlxOAuthProviderRepository {
.bind(&record.frontend_callback_url)
.bind(record.attribute_mapping.as_ref())
.bind(record.extra_config.as_ref())
.bind(record.icon_url.as_deref())
.bind(record.is_enabled)
.fetch_one(&self.pool)
.await
@@ -330,6 +336,7 @@ fn map_oauth_provider_row(row: &PgRow) -> Result<StoredOAuthProviderConfig, Data
parse_scopes(row.try_get("scopes").map_postgres_err()?)?,
row.try_get("attribute_mapping").map_postgres_err()?,
row.try_get("extra_config").map_postgres_err()?,
row.try_get("icon_url").map_postgres_err()?,
row.try_get("is_enabled").map_postgres_err()?,
)
.with_timestamps(
crates/aether-data/src/repository/oauth_providers/sqlite.rs
+7 -1
View File
@@ -56,6 +56,7 @@ SELECT
frontend_callback_url,
attribute_mapping,
extra_config,
icon_url,
is_enabled,
created_at AS created_at_unix_ms,
updated_at AS updated_at_unix_secs
@@ -162,13 +163,14 @@ INSERT INTO oauth_providers (
frontend_callback_url,
attribute_mapping,
extra_config,
icon_url,
is_enabled,
created_at,
updated_at
) VALUES (
?, ?, ?,
CASE ? WHEN 'set' THEN ? WHEN 'clear' THEN NULL ELSE NULL END,
?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?
?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?
)
ON CONFLICT(provider_type) DO UPDATE SET
display_name = excluded.display_name,
@@ -186,6 +188,7 @@ ON CONFLICT(provider_type) DO UPDATE SET
frontend_callback_url = excluded.frontend_callback_url,
attribute_mapping = excluded.attribute_mapping,
extra_config = excluded.extra_config,
icon_url = excluded.icon_url,
is_enabled = excluded.is_enabled,
updated_at = excluded.updated_at
"#,
@@ -203,6 +206,7 @@ ON CONFLICT(provider_type) DO UPDATE SET
.bind(&record.frontend_callback_url)
.bind(json_to_string(record.attribute_mapping.as_ref())?)
.bind(json_to_string(record.extra_config.as_ref())?)
.bind(record.icon_url.as_deref())
.bind(record.is_enabled)
.bind(now as i64)
.bind(now as i64)
@@ -350,6 +354,7 @@ fn map_oauth_provider_row(row: &SqliteRow) -> Result<StoredOAuthProviderConfig,
row.try_get("extra_config").map_sql_err()?,
"oauth_providers.extra_config",
)?,
row.try_get("icon_url").map_sql_err()?,
row.try_get("is_enabled").map_sql_err()?,
)
.with_timestamps(
@@ -381,6 +386,7 @@ mod tests {
frontend_callback_url: "https://frontend.example.com/auth/callback".to_string(),
attribute_mapping: Some(serde_json::json!({"email": "email"})),
extra_config: Some(serde_json::json!({"team": true})),
icon_url: None,
is_enabled: true,
}
}
crates/aether-data/src/repository/oauth_providers/types.rs
+5
View File
@@ -14,6 +14,7 @@ pub struct StoredOAuthProviderConfig {
pub frontend_callback_url: String,
pub attribute_mapping: Option<serde_json::Value>,
pub extra_config: Option<serde_json::Value>,
pub icon_url: Option<String>,
pub is_enabled: bool,
pub created_at_unix_ms: Option<u64>,
pub updated_at_unix_secs: Option<u64>,
@@ -66,6 +67,7 @@ impl StoredOAuthProviderConfig {
frontend_callback_url,
attribute_mapping: None,
extra_config: None,
icon_url: None,
is_enabled: false,
created_at_unix_ms: None,
updated_at_unix_secs: None,
@@ -82,6 +84,7 @@ impl StoredOAuthProviderConfig {
scopes: Option<Vec<String>>,
attribute_mapping: Option<serde_json::Value>,
extra_config: Option<serde_json::Value>,
icon_url: Option<String>,
is_enabled: bool,
) -> Self {
self.client_secret_encrypted = client_secret_encrypted;
@@ -91,6 +94,7 @@ impl StoredOAuthProviderConfig {
self.scopes = scopes;
self.attribute_mapping = attribute_mapping;
self.extra_config = extra_config;
self.icon_url = icon_url;
self.is_enabled = is_enabled;
self
}
@@ -145,6 +149,7 @@ pub struct UpsertOAuthProviderConfigRecord {
pub frontend_callback_url: String,
pub attribute_mapping: Option<serde_json::Value>,
pub extra_config: Option<serde_json::Value>,
pub icon_url: Option<String>,
pub is_enabled: bool,
}
crates/aether-data/src/repository/provider_catalog/postgres.rs
+12 -12
View File
@@ -289,12 +289,12 @@ SELECT
NULL::jsonb AS utilization_samples,
NULL::bigint AS last_probe_increase_at_unix_secs,
NULL::integer AS last_rpm_peak,
NULL::integer AS request_count,
NULL::bigint AS request_count,
0::bigint AS total_tokens,
0::double precision AS total_cost_usd,
NULL::integer AS success_count,
NULL::integer AS error_count,
NULL::integer AS total_response_time_ms,
NULL::bigint AS success_count,
NULL::bigint AS error_count,
NULL::bigint AS total_response_time_ms,
NULL::bigint AS last_used_at_unix_secs,
FALSE AS auto_fetch_models,
NULL::bigint AS last_models_fetch_at_unix_secs,
@@ -1273,7 +1273,7 @@ INSERT INTO provider_api_keys (
.map(|value| value as f64),
)
.bind(key.last_rpm_peak.map(|value| value as i32))
.bind(key.request_count.map(|value| value as i32))
.bind(key.request_count.map(i64::from))
.bind(Some(i64::try_from(key.total_tokens).map_err(|_| {
DataLayerError::InvalidInput(format!(
"provider catalog key.total_tokens exceeds i64: {}",
@@ -1281,9 +1281,9 @@ INSERT INTO provider_api_keys (
))
})?))
.bind(key.total_cost_usd)
.bind(key.success_count.map(|value| value as i32))
.bind(key.error_count.map(|value| value as i32))
.bind(key.total_response_time_ms.map(|value| value as i32))
.bind(key.success_count.map(i64::from))
.bind(key.error_count.map(i64::from))
.bind(key.total_response_time_ms.map(i64::from))
.bind(key.last_used_at_unix_secs.map(|value| value as f64))
.bind(key.last_models_fetch_at_unix_secs.map(|value| value as f64))
.bind(&key.last_models_fetch_error)
@@ -2293,7 +2293,7 @@ fn map_key_row(row: &PgRow) -> Result<StoredProviderCatalogKey, DataLayerError>
})
})
.transpose()?;
let request_count = row_get::<Option<i32>>(row, "request_count")?
let request_count = row_get::<Option<i64>>(row, "request_count")?
.map(|value| {
u32::try_from(value).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
@@ -2314,7 +2314,7 @@ fn map_key_row(row: &PgRow) -> Result<StoredProviderCatalogKey, DataLayerError>
"invalid provider_api_keys.total_cost_usd".to_string(),
));
}
let success_count = row_get::<Option<i32>>(row, "success_count")?
let success_count = row_get::<Option<i64>>(row, "success_count")?
.map(|value| {
u32::try_from(value).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
@@ -2323,7 +2323,7 @@ fn map_key_row(row: &PgRow) -> Result<StoredProviderCatalogKey, DataLayerError>
})
})
.transpose()?;
let error_count = row_get::<Option<i32>>(row, "error_count")?
let error_count = row_get::<Option<i64>>(row, "error_count")?
.map(|value| {
u32::try_from(value).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
@@ -2332,7 +2332,7 @@ fn map_key_row(row: &PgRow) -> Result<StoredProviderCatalogKey, DataLayerError>
})
})
.transpose()?;
let total_response_time_ms = row_get::<Option<i32>>(row, "total_response_time_ms")?
let total_response_time_ms = row_get::<Option<i64>>(row, "total_response_time_ms")?
.map(|value| {
u32::try_from(value).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
crates/aether-data/src/repository/proxy_nodes/memory.rs
+2 -2
View File
@@ -561,7 +561,7 @@ impl ProxyNodeWriteRepository for InMemoryProxyNodeRepository {
};
if !node.tunnel_mode {
return Err(DataLayerError::InvalidInput(
"non-tunnel mode is no longer supported, please upgrade aether-proxy to use tunnel mode"
"non-tunnel mode is no longer supported, please upgrade aether-tunnel to use tunnel mode"
.to_string(),
));
}
@@ -1163,7 +1163,7 @@ mod tests {
dns_failures_delta: Some(0),
stream_errors_delta: Some(0),
proxy_metadata: Some(json!({"arch": "arm64"})),
proxy_version: Some("proxy-v2.2.0".to_string()),
proxy_version: Some("tunnel-v2.2.0".to_string()),
})
.await
.expect("heartbeat should succeed")
crates/aether-data/src/repository/proxy_nodes/mysql.rs
+1 -1
View File
@@ -755,7 +755,7 @@ WHERE is_manual = 0
};
if !node.tunnel_mode {
return Err(DataLayerError::InvalidInput(
"non-tunnel mode is no longer supported, please upgrade aether-proxy to use tunnel mode"
"non-tunnel mode is no longer supported, please upgrade aether-tunnel to use tunnel mode"
.to_string(),
));
}
crates/aether-data/src/repository/proxy_nodes/postgres.rs
+1 -1
View File
@@ -1165,7 +1165,7 @@ impl ProxyNodeWriteRepository for SqlxProxyNodeRepository {
};
if !existing.tunnel_mode {
return Err(DataLayerError::InvalidInput(
"non-tunnel mode is no longer supported, please upgrade aether-proxy to use tunnel mode"
"non-tunnel mode is no longer supported, please upgrade aether-tunnel to use tunnel mode"
.to_string(),
));
}
crates/aether-data/src/repository/proxy_nodes/sqlite.rs
+2 -2
View File
@@ -746,7 +746,7 @@ WHERE is_manual = 0
};
if !node.tunnel_mode {
return Err(DataLayerError::InvalidInput(
"non-tunnel mode is no longer supported, please upgrade aether-proxy to use tunnel mode"
"non-tunnel mode is no longer supported, please upgrade aether-tunnel to use tunnel mode"
.to_string(),
));
}
@@ -1431,7 +1431,7 @@ VALUES ('node-1', 'registered', 'ok', 3)
log_level: Some("debug".to_string()),
heartbeat_interval: Some(45),
scheduling_state: Some(Some("draining".to_string())),
upgrade_to: Some(Some("proxy-v2.0.0".to_string())),
upgrade_to: Some(Some("tunnel-v2.0.0".to_string())),
})
.await
.expect("remote config should update")
crates/aether-data/src/repository/proxy_nodes/types.rs
+4 -3
View File
@@ -570,7 +570,8 @@ fn normalize_proxy_version_label(value: &str) -> Option<String> {
}
Some(
trimmed
.strip_prefix("proxy-v")
.strip_prefix("tunnel-v")
.or_else(|| trimmed.strip_prefix("proxy-v"))
.unwrap_or(trimmed)
.to_ascii_lowercase(),
)
@@ -781,7 +782,7 @@ mod tests {
fn normalizes_reported_versions_and_clears_completed_upgrade_targets() {
let remote_config = json!({
"node_name": "edge-1",
"upgrade_to": "proxy-v2.0.0",
"upgrade_to": "tunnel-v2.0.0",
});
let proxy_metadata = json!({
"version": "2.0.0",
@@ -798,7 +799,7 @@ mod tests {
);
let reconciled =
reconcile_remote_config_after_heartbeat(Some(&remote_config), Some("proxy-v2.0.0"))
reconcile_remote_config_after_heartbeat(Some(&remote_config), Some("tunnel-v2.0.0"))
.expect("reconciled config should remain an object");
assert_eq!(reconciled.get("upgrade_to"), None);
assert_eq!(reconciled.get("node_name"), Some(&json!("edge-1")));
crates/aether-data/src/repository/settlement/memory.rs
+39 -22
View File
@@ -4,8 +4,8 @@ use std::sync::{Arc, RwLock};
use async_trait::async_trait;
use super::{
plan_finite_wallet_debit, SettlementWriteRepository, StoredUsageSettlement,
UsageSettlementInput, SETTLEMENT_EPSILON_USD,
plan_finite_wallet_debit, settlement_billing_status_for_usage_status,
SettlementWriteRepository, StoredUsageSettlement, UsageSettlementInput, SETTLEMENT_EPSILON_USD,
};
use crate::repository::wallet::{InMemoryWalletRepository, StoredWalletSnapshot};
use crate::DataLayerError;
@@ -102,11 +102,8 @@ impl SettlementWriteRepository for InMemorySettlementRepository {
})));
}
let mut final_billing_status = if input.status == "completed" {
"settled".to_string()
} else {
"void".to_string()
};
let mut final_billing_status =
settlement_billing_status_for_usage_status(&input.status).to_string();
let mut settlement = self.wallets.with_mut(|wallets| {
let wallet_id = input
.api_key_id
@@ -164,15 +161,9 @@ impl SettlementWriteRepository for InMemorySettlementRepository {
before_gift,
input.total_cost_usd,
);
if debit_plan.covered_usd() + SETTLEMENT_EPSILON_USD < input.total_cost_usd
{
final_billing_status = "insufficient_quota".to_string();
settlement.billing_status = final_billing_status.clone();
} else {
wallet.balance = before_recharge - debit_plan.recharge_deduction;
wallet.gift_balance = before_gift - debit_plan.gift_deduction;
wallet.total_consumed += input.total_cost_usd;
}
(wallet.balance, wallet.gift_balance) =
debit_plan.after_balances(before_recharge, before_gift);
wallet.total_consumed += input.total_cost_usd;
}
}
@@ -306,6 +297,32 @@ mod tests {
assert_eq!(settlement.wallet_balance_after, Some(9.0));
}
#[tokio::test]
async fn settles_cancelled_usage_against_wallet_and_provider_quota() {
let repository = InMemorySettlementRepository::seed(vec![sample_wallet()]);
let settlement = repository
.settle_usage(UsageSettlementInput {
request_id: "req-cancelled".to_string(),
user_id: Some("user-1".to_string()),
api_key_id: Some("key-1".to_string()),
api_key_is_standalone: false,
provider_id: Some("provider-1".to_string()),
status: "cancelled".to_string(),
billing_status: "pending".to_string(),
total_cost_usd: 3.0,
actual_total_cost_usd: 1.5,
finalized_at_unix_secs: Some(200),
})
.await
.expect("settlement should succeed")
.expect("settlement should exist");
assert_eq!(settlement.billing_status, "settled");
assert_eq!(settlement.wallet_balance_before, Some(12.0));
assert_eq!(settlement.wallet_balance_after, Some(9.0));
assert_eq!(settlement.provider_monthly_used_usd, Some(1.5));
}
#[tokio::test]
async fn standalone_key_settlement_never_falls_back_to_owner_wallet() {
let repository = InMemorySettlementRepository::seed(vec![sample_user_wallet(
@@ -336,7 +353,7 @@ mod tests {
}
#[tokio::test]
async fn finite_wallet_insufficient_balance_does_not_overdraw() {
async fn finite_wallet_insufficient_balance_overdraws_and_settles() {
let repository = InMemorySettlementRepository::seed(vec![sample_wallet()]);
let settlement = repository
.settle_usage(UsageSettlementInput {
@@ -355,12 +372,12 @@ mod tests {
.expect("settlement should succeed")
.expect("settlement should exist");
assert_eq!(settlement.billing_status, "insufficient_quota");
assert_eq!(settlement.billing_status, "settled");
assert_eq!(settlement.wallet_balance_before, Some(12.0));
assert_eq!(settlement.wallet_balance_after, Some(12.0));
assert_eq!(settlement.wallet_recharge_balance_after, Some(10.0));
assert_eq!(settlement.wallet_gift_balance_after, Some(2.0));
assert_eq!(settlement.provider_monthly_used_usd, None);
assert_eq!(settlement.wallet_balance_after, Some(-3.0));
assert_eq!(settlement.wallet_recharge_balance_after, Some(-3.0));
assert_eq!(settlement.wallet_gift_balance_after, Some(0.0));
assert_eq!(settlement.provider_monthly_used_usd, Some(7.5));
}
#[tokio::test]
crates/aether-data/src/repository/settlement/mod.rs
+37 -6
View File
@@ -9,11 +9,15 @@ const SETTLEMENT_EPSILON_USD: f64 = 0.000_000_01;
struct WalletDebitPlan {
recharge_deduction: f64,
gift_deduction: f64,
recharge_overdraft: f64,
}
impl WalletDebitPlan {
fn covered_usd(self) -> f64 {
self.recharge_deduction + self.gift_deduction
fn after_balances(self, recharge_balance: f64, gift_balance: f64) -> (f64, f64) {
(
recharge_balance - self.recharge_deduction - self.recharge_overdraft,
gift_balance - self.gift_deduction,
)
}
}
@@ -26,13 +30,22 @@ fn plan_finite_wallet_debit(
gift_balance: f64,
requested_usd: f64,
) -> WalletDebitPlan {
let recharge_deduction = recharge_balance.max(0.0).min(requested_usd.max(0.0));
let gift_deduction = gift_balance
.max(0.0)
.min((requested_usd - recharge_deduction).max(0.0));
let requested_usd = requested_usd.max(0.0);
let recharge_deduction = recharge_balance.max(0.0).min(requested_usd);
let after_recharge_remaining = (requested_usd - recharge_deduction).max(0.0);
let gift_deduction = gift_balance.max(0.0).min(after_recharge_remaining);
let recharge_overdraft = (after_recharge_remaining - gift_deduction).max(0.0);
WalletDebitPlan {
recharge_deduction,
gift_deduction,
recharge_overdraft,
}
}
fn settlement_billing_status_for_usage_status(status: &str) -> &'static str {
match status {
"completed" | "cancelled" => "settled",
_ => "void",
}
}
@@ -44,3 +57,21 @@ pub use memory::InMemorySettlementRepository;
pub use mysql::MysqlSettlementRepository;
pub use postgres::SqlxSettlementRepository;
pub use sqlite::SqliteSettlementRepository;
#[cfg(test)]
mod tests {
use super::settlement_billing_status_for_usage_status;
#[test]
fn cancelled_usage_status_is_billable() {
assert_eq!(
settlement_billing_status_for_usage_status("completed"),
"settled"
);
assert_eq!(
settlement_billing_status_for_usage_status("cancelled"),
"settled"
);
assert_eq!(settlement_billing_status_for_usage_status("failed"), "void");
}
}
crates/aether-data/src/repository/settlement/mysql.rs
+11 -15
View File
@@ -2,8 +2,9 @@ use async_trait::async_trait;
use sqlx::{mysql::MySqlRow, Row};
use super::{
finite_wallet_available_usd, plan_finite_wallet_debit, SettlementWriteRepository,
StoredUsageSettlement, UsageSettlementInput, SETTLEMENT_EPSILON_USD,
finite_wallet_available_usd, plan_finite_wallet_debit,
settlement_billing_status_for_usage_status, SettlementWriteRepository, StoredUsageSettlement,
UsageSettlementInput, SETTLEMENT_EPSILON_USD,
};
use crate::driver::mysql::MysqlPool;
use crate::error::SqlResultExt;
@@ -206,6 +207,7 @@ async fn consume_daily_quota_mysql(
request_id: &str,
total_cost_usd: f64,
wallet_available_usd: Option<f64>,
wallet_can_overdraft: bool,
now_unix_secs: i64,
) -> Result<DailyQuotaDebitResult, DataLayerError> {
if total_cost_usd <= 0.0 {
@@ -279,6 +281,7 @@ WHERE user_entitlement_id = ?
});
}
if allow_wallet_overage
&& !wallet_can_overdraft
&& wallet_available_usd.is_some_and(|available| {
total_remaining + available + SETTLEMENT_EPSILON_USD < total_cost_usd
})
@@ -364,11 +367,8 @@ impl SettlementWriteRepository for MysqlSettlementRepository {
return Ok(Some(settlement));
}
let mut final_billing_status = if input.status == "completed" {
"settled".to_string()
} else {
"void".to_string()
};
let mut final_billing_status =
settlement_billing_status_for_usage_status(&input.status).to_string();
let mut settlement = StoredUsageSettlement {
request_id: input.request_id.clone(),
wallet_id: None,
@@ -450,6 +450,7 @@ FOR UPDATE
None
};
let wallet_can_overdraft = wallet_row.is_some();
let wallet_available_usd = match wallet_row.as_ref() {
Some(row) => {
let limit_mode: String = row.try_get("limit_mode").map_sql_err()?;
@@ -486,6 +487,7 @@ FOR UPDATE
&input.request_id,
input.total_cost_usd,
wallet_available_usd,
wallet_can_overdraft,
updated_at,
)
.await?;
@@ -546,14 +548,8 @@ FOR UPDATE
before_gift,
wallet_debit_cost_usd,
);
if debit_plan.covered_usd() + SETTLEMENT_EPSILON_USD < wallet_debit_cost_usd
{
final_billing_status = "insufficient_quota".to_string();
settlement.billing_status = final_billing_status.clone();
} else {
after_recharge = before_recharge - debit_plan.recharge_deduction;
after_gift = before_gift - debit_plan.gift_deduction;
}
(after_recharge, after_gift) =
debit_plan.after_balances(before_recharge, before_gift);
}
if final_billing_status == "settled" {
sqlx::query(
crates/aether-data/src/repository/settlement/postgres.rs
+72 -33
View File
@@ -2,8 +2,9 @@ use async_trait::async_trait;
use sqlx::{PgPool, Row};
use super::{
finite_wallet_available_usd, plan_finite_wallet_debit, SettlementWriteRepository,
StoredUsageSettlement, UsageSettlementInput, SETTLEMENT_EPSILON_USD,
finite_wallet_available_usd, plan_finite_wallet_debit,
settlement_billing_status_for_usage_status, SettlementWriteRepository, StoredUsageSettlement,
UsageSettlementInput, SETTLEMENT_EPSILON_USD,
};
use crate::driver::postgres::PostgresTransactionRunner;
use crate::error::SqlxResultExt;
@@ -125,6 +126,22 @@ DO UPDATE SET
updated_at = NOW()
"#;
const ENQUEUE_PROVIDER_MONTHLY_USAGE_DELTA_SQL: &str = r#"
INSERT INTO usage_counter_deltas (
id,
request_id,
kind,
target_id,
total_cost_usd_delta
) VALUES (
$1,
$2,
'provider_monthly',
$3,
$4
)
"#;
#[derive(Debug, Clone)]
pub struct SqlxSettlementRepository {
tx_runner: PostgresTransactionRunner,
@@ -193,6 +210,37 @@ where
Ok(())
}
async fn enqueue_provider_monthly_usage_delta<'e, E>(
executor: E,
request_id: &str,
provider_id: &str,
total_cost_usd_delta: f64,
) -> Result<(), DataLayerError>
where
E: sqlx::Executor<'e, Database = sqlx::Postgres>,
{
let request_id = request_id.trim();
let provider_id = provider_id.trim();
if request_id.is_empty() || provider_id.is_empty() || total_cost_usd_delta == 0.0 {
return Ok(());
}
if !total_cost_usd_delta.is_finite() {
return Err(DataLayerError::UnexpectedValue(format!(
"provider monthly usage delta is not finite for {provider_id}"
)));
}
sqlx::query(ENQUEUE_PROVIDER_MONTHLY_USAGE_DELTA_SQL)
.bind(uuid::Uuid::new_v4().to_string())
.bind(request_id)
.bind(provider_id)
.bind(total_cost_usd_delta)
.execute(executor)
.await
.map_postgres_err()?;
Ok(())
}
#[derive(Debug, Default)]
struct DailyQuotaDebitResult {
debited_usd: f64,
@@ -264,6 +312,7 @@ async fn consume_daily_quota_postgres(
request_id: &str,
total_cost_usd: f64,
wallet_available_usd: Option<f64>,
wallet_can_overdraft: bool,
) -> Result<DailyQuotaDebitResult, DataLayerError> {
if total_cost_usd <= 0.0 {
return Ok(DailyQuotaDebitResult::default());
@@ -331,6 +380,7 @@ WHERE user_entitlement_id = $1
});
}
if allow_wallet_overage
&& !wallet_can_overdraft
&& wallet_available_usd.is_some_and(|available| {
total_remaining + available + SETTLEMENT_EPSILON_USD < total_cost_usd
})
@@ -408,11 +458,8 @@ impl SettlementWriteRepository for SqlxSettlementRepository {
return settlement_from_row(&usage_row).map(Some);
}
let mut final_billing_status = if input.status == "completed" {
"settled".to_string()
} else {
"void".to_string()
};
let mut final_billing_status =
settlement_billing_status_for_usage_status(&input.status).to_string();
let finalized_at =
i64::try_from(input.finalized_at_unix_secs.unwrap_or_else(|| {
std::time::SystemTime::now()
@@ -515,6 +562,7 @@ LIMIT 1
None
};
let wallet_can_overdraft = wallet_row.is_some();
let wallet_available_usd = match wallet_row.as_ref() {
Some(row) => {
let limit_mode: String =
@@ -555,6 +603,7 @@ LIMIT 1
&input.request_id,
input.total_cost_usd,
wallet_available_usd,
wallet_can_overdraft,
)
.await?;
if quota.insufficient {
@@ -601,16 +650,8 @@ LIMIT 1
before_gift,
wallet_debit_cost_usd,
);
if debit_plan.covered_usd() + SETTLEMENT_EPSILON_USD
< wallet_debit_cost_usd
{
final_billing_status = "insufficient_quota".to_string();
settlement.billing_status = final_billing_status.clone();
} else {
after_recharge =
before_recharge - debit_plan.recharge_deduction;
after_gift = before_gift - debit_plan.gift_deduction;
}
(after_recharge, after_gift) =
debit_plan.after_balances(before_recharge, before_gift);
}
if final_billing_status == "settled" {
sqlx::query(
@@ -663,23 +704,13 @@ WHERE id = $1
.as_deref()
.filter(|value| !value.is_empty())
{
let quota_row = sqlx::query(
r#"
UPDATE providers
SET
monthly_used_usd = COALESCE(monthly_used_usd, 0) + $2,
updated_at = NOW()
WHERE id = $1
RETURNING CAST(monthly_used_usd AS DOUBLE PRECISION) AS monthly_used_usd
"#,
enqueue_provider_monthly_usage_delta(
&mut **tx,
&input.request_id,
provider_id,
input.actual_total_cost_usd,
)
.bind(provider_id)
.bind(input.actual_total_cost_usd)
.fetch_optional(&mut **tx)
.await
.map_postgres_err()?;
settlement.provider_monthly_used_usd =
quota_row.and_then(|row| row.try_get("monthly_used_usd").ok());
.await?;
}
}
@@ -732,6 +763,14 @@ mod tests {
assert!(!source.contains("UPDATE \"usage\"\nSET\n wallet_id = $2"));
}
#[test]
fn settlement_sql_enqueues_provider_monthly_usage_delta() {
let source = include_str!("postgres.rs");
assert!(super::ENQUEUE_PROVIDER_MONTHLY_USAGE_DELTA_SQL.contains("usage_counter_deltas"));
assert!(super::ENQUEUE_PROVIDER_MONTHLY_USAGE_DELTA_SQL.contains("'provider_monthly'"));
assert!(!source.contains("UPDATE providers\nSET\n monthly_used_usd"));
}
#[test]
fn settlement_sql_blocks_standalone_key_owner_wallet_fallback() {
let source = include_str!("postgres.rs");
crates/aether-data/src/repository/settlement/sqlite.rs
+62 -16
View File
@@ -2,8 +2,9 @@ use async_trait::async_trait;
use sqlx::{sqlite::SqliteRow, Row};
use super::{
finite_wallet_available_usd, plan_finite_wallet_debit, SettlementWriteRepository,
StoredUsageSettlement, UsageSettlementInput, SETTLEMENT_EPSILON_USD,
finite_wallet_available_usd, plan_finite_wallet_debit,
settlement_billing_status_for_usage_status, SettlementWriteRepository, StoredUsageSettlement,
UsageSettlementInput, SETTLEMENT_EPSILON_USD,
};
use crate::driver::sqlite::{sqlite_optional_real, sqlite_real, SqlitePool};
use crate::error::SqlResultExt;
@@ -220,6 +221,7 @@ async fn consume_daily_quota_sqlite(
request_id: &str,
total_cost_usd: f64,
wallet_available_usd: Option<f64>,
wallet_can_overdraft: bool,
now_unix_secs: i64,
) -> Result<DailyQuotaDebitResult, DataLayerError> {
if total_cost_usd <= 0.0 {
@@ -292,6 +294,7 @@ WHERE user_entitlement_id = ?
});
}
if allow_wallet_overage
&& !wallet_can_overdraft
&& wallet_available_usd.is_some_and(|available| {
total_remaining + available + SETTLEMENT_EPSILON_USD < total_cost_usd
})
@@ -377,11 +380,8 @@ impl SettlementWriteRepository for SqliteSettlementRepository {
return Ok(Some(settlement));
}
let mut final_billing_status = if input.status == "completed" {
"settled".to_string()
} else {
"void".to_string()
};
let mut final_billing_status =
settlement_billing_status_for_usage_status(&input.status).to_string();
let mut settlement = StoredUsageSettlement {
request_id: input.request_id.clone(),
wallet_id: None,
@@ -461,6 +461,7 @@ LIMIT 1
None
};
let wallet_can_overdraft = wallet_row.is_some();
let wallet_available_usd = match wallet_row.as_ref() {
Some(row) => {
let limit_mode: String = row.try_get("limit_mode").map_sql_err()?;
@@ -497,6 +498,7 @@ LIMIT 1
&input.request_id,
input.total_cost_usd,
wallet_available_usd,
wallet_can_overdraft,
updated_at,
)
.await?;
@@ -557,14 +559,8 @@ LIMIT 1
before_gift,
wallet_debit_cost_usd,
);
if debit_plan.covered_usd() + SETTLEMENT_EPSILON_USD < wallet_debit_cost_usd
{
final_billing_status = "insufficient_quota".to_string();
settlement.billing_status = final_billing_status.clone();
} else {
after_recharge = before_recharge - debit_plan.recharge_deduction;
after_gift = before_gift - debit_plan.gift_deduction;
}
(after_recharge, after_gift) =
debit_plan.after_balances(before_recharge, before_gift);
}
if final_billing_status == "settled" {
sqlx::query(
@@ -815,6 +811,55 @@ mod tests {
assert_eq!(wallet_total, 12.0);
}
#[tokio::test]
async fn sqlite_repository_overdraws_finite_wallet_and_settles_usage() {
let pool = sqlx::sqlite::SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.expect("sqlite pool should connect");
run_sqlite_migrations(&pool)
.await
.expect("sqlite migrations should run");
seed_settlement_rows(&pool).await;
let repository = SqliteSettlementRepository::new(pool.clone());
let settlement = repository
.settle_usage(UsageSettlementInput {
request_id: "request-overdraw".to_string(),
user_id: Some("user-1".to_string()),
api_key_id: None,
api_key_is_standalone: false,
provider_id: Some("provider-1".to_string()),
status: "completed".to_string(),
billing_status: "pending".to_string(),
total_cost_usd: 15.0,
actual_total_cost_usd: 7.5,
finalized_at_unix_secs: Some(1_236),
})
.await
.expect("settlement should run")
.expect("usage should exist");
assert_eq!(settlement.billing_status, "settled");
assert_eq!(settlement.wallet_id.as_deref(), Some("wallet-1"));
assert_eq!(settlement.wallet_balance_before, Some(12.0));
assert_eq!(settlement.wallet_balance_after, Some(-3.0));
assert_eq!(settlement.wallet_recharge_balance_after, Some(-3.0));
assert_eq!(settlement.wallet_gift_balance_after, Some(0.0));
assert_eq!(settlement.provider_monthly_used_usd, Some(12.5));
let wallet = sqlx::query(
"SELECT balance, gift_balance, total_consumed FROM wallets WHERE id = 'wallet-1'",
)
.fetch_one(&pool)
.await
.expect("wallet should load");
assert_eq!(wallet.try_get::<f64, _>("balance").unwrap(), -3.0);
assert_eq!(wallet.try_get::<f64, _>("gift_balance").unwrap(), 0.0);
assert_eq!(wallet.try_get::<f64, _>("total_consumed").unwrap(), 15.0);
}
#[tokio::test]
async fn sqlite_repository_records_wallet_for_quota_covered_user_usage() {
let pool = sqlx::sqlite::SqlitePoolOptions::new()
@@ -885,7 +930,8 @@ INSERT INTO "usage" (
)
VALUES
('request-1', 'user-1', 'provider-1', 'completed', 'pending', 3.0, 2.0),
('request-2', 'user-1', 'provider-1', 'failed', 'pending', 3.0, 2.0);
('request-2', 'user-1', 'provider-1', 'failed', 'pending', 3.0, 2.0),
('request-overdraw', 'user-1', 'provider-1', 'completed', 'pending', 15.0, 7.5);
"#,
)
.execute(pool)
crates/aether-data/src/repository/usage/memory.rs
+23 -5
View File
@@ -29,11 +29,12 @@ use serde_json::Value;
use super::{
api_key_usage_contribution, provider_api_key_usage_contribution,
strip_deprecated_usage_display_fields, usage_can_recover_terminal_failure,
ApiKeyUsageContribution, ApiKeyUsageDelta, ProviderApiKeyUsageContribution,
ProviderApiKeyUsageDelta, ProviderApiKeyWindowUsageRequest, StoredProviderApiKeyUsageSummary,
StoredProviderApiKeyWindowUsageSummary, StoredProviderUsageSummary, StoredProviderUsageWindow,
StoredRequestUsageAudit, StoredUsageDailySummary, UpsertUsageRecord, UsageAuditListQuery,
UsageDailyHeatmapQuery, UsageReadRepository, UsageWriteRepository,
usage_request_metadata_client_family, ApiKeyUsageContribution, ApiKeyUsageDelta,
ProviderApiKeyUsageContribution, ProviderApiKeyUsageDelta, ProviderApiKeyWindowUsageRequest,
StoredProviderApiKeyUsageSummary, StoredProviderApiKeyWindowUsageSummary,
StoredProviderUsageSummary, StoredProviderUsageWindow, StoredRequestUsageAudit,
StoredUsageDailySummary, UpsertUsageRecord, UsageAuditListQuery, UsageDailyHeatmapQuery,
UsageReadRepository, UsageWriteRepository,
};
use crate::repository::auth::InMemoryAuthApiKeySnapshotRepository;
use crate::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
@@ -56,6 +57,7 @@ impl InMemoryUsageReadRepository {
let mut by_request_id = BTreeMap::new();
for mut item in items {
hydrate_legacy_body_refs(&mut item);
hydrate_client_family(&mut item);
by_request_id.insert(item.request_id.clone(), item);
}
Self {
@@ -75,6 +77,7 @@ impl InMemoryUsageReadRepository {
let mut detached_bodies = BTreeMap::new();
for mut item in items {
hydrate_legacy_body_refs(&mut item);
hydrate_client_family(&mut item);
let request_id = item.request_id.clone();
if let Some(body_ref) = detach_usage_body(
&request_id,
@@ -2546,6 +2549,13 @@ fn hydrate_legacy_body_refs(item: &mut StoredRequestUsageAudit) {
}
}
fn hydrate_client_family(item: &mut StoredRequestUsageAudit) {
if item.client_family.is_none() {
item.client_family = usage_request_metadata_client_family(item.request_metadata.as_ref())
.map(ToOwned::to_owned);
}
}
fn persisted_usage_body_ref(
incoming_ref: Option<&str>,
incoming_body: Option<&Value>,
@@ -2851,6 +2861,13 @@ impl UsageWriteRepository for InMemoryUsageReadRepository {
})
},
),
client_family: usage_request_metadata_client_family(request_metadata.as_ref())
.map(ToOwned::to_owned)
.or_else(|| {
existing
.as_ref()
.and_then(|existing| existing.client_family.clone())
}),
request_metadata,
created_at_unix_ms,
updated_at_unix_secs: usage.updated_at_unix_secs,
@@ -4399,6 +4416,7 @@ mod tests {
provider_endpoint_kind: Some("chat".to_string()),
has_format_conversion: false,
is_stream: false,
client_family: None,
input_tokens: 10,
output_tokens: 20,
total_tokens: 30,
crates/aether-data/src/repository/usage/mod.rs
+125 -13
View File
@@ -363,14 +363,15 @@ mod sqlite;
#[allow(unused_imports)]
pub(crate) use aether_data_contracts::repository::usage::{
PendingUsageCleanupSummary, ProviderApiKeyWindowUsageRequest, StoredProviderApiKeyUsageSummary,
StoredProviderApiKeyWindowUsageSummary, StoredProviderUsageSummary, StoredProviderUsageWindow,
StoredRequestUsageAudit, StoredUsageAuditAggregation, StoredUsageAuditSummary,
StoredUsageBreakdownSummaryRow, StoredUsageCacheAffinityHitSummary,
StoredUsageCacheAffinityIntervalRow, StoredUsageCacheHitSummary, StoredUsageCostSavingsSummary,
StoredUsageDailySummary, StoredUsageDashboardDailyBreakdownRow,
StoredUsageDashboardProviderCount, StoredUsageDashboardSummary,
StoredUsageErrorDistributionRow, StoredUsageLeaderboardSummary,
usage_request_metadata_client_family, ApiKeyLastUsedDelta, ManagementTokenCounterDelta,
PendingUsageCleanupSummary, ProviderApiKeyWindowUsageRequest, ProxyNodeCounterDelta,
StoredProviderApiKeyUsageSummary, StoredProviderApiKeyWindowUsageSummary,
StoredProviderUsageSummary, StoredProviderUsageWindow, StoredRequestUsageAudit,
StoredUsageAuditAggregation, StoredUsageAuditSummary, StoredUsageBreakdownSummaryRow,
StoredUsageCacheAffinityHitSummary, StoredUsageCacheAffinityIntervalRow,
StoredUsageCacheHitSummary, StoredUsageCostSavingsSummary, StoredUsageDailySummary,
StoredUsageDashboardDailyBreakdownRow, StoredUsageDashboardProviderCount,
StoredUsageDashboardSummary, StoredUsageErrorDistributionRow, StoredUsageLeaderboardSummary,
StoredUsagePerformancePercentilesRow, StoredUsageProviderPerformance,
StoredUsageProviderPerformanceProviderRow, StoredUsageProviderPerformanceSummary,
StoredUsageProviderPerformanceTimelineRow, StoredUsageSettledCostSummary,
@@ -379,7 +380,8 @@ pub(crate) use aether_data_contracts::repository::usage::{
UsageAuditListQuery, UsageAuditSummaryQuery, UsageBreakdownGroupBy, UsageBreakdownSummaryQuery,
UsageCacheAffinityHitSummaryQuery, UsageCacheAffinityIntervalGroupBy,
UsageCacheAffinityIntervalQuery, UsageCacheHitSummaryQuery, UsageCleanupPreviewCounts,
UsageCleanupSummary, UsageCleanupWindow, UsageCostSavingsSummaryQuery, UsageDailyHeatmapQuery,
UsageCleanupSummary, UsageCleanupWindow, UsageCostSavingsSummaryQuery,
UsageCounterFlushSummary, UsageCounterHealthSnapshot, UsageDailyHeatmapQuery,
UsageDashboardDailyBreakdownQuery, UsageDashboardProviderCountsQuery,
UsageDashboardSummaryQuery, UsageErrorDistributionQuery, UsageLeaderboardGroupBy,
UsageLeaderboardQuery, UsageMonitoringErrorCountQuery, UsageMonitoringErrorListQuery,
@@ -422,7 +424,10 @@ impl ApiKeyUsageDelta {
total_requests: after.total_requests - before.total_requests,
total_tokens: after.total_tokens - before.total_tokens,
total_cost_usd: after.total_cost_usd - before.total_cost_usd,
candidate_last_used_at_unix_secs: after.last_used_at_unix_secs,
candidate_last_used_at_unix_secs: newer_last_used_at(
before.last_used_at_unix_secs,
after.last_used_at_unix_secs,
),
removed_last_used_at_unix_secs: None,
}
}
@@ -529,7 +534,10 @@ impl ProviderApiKeyUsageDelta {
total_tokens: after.total_tokens - before.total_tokens,
total_cost_usd: after.total_cost_usd - before.total_cost_usd,
total_response_time_ms: after.total_response_time_ms - before.total_response_time_ms,
candidate_last_used_at_unix_secs: after.last_used_at_unix_secs,
candidate_last_used_at_unix_secs: newer_last_used_at(
before.last_used_at_unix_secs,
after.last_used_at_unix_secs,
),
removed_last_used_at_unix_secs: None,
usage_created_at_unix_secs: after.usage_created_at_unix_secs,
}
@@ -632,6 +640,9 @@ pub(crate) fn provider_api_key_usage_is_error(
pub(crate) fn provider_api_key_usage_contribution(
usage: &StoredRequestUsageAudit,
) -> Option<ProviderApiKeyUsageContribution> {
if matches!(usage.status.as_str(), "pending" | "streaming") {
return None;
}
let key_id = usage
.provider_api_key_id
.as_deref()
@@ -693,6 +704,9 @@ pub(crate) fn model_usage_contribution(
pub(crate) fn api_key_usage_contribution(
usage: &StoredRequestUsageAudit,
) -> Option<ApiKeyUsageContribution> {
if matches!(usage.status.as_str(), "pending" | "streaming") {
return None;
}
let api_key_id = usage
.api_key_id
.as_deref()
@@ -713,14 +727,23 @@ pub(crate) fn api_key_usage_contribution(
})
}
fn newer_last_used_at(before: Option<u64>, after: Option<u64>) -> Option<u64> {
match (before, after) {
(Some(before), Some(after)) if after > before => Some(after),
(None, Some(after)) => Some(after),
_ => None,
}
}
#[cfg(test)]
mod tests {
use super::{
api_key_usage_contribution, incoming_usage_can_recover_terminal_failure,
model_usage_contribution, provider_api_key_usage_contribution,
provider_api_key_usage_is_error, provider_api_key_usage_is_success,
strip_deprecated_usage_display_fields, usage_can_recover_terminal_failure, ModelUsageDelta,
StoredRequestUsageAudit, UpsertUsageRecord,
strip_deprecated_usage_display_fields, usage_can_recover_terminal_failure,
ApiKeyUsageDelta, ModelUsageDelta, ProviderApiKeyUsageDelta, StoredRequestUsageAudit,
UpsertUsageRecord,
};
#[test]
@@ -995,6 +1018,95 @@ mod tests {
assert_eq!(contribution.total_tokens, 20);
assert_eq!(contribution.total_cost_usd, 0.25);
assert_eq!(contribution.last_used_at_unix_secs, Some(123));
let mut streaming = usage.clone();
streaming.status = "streaming".to_string();
assert!(api_key_usage_contribution(&streaming).is_none());
let mut pending = usage;
pending.status = "pending".to_string();
assert!(api_key_usage_contribution(&pending).is_none());
}
#[test]
fn provider_api_key_usage_contribution_tracks_terminal_requests_only() {
let usage = StoredRequestUsageAudit::new(
"usage-1".to_string(),
"request-1".to_string(),
Some("user-1".to_string()),
Some("api-key-1".to_string()),
None,
None,
"OpenAI".to_string(),
"gpt-5".to_string(),
None,
Some("provider-1".to_string()),
None,
Some("provider-key-1".to_string()),
None,
None,
None,
None,
None,
None,
None,
false,
false,
12,
8,
20,
0.25,
0.25,
Some(200),
None,
None,
Some(120),
None,
"completed".to_string(),
"settled".to_string(),
123,
124,
Some(125),
)
.expect("usage should build");
assert!(provider_api_key_usage_contribution(&usage).is_some());
let mut streaming = usage.clone();
streaming.status = "streaming".to_string();
assert!(provider_api_key_usage_contribution(&streaming).is_none());
let mut pending = usage;
pending.status = "pending".to_string();
assert!(provider_api_key_usage_contribution(&pending).is_none());
}
#[test]
fn usage_delta_between_does_not_emit_duplicate_last_used_candidate() {
let api_key_contribution = super::ApiKeyUsageContribution {
api_key_id: "api-key-1".to_string(),
total_requests: 1,
total_tokens: 20,
total_cost_usd: 0.25,
last_used_at_unix_secs: Some(123),
};
assert!(ApiKeyUsageDelta::between(&api_key_contribution, &api_key_contribution).is_noop());
let provider_contribution = super::ProviderApiKeyUsageContribution {
key_id: "provider-key-1".to_string(),
request_count: 1,
success_count: 1,
error_count: 0,
total_tokens: 20,
total_cost_usd: 0.25,
total_response_time_ms: 120,
last_used_at_unix_secs: Some(123),
usage_created_at_unix_secs: Some(123),
};
assert!(
ProviderApiKeyUsageDelta::between(&provider_contribution, &provider_contribution,)
.is_noop()
);
}
#[test]
crates/aether-data/src/repository/usage/mysql.rs
+4 -2
View File
@@ -6,8 +6,8 @@ use sqlx::{mysql::MySqlRow, Row};
use super::{
provider_api_key_usage_is_error, provider_api_key_usage_is_success,
strip_deprecated_usage_display_fields, usage_can_recover_terminal_failure,
InMemoryUsageReadRepository, PendingUsageCleanupSummary, StoredRequestUsageAudit,
UpsertUsageRecord, UsageWriteRepository,
usage_request_metadata_client_family, InMemoryUsageReadRepository, PendingUsageCleanupSummary,
StoredRequestUsageAudit, UpsertUsageRecord, UsageWriteRepository,
};
use crate::driver::mysql::MysqlPool;
use crate::error::SqlResultExt;
@@ -850,6 +850,8 @@ fn map_usage_row(row: &MySqlRow) -> Result<StoredRequestUsageAudit, DataLayerErr
.map(|raw| serde_json::from_str(&raw))
.transpose()
.map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?;
audit.client_family = usage_request_metadata_client_family(audit.request_metadata.as_ref())
.map(ToOwned::to_owned);
let upstream_is_stream = row
.try_get::<Option<bool>, _>("upstream_is_stream")
.map_sql_err()?;
crates/aether-data/src/repository/usage/postgres/mod.rs
+1071 -75
View File
File diff suppressed because it is too large Load Diff
crates/aether-data/src/repository/usage/postgres/queries/apply_api_key_usage_delta_sql.sql
+1
View File
@@ -17,6 +17,7 @@ SET
SELECT MAX(created_at)
FROM "usage"
WHERE api_key_id = $1
AND status NOT IN ('pending', 'streaming')
)
ELSE last_used_at
END
crates/aether-data/src/repository/usage/postgres/queries/apply_provider_api_key_codex_window_usage_delta_sql.sql
-170
View File
@@ -1,170 +0,0 @@
WITH target_key AS (
SELECT
id,
COALESCE(status_snapshot::jsonb, '{}'::jsonb) AS snapshot
FROM provider_api_keys
WHERE id = $1
AND jsonb_typeof((status_snapshot::jsonb) -> 'quota' -> 'windows') = 'array'
AND lower(BTRIM(COALESCE((status_snapshot::jsonb) -> 'quota' ->> 'provider_type', ''))) = 'codex'
FOR UPDATE
),
window_items AS (
SELECT
target_key.id,
window_rows.window_item,
window_rows.window_ordinality
FROM target_key
CROSS JOIN LATERAL jsonb_array_elements(target_key.snapshot -> 'quota' -> 'windows')
WITH ORDINALITY AS window_rows(window_item, window_ordinality)
),
parsed_windows AS (
SELECT
window_items.id,
window_items.window_item,
window_items.window_ordinality,
lower(BTRIM(COALESCE(window_items.window_item ->> 'code', ''))) AS window_code,
CASE
WHEN text_values.reset_at_text ~ '^[0-9]+$'
AND (
length(text_values.reset_at_text) < 19
OR (
length(text_values.reset_at_text) = 19
AND text_values.reset_at_text <= '9223372036854775807'
)
)
THEN text_values.reset_at_text::BIGINT
ELSE NULL
END AS reset_at,
CASE
WHEN text_values.window_minutes_text ~ '^[0-9]+$'
AND (
length(text_values.window_minutes_text) < 19
OR (
length(text_values.window_minutes_text) = 19
AND text_values.window_minutes_text <= '9223372036854775807'
)
)
THEN text_values.window_minutes_text::BIGINT
ELSE CASE lower(BTRIM(COALESCE(window_items.window_item ->> 'code', '')))
WHEN '5h' THEN 300
WHEN 'weekly' THEN 10080
ELSE NULL
END
END AS window_minutes,
CASE
WHEN text_values.usage_reset_at_text ~ '^[0-9]+$'
AND (
length(text_values.usage_reset_at_text) < 19
OR (
length(text_values.usage_reset_at_text) = 19
AND text_values.usage_reset_at_text <= '9223372036854775807'
)
)
THEN text_values.usage_reset_at_text::BIGINT
ELSE NULL
END AS usage_reset_at,
CASE
WHEN text_values.request_count_text ~ '^[0-9]+$'
AND (
length(text_values.request_count_text) < 19
OR (
length(text_values.request_count_text) = 19
AND text_values.request_count_text <= '9223372036854775807'
)
)
THEN text_values.request_count_text::BIGINT
ELSE 0
END AS current_request_count,
CASE
WHEN text_values.total_tokens_text ~ '^[0-9]+$'
AND (
length(text_values.total_tokens_text) < 19
OR (
length(text_values.total_tokens_text) = 19
AND text_values.total_tokens_text <= '9223372036854775807'
)
)
THEN text_values.total_tokens_text::BIGINT
ELSE 0
END AS current_total_tokens,
CASE
WHEN text_values.total_cost_usd_text ~ '^[-+]?[0-9]+([.][0-9]+)?$'
THEN text_values.total_cost_usd_text::DOUBLE PRECISION
ELSE 0
END AS current_total_cost_usd
FROM window_items
CROSS JOIN LATERAL (
SELECT
BTRIM(COALESCE(window_items.window_item ->> 'reset_at', '')) AS reset_at_text,
BTRIM(COALESCE(window_items.window_item ->> 'window_minutes', '')) AS window_minutes_text,
BTRIM(COALESCE(window_items.window_item ->> 'usage_reset_at', '')) AS usage_reset_at_text,
BTRIM(COALESCE(window_items.window_item -> 'usage' ->> 'request_count', '')) AS request_count_text,
BTRIM(COALESCE(window_items.window_item -> 'usage' ->> 'total_tokens', '')) AS total_tokens_text,
BTRIM(COALESCE(window_items.window_item -> 'usage' ->> 'total_cost_usd', '')) AS total_cost_usd_text
) AS text_values
),
window_usage AS (
SELECT
parsed_windows.*,
CASE
WHEN parsed_windows.window_minutes BETWEEN 0 AND 153722867280912930
THEN parsed_windows.window_minutes * 60
ELSE NULL
END AS window_seconds
FROM parsed_windows
),
updated_windows AS (
SELECT
window_usage.id,
jsonb_agg(
CASE
WHEN window_usage.window_code IN ('5h', 'weekly')
AND window_usage.reset_at IS NOT NULL
AND window_usage.window_seconds IS NOT NULL
AND window_usage.reset_at >= window_usage.window_seconds
AND window_usage.reset_at > $2
AND GREATEST(
window_usage.reset_at - window_usage.window_seconds,
COALESCE(window_usage.usage_reset_at, 0)
) <= $2
THEN jsonb_set(
window_usage.window_item,
'{usage}',
jsonb_build_object(
'request_count',
LEAST(
GREATEST(window_usage.current_request_count::NUMERIC + $3::NUMERIC, 0),
9223372036854775807
)::BIGINT,
'total_tokens',
LEAST(
GREATEST(window_usage.current_total_tokens::NUMERIC + $4::NUMERIC, 0),
9223372036854775807
)::BIGINT,
'total_cost_usd',
to_char(
GREATEST(COALESCE(window_usage.current_total_cost_usd, 0) + $5, 0),
'FM999999999999999990.00000000'
)
),
true
)
ELSE window_usage.window_item
END
ORDER BY window_usage.window_ordinality
) AS windows
FROM window_usage
GROUP BY window_usage.id
)
UPDATE provider_api_keys AS keys
SET
status_snapshot = jsonb_set(
target_key.snapshot,
'{quota,windows}',
updated_windows.windows,
true
)::json,
updated_at = NOW()
FROM target_key
JOIN updated_windows ON updated_windows.id = target_key.id
WHERE keys.id = target_key.id
crates/aether-data/src/repository/usage/postgres/queries/apply_provider_api_key_usage_delta_sql.sql
+1
View File
@@ -20,6 +20,7 @@ SET
SELECT MAX(created_at)
FROM "usage"
WHERE provider_api_key_id = $1
AND status NOT IN ('pending', 'streaming')
)
ELSE last_used_at
END
crates/aether-data/src/repository/usage/postgres/queries/find_by_id_sql.sql
+4
View File
@@ -93,6 +93,10 @@ SELECT
"usage".first_byte_time_ms,
"usage".status,
COALESCE(usage_settlement_snapshots.billing_status, "usage".billing_status) AS billing_status,
COALESCE(
NULLIF(BTRIM("usage".request_metadata->'client_session_affinity'->>'client_family'), ''),
NULLIF(BTRIM("usage".request_metadata->>'client_family'), '')
) AS client_family,
COALESCE(usage_http_audits.request_headers, "usage".request_headers) AS request_headers,
"usage".request_body,
"usage".request_body_compressed,
crates/aether-data/src/repository/usage/postgres/queries/find_by_request_id_sql.sql
+4
View File
@@ -93,6 +93,10 @@ SELECT
"usage".first_byte_time_ms,
"usage".status,
COALESCE(usage_settlement_snapshots.billing_status, "usage".billing_status) AS billing_status,
COALESCE(
NULLIF(BTRIM("usage".request_metadata->'client_session_affinity'->>'client_family'), ''),
NULLIF(BTRIM("usage".request_metadata->>'client_family'), '')
) AS client_family,
COALESCE(usage_http_audits.request_headers, "usage".request_headers) AS request_headers,
"usage".request_body,
"usage".request_body_compressed,
crates/aether-data/src/repository/usage/postgres/queries/list_recent_usage_audits_prefix.sql
+19 -3
View File
@@ -93,6 +93,10 @@ SELECT
"usage".first_byte_time_ms,
"usage".status,
COALESCE(usage_settlement_snapshots.billing_status, "usage".billing_status) AS billing_status,
COALESCE(
NULLIF(BTRIM("usage".request_metadata->'client_session_affinity'->>'client_family'), ''),
NULLIF(BTRIM("usage".request_metadata->>'client_family'), '')
) AS client_family,
NULL::json AS request_headers,
NULL::json AS request_body,
NULL::bytea AS request_body_compressed,
@@ -106,9 +110,21 @@ SELECT
NULL::json AS client_response_body,
NULL::bytea AS client_response_body_compressed,
CASE
WHEN ("usage".request_metadata->>'client_requested_stream') IN ('true', 'false')
WHEN NULLIF(BTRIM("usage".request_metadata->>'client_ip'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'user_agent'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'request_path'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'request_path_and_query'), '') IS NOT NULL
OR ("usage".request_metadata->>'client_requested_stream') IN ('true', 'false')
OR ("usage".request_metadata->>'upstream_is_stream') IN ('true', 'false')
THEN jsonb_build_object(
THEN jsonb_strip_nulls(jsonb_build_object(
'client_ip',
NULLIF(BTRIM("usage".request_metadata->>'client_ip'), ''),
'user_agent',
NULLIF(BTRIM("usage".request_metadata->>'user_agent'), ''),
'request_path',
NULLIF(BTRIM("usage".request_metadata->>'request_path'), ''),
'request_path_and_query',
NULLIF(BTRIM("usage".request_metadata->>'request_path_and_query'), ''),
'client_requested_stream',
CASE
WHEN ("usage".request_metadata->>'client_requested_stream') IN ('true', 'false')
@@ -121,7 +137,7 @@ SELECT
THEN ("usage".request_metadata->>'upstream_is_stream')::boolean
ELSE NULL
END
)::json
))::json
ELSE NULL::json
END AS request_metadata,
NULL::varchar AS http_request_body_ref,
crates/aether-data/src/repository/usage/postgres/queries/list_usage_audits_prefix.sql
+19 -3
View File
@@ -93,6 +93,10 @@ SELECT
"usage".first_byte_time_ms,
"usage".status,
COALESCE(usage_settlement_snapshots.billing_status, "usage".billing_status) AS billing_status,
COALESCE(
NULLIF(BTRIM("usage".request_metadata->'client_session_affinity'->>'client_family'), ''),
NULLIF(BTRIM("usage".request_metadata->>'client_family'), '')
) AS client_family,
NULL::json AS request_headers,
NULL::json AS request_body,
NULL::bytea AS request_body_compressed,
@@ -106,9 +110,21 @@ SELECT
NULL::json AS client_response_body,
NULL::bytea AS client_response_body_compressed,
CASE
WHEN ("usage".request_metadata->>'client_requested_stream') IN ('true', 'false')
WHEN NULLIF(BTRIM("usage".request_metadata->>'client_ip'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'user_agent'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'request_path'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'request_path_and_query'), '') IS NOT NULL
OR ("usage".request_metadata->>'client_requested_stream') IN ('true', 'false')
OR ("usage".request_metadata->>'upstream_is_stream') IN ('true', 'false')
THEN jsonb_build_object(
THEN jsonb_strip_nulls(jsonb_build_object(
'client_ip',
NULLIF(BTRIM("usage".request_metadata->>'client_ip'), ''),
'user_agent',
NULLIF(BTRIM("usage".request_metadata->>'user_agent'), ''),
'request_path',
NULLIF(BTRIM("usage".request_metadata->>'request_path'), ''),
'request_path_and_query',
NULLIF(BTRIM("usage".request_metadata->>'request_path_and_query'), ''),
'client_requested_stream',
CASE
WHEN ("usage".request_metadata->>'client_requested_stream') IN ('true', 'false')
@@ -121,7 +137,7 @@ SELECT
THEN ("usage".request_metadata->>'upstream_is_stream')::boolean
ELSE NULL
END
)::json
))::json
ELSE NULL::json
END AS request_metadata,
NULL::varchar AS http_request_body_ref,
crates/aether-data/src/repository/usage/postgres/queries/rebuild_api_key_usage_stats_sql.sql
+2 -1
View File
@@ -1,7 +1,7 @@
WITH aggregated AS (
SELECT
api_key_id,
COUNT(*)::INTEGER AS total_requests,
COUNT(*)::BIGINT AS total_requests,
COALESCE(SUM(
GREATEST(
COALESCE(
@@ -16,6 +16,7 @@ WITH aggregated AS (
FROM usage_billing_facts AS "usage"
WHERE api_key_id IS NOT NULL
AND BTRIM(api_key_id) <> ''
AND status NOT IN ('pending', 'streaming')
GROUP BY api_key_id
)
UPDATE api_keys
crates/aether-data/src/repository/usage/postgres/queries/rebuild_provider_api_key_usage_stats_sql.sql
+5 -4
View File
@@ -1,7 +1,7 @@
WITH aggregated AS (
SELECT
provider_api_key_id,
COUNT(*)::INTEGER AS request_count,
COUNT(*)::BIGINT AS request_count,
COALESCE(SUM(
CASE
WHEN status IN ('completed', 'success', 'ok', 'billed', 'settled')
@@ -10,7 +10,7 @@ WITH aggregated AS (
THEN 1
ELSE 0
END
), 0)::INTEGER AS success_count,
), 0)::BIGINT AS success_count,
COALESCE(SUM(
CASE
WHEN status NOT IN ('pending', 'streaming')
@@ -22,7 +22,7 @@ WITH aggregated AS (
THEN 1
ELSE 0
END
), 0)::INTEGER AS error_count,
), 0)::BIGINT AS error_count,
COALESCE(SUM(
GREATEST(
COALESCE(
@@ -42,11 +42,12 @@ WITH aggregated AS (
THEN GREATEST(response_time_ms, 0)
ELSE 0
END
), 0)::INTEGER AS total_response_time_ms,
), 0)::BIGINT AS total_response_time_ms,
MAX(created_at) AS last_used_at
FROM usage_billing_facts AS "usage"
WHERE provider_api_key_id IS NOT NULL
AND BTRIM(provider_api_key_id) <> ''
AND status NOT IN ('pending', 'streaming')
GROUP BY provider_api_key_id
)
UPDATE provider_api_keys
crates/aether-data/src/repository/usage/postgres/queries/upsert_sql.sql
+26 -1
View File
@@ -24,6 +24,8 @@ INSERT INTO "usage" (
input_tokens,
output_tokens,
total_tokens,
input_output_total_tokens,
input_context_tokens,
cache_creation_input_tokens,
cache_creation_input_tokens_5m,
cache_creation_input_tokens_1h,
@@ -87,7 +89,28 @@ INSERT INTO "usage" (
),
COALESCE($22, 0),
COALESCE($23, 0),
COALESCE($24, COALESCE($22, 0) + COALESCE($23, 0)),
COALESCE(
$24,
COALESCE($22, 0)
+ COALESCE($23, 0)
+ COALESCE(
NULLIF(COALESCE($25, 0), 0),
COALESCE($26, 0) + COALESCE($27, 0),
0
)
+ COALESCE($28, 0)
),
COALESCE($22, 0) + COALESCE($23, 0),
COALESCE(
COALESCE($22, 0)
+ COALESCE(
NULLIF(COALESCE($25, 0), 0),
COALESCE($26, 0) + COALESCE($27, 0),
0
)
+ COALESCE($28, 0),
0
),
COALESCE($25, 0),
COALESCE($26, 0),
COALESCE($27, 0),
@@ -148,6 +171,8 @@ DO UPDATE SET
input_tokens = CASE WHEN "usage".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST("usage".input_tokens, EXCLUDED.input_tokens) ELSE "usage".input_tokens END,
output_tokens = CASE WHEN "usage".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST("usage".output_tokens, EXCLUDED.output_tokens) ELSE "usage".output_tokens END,
total_tokens = CASE WHEN "usage".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST("usage".total_tokens, EXCLUDED.total_tokens) ELSE "usage".total_tokens END,
input_output_total_tokens = CASE WHEN "usage".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST("usage".input_output_total_tokens, EXCLUDED.input_output_total_tokens) ELSE "usage".input_output_total_tokens END,
input_context_tokens = CASE WHEN "usage".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST("usage".input_context_tokens, EXCLUDED.input_context_tokens) ELSE "usage".input_context_tokens END,
cache_creation_input_tokens = CASE WHEN "usage".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST("usage".cache_creation_input_tokens, EXCLUDED.cache_creation_input_tokens) ELSE "usage".cache_creation_input_tokens END,
cache_creation_input_tokens_5m = CASE WHEN "usage".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST("usage".cache_creation_input_tokens_5m, EXCLUDED.cache_creation_input_tokens_5m) ELSE "usage".cache_creation_input_tokens_5m END,
cache_creation_input_tokens_1h = CASE WHEN "usage".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST("usage".cache_creation_input_tokens_1h, EXCLUDED.cache_creation_input_tokens_1h) ELSE "usage".cache_creation_input_tokens_1h END,
crates/aether-data/src/repository/usage/postgres/tests.rs
+42 -22
View File
@@ -234,25 +234,6 @@ fn usage_sql_summarizes_usage_by_provider_api_key_ids_in_database() {
assert!(super::SUMMARIZE_USAGE_BY_PROVIDER_API_KEY_IDS_SQL.contains("ANY($1::TEXT[])"));
}
#[test]
fn usage_sql_materializes_provider_key_window_usage_in_status_snapshot() {
assert!(super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL
.contains("UPDATE provider_api_keys AS keys"));
assert!(super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL.contains("jsonb_set"));
assert!(
super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL.contains("'{quota,windows}'")
);
assert!(super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL.contains("'usage'"));
assert!(
super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL.contains("WHEN '5h' THEN 300")
);
assert!(super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL
.contains("WHEN 'weekly' THEN 10080"));
assert!(
!super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL.contains("usage_billing_facts")
);
}
#[test]
fn usage_sql_rebuilds_provider_key_window_usage_into_status_snapshot() {
assert!(super::REBUILD_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_STATS_SQL
@@ -280,18 +261,44 @@ fn usage_sql_serializes_request_id_upserts_before_reading_previous_usage() {
.contains("lock_usage_request_id_in_tx(tx, &usage.request_id).await?;"));
}
#[test]
fn usage_sql_moves_shared_counter_updates_behind_outbox() {
let source = include_str!("mod.rs");
assert!(super::INSERT_USAGE_COUNTER_DELTA_SQL.contains("usage_counter_deltas"));
assert!(super::CLAIM_USAGE_COUNTER_DELTAS_SQL.contains("FOR UPDATE SKIP LOCKED"));
assert!(super::MARK_USAGE_COUNTER_DELTAS_PROCESSED_SQL.contains("processed_at = NOW()"));
assert!(super::TRY_LOCK_USAGE_COUNTER_FLUSH_SQL.contains("pg_try_advisory_xact_lock"));
assert!(source.contains("enqueue_api_key_usage_delta_in_tx("));
assert!(source.contains("enqueue_provider_api_key_usage_delta_in_tx("));
assert!(source.contains("enqueue_model_usage_delta_in_tx("));
assert!(source.contains("apply_provider_api_key_main_usage_delta_in_tx("));
assert!(source.contains("USAGE_COUNTER_KIND_PROVIDER_MONTHLY"));
assert!(source.contains("apply_provider_monthly_usage_delta_in_tx("));
}
#[test]
fn usage_sql_exposes_counter_outbox_health() {
assert!(super::READ_USAGE_COUNTER_HEALTH_SQL.contains("pending_rows"));
assert!(super::READ_USAGE_COUNTER_HEALTH_SQL.contains("oldest_pending_created_at_unix_secs"));
assert!(super::READ_USAGE_COUNTER_HEALTH_SQL.contains("latest_processed_at_unix_secs"));
assert!(super::READ_PENDING_USAGE_COUNTER_DELTAS_BY_KIND_SQL.contains("GROUP BY kind"));
}
#[test]
fn usage_sql_rebuild_matches_online_api_key_usage_semantics() {
assert!(super::REBUILD_API_KEY_USAGE_STATS_SQL.contains("COUNT(*)::INTEGER"));
assert!(super::REBUILD_API_KEY_USAGE_STATS_SQL.contains("COUNT(*)::BIGINT"));
assert!(super::REBUILD_API_KEY_USAGE_STATS_SQL.contains("COALESCE("));
assert!(super::REBUILD_API_KEY_USAGE_STATS_SQL.contains("total_tokens,"));
assert!(super::REBUILD_API_KEY_USAGE_STATS_SQL
.contains("COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)"));
assert!(super::REBUILD_API_KEY_USAGE_STATS_SQL.contains("AND BTRIM(api_key_id) <> ''"));
assert!(super::REBUILD_API_KEY_USAGE_STATS_SQL
.contains("AND status NOT IN ('pending', 'streaming')"));
}
#[test]
fn usage_sql_rebuild_matches_online_provider_key_usage_semantics() {
assert!(super::REBUILD_PROVIDER_API_KEY_USAGE_STATS_SQL.contains("COUNT(*)::BIGINT"));
assert!(super::REBUILD_PROVIDER_API_KEY_USAGE_STATS_SQL
.contains("NULLIF(BTRIM(error_message), '') IS NULL"));
assert!(super::REBUILD_PROVIDER_API_KEY_USAGE_STATS_SQL.contains("COALESCE("));
@@ -300,6 +307,8 @@ fn usage_sql_rebuild_matches_online_provider_key_usage_semantics() {
.contains("COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)"));
assert!(super::REBUILD_PROVIDER_API_KEY_USAGE_STATS_SQL
.contains("AND BTRIM(provider_api_key_id) <> ''"));
assert!(super::REBUILD_PROVIDER_API_KEY_USAGE_STATS_SQL
.contains("AND status NOT IN ('pending', 'streaming')"));
}
#[test]
@@ -503,8 +512,7 @@ fn usage_sql_raw_aggregates_use_canonical_billing_facts() {
.contains("FROM usage_billing_facts AS \"usage\""));
assert!(super::SUMMARIZE_USAGE_TOTALS_BY_USER_IDS_SQL
.contains("FROM usage_billing_facts AS \"usage\""));
assert!(!super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL
.contains("usage_billing_facts AS \"usage\""));
assert!(!source.contains("apply_provider_api_key_codex_window_usage_delta_in_tx"));
}
#[test]
@@ -633,6 +641,14 @@ fn usage_sql_uses_json_null_placeholders_for_usage_payload_columns() {
super::LIST_USAGE_AUDITS_PREFIX,
super::LIST_RECENT_USAGE_AUDITS_PREFIX,
] {
assert!(sql.contains("jsonb_strip_nulls(jsonb_build_object("));
assert!(sql.contains("'client_ip'"));
assert!(sql.contains("request_metadata->>'client_ip'"));
assert!(sql.contains("'user_agent'"));
assert!(sql.contains("request_metadata->>'user_agent'"));
assert!(sql.contains("AS client_family"));
assert!(sql.contains("request_metadata->'client_session_affinity'->>'client_family'"));
assert!(sql.contains("request_metadata->>'client_family'"));
assert!(sql.contains("CAST(\"usage\".input_tokens AS INTEGER) AS input_tokens"));
assert!(sql.contains(
"usage_settlement_snapshots.billing_input_tokens AS settlement_billing_input_tokens"
@@ -690,6 +706,8 @@ fn usage_sql_upsert_returning_includes_routing_placeholders() {
super::UPSERT_SQL.contains("NULL::varchar AS settlement_billing_snapshot_schema_version")
);
assert!(super::UPSERT_SQL.contains("NULL::double precision AS settlement_input_price_per_1m"));
assert!(super::UPSERT_SQL.contains("input_output_total_tokens"));
assert!(super::UPSERT_SQL.contains("input_context_tokens"));
}
#[test]
@@ -721,6 +739,8 @@ fn usage_sql_updates_usage_mirror_columns_from_terminal_events_only() {
"input_tokens = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".input_tokens, EXCLUDED.input_tokens) ELSE \"usage\".input_tokens END",
"output_tokens = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".output_tokens, EXCLUDED.output_tokens) ELSE \"usage\".output_tokens END",
"total_tokens = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".total_tokens, EXCLUDED.total_tokens) ELSE \"usage\".total_tokens END",
"input_output_total_tokens = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".input_output_total_tokens, EXCLUDED.input_output_total_tokens) ELSE \"usage\".input_output_total_tokens END",
"input_context_tokens = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".input_context_tokens, EXCLUDED.input_context_tokens) ELSE \"usage\".input_context_tokens END",
"cache_creation_input_tokens = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".cache_creation_input_tokens, EXCLUDED.cache_creation_input_tokens) ELSE \"usage\".cache_creation_input_tokens END",
"cache_creation_input_tokens_5m = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".cache_creation_input_tokens_5m, EXCLUDED.cache_creation_input_tokens_5m) ELSE \"usage\".cache_creation_input_tokens_5m END",
"cache_creation_input_tokens_1h = CASE WHEN \"usage\".billing_status = 'pending' AND EXCLUDED.status IN ('completed', 'failed', 'cancelled') THEN GREATEST(\"usage\".cache_creation_input_tokens_1h, EXCLUDED.cache_creation_input_tokens_1h) ELSE \"usage\".cache_creation_input_tokens_1h END",
crates/aether-data/src/repository/usage/sqlite.rs
+4 -1
View File
@@ -8,7 +8,8 @@ use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite};
use super::{
strip_deprecated_usage_display_fields, usage_can_recover_terminal_failure,
PendingUsageCleanupSummary, ProviderApiKeyWindowUsageRequest, StoredProviderApiKeyUsageSummary,
usage_request_metadata_client_family, PendingUsageCleanupSummary,
ProviderApiKeyWindowUsageRequest, StoredProviderApiKeyUsageSummary,
StoredProviderApiKeyWindowUsageSummary, StoredProviderUsageSummary, StoredRequestUsageAudit,
StoredUsageAuditAggregation, StoredUsageAuditSummary, StoredUsageBreakdownSummaryRow,
StoredUsageCacheAffinityHitSummary, StoredUsageCacheAffinityIntervalRow,
@@ -3812,6 +3813,8 @@ fn map_usage_row(row: &SqliteRow) -> Result<StoredRequestUsageAudit, DataLayerEr
.map(|raw| serde_json::from_str(&raw))
.transpose()
.map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?;
audit.client_family = usage_request_metadata_client_family(audit.request_metadata.as_ref())
.map(ToOwned::to_owned);
let upstream_is_stream = row
.try_get::<Option<i64>, _>("upstream_is_stream")
.map_sql_err()?
crates/aether-data/src/repository/users/postgres.rs
+4 -1
View File
@@ -336,8 +336,11 @@ SELECT
users.role::text AS role,
users.auth_source::text AS auth_source,
users.allowed_providers,
users.allowed_providers_mode,
users.allowed_api_formats,
users.allowed_api_formats_mode,
users.allowed_models,
users.allowed_models_mode,
users.is_active,
users.is_deleted,
users.created_at,
@@ -353,7 +356,7 @@ const TOUCH_OAUTH_LINK_SQL: &str = r#"
UPDATE user_oauth_links
SET provider_username = COALESCE($3, provider_username),
provider_email = COALESCE($4, provider_email),
extra_data = COALESCE($5, extra_data),
extra_data = COALESCE($5::json, extra_data),
last_login_at = $6
WHERE provider_type = $1
AND provider_user_id = $2