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

Merge branch 'fawney19:main' into main

Browse Source
This commit is contained in:
ZheFox
2026-05-20 22:49:49 +08:00
committed by GitHub
parent 5a9f19cbf2 d49b0a8a45
commit 779877acd0
38 changed files with 1523 additions and 172 deletions
Download Patch File Download Diff File Expand all files Collapse all files

128
crates/aether-ai-formats/src/formats/openai/responses/codex.rs
View File

@@ -375,6 +375,51 @@ fn ensure_codex_chat_reasoning_defaults(
.or_insert_with(|| json!(CODEX_DEFAULT_REASONING_SUMMARY));
}
fn codex_tool_type_rejects_top_level_name(tool_type: &str) -> bool {
let normalized = tool_type.trim().to_ascii_lowercase();
!normalized.is_empty()
&& normalized != "function"
&& normalized != "custom"
&& normalized != "namespace"
}
fn strip_codex_hosted_tool_names_for_backend(body_object: &mut serde_json::Map<String, Value>) {
let Some(tools) = body_object.get_mut("tools").and_then(Value::as_array_mut) else {
return;
};
for tool in tools {
let Some(tool_object) = tool.as_object_mut() else {
continue;
};
if tool_object
.get("type")
.and_then(Value::as_str)
.is_some_and(codex_tool_type_rejects_top_level_name)
{
tool_object.remove("name");
}
}
}
fn strip_codex_hosted_tool_choice_name_for_backend(
body_object: &mut serde_json::Map<String, Value>,
) {
let Some(tool_choice_object) = body_object
.get_mut("tool_choice")
.and_then(Value::as_object_mut)
else {
return;
};
if tool_choice_object
.get("type")
.and_then(Value::as_str)
.is_some_and(codex_tool_type_rejects_top_level_name)
{
tool_choice_object.remove("name");
}
}
pub fn apply_codex_openai_responses_special_body_edits(
provider_request_body: &mut Value,
provider_type: &str,
@@ -420,6 +465,8 @@ pub fn apply_codex_openai_responses_special_body_edits(
{
body_object.insert("instructions".to_string(), json!(""));
}
strip_codex_hosted_tool_names_for_backend(body_object);
strip_codex_hosted_tool_choice_name_for_backend(body_object);
if is_openai_image_request(provider_api_format)
|| codex_openai_responses_tool_choice_references_image_generation(body_object)
{
@@ -620,6 +667,87 @@ mod tests {
assert_eq!(provider_request_body["parallel_tool_calls"], json!(false));
}
#[test]
fn codex_responses_body_edits_preserve_function_tools_for_codex_backend() {
let mut provider_request_body = json!({
"input": [],
"model": "gpt-5.4",
"tools": [{
"type": "function",
"name": "lookup_account",
"description": "Lookup an account by id.",
"parameters": {
"type": "object",
"properties": {
"account_id": {
"type": "string"
}
},
"required": ["account_id"],
"additionalProperties": false
},
"strict": true
}],
"tool_choice": {
"type": "function",
"name": "lookup_account"
}
});
apply_codex_openai_responses_special_body_edits(
&mut provider_request_body,
"codex",
"openai:responses",
None,
None,
);
assert_eq!(
provider_request_body["tools"][0]["name"],
json!("lookup_account")
);
assert_eq!(
provider_request_body["tools"][0]["parameters"]["properties"]["account_id"]["type"],
json!("string")
);
assert_eq!(
provider_request_body["tool_choice"]["name"],
json!("lookup_account")
);
assert!(provider_request_body["tools"][0].get("function").is_none());
}
#[test]
fn codex_responses_body_edits_strip_name_from_hosted_web_search_tool() {
let mut provider_request_body = json!({
"input": [],
"model": "gpt-5.4",
"tools": [{
"type": "web_search",
"name": "web_search"
}],
"tool_choice": {
"type": "web_search",
"name": "web_search"
}
});
apply_codex_openai_responses_special_body_edits(
&mut provider_request_body,
"codex",
"openai:responses",
None,
None,
);
assert!(provider_request_body["tools"][0].get("name").is_none());
assert!(provider_request_body["tool_choice"].get("name").is_none());
assert_eq!(
provider_request_body["tool_choice"]["type"],
json!("web_search")
);
}
#[test]
fn compact_body_edits_strip_include_store_and_stream() {
let mut provider_request_body = json!({

10
crates/aether-data/migrations/mysql/20260520010000_preserve_usage_provider_identity.sql Normal file
View File

@@ -0,0 +1,10 @@
-- Usage is a historical fact table. Backfill nullable provider_id snapshots
-- from the unique provider name where the catalog row still exists.
UPDATE `usage` AS usage_rows
JOIN providers
ON providers.name = TRIM(usage_rows.provider_name)
SET usage_rows.provider_id = providers.id
WHERE usage_rows.provider_id IS NULL
AND TRIM(COALESCE(usage_rows.provider_name, '')) <> ''
AND LOWER(TRIM(COALESCE(usage_rows.provider_name, ''))) NOT IN ('unknown', 'unknow', 'pending');

15
crates/aether-data/migrations/postgres/20260403000000_baseline.sql
View File

@@ -4427,21 +4427,6 @@ END $mig$;
--
-- Name: usage usage_provider_id_fkey; Type: FK CONSTRAINT; Schema: public; Owner: -
--
DO $mig$ BEGIN
ALTER TABLE ONLY public.usage
ADD CONSTRAINT usage_provider_id_fkey FOREIGN KEY (provider_id) REFERENCES public.providers(id) ON DELETE SET NULL;
EXCEPTION
WHEN duplicate_object THEN NULL;
WHEN duplicate_table THEN NULL;
WHEN invalid_table_definition THEN NULL;
END $mig$;
--
-- Name: usage usage_user_id_fkey; Type: FK CONSTRAINT; Schema: public; Owner: -
--

14
crates/aether-data/migrations/postgres/20260520010000_preserve_usage_provider_identity.sql Normal file
View File

@@ -0,0 +1,14 @@
-- Usage is a historical fact table. Keep the provider_id snapshot even if the
-- provider catalog row is deleted, and backfill rows that can still be matched
-- by the unique provider name.
UPDATE public.usage AS usage_rows
SET provider_id = providers.id
FROM public.providers AS providers
WHERE usage_rows.provider_id IS NULL
AND BTRIM(COALESCE(usage_rows.provider_name, '')) <> ''
AND lower(BTRIM(COALESCE(usage_rows.provider_name, ''))) NOT IN ('unknown', 'unknow', 'pending')
AND providers.name = BTRIM(usage_rows.provider_name);
ALTER TABLE ONLY public.usage
DROP CONSTRAINT IF EXISTS usage_provider_id_fkey;

18
crates/aether-data/migrations/sqlite/20260520010000_preserve_usage_provider_identity.sql Normal file
View File

@@ -0,0 +1,18 @@
-- Usage is a historical fact table. Backfill nullable provider_id snapshots
-- from the unique provider name where the catalog row still exists.
UPDATE "usage"
SET provider_id = (
SELECT providers.id
FROM providers
WHERE providers.name = TRIM("usage".provider_name)
LIMIT 1
)
WHERE provider_id IS NULL
AND TRIM(COALESCE(provider_name, '')) <> ''
AND LOWER(TRIM(COALESCE(provider_name, ''))) NOT IN ('unknown', 'unknow', 'pending')
AND EXISTS (
SELECT 1
FROM providers
WHERE providers.name = TRIM("usage".provider_name)
);

15
crates/aether-data/schema/bootstrap/postgres/005_foreign_keys.sql
View File

@@ -627,21 +627,6 @@ END $mig$;
--
-- Name: usage usage_provider_id_fkey; Type: FK CONSTRAINT; Schema: public; Owner: -
--
DO $mig$ BEGIN
ALTER TABLE ONLY public.usage
ADD CONSTRAINT usage_provider_id_fkey FOREIGN KEY (provider_id) REFERENCES public.providers(id) ON DELETE SET NULL;
EXCEPTION
WHEN duplicate_object THEN NULL;
WHEN duplicate_table THEN NULL;
WHEN invalid_table_definition THEN NULL;
END $mig$;
--
-- Name: usage usage_user_id_fkey; Type: FK CONSTRAINT; Schema: public; Owner: -
--

15
crates/aether-data/schema/drivers/postgres/baseline/005_foreign_keys.sql
View File

@@ -552,21 +552,6 @@ END $mig$;
--
-- Name: usage usage_provider_id_fkey; Type: FK CONSTRAINT; Schema: public; Owner: -
--
DO $mig$ BEGIN
ALTER TABLE ONLY public.usage
ADD CONSTRAINT usage_provider_id_fkey FOREIGN KEY (provider_id) REFERENCES public.providers(id) ON DELETE SET NULL;
EXCEPTION
WHEN duplicate_object THEN NULL;
WHEN duplicate_table THEN NULL;
WHEN invalid_table_definition THEN NULL;
END $mig$;
--
-- Name: usage usage_user_id_fkey; Type: FK CONSTRAINT; Schema: public; Owner: -
--

2
crates/aether-data/src/lifecycle/bootstrap/postgres.rs
View File

@@ -7,7 +7,7 @@ use tracing::info;
// Generated by build.rs from schema/bootstrap/postgres.
pub(crate) static EMPTY_DATABASE_SNAPSHOT_SQL: &str =
include_str!(concat!(env!("OUT_DIR"), "/empty_database_snapshot.sql"));
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260520000000;
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260520010000;
const PUBLIC_BASE_TABLE_COUNT_SQL: &str = r#"
SELECT COUNT(*)::BIGINT

4
crates/aether-data/src/lifecycle/migrate/tests.rs
View File

@@ -311,6 +311,7 @@ fn empty_database_snapshot_covers_current_cutoff_versions() {
20260519120000,
20260519130000,
20260520000000,
20260520010000,
]
);
}
@@ -631,6 +632,7 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() {
20260519120000,
20260519130000,
20260520000000,
20260520010000,
]
);
assert_eq!(
@@ -653,6 +655,7 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() {
20260519120000,
20260519130000,
20260520000000,
20260520010000,
]
);
}
@@ -1175,6 +1178,7 @@ fn pending_migrations_from_applied_skips_versions_already_applied() {
20260519120000,
20260519130000,
20260520000000,
20260520010000,
]
);
}

103
crates/aether-data/src/repository/usage/memory.rs
View File

@@ -983,6 +983,30 @@ fn usage_provider_display_name(item: &StoredRequestUsageAudit) -> Option<String>
}
}
fn usage_provider_id(item: &StoredRequestUsageAudit) -> Option<String> {
let provider_id = item.provider_id.as_deref()?.trim();
if provider_id.is_empty() || usage_reserved_provider_label(provider_id) {
None
} else {
Some(provider_id.to_string())
}
}
fn usage_provider_aggregation_identity(
item: &StoredRequestUsageAudit,
) -> Option<(String, Option<String>, String)> {
let display_name = usage_provider_display_name(item);
if let Some(provider_id) = usage_provider_id(item) {
return Some((provider_id, display_name, "provider_id".to_string()));
}
let display_name = display_name?;
Some((
display_name.clone(),
Some(display_name),
"legacy_name".to_string(),
))
}
#[async_trait]
impl UsageReadRepository for InMemoryUsageReadRepository {
async fn find_by_id(
@@ -1167,14 +1191,14 @@ impl UsageReadRepository for InMemoryUsageReadRepository {
|| item.created_at_unix_ms >= query.created_until_unix_secs
|| matches!(item.status.as_str(), "pending" | "streaming")
|| (query.exclude_reserved_provider_labels
&& usage_provider_display_name(item).is_none())
&& usage_provider_aggregation_identity(item).is_none())
{
continue;
}
let provider_display_name =
let provider_identity =
if matches!(query.group_by, UsageAuditAggregationGroupBy::Provider) {
match usage_provider_display_name(item) {
match usage_provider_aggregation_identity(item) {
Some(value) => Some(value),
None => continue,
}
@@ -1184,19 +1208,11 @@ impl UsageReadRepository for InMemoryUsageReadRepository {
let group_key = match query.group_by {
UsageAuditAggregationGroupBy::Model => item.model.clone(),
UsageAuditAggregationGroupBy::Provider => {
let display_name = provider_display_name
.as_deref()
.expect("provider display name is set for provider aggregation");
item.provider_id
.as_deref()
.map(str::trim)
.filter(|provider_id| {
!provider_id.is_empty() && !usage_reserved_provider_label(provider_id)
})
.unwrap_or(display_name)
.to_string()
}
UsageAuditAggregationGroupBy::Provider => provider_identity
.as_ref()
.expect("provider identity is set for provider aggregation")
.0
.clone(),
UsageAuditAggregationGroupBy::ApiFormat => item
.api_format
.clone()
@@ -1211,7 +1227,17 @@ impl UsageReadRepository for InMemoryUsageReadRepository {
&& (bucket.display_name.is_none()
|| bucket.display_name.as_deref() == Some("Unknown"))
{
bucket.display_name = provider_display_name;
bucket.display_name = provider_identity
.as_ref()
.and_then(|(_, display_name, _)| display_name.clone());
}
if matches!(query.group_by, UsageAuditAggregationGroupBy::Provider)
&& (bucket.secondary_name.is_none()
|| bucket.secondary_name.as_deref() == Some("legacy_name"))
{
bucket.secondary_name = provider_identity
.as_ref()
.map(|(_, _, identity_source)| identity_source.clone());
}
bucket.request_count = bucket.request_count.saturating_add(1);
bucket.total_tokens = bucket.total_tokens.saturating_add(item.total_tokens);
@@ -3149,6 +3175,12 @@ mod tests {
#[tokio::test]
async fn provider_aggregation_skips_unknown_provider_labels() {
let valid_provider = sample_usage("req-valid-provider", 300);
let mut legacy_provider = sample_usage("req-legacy-provider", 250);
legacy_provider.provider_id = None;
legacy_provider.provider_name = "Legacy Provider".to_string();
let mut unknown = sample_usage("req-unknown-provider", 100);
unknown.provider_id = None;
unknown.provider_name = "unknown".to_string();
@@ -3158,7 +3190,8 @@ mod tests {
typo_unknown.provider_name = "unknow".to_string();
let repository = InMemoryUsageReadRepository::seed(vec![
sample_usage("req-valid-provider", 300),
valid_provider,
legacy_provider,
unknown,
typo_unknown,
]);
@@ -3174,9 +3207,29 @@ mod tests {
.await
.expect("aggregation should succeed");
assert_eq!(rows.len(), 1);
assert_eq!(rows[0].group_key, "provider-1");
assert_eq!(rows[0].display_name.as_deref(), Some("OpenAI"));
assert_eq!(rows.len(), 2);
let provider_id_row = rows
.iter()
.find(|row| row.group_key == "provider-1")
.expect("provider_id row should be present");
assert_eq!(provider_id_row.display_name.as_deref(), Some("OpenAI"));
assert_eq!(
provider_id_row.secondary_name.as_deref(),
Some("provider_id")
);
let legacy_name_row = rows
.iter()
.find(|row| row.group_key == "Legacy Provider")
.expect("legacy provider name row should be present");
assert_eq!(
legacy_name_row.display_name.as_deref(),
Some("Legacy Provider")
);
assert_eq!(
legacy_name_row.secondary_name.as_deref(),
Some("legacy_name")
);
}
#[tokio::test]
@@ -3193,11 +3246,15 @@ mod tests {
pending_provider.provider_id = None;
pending_provider.provider_name = "pending".to_string();
let mut id_only_provider = sample_usage("req-id-only-provider", 350);
id_only_provider.provider_name = "unknown".to_string();
let repository = InMemoryUsageReadRepository::seed(vec![
sample_usage("req-valid-provider", 400),
unknown,
typo_unknown,
pending_provider,
id_only_provider,
]);
let model_rows = repository
@@ -3212,7 +3269,7 @@ mod tests {
.expect("model aggregation should succeed");
assert_eq!(model_rows.len(), 1);
assert_eq!(model_rows[0].group_key, "gpt-4.1");
assert_eq!(model_rows[0].request_count, 1);
assert_eq!(model_rows[0].request_count, 2);
let api_format_rows = repository
.aggregate_usage_audits(&UsageAuditAggregationQuery {
@@ -3226,7 +3283,7 @@ mod tests {
.expect("api format aggregation should succeed");
assert_eq!(api_format_rows.len(), 1);
assert_eq!(api_format_rows[0].group_key, "openai:chat");
assert_eq!(api_format_rows[0].request_count, 1);
assert_eq!(api_format_rows[0].request_count, 2);
}
#[tokio::test]

46
crates/aether-data/src/repository/usage/postgres/mod.rs
View File

@@ -1488,8 +1488,16 @@ const REBUILD_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_STATS_SQL: &str =
include_str!("queries/rebuild_provider_api_key_codex_window_usage_stats_sql.sql");
const LIST_USAGE_AUDITS_PREFIX: &str = include_str!("queries/list_usage_audits_prefix.sql");
const USAGE_RESERVED_PROVIDER_LABELS_FILTER_SQL: &str = " AND BTRIM(COALESCE(\"usage\".provider_name, '')) <> '' AND lower(BTRIM(COALESCE(\"usage\".provider_name, ''))) NOT IN ('unknown', 'unknow', 'pending')";
const USAGE_PROVIDER_IDENTITY_FILTER_SQL: &str = " AND BTRIM(COALESCE(\"usage\".provider_id, '')) <> '' AND lower(BTRIM(COALESCE(\"usage\".provider_id, ''))) NOT IN ('unknown', 'unknow', 'pending')";
const USAGE_PROVIDER_IDENTITY_FILTER_SQL: &str = r#" AND (
(
BTRIM(COALESCE("usage".provider_id, '')) <> ''
AND lower(BTRIM(COALESCE("usage".provider_id, ''))) NOT IN ('unknown', 'unknow', 'pending')
)
OR (
BTRIM(COALESCE("usage".provider_name, '')) <> ''
AND lower(BTRIM(COALESCE("usage".provider_name, ''))) NOT IN ('unknown', 'unknow', 'pending')
)
)"#;
const USAGE_RAW_PROVIDER_GROUP_KEY_SQL: &str = r#"CASE
WHEN BTRIM(COALESCE("usage".provider_id, '')) = ''
OR lower(BTRIM(COALESCE("usage".provider_id, ''))) IN ('unknown', 'unknow', 'pending')
@@ -1502,13 +1510,33 @@ const USAGE_RAW_PROVIDER_DISPLAY_NAME_SQL: &str = r#"CASE
THEN NULL
ELSE BTRIM("usage".provider_name)
END"#;
const USAGE_PROVIDER_IDENTITY_SOURCE_SQL: &str = r#"CASE
WHEN BTRIM(COALESCE("usage".provider_id, '')) <> ''
AND lower(BTRIM(COALESCE("usage".provider_id, ''))) NOT IN ('unknown', 'unknow', 'pending')
THEN 'provider_id'
WHEN BTRIM(COALESCE("usage".provider_name, '')) <> ''
AND lower(BTRIM(COALESCE("usage".provider_name, ''))) NOT IN ('unknown', 'unknow', 'pending')
THEN 'legacy_name'
ELSE NULL
END"#;
const USAGE_PROVIDER_IDENTITY_JOIN_SQL: &str = r#" LEFT JOIN providers AS provider_by_id
ON BTRIM(COALESCE("usage".provider_id, '')) <> ''
AND lower(BTRIM(COALESCE("usage".provider_id, ''))) NOT IN ('unknown', 'unknow', 'pending')
AND provider_by_id.id = BTRIM("usage".provider_id)"#;
const USAGE_RESOLVED_PROVIDER_GROUP_KEY_SQL: &str = r#"COALESCE(
provider_by_id.id,
BTRIM("usage".provider_id)
CASE
WHEN BTRIM(COALESCE("usage".provider_id, '')) = ''
OR lower(BTRIM(COALESCE("usage".provider_id, ''))) IN ('unknown', 'unknow', 'pending')
THEN NULL
ELSE BTRIM("usage".provider_id)
END,
CASE
WHEN BTRIM(COALESCE("usage".provider_name, '')) = ''
OR lower(BTRIM(COALESCE("usage".provider_name, ''))) IN ('unknown', 'unknow', 'pending')
THEN NULL
ELSE BTRIM("usage".provider_name)
END
)"#;
const USAGE_RESOLVED_PROVIDER_DISPLAY_NAME_SQL: &str = r#"COALESCE(
provider_by_id.name,
@@ -1558,9 +1586,9 @@ fn usage_audit_aggregation_sql_fragments(
filtered_extra_where: "",
group_key_expr: "provider_group_key",
display_name_expr: "provider_display_name",
secondary_name_expr: "NULL::varchar",
secondary_name_expr: "provider_identity_source",
aggregate_display_name_expr: "MAX(display_name)",
aggregate_secondary_name_expr: "NULL::varchar",
aggregate_secondary_name_expr: "CASE WHEN COUNT(*) FILTER (WHERE secondary_name = 'provider_id') > 0 THEN 'provider_id' WHEN COUNT(*) FILTER (WHERE secondary_name = 'legacy_name') > 0 THEN 'legacy_name' ELSE NULL END",
avg_response_time_expr: "AVG(response_time_ms::DOUBLE PRECISION)",
success_count_expr: "COALESCE(SUM(success_flag), 0)::BIGINT",
},
@@ -6687,10 +6715,10 @@ ORDER BY request_count DESC, group_key ASC
) -> Result<Vec<StoredUsageAuditAggregation>, DataLayerError> {
let fragments = usage_audit_aggregation_sql_fragments(query.group_by);
let provider_extra_where =
if matches!(query.group_by, UsageAuditAggregationGroupBy::Provider) {
if matches!(query.group_by, UsageAuditAggregationGroupBy::Provider)
|| query.exclude_reserved_provider_labels
{
USAGE_PROVIDER_IDENTITY_FILTER_SQL
} else if query.exclude_reserved_provider_labels {
USAGE_RESERVED_PROVIDER_LABELS_FILTER_SQL
} else {
""
};
@@ -6702,6 +6730,7 @@ WITH filtered_usage AS (
"usage".user_id AS user_id,
{provider_group_key_expr} AS provider_group_key,
{provider_display_name_expr} AS provider_display_name,
{provider_identity_source_expr} AS provider_identity_source,
COALESCE("usage".api_format, 'unknown') AS api_format_group_key,
GREATEST(COALESCE("usage".input_tokens, 0), 0) AS input_tokens,
GREATEST(COALESCE("usage".output_tokens, 0), 0) AS output_tokens,
@@ -6835,6 +6864,7 @@ LIMIT $3
provider_identity_join = fragments.provider_identity_join,
provider_group_key_expr = fragments.provider_group_key_expr,
provider_display_name_expr = fragments.provider_display_name_expr,
provider_identity_source_expr = USAGE_PROVIDER_IDENTITY_SOURCE_SQL,
group_key_expr = fragments.group_key_expr,
display_name_expr = fragments.display_name_expr,
secondary_name_expr = fragments.secondary_name_expr,

16
crates/aether-data/src/repository/usage/postgres/tests.rs
View File

@@ -464,16 +464,24 @@ fn usage_sql_aggregate_usage_audits_supports_daily_model_and_provider_aggregates
#[test]
fn usage_sql_provider_aggregation_excludes_unknown_provider_labels() {
let source = include_str!("mod.rs");
assert!(source.contains(
r#"const USAGE_PROVIDER_IDENTITY_FILTER_SQL: &str = " AND BTRIM(COALESCE(\"usage\".provider_id, '')) <> ''"#
));
assert!(source.contains("const USAGE_PROVIDER_IDENTITY_FILTER_SQL"));
assert!(source.contains("const USAGE_PROVIDER_IDENTITY_SOURCE_SQL"));
assert!(source.contains(r#"BTRIM(COALESCE("usage".provider_id, '')) <> ''"#));
assert!(source.contains(r#"BTRIM(COALESCE("usage".provider_name, '')) <> ''"#));
assert!(source.contains("LEFT JOIN providers AS provider_by_id"));
assert!(source.contains("provider_by_id.id = BTRIM(\"usage\".provider_id)"));
assert!(source.contains("COALESCE(\n provider_by_id.id,\n CASE"));
assert!(
source.contains("COALESCE(\n provider_by_id.id,\n BTRIM(\"usage\".provider_id)")
source.contains(
"ELSE BTRIM(\"usage\".provider_id)\n END,\n CASE\n WHEN BTRIM(COALESCE(\"usage\".provider_name, ''))"
)
);
assert!(source.contains("COALESCE(\n provider_by_id.name,"));
assert!(!source.contains("provider_by_name.name = BTRIM(\"usage\".provider_name)"));
assert!(source.contains("{provider_identity_source_expr} AS provider_identity_source"));
assert!(source.contains(r#"secondary_name_expr: "provider_identity_source""#));
assert!(source
.contains("COUNT(*) FILTER (WHERE secondary_name = 'provider_id') > 0 THEN 'provider_id'"));
assert!(source.contains(
"if matches!(query.group_by, UsageAuditAggregationGroupBy::Provider) {\n return self.aggregate_usage_audits_raw(query).await;"
));

43
crates/aether-data/src/repository/usage/sqlite.rs
View File

@@ -223,10 +223,19 @@ WHERE request_id = ?
AND status IN ('streaming', 'success')
"#;
const SQLITE_PROVIDER_NAME_IS_NOT_RESERVED: &str = r#"
provider_name IS NOT NULL
AND TRIM(provider_name) <> ''
AND LOWER(TRIM(provider_name)) NOT IN ('unknown', 'unknow', 'pending')
const SQLITE_PROVIDER_IDENTITY_IS_NOT_RESERVED: &str = r#"
(
(
provider_id IS NOT NULL
AND TRIM(provider_id) <> ''
AND LOWER(TRIM(provider_id)) NOT IN ('unknown', 'unknow', 'pending')
)
OR (
provider_name IS NOT NULL
AND TRIM(provider_name) <> ''
AND LOWER(TRIM(provider_name)) NOT IN ('unknown', 'unknow', 'pending')
)
)
"#;
const SQLITE_USAGE_CACHE_CREATION_TOKENS_EXPR: &str = r#"
@@ -617,6 +626,23 @@ THEN TRIM(provider_id) ELSE TRIM(provider_name) END"
}
}
fn sqlite_usage_aggregation_secondary_expr(group_by: UsageAuditAggregationGroupBy) -> &'static str {
match group_by {
UsageAuditAggregationGroupBy::Provider => {
"CASE WHEN SUM(CASE WHEN provider_id IS NOT NULL \
AND TRIM(provider_id) <> '' \
AND LOWER(TRIM(provider_id)) NOT IN ('unknown', 'unknow', 'pending') \
THEN 1 ELSE 0 END) > 0 THEN 'provider_id' \
WHEN SUM(CASE WHEN provider_name IS NOT NULL \
AND TRIM(provider_name) <> '' \
AND LOWER(TRIM(provider_name)) NOT IN ('unknown', 'unknow', 'pending') \
THEN 1 ELSE 0 END) > 0 THEN 'legacy_name' \
ELSE NULL END"
}
_ => "NULL",
}
}
fn sqlite_aggregate_u64(row: &SqliteRow, field: &str) -> Result<u64, DataLayerError> {
Ok(row.try_get::<i64, _>(field).map_sql_err()?.max(0) as u64)
}
@@ -682,7 +708,7 @@ fn decode_sqlite_usage_aggregation_row(
Ok(StoredUsageAuditAggregation {
group_key: row.try_get::<String, _>("group_key").map_sql_err()?,
display_name: row.try_get("display_name").map_sql_err()?,
secondary_name: None,
secondary_name: row.try_get("secondary_name").map_sql_err()?,
request_count: sqlite_aggregate_u64(row, "request_count")?,
total_tokens: sqlite_aggregate_u64(row, "total_tokens")?,
output_tokens: sqlite_aggregate_u64(row, "output_tokens")?,
@@ -1323,8 +1349,9 @@ FROM "usage"
}
let group_expr = sqlite_usage_aggregation_group_expr(query.group_by);
let secondary_expr = sqlite_usage_aggregation_secondary_expr(query.group_by);
let display_expr = if matches!(query.group_by, UsageAuditAggregationGroupBy::Provider) {
"NULLIF(TRIM(provider_name), '')"
"CASE WHEN provider_name IS NOT NULL AND TRIM(provider_name) <> '' AND LOWER(TRIM(provider_name)) NOT IN ('unknown', 'unknow', 'pending') THEN TRIM(provider_name) ELSE NULL END"
} else {
"NULL"
};
@@ -1348,6 +1375,7 @@ FROM "usage"
SELECT
{group_expr} AS group_key,
{display_expr} AS display_name,
{secondary_expr} AS secondary_name,
COUNT(*) AS request_count,
COALESCE(SUM(MAX(COALESCE(total_tokens, 0), 0)), 0) AS total_tokens,
COALESCE(SUM(MAX(COALESCE(output_tokens, 0), 0)), 0) AS output_tokens,
@@ -1368,6 +1396,7 @@ FROM "usage"
"#,
effective_input_expr = SQLITE_USAGE_EFFECTIVE_INPUT_TOKENS_EXPR,
total_input_context_expr = SQLITE_USAGE_TOTAL_INPUT_CONTEXT_EXPR,
secondary_expr = secondary_expr,
cache_creation_expr = SQLITE_USAGE_CACHE_CREATION_TOKENS_EXPR
));
let mut has_where = false;
@@ -1383,7 +1412,7 @@ FROM "usage"
builder.push("status NOT IN ('pending', 'streaming')");
if query.exclude_reserved_provider_labels {
push_sqlite_usage_where(&mut builder, &mut has_where);
builder.push(SQLITE_PROVIDER_NAME_IS_NOT_RESERVED);
builder.push(SQLITE_PROVIDER_IDENTITY_IS_NOT_RESERVED);
}
if matches!(query.group_by, UsageAuditAggregationGroupBy::User) {
push_sqlite_usage_where(&mut builder, &mut has_where);

75
crates/aether-provider-pool/src/lib.rs
View File

@@ -380,6 +380,81 @@ mod tests {
));
}
#[test]
fn provider_quota_exhaustion_snapshot_expires_after_reset_at() {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.expect("system time should be after unix epoch")
.as_secs();
let mut expired = sample_key(None);
expired.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "codex",
"code": "exhausted",
"exhausted": true,
"updated_at": now.saturating_sub(600),
"windows": [{
"code": "5h",
"used_ratio": 1.0,
"reset_at": now.saturating_sub(60),
"is_exhausted": true
}]
}
}));
assert!(!provider_pool_key_account_quota_exhausted(
&expired, "codex"
));
let mut active = sample_key(None);
active.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "codex",
"code": "exhausted",
"exhausted": true,
"updated_at": now,
"windows": [{
"code": "5h",
"used_ratio": 1.0,
"reset_at": now.saturating_add(3600),
"is_exhausted": true
}]
}
}));
assert!(provider_pool_key_account_quota_exhausted(&active, "codex"));
}
#[test]
fn provider_quota_exhaustion_metadata_expires_after_reset_at() {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.expect("system time should be after unix epoch")
.as_secs();
assert!(!provider_pool_key_account_quota_exhausted(
&sample_key(Some(json!({
"codex": {
"updated_at": now.saturating_sub(600),
"primary_used_percent": 100.0,
"primary_reset_at": now.saturating_sub(60)
}
}))),
"codex",
));
assert!(provider_pool_key_account_quota_exhausted(
&sample_key(Some(json!({
"codex": {
"updated_at": now,
"primary_used_percent": 100.0,
"primary_reset_at": now.saturating_add(3600)
}
}))),
"codex",
));
}
#[test]
fn grok_quota_tier_boundaries_match_pool_modes() {
assert_eq!(

18
crates/aether-provider-pool/src/providers/chatgpt_web.rs
View File

@@ -10,8 +10,9 @@ use crate::provider::{
ProviderPoolMemberInput,
};
use crate::quota::{
provider_pool_json_bool, provider_pool_json_f64, provider_pool_metadata_bucket,
provider_pool_quota_snapshot_exhausted_decision,
provider_pool_current_unix_secs, provider_pool_json_bool, provider_pool_json_f64,
provider_pool_metadata_bucket, provider_pool_quota_snapshot_exhausted_decision,
provider_pool_reset_deadline_elapsed, provider_pool_timestamp_unix_secs,
};
use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
@@ -255,6 +256,19 @@ fn infer_chatgpt_web_image_quota_limit(
}
pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
if provider_pool_current_unix_secs().is_some_and(|now| {
let mut image_quota = Map::new();
if let Some(value) = bucket.get("image_quota_reset_at") {
image_quota.insert("reset_at".to_string(), value.clone());
}
provider_pool_reset_deadline_elapsed(
&image_quota,
provider_pool_timestamp_unix_secs(bucket.get("updated_at")),
now,
)
}) {
return false;
}
if provider_pool_json_bool(bucket.get("image_quota_blocked")) == Some(true) {
return true;
}

41
crates/aether-provider-pool/src/providers/codex.rs
View File

@@ -10,8 +10,9 @@ use crate::provider::{
ProviderPoolMemberInput,
};
use crate::quota::{
provider_pool_json_bool, provider_pool_json_f64, provider_pool_metadata_bucket,
provider_pool_quota_snapshot_exhausted_decision,
provider_pool_current_unix_secs, provider_pool_json_bool, provider_pool_json_f64,
provider_pool_metadata_bucket, provider_pool_quota_snapshot_exhausted_decision,
provider_pool_reset_deadline_elapsed, provider_pool_timestamp_unix_secs,
};
use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
@@ -121,6 +122,37 @@ pub fn build_codex_pool_quota_request(
})
}
fn codex_window_reset_elapsed(bucket: &Map<String, Value>, prefix: &str) -> bool {
let Some(now_unix_secs) = provider_pool_current_unix_secs() else {
return false;
};
let mut window = Map::new();
for (target, source) in [
("reset_at", format!("{prefix}_reset_at")),
("next_reset_at", format!("{prefix}_next_reset_at")),
("reset_seconds", format!("{prefix}_reset_seconds")),
(
"reset_after_seconds",
format!("{prefix}_reset_after_seconds"),
),
] {
if let Some(value) = bucket.get(source.as_str()) {
window.insert(target.to_string(), value.clone());
}
}
provider_pool_reset_deadline_elapsed(
&window,
provider_pool_timestamp_unix_secs(bucket.get("updated_at")),
now_unix_secs,
)
}
fn codex_window_used_percent_exhausted(bucket: &Map<String, Value>, prefix: &str) -> bool {
let used_percent_key = format!("{prefix}_used_percent");
provider_pool_json_f64(bucket.get(used_percent_key.as_str()))
.is_some_and(|value| value >= 100.0 && !codex_window_reset_elapsed(bucket, prefix))
}
pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
if provider_pool_json_bool(bucket.get("credits_unlimited")) == Some(true) {
return false;
@@ -130,7 +162,6 @@ pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
if !has_window_data && provider_pool_json_bool(bucket.get("has_credits")) == Some(false) {
return true;
}
provider_pool_json_f64(bucket.get("primary_used_percent")).is_some_and(|value| value >= 100.0)
|| provider_pool_json_f64(bucket.get("secondary_used_percent"))
.is_some_and(|value| value >= 100.0)
codex_window_used_percent_exhausted(bucket, "primary")
|| codex_window_used_percent_exhausted(bucket, "secondary")
}

15
crates/aether-provider-pool/src/providers/grok.rs
View File

@@ -7,8 +7,9 @@ use crate::provider::{
ProviderPoolMemberInput,
};
use crate::quota::{
provider_pool_json_bool, provider_pool_json_f64, provider_pool_metadata_bucket,
provider_pool_quota_snapshot_exhausted_decision,
provider_pool_current_unix_secs, provider_pool_json_bool, provider_pool_json_f64,
provider_pool_metadata_bucket, provider_pool_quota_snapshot_exhausted_decision,
provider_pool_reset_deadline_elapsed, provider_pool_timestamp_unix_secs,
};
pub const GROK_QUOTA_WINDOWS_BASIC: &[(&str, &str)] = &[("quota_fast", "fast")];
@@ -191,6 +192,8 @@ pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
let mut model_count = 0usize;
let mut exhausted_count = 0usize;
let now_unix_secs = provider_pool_current_unix_secs();
let bucket_updated_at = provider_pool_timestamp_unix_secs(bucket.get("updated_at"));
for (model_key, item) in models.iter() {
if !supported_mode_keys.is_empty() && !supported_mode_keys.contains(&model_key.as_str()) {
continue;
@@ -207,12 +210,14 @@ pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
continue;
}
model_count += 1;
if provider_pool_json_bool(item.get("is_exhausted")) == Some(true)
let quota_exhausted = provider_pool_json_bool(item.get("is_exhausted")) == Some(true)
|| provider_pool_json_f64(item.get("used_percent")).is_some_and(|value| value >= 100.0)
|| provider_pool_json_f64(item.get("remaining")).is_some_and(|value| value <= 0.0)
|| provider_pool_json_f64(item.get("remaining_fraction"))
.is_some_and(|value| value <= 0.0)
{
.is_some_and(|value| value <= 0.0);
let reset_elapsed = now_unix_secs
.is_some_and(|now| provider_pool_reset_deadline_elapsed(item, bucket_updated_at, now));
if quota_exhausted && !reset_elapsed {
exhausted_count += 1;
}
}

14
crates/aether-provider-pool/src/providers/kiro.rs
View File

@@ -11,8 +11,9 @@ use crate::provider::{
ProviderPoolMemberInput,
};
use crate::quota::{
provider_pool_json_f64, provider_pool_metadata_bucket,
provider_pool_quota_snapshot_exhausted_decision,
provider_pool_current_unix_secs, provider_pool_json_f64, provider_pool_metadata_bucket,
provider_pool_quota_snapshot_exhausted_decision, provider_pool_reset_deadline_elapsed,
provider_pool_timestamp_unix_secs,
};
use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
@@ -151,6 +152,15 @@ fn normalize_kiro_version(value: &str) -> &str {
}
pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
if provider_pool_current_unix_secs().is_some_and(|now| {
provider_pool_reset_deadline_elapsed(
bucket,
provider_pool_timestamp_unix_secs(bucket.get("updated_at")),
now,
)
}) {
return false;
}
if provider_pool_json_f64(bucket.get("remaining")).is_some_and(|value| value <= 0.0) {
return true;
}

95
crates/aether-provider-pool/src/quota.rs
View File

@@ -1,3 +1,5 @@
use std::time::{SystemTime, UNIX_EPOCH};
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use serde_json::{json, Map, Value};
@@ -135,7 +137,7 @@ pub(crate) fn provider_pool_json_f64(value: Option<&Value>) -> Option<f64> {
.filter(|value| value.is_finite())
}
fn provider_pool_timestamp_unix_secs(value: Option<&Value>) -> Option<u64> {
pub(crate) fn provider_pool_timestamp_unix_secs(value: Option<&Value>) -> Option<u64> {
let mut timestamp = provider_pool_json_f64(value)?;
if timestamp <= 0.0 {
return None;
@@ -146,6 +148,49 @@ fn provider_pool_timestamp_unix_secs(value: Option<&Value>) -> Option<u64> {
Some(timestamp as u64)
}
pub(crate) fn provider_pool_current_unix_secs() -> Option<u64> {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.ok()
.map(|duration| duration.as_secs())
}
fn provider_pool_reset_deadline_unix_secs(
item: &Map<String, Value>,
fallback_observed_at: Option<u64>,
) -> Option<u64> {
provider_pool_timestamp_unix_secs(item.get("reset_at"))
.or_else(|| provider_pool_timestamp_unix_secs(item.get("next_reset_at")))
.or_else(|| {
let reset_seconds = provider_pool_json_f64(item.get("reset_seconds"))
.or_else(|| provider_pool_json_f64(item.get("reset_after_seconds")))?;
if reset_seconds < 0.0 {
return None;
}
let base = provider_pool_timestamp_unix_secs(item.get("observed_at"))
.or_else(|| provider_pool_timestamp_unix_secs(item.get("updated_at")))
.or(fallback_observed_at)?;
Some(base.saturating_add(reset_seconds.ceil() as u64))
})
}
pub(crate) fn provider_pool_reset_deadline_elapsed(
item: &Map<String, Value>,
fallback_observed_at: Option<u64>,
now_unix_secs: u64,
) -> bool {
provider_pool_reset_deadline_unix_secs(item, fallback_observed_at)
.is_some_and(|reset_at| reset_at <= now_unix_secs)
}
fn provider_pool_quota_window_is_exhausted(window: &Map<String, Value>) -> bool {
provider_pool_json_bool(window.get("is_exhausted"))
.or_else(|| {
provider_pool_json_f64(window.get("used_ratio")).map(|value| value >= 1.0 - 1e-6)
})
.unwrap_or(false)
}
fn provider_pool_quota_snapshot_matches_provider(
quota_snapshot: &Map<String, Value>,
provider_type: &str,
@@ -203,19 +248,45 @@ pub(crate) fn provider_pool_quota_snapshot_exhausted_decision(
}
let exhausted = provider_pool_json_bool(quota_snapshot.get("exhausted"))?;
if exhausted {
let windows_max_ratio = quota_snapshot
let now_unix_secs = provider_pool_current_unix_secs();
let snapshot_observed_at =
provider_pool_timestamp_unix_secs(quota_snapshot.get("observed_at"))
.or_else(|| provider_pool_timestamp_unix_secs(quota_snapshot.get("updated_at")));
if let Some(windows) = quota_snapshot
.get("windows")
.and_then(Value::as_array)
.filter(|w| !w.is_empty())
.and_then(|windows| {
windows
.iter()
.filter_map(Value::as_object)
.filter_map(|w| w.get("used_ratio"))
.filter_map(Value::as_f64)
.max_by(f64::total_cmp)
});
if windows_max_ratio.is_some_and(|ratio| ratio < 1.0 - 1e-6) {
.filter(|windows| !windows.is_empty())
{
let mut saw_exhausted_window = false;
let mut saw_active_exhausted_window = false;
let mut windows_max_ratio = None::<f64>;
for window in windows.iter().filter_map(Value::as_object) {
if let Some(ratio) = provider_pool_json_f64(window.get("used_ratio")) {
windows_max_ratio =
Some(windows_max_ratio.map_or(ratio, |current| current.max(ratio)));
}
if provider_pool_quota_window_is_exhausted(window) {
saw_exhausted_window = true;
let reset_elapsed = now_unix_secs.is_some_and(|now| {
provider_pool_reset_deadline_elapsed(window, snapshot_observed_at, now)
});
if !reset_elapsed {
saw_active_exhausted_window = true;
}
}
}
if saw_exhausted_window {
return Some(saw_active_exhausted_window);
}
if windows_max_ratio.is_some_and(|ratio| ratio < 1.0 - 1e-6) {
return Some(false);
}
} else if now_unix_secs.is_some_and(|now| {
provider_pool_reset_deadline_elapsed(quota_snapshot, snapshot_observed_at, now)
}) {
return Some(false);
}
}

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

@@ -488,6 +488,7 @@ pub fn build_local_request_candidate_status_record(
demoted_by: metadata.demoted_by.clone(),
routing_trace: metadata.routing_trace.clone(),
});
let extra_data = mark_request_candidate_stream_completed_if_success(status, extra_data);
let created_at_unix_ms = started_at_unix_ms.or(finished_at_unix_ms);
Some(UpsertRequestCandidateRecord {
@@ -568,7 +569,7 @@ pub fn build_report_request_candidate_status_record(
error_message,
latency_ms,
concurrent_requests: None,
extra_data: slot.extra_data,
extra_data: mark_request_candidate_stream_completed_if_success(status, slot.extra_data),
required_capabilities: None,
created_at_unix_ms: Some(created_at_unix_ms),
started_at_unix_ms,
@@ -829,6 +830,23 @@ fn build_report_candidate_extra_data(input: ReportCandidateExtraDataInput) -> Op
(!extra_data.is_empty()).then_some(Value::Object(extra_data))
}
fn mark_request_candidate_stream_completed_if_success(
status: RequestCandidateStatus,
extra_data: Option<Value>,
) -> Option<Value> {
if status != RequestCandidateStatus::Success {
return extra_data;
}
let mut object = match extra_data {
Some(Value::Object(object)) => object,
Some(other) => return Some(other),
None => Map::new(),
};
object.insert("stream_completed".to_string(), Value::Bool(true));
Some(Value::Object(object))
}
fn merge_request_candidate_extra_data(
existing: Option<Value>,
overlay: Option<Value>,
@@ -1295,6 +1313,53 @@ mod tests {
assert_eq!(record.finished_at_unix_ms, Some(123));
assert_eq!(record.created_at_unix_ms, Some(123));
assert_eq!(record.status, RequestCandidateStatus::Success);
assert_eq!(
record
.extra_data
.as_ref()
.and_then(|value| value.get("stream_completed")),
Some(&json!(true))
);
}
#[test]
fn local_success_status_marks_stream_completed_for_pending_cleanup_recovery() {
let mut plan = sample_plan();
plan.candidate_id = Some("cand-1".to_string());
let report_context = json!({
"request_id": "req-1",
"candidate_id": "cand-1",
"candidate_index": 0,
"retry_index": 0,
"user_id": "user-1",
"api_key_id": "api-key-1",
"client_api_format": "openai:responses",
"provider_api_format": "openai:responses",
});
let record =
build_local_request_candidate_status_record(LocalRequestCandidateStatusRecordInput {
plan: &plan,
report_context: Some(&report_context),
status_update: SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Success,
status_code: Some(200),
error_type: None,
error_message: None,
latency_ms: Some(25),
started_at_unix_ms: Some(1_000),
finished_at_unix_ms: Some(1_025),
},
})
.expect("success status record should build");
assert_eq!(
record
.extra_data
.as_ref()
.and_then(|value| value.get("stream_completed")),
Some(&json!(true))
);
}
#[test]