This commit is contained in:
zhefox
2026-05-25 13:38:38 +08:00
65 changed files with 6848 additions and 3556 deletions
+1 -4
View File
@@ -261,8 +261,7 @@ jobs:
root="package/${bundle}"
mkdir -p \
"${root}/bin" \
"${root}/frontend" \
"${root}/scripts"
"${root}/frontend"
install -m 0755 "artifacts/aether-gateway-${platform}-${arch}/aether-gateway" "${root}/bin/aether-gateway"
cp -R artifacts/frontend-dist/. "${root}/frontend/"
@@ -274,8 +273,6 @@ jobs:
install -m 0755 update.sh "${root}/update.sh"
install -m 0644 docker-compose.yml "${root}/docker-compose.yml"
install -m 0644 docker-compose.single-node.yml "${root}/docker-compose.single-node.yml"
install -m 0755 scripts/migrate-pg-compose-to-single-node.sh "${root}/scripts/migrate-pg-compose-to-single-node.sh"
install -m 0755 scripts/migrate-pg-to-single-node.sh "${root}/scripts/migrate-pg-to-single-node.sh"
install -m 0644 .env.example "${root}/.env.example"
install -m 0755 generate_keys.sh "${root}/generate_keys.sh"
install -m 0644 README.md "${root}/README.md"
Generated
+1
View File
@@ -432,6 +432,7 @@ dependencies = [
"aether-contracts",
"aether-data-contracts",
"aether-wallet",
"chrono",
"regex",
"serde",
"serde_json",
@@ -6,6 +6,7 @@ use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKe
use aether_pool_core::{
score_pool_member_with_rules, PoolMemberScoreInput, PoolMemberScoreRules, POOL_SCORE_VERSION,
};
use aether_scheduler_core::any_provider_key_circuit_open_at;
use serde_json::Value;
use crate::handlers::shared::{provider_key_health_summary, provider_key_status_snapshot_payload};
@@ -98,7 +99,8 @@ fn provider_key_score_input(
.as_object()
.and_then(|snapshot| snapshot.get("account"))
.and_then(Value::as_object);
let (health_score, _, _, any_circuit_open, _) = provider_key_health_summary(key);
let (health_score, _, _, _, _) = provider_key_health_summary(key);
let active_circuit_open = any_provider_key_circuit_open_at(key, now_unix_secs);
let health_score = key
.health_by_format
.as_ref()
@@ -125,7 +127,7 @@ fn provider_key_score_input(
.and_then(Value::as_bool)
.unwrap_or(false),
oauth_invalid_reason: key.oauth_invalid_reason.clone(),
circuit_open: any_circuit_open,
circuit_open: active_circuit_open,
success_count: key.success_count.unwrap_or(0).into(),
error_count: key.error_count.unwrap_or(0).into(),
total_response_time_ms: key.total_response_time_ms.unwrap_or(0).into(),
@@ -159,3 +161,70 @@ fn stable_hash(bytes: &[u8]) -> u64 {
}
hash
}
#[cfg(test)]
mod tests {
use super::*;
use aether_data_contracts::repository::pool_scores::PoolMemberHardState;
use serde_json::json;
fn sample_key_with_circuit_next_probe(
next_probe_at_unix_secs: u64,
) -> StoredProviderCatalogKey {
let mut key = StoredProviderCatalogKey::new(
"key-gemini-5".to_string(),
"provider-google-api".to_string(),
"5".to_string(),
"api_key".to_string(),
None,
true,
)
.expect("sample key should be valid");
key.health_by_format = Some(json!({
"gemini:generate_content": {
"health_score": 0.2,
"consecutive_failures": 8
}
}));
key.circuit_breaker_by_format = Some(json!({
"gemini:generate_content": {
"open": true,
"reason": "consecutive_failures_8",
"next_probe_at_unix_secs": next_probe_at_unix_secs
}
}));
key
}
#[test]
fn expired_circuit_probe_deadline_does_not_leave_pool_score_in_cooldown() {
let now_unix_secs = 1_000;
let key = sample_key_with_circuit_next_probe(900);
let score = build_provider_key_pool_score_upsert(
&key,
"custom",
None,
now_unix_secs,
PoolMemberScoreRules::default(),
);
assert_eq!(score.hard_state, PoolMemberHardState::Available);
}
#[test]
fn future_circuit_probe_deadline_keeps_pool_score_in_cooldown() {
let now_unix_secs = 1_000;
let key = sample_key_with_circuit_next_probe(1_100);
let score = build_provider_key_pool_score_upsert(
&key,
"custom",
None,
now_unix_secs,
PoolMemberScoreRules::default(),
);
assert_eq!(score.hard_state, PoolMemberHardState::Cooldown);
}
}
@@ -1073,22 +1073,42 @@ fn build_chatgpt_web_image_provider_body_from_openai_responses_body(
.unwrap_or("gpt-5-5-thinking");
let image_urls = openai_image_inputs_as_urls(&images);
let body = json!({
let mut body = json!({
"operation": operation,
"model": if model.is_empty() { "gpt-image-2" } else { model },
"web_model": web_model,
"prompt": prompt,
"size": size,
"ratio": chatgpt_web_ratio_for_size(size),
"quality": quality,
"output_format": output_format,
"images": image_urls,
});
let summary = json!({
if let Some(partial_images) = tool
.as_ref()
.and_then(|tool| tool.get("partial_images"))
.or_else(|| object.get("partial_images"))
.cloned()
{
body.as_object_mut()?
.insert("partial_images".to_string(), partial_images);
}
let mut summary = json!({
"operation": operation,
"output_format": output_format,
"size": size,
"quality": quality,
});
if let Some(partial_images) = tool
.as_ref()
.and_then(|tool| tool.get("partial_images"))
.or_else(|| object.get("partial_images"))
.cloned()
{
summary
.as_object_mut()?
.insert("partial_images".to_string(), partial_images);
}
Some((body, summary))
}
@@ -1426,4 +1446,36 @@ mod tests {
assert_eq!(summary["operation"], "generate");
assert_eq!(summary["output_format"], "png");
}
#[test]
fn chatgpt_web_responses_image_body_preserves_usage_options() {
let body_json = json!({
"model": "gpt-image-2",
"input": "Draw a glass city",
"tools": [
{
"type": "image_generation",
"size": "1024x1024",
"quality": "high",
"output_format": "png",
"partial_images": 2
}
],
"tool_choice": {
"type": "image_generation"
}
});
let (provider_body, summary) =
build_chatgpt_web_image_provider_body_from_openai_responses_body(
&body_json,
"gpt-image-2",
)
.expect("responses image body should convert");
assert_eq!(provider_body["quality"], "high");
assert_eq!(provider_body["partial_images"], 2);
assert_eq!(summary["quality"], "high");
assert_eq!(summary["partial_images"], 2);
}
}
@@ -1684,6 +1684,28 @@ impl GatewayDataState {
}
}
pub(crate) async fn set_api_key_usage_totals(
&self,
api_key_id: &str,
total_requests: u64,
total_tokens: u64,
total_cost_usd: f64,
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
match &self.auth_api_key_writer {
Some(repository) => {
repository
.set_api_key_usage_totals(
api_key_id,
total_requests,
total_tokens,
total_cost_usd,
)
.await
}
None => Ok(None),
}
}
pub(crate) async fn set_standalone_api_key_feature_settings(
&self,
api_key_id: &str,
@@ -511,6 +511,45 @@ impl GatewayDataState {
}
}
pub(crate) async fn export_admin_system_usage_aggregates(
&self,
) -> Result<aether_data::repository::system::AdminSystemUsageAggregateSnapshot, DataLayerError>
{
match self.backends.as_ref() {
Some(backends) => backends.export_admin_system_usage_aggregates().await,
None => {
Ok(aether_data::repository::system::AdminSystemUsageAggregateSnapshot::default())
}
}
}
pub(crate) async fn import_admin_system_usage_aggregates(
&self,
snapshot: &aether_data::repository::system::AdminSystemUsageAggregateSnapshot,
user_id_map: &std::collections::BTreeMap<String, String>,
api_key_id_map: &std::collections::BTreeMap<String, String>,
mode: aether_data::repository::system::AdminSystemUsageAggregateImportMode,
) -> Result<
aether_data::repository::system::AdminSystemUsageAggregateImportSummary,
DataLayerError,
> {
match self.backends.as_ref() {
Some(backends) => {
backends
.import_admin_system_usage_aggregates(
snapshot,
user_id_map,
api_key_id_map,
mode,
)
.await
}
None => Ok(
aether_data::repository::system::AdminSystemUsageAggregateImportSummary::default(),
),
}
}
pub(crate) async fn purge_admin_request_bodies_batch(
&self,
batch_size: usize,
File diff suppressed because it is too large Load Diff
@@ -1,6 +1,9 @@
use crate::handlers::admin::request::AdminAppState;
use crate::provider_key_auth::provider_key_effective_api_formats;
use aether_scheduler_core::count_recent_rpm_requests_for_provider_key_since;
use aether_scheduler_core::{
count_recent_rpm_requests_for_provider_key_since,
provider_key_circuit_payload_is_active_open_at,
};
use serde_json::json;
use std::time::{SystemTime, UNIX_EPOCH};
@@ -18,6 +21,10 @@ pub(crate) async fn build_admin_key_health_payload(
.await
.ok()
.and_then(|mut keys| keys.drain(..).next())?;
let now_unix_secs = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_secs())
.unwrap_or_default();
let provider = state
.read_provider_catalog_providers_by_ids(std::slice::from_ref(&key.provider_id))
.await
@@ -81,10 +88,10 @@ pub(crate) async fn build_admin_key_health_payload(
.and_then(|value| value.get("last_failure_at"))
.cloned()
.unwrap_or(serde_json::Value::Null);
payload["circuit_breaker_open"] = json!(circuit_data
.and_then(|value| value.get("open"))
.and_then(serde_json::Value::as_bool)
.unwrap_or(false));
payload["circuit_breaker_open"] =
json!(circuit_data.is_some_and(
|value| provider_key_circuit_payload_is_active_open_at(value, now_unix_secs)
));
payload["circuit_breaker_open_at"] = circuit_data
.and_then(|value| value.get("open_at"))
.cloned()
@@ -107,15 +114,16 @@ pub(crate) async fn build_admin_key_health_payload(
.unwrap_or(0));
} else {
let mut formats_payload = serde_json::Map::new();
let mut any_circuit_open = false;
for format_name in
provider_key_effective_api_formats(&key, &provider.provider_type, &endpoints)
{
let health_data = health_by_format.and_then(|formats| formats.get(&format_name));
let circuit_data = circuit_by_format.and_then(|formats| formats.get(&format_name));
let is_open = circuit_data
.and_then(|value| value.get("open"))
.and_then(serde_json::Value::as_bool)
.unwrap_or(false);
let active_open = circuit_data.is_some_and(|value| {
provider_key_circuit_payload_is_active_open_at(value, now_unix_secs)
});
any_circuit_open |= active_open;
formats_payload.insert(
format_name.clone(),
json!({
@@ -134,8 +142,8 @@ pub(crate) async fn build_admin_key_health_payload(
.cloned()
.unwrap_or(serde_json::Value::Null),
"circuit_breaker": {
"state": if is_open { "open" } else { "closed" },
"open": is_open,
"state": if active_open { "open" } else { "closed" },
"open": active_open,
"open_at": circuit_data
.and_then(|value| value.get("open_at"))
.cloned()
@@ -167,13 +175,6 @@ pub(crate) async fn build_admin_key_health_payload(
.filter_map(serde_json::Value::as_f64)
.reduce(f64::min)
.unwrap_or(1.0);
let any_circuit_open = formats_payload.values().any(|value| {
value
.get("circuit_breaker")
.and_then(|circuit| circuit.get("open"))
.and_then(serde_json::Value::as_bool)
.unwrap_or(false)
});
payload["key_health_score"] = json!(key_health_score);
payload["any_circuit_open"] = json!(any_circuit_open);
@@ -4,7 +4,9 @@ use crate::handlers::public::{api_format_display_name, build_public_health_timel
use crate::handlers::shared::unix_ms_to_rfc3339;
use crate::provider_key_auth::provider_key_effective_api_formats;
use aether_data_contracts::repository::candidates::PublicHealthTimelineBucket;
use aether_scheduler_core::{is_provider_key_circuit_open, provider_key_health_score};
use aether_scheduler_core::{
any_provider_key_circuit_open_at, is_provider_key_circuit_open_at, provider_key_health_score,
};
use serde_json::json;
use std::collections::{BTreeMap, BTreeSet};
use std::time::{SystemTime, UNIX_EPOCH};
@@ -99,7 +101,9 @@ pub(crate) async fn build_admin_endpoint_health_status_payload(
.entry(api_format.clone())
.or_default()
.insert(key.id.clone());
if key.is_active && !is_provider_key_circuit_open(&key, &api_format) {
if key.is_active
&& !is_provider_key_circuit_open_at(&key, &api_format, now_unix_secs)
{
let key_health_score =
provider_key_health_score(&key, &api_format).unwrap_or(1.0);
active_keys_by_format
@@ -229,6 +233,10 @@ pub(crate) async fn build_admin_health_summary_payload(
return None;
}
let now_unix_secs = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_secs())
.unwrap_or_default();
let providers = state
.list_provider_catalog_providers(false)
.await
@@ -286,20 +294,7 @@ pub(crate) async fn build_admin_health_summary_payload(
.count();
let circuit_open_keys = keys
.iter()
.filter(|key| {
key.circuit_breaker_by_format
.as_ref()
.and_then(serde_json::Value::as_object)
.map(|formats| {
formats.values().any(|circuit| {
circuit
.get("open")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false)
})
})
.unwrap_or(false)
})
.filter(|key| any_provider_key_circuit_open_at(key, now_unix_secs))
.count();
Some(json!({
@@ -8,10 +8,12 @@ use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
};
use aether_scheduler_core::{
is_provider_key_circuit_open, matches_model_mapping, provider_key_health_score,
is_provider_key_circuit_open_at, matches_model_mapping,
provider_key_circuit_payload_is_active_open_at, provider_key_health_score,
};
use serde_json::json;
use std::collections::BTreeMap;
use std::time::{SystemTime, UNIX_EPOCH};
use uuid::Uuid;
pub(crate) async fn build_admin_global_model_routing_payload(
@@ -86,6 +88,10 @@ pub(crate) async fn build_admin_global_model_routing_payload(
.flatten()
.and_then(|value| value.as_bool())
.unwrap_or(false);
let now_unix_secs = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_secs())
.unwrap_or_default();
let global_model_mappings = global_model
.config
@@ -163,9 +169,11 @@ pub(crate) async fn build_admin_global_model_routing_payload(
entries
.iter()
.filter_map(|(api_format, value)| {
value.get("open")
.and_then(serde_json::Value::as_bool)
.filter(|is_open| *is_open)
provider_key_circuit_payload_is_active_open_at(
value,
now_unix_secs,
)
.then_some(())
.map(|_| api_format.clone())
})
.collect::<Vec<_>>()
@@ -188,7 +196,7 @@ pub(crate) async fn build_admin_global_model_routing_payload(
"effective_rpm": effective_rpm,
"allowed_models": allowed_models,
"health_score": provider_key_health_score(key, &endpoint.api_format),
"circuit_breaker_open": is_provider_key_circuit_open(key, &endpoint.api_format),
"circuit_breaker_open": is_provider_key_circuit_open_at(key, &endpoint.api_format, now_unix_secs),
"circuit_breaker_formats": circuit_breaker_formats,
"next_probe_at": next_probe_at,
});
@@ -1,6 +1,6 @@
use super::super::usage_helpers::admin_monitoring_usage_is_error;
use crate::handlers::admin::request::AdminAppState;
use crate::handlers::admin::shared::{provider_key_health_summary, unix_secs_to_rfc3339};
use crate::handlers::admin::shared::{provider_key_health_summary_at, unix_secs_to_rfc3339};
use crate::GatewayError;
use aether_data_contracts::repository::{
provider_catalog::StoredProviderCatalogKey, usage::UsageMonitoringErrorListQuery,
@@ -99,7 +99,7 @@ pub(super) async fn build_admin_monitoring_resilience_snapshot(
last_failure_at,
circuit_breaker_open,
circuit_by_format,
) = provider_key_health_summary(key);
) = provider_key_health_summary_at(key, now.timestamp().max(0) as u64);
if health_score < 0.8 {
degraded_keys += 1;
}
@@ -110,11 +110,12 @@ pub(super) async fn build_admin_monitoring_resilience_snapshot(
let open_formats = circuit_by_format
.iter()
.filter_map(|(api_format, value)| {
value
.get("open")
.and_then(serde_json::Value::as_bool)
.filter(|open| *open)
.map(|_| api_format.clone())
aether_scheduler_core::provider_key_circuit_payload_is_active_open_at(
value,
now.timestamp().max(0) as u64,
)
.then_some(())
.map(|_| api_format.clone())
})
.collect::<Vec<_>>();
@@ -1237,7 +1237,8 @@ async fn admin_monitoring_circuit_history_returns_local_payload() {
"openai:chat": {
"open": true,
"open_at": "2026-03-30T12:00:00+00:00",
"next_probe_at": "2026-03-30T12:05:00+00:00",
"next_probe_at": "2099-03-30T12:05:00+00:00",
"recovery_seconds": 300,
"reason": "错误率过高"
}
})),
@@ -5,12 +5,15 @@ use super::shared::{
quota_key_auto_removed, quota_refresh_success_invalid_state, ProviderQuotaExecutionOutcome,
};
use crate::handlers::admin::provider::shared::payloads::{
OAUTH_ACCOUNT_BLOCK_PREFIX, OAUTH_EXPIRED_PREFIX,
OAUTH_ACCOUNT_BLOCK_PREFIX, OAUTH_EXPIRED_PREFIX, OAUTH_REFRESH_FAILED_PREFIX,
};
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
use crate::GatewayError;
use aether_admin::provider::quota::parse_chatgpt_web_conversation_init_response;
use aether_contracts::ProxySnapshot;
use aether_contracts::{
ExecutionResult, ProxySnapshot, ResolvedTransportProfile, TRANSPORT_BACKEND_BROWSER_WREQ,
TRANSPORT_HTTP_MODE_AUTO, TRANSPORT_POOL_SCOPE_KEY,
};
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
@@ -18,10 +21,12 @@ use aether_provider_pool::{
build_chatgpt_web_pool_quota_request, enrich_chatgpt_web_quota_metadata,
normalize_chatgpt_web_image_quota_limit,
};
use base64::Engine as _;
use serde_json::json;
use std::time::{SystemTime, UNIX_EPOCH};
const PLACEHOLDER_API_KEY: &str = "__placeholder__";
const CHATGPT_WEB_BROWSER_PROFILE: &str = "chrome143";
fn chatgpt_web_auth_config(
transport: &AdminGatewayProviderTransportSnapshot,
@@ -74,19 +79,99 @@ async fn execute_chatgpt_web_quota_plan(
)));
let spec =
build_chatgpt_web_pool_quota_request(&transport.key.id, &endpoint.base_url, authorization);
let resolved_transport_profile = state.resolve_transport_profile(transport);
let plan = super::shared::build_provider_quota_execution_plan(
transport,
spec,
proxy,
state.resolve_transport_profile(transport),
chatgpt_web_quota_transport_profile(resolved_transport_profile.as_ref()),
timeouts,
);
execute_provider_quota_plan(state, transport, plan, "chatgpt_web").await
}
fn chatgpt_web_quota_transport_profile(
transport_profile: Option<&ResolvedTransportProfile>,
) -> Option<ResolvedTransportProfile> {
match transport_profile {
Some(profile)
if profile
.backend
.trim()
.eq_ignore_ascii_case(TRANSPORT_BACKEND_BROWSER_WREQ) =>
{
Some(profile.clone())
}
_ => Some(default_chatgpt_web_quota_transport_profile()),
}
}
fn default_chatgpt_web_quota_transport_profile() -> ResolvedTransportProfile {
ResolvedTransportProfile {
profile_id: CHATGPT_WEB_BROWSER_PROFILE.to_string(),
backend: TRANSPORT_BACKEND_BROWSER_WREQ.to_string(),
http_mode: TRANSPORT_HTTP_MODE_AUTO.to_string(),
pool_scope: TRANSPORT_POOL_SCOPE_KEY.to_string(),
header_fingerprint: None,
extra: Some(json!({
"browser_profile": CHATGPT_WEB_BROWSER_PROFILE,
"source": "chatgpt_web_quota_default",
})),
}
}
fn chatgpt_web_quota_error_detail(result: &ExecutionResult) -> Option<String> {
extract_execution_error_message(result).or_else(|| {
let body = result.body.as_ref()?.body_bytes_b64.as_deref()?;
let decoded = base64::engine::general_purpose::STANDARD
.decode(body)
.ok()?;
let text = String::from_utf8_lossy(&decoded).trim().to_string();
(!text.is_empty()).then_some(text)
})
}
fn chatgpt_web_is_structured_account_block(message: &str) -> bool {
let lowered = message.to_ascii_lowercase();
[
"account has been disabled",
"account disabled",
"account has been deactivated",
"account_deactivated",
"account deactivated",
"organization has been disabled",
"organization_disabled",
"deactivated_workspace",
"account suspended",
"account banned",
"account_block",
"account blocked",
"访问被禁止",
"账户访问被禁止",
"账户已封禁",
"封禁",
"封号",
"被封",
]
.iter()
.any(|keyword| lowered.contains(keyword))
}
fn chatgpt_web_quota_403_refresh_failed_reason(message: Option<&str>) -> String {
let detail = message
.map(str::trim)
.filter(|value| !value.is_empty())
.filter(|value| !value.contains('<'))
.unwrap_or("ChatGPT Web 访问验证失败,请检查浏览器指纹、Cloudflare 验证或代理/地区限制");
format!("{OAUTH_REFRESH_FAILED_PREFIX}{detail}")
}
fn chatgpt_web_quota_invalid_reason(status_code: u16, upstream_message: Option<&str>) -> String {
let message = upstream_message.unwrap_or_default().trim();
if status_code == 403 && !chatgpt_web_is_structured_account_block(message) {
return chatgpt_web_quota_403_refresh_failed_reason(upstream_message);
}
let detail = if message.is_empty() {
match status_code {
401 => "ChatGPT Web Token 无效或已过期",
@@ -103,6 +188,19 @@ fn chatgpt_web_quota_invalid_reason(status_code: u16, upstream_message: Option<&
}
}
fn chatgpt_web_quota_result_message(reason: &str) -> String {
for prefix in [
OAUTH_REFRESH_FAILED_PREFIX,
OAUTH_EXPIRED_PREFIX,
OAUTH_ACCOUNT_BLOCK_PREFIX,
] {
if let Some(message) = reason.strip_prefix(prefix) {
return message.trim().to_string();
}
}
reason.trim().to_string()
}
pub(crate) async fn refresh_chatgpt_web_provider_quota_locally(
state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
@@ -216,8 +314,20 @@ pub(crate) async fn refresh_chatgpt_web_provider_quota_locally(
message = Some("响应中未包含 ChatGPT Web 生图限额信息".to_string());
}
} else {
let err_msg = extract_execution_error_message(&result);
message = Some(match err_msg.as_deref() {
let err_msg = chatgpt_web_quota_error_detail(&result);
let invalid_reason = if matches!(result.status_code, 401 | 403) {
Some(chatgpt_web_quota_invalid_reason(
result.status_code,
err_msg.as_deref(),
))
} else {
None
};
let display_detail = invalid_reason
.as_deref()
.map(chatgpt_web_quota_result_message)
.or_else(|| err_msg.clone());
message = Some(match display_detail.as_deref() {
Some(detail) if !detail.is_empty() => {
format!(
"conversation/init 返回状态码 {}: {}",
@@ -229,12 +339,14 @@ pub(crate) async fn refresh_chatgpt_web_provider_quota_locally(
if matches!(result.status_code, 401 | 403) {
oauth_invalid_at_unix_secs = Some(now_unix_secs);
oauth_invalid_reason = Some(chatgpt_web_quota_invalid_reason(
result.status_code,
err_msg.as_deref(),
));
oauth_invalid_reason = invalid_reason;
status = if result.status_code == 401 {
"auth_invalid".to_string()
} else if oauth_invalid_reason
.as_deref()
.is_some_and(|reason| reason.starts_with(OAUTH_REFRESH_FAILED_PREFIX))
{
"refresh_failed".to_string()
} else {
"forbidden".to_string()
};
@@ -303,3 +415,81 @@ pub(crate) async fn refresh_chatgpt_web_provider_quota_locally(
"auto_removed": auto_removed_count,
})))
}
#[cfg(test)]
mod tests {
use super::*;
use aether_contracts::{ResponseBody, TRANSPORT_BACKEND_REQWEST_RUSTLS};
use base64::Engine as _;
use std::collections::BTreeMap;
#[test]
fn quota_refresh_defaults_to_browser_wreq_transport() {
let profile = chatgpt_web_quota_transport_profile(None).expect("transport profile");
assert_eq!(profile.backend, TRANSPORT_BACKEND_BROWSER_WREQ);
assert_eq!(profile.profile_id, CHATGPT_WEB_BROWSER_PROFILE);
assert_eq!(profile.http_mode, TRANSPORT_HTTP_MODE_AUTO);
assert_eq!(profile.pool_scope, TRANSPORT_POOL_SCOPE_KEY);
assert_eq!(
profile
.extra
.as_ref()
.and_then(|value| value.get("browser_profile"))
.and_then(serde_json::Value::as_str),
Some(CHATGPT_WEB_BROWSER_PROFILE)
);
}
#[test]
fn quota_refresh_overrides_non_browser_transport() {
let reqwest_profile = ResolvedTransportProfile {
profile_id: "chrome_136".to_string(),
backend: TRANSPORT_BACKEND_REQWEST_RUSTLS.to_string(),
http_mode: TRANSPORT_HTTP_MODE_AUTO.to_string(),
pool_scope: TRANSPORT_POOL_SCOPE_KEY.to_string(),
header_fingerprint: None,
extra: None,
};
let profile =
chatgpt_web_quota_transport_profile(Some(&reqwest_profile)).expect("transport profile");
assert_eq!(profile.backend, TRANSPORT_BACKEND_BROWSER_WREQ);
assert_eq!(profile.profile_id, CHATGPT_WEB_BROWSER_PROFILE);
}
#[test]
fn browser_challenge_403_is_not_account_block() {
let body = "<!DOCTYPE html><html><head><title>Just a moment...</title></head><body>Cloudflare</body></html>";
let result = ExecutionResult {
request_id: "chatgpt-web-quota:test".to_string(),
candidate_id: None,
status_code: 403,
headers: BTreeMap::new(),
body: Some(ResponseBody {
json_body: None,
body_bytes_b64: Some(base64::engine::general_purpose::STANDARD.encode(body)),
}),
telemetry: None,
error: None,
};
let detail = chatgpt_web_quota_error_detail(&result).expect("html body should decode");
let reason = chatgpt_web_quota_invalid_reason(result.status_code, Some(&detail));
assert!(reason.starts_with(OAUTH_REFRESH_FAILED_PREFIX));
assert!(!reason.starts_with(OAUTH_ACCOUNT_BLOCK_PREFIX));
assert_eq!(
chatgpt_web_quota_result_message(&reason),
"ChatGPT Web 访问验证失败,请检查浏览器指纹、Cloudflare 验证或代理/地区限制"
);
}
#[test]
fn explicit_account_block_403_remains_account_block() {
let reason = chatgpt_web_quota_invalid_reason(403, Some("account has been deactivated"));
assert!(reason.starts_with(OAUTH_ACCOUNT_BLOCK_PREFIX));
}
}
@@ -13,6 +13,7 @@ use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
};
use aether_data_contracts::repository::usage::StoredProviderApiKeyWindowUsageSummary;
use aether_scheduler_core::provider_key_circuit_payload_is_active_open_at;
use serde_json::json;
use std::collections::BTreeMap;
use std::time::{SystemTime, UNIX_EPOCH};
@@ -921,19 +922,14 @@ fn admin_pool_health_score(key: &StoredProviderCatalogKey) -> f64 {
}
}
fn admin_pool_circuit_breaker_open(key: &StoredProviderCatalogKey) -> bool {
fn admin_pool_circuit_breaker_open(key: &StoredProviderCatalogKey, now_unix_secs: u64) -> bool {
key.circuit_breaker_by_format
.as_ref()
.and_then(serde_json::Value::as_object)
.map(|formats| {
formats
.values()
.filter_map(serde_json::Value::as_object)
.any(|item| {
item.get("open")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false)
})
.any(|item| provider_key_circuit_payload_is_active_open_at(item, now_unix_secs))
})
.unwrap_or(false)
}
@@ -1032,7 +1028,7 @@ pub(super) fn build_admin_pool_key_payload(
.as_ref()
.and_then(|_| runtime.cooldown_ttl_by_key.get(&key.id).copied());
let health_score = admin_pool_health_score(key);
let circuit_breaker_open = admin_pool_circuit_breaker_open(key);
let circuit_breaker_open = admin_pool_circuit_breaker_open(key, now_unix_secs);
let auth_semantics = provider_key_auth_semantics(key, provider_type);
let account_quota_exhausted = pool_config
.as_ref()
@@ -68,6 +68,7 @@ use aether_model_fetch::{
aggregate_models_for_cache, fetch_models_from_transports, json_string_list,
preset_models_for_provider, selected_models_fetch_endpoints,
};
use aether_scheduler_core::provider_key_circuit_payload_is_active_open_at;
use axum::{
body::{to_bytes, Body},
http::{self, HeaderMap, HeaderName, HeaderValue},
@@ -920,6 +921,7 @@ fn provider_query_test_key_sort_key(
provider_type: &str,
key: &StoredProviderCatalogKey,
endpoint_api_format: &str,
now_unix_secs: u64,
) -> (u8, u8, i32, u64, i32) {
let quota_exhausted =
admin_provider_pool_pure::admin_pool_key_account_quota_exhausted(key, provider_type);
@@ -928,10 +930,7 @@ fn provider_query_test_key_sort_key(
.as_ref()
.and_then(Value::as_object)
.and_then(|value| value.get(endpoint_api_format))
.and_then(Value::as_object)
.and_then(|value| value.get("open"))
.and_then(Value::as_bool)
.unwrap_or(false);
.is_some_and(|value| provider_key_circuit_payload_is_active_open_at(value, now_unix_secs));
let health_score = key
.health_by_format
.as_ref()
@@ -1390,6 +1389,7 @@ async fn provider_query_build_kiro_test_candidates(
provider_query_key_supports_endpoint(key, &provider.provider_type, &endpoint.api_format)
})
.collect::<Vec<_>>();
let now_unix_secs = current_unix_ms() / 1000;
let candidates = if test_mode.eq_ignore_ascii_case("pool") {
if let Some(pool_config) =
@@ -1411,6 +1411,7 @@ async fn provider_query_build_kiro_test_candidates(
provider.provider_type.as_str(),
key,
&endpoint.api_format,
now_unix_secs,
)
});
keys.into_iter()
@@ -1428,6 +1429,7 @@ async fn provider_query_build_kiro_test_candidates(
provider.provider_type.as_str(),
key,
&endpoint.api_format,
now_unix_secs,
)
});
keys.into_iter()
@@ -185,6 +185,7 @@ impl<'a> AdminAppState<'a> {
let standalone_wallets = self
.list_wallet_snapshots_by_api_key_ids(&standalone_api_key_ids)
.await?;
let usage_aggregates = self.export_admin_system_usage_aggregates().await?;
let wallets_by_user_id = user_wallets
.into_iter()
@@ -262,6 +263,7 @@ impl<'a> AdminAppState<'a> {
.collect::<Vec<_>>();
json!({
"id": user.id.clone(),
"email": user.email.clone(),
"email_verified": user.email_verified,
"username": user.username.clone(),
@@ -306,6 +308,7 @@ impl<'a> AdminAppState<'a> {
"user_groups": user_groups_data,
"users": users_data,
"standalone_keys": standalone_keys_data,
"usage_aggregates": usage_aggregates,
}))
}
@@ -330,6 +333,7 @@ impl<'a> AdminAppState<'a> {
include_is_standalone: bool,
) -> serde_json::Value {
let mut payload = serde_json::Map::from_iter([
("api_key_id".to_string(), json!(key.api_key_id.clone())),
("key_hash".to_string(), json!(key.key_hash.clone())),
("name".to_string(), json!(key.name.clone())),
(
@@ -38,6 +38,10 @@ use aether_data::repository::auth_modules::StoredLdapModuleConfig;
use aether_data::repository::oauth_providers::{
EncryptedSecretUpdate, UpsertOAuthProviderConfigRecord,
};
use aether_data::repository::system::{
AdminSystemUsageAggregateImportMode, AdminSystemUsageAggregateImportSummary,
AdminSystemUsageAggregateSnapshot,
};
use aether_data::repository::wallet::WalletLookupKey;
use aether_data_contracts::repository::global_models::{
AdminGlobalModelListQuery, AdminProviderModelListQuery, CreateAdminGlobalModelRecord,
@@ -257,6 +261,68 @@ fn build_import_key_match_name(item: &ImportedProviderKey) -> Option<String> {
.map(ToOwned::to_owned)
}
fn normalize_selected_import_key_format(
value: &str,
allowed_formats: &BTreeSet<String>,
) -> Option<String> {
let normalized = normalize_import_endpoint_format(value).ok()?;
allowed_formats.contains(&normalized).then_some(normalized)
}
fn normalize_import_key_format_scoped_list(
value: Option<&Value>,
normalized_api_formats: &[String],
) -> Option<Value> {
let value = value?;
let Value::Array(items) = value else {
return Some(value.clone());
};
let allowed_formats = normalized_api_formats
.iter()
.cloned()
.collect::<BTreeSet<_>>();
let mut seen = BTreeSet::new();
let mut normalized = Vec::new();
for item in items {
let Some(raw) = item
.as_str()
.map(str::trim)
.filter(|value| !value.is_empty())
else {
continue;
};
let Some(api_format) = normalize_selected_import_key_format(raw, &allowed_formats) else {
continue;
};
if seen.insert(api_format.clone()) {
normalized.push(json!(api_format));
}
}
Some(Value::Array(normalized))
}
fn normalize_import_key_format_scoped_object(
value: Option<&Value>,
normalized_api_formats: &[String],
) -> Option<Value> {
let value = value?;
let Value::Object(map) = value else {
return Some(value.clone());
};
let allowed_formats = normalized_api_formats
.iter()
.cloned()
.collect::<BTreeSet<_>>();
let mut normalized = Map::new();
for (key, value) in map {
let Some(api_format) = normalize_selected_import_key_format(key, &allowed_formats) else {
continue;
};
normalized.insert(api_format, value.clone());
}
Some(Value::Object(normalized))
}
fn normalize_import_key_raw_payload(
raw_key: &Map<String, Value>,
auth_type: &str,
@@ -268,6 +334,21 @@ fn normalize_import_key_raw_payload(
payload.remove("api_key");
}
payload.insert("api_formats".to_string(), json!(normalized_api_formats));
if let Some(auth_type_by_format) = normalize_import_key_format_scoped_object(
raw_key.get("auth_type_by_format"),
normalized_api_formats,
) {
payload.insert("auth_type_by_format".to_string(), auth_type_by_format);
}
if let Some(allow_auth_channel_mismatch_formats) = normalize_import_key_format_scoped_list(
raw_key.get("allow_auth_channel_mismatch_formats"),
normalized_api_formats,
) {
payload.insert(
"allow_auth_channel_mismatch_formats".to_string(),
allow_auth_channel_mismatch_formats,
);
}
if let Some(auth_config) = normalized_auth_config {
payload.insert("auth_config".to_string(), auth_config);
} else if raw_key.contains_key("auth_config") {
@@ -344,54 +425,6 @@ fn imported_oauth_expires_at_unix_secs(normalized_auth_config: Option<&Value>) -
None
}
fn imported_oauth_has_refresh_token(normalized_auth_config: Option<&Value>) -> bool {
normalized_auth_config
.and_then(Value::as_object)
.and_then(|object| object.get("refresh_token"))
.and_then(Value::as_str)
.map(str::trim)
.is_some_and(|value| !value.is_empty())
}
async fn refresh_imported_oauth_key_after_persist(
state: &AdminAppState<'_>,
provider: &aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider,
key_id: &str,
) -> Result<(), GatewayError> {
let Some(endpoint) =
crate::handlers::admin::provider::oauth::runtime::resolve_provider_oauth_runtime_endpoints(
state,
provider,
provider.provider_type.as_str(),
)
.await?
.runtime_endpoint
else {
return Ok(());
};
let Some(transport) = state
.read_provider_transport_snapshot(&provider.id, &endpoint.id, key_id)
.await?
else {
return Ok(());
};
if !crate::provider_transport::supports_local_oauth_request_auth_resolution(&transport) {
return Ok(());
}
if let Err(error) = state.force_local_oauth_refresh_entry(&transport).await {
tracing::warn!(
provider_id = %provider.id,
provider_type = %provider.provider_type,
key_id = %key_id,
error = ?error,
"admin system import oauth refresh after credential import failed"
);
}
Ok(())
}
fn build_import_provider_model_record(
provider_id: &str,
existing_id: Option<&str>,
@@ -435,6 +468,8 @@ struct AdminSystemUsersImportStats {
users: AdminSystemConfigImportCounter,
api_keys: AdminSystemConfigImportCounter,
standalone_keys: AdminSystemConfigImportCounter,
#[serde(skip_serializing_if = "Option::is_none")]
usage_aggregates: Option<AdminSystemUsageAggregateImportSummary>,
errors: Vec<String>,
}
@@ -488,6 +523,16 @@ fn validate_imported_system_users_export_version(version: Option<&Value>) -> Res
Ok(())
}
fn usage_aggregate_import_mode(
merge_mode: AdminImportMergeMode,
) -> AdminSystemUsageAggregateImportMode {
match merge_mode {
AdminImportMergeMode::Skip => AdminSystemUsageAggregateImportMode::Skip,
AdminImportMergeMode::Overwrite => AdminSystemUsageAggregateImportMode::Overwrite,
AdminImportMergeMode::Error => AdminSystemUsageAggregateImportMode::Error,
}
}
fn imported_object_field<'a>(
value: &'a Value,
field_name: &str,
@@ -1560,16 +1605,6 @@ impl<'a> AdminAppState<'a> {
"更新 Provider '{provider_name}' 的 Key 失败"
))));
};
if auth_type == "oauth"
&& imported_oauth_has_refresh_token(normalized_auth_config.as_ref())
{
refresh_imported_oauth_key_after_persist(
self,
&provider,
&persisted.id,
)
.await?;
}
existing_keys[existing_index] = persisted;
stats.keys.updated += 1;
}
@@ -1610,11 +1645,6 @@ impl<'a> AdminAppState<'a> {
"创建 Provider '{provider_name}' 的 Key 失败"
))));
};
if auth_type == "oauth"
&& imported_oauth_has_refresh_token(normalized_auth_config.as_ref())
{
refresh_imported_oauth_key_after_persist(self, &provider, &created.id).await?;
}
existing_keys.push(created);
stats.keys.created += 1;
}
@@ -2052,6 +2082,8 @@ impl<'a> AdminAppState<'a> {
));
let mut stats = AdminSystemUsersImportStats::default();
let mut imported_user_id_map = BTreeMap::<String, String>::new();
let mut imported_api_key_id_map = BTreeMap::<String, String>::new();
let default_group_id = self.effective_default_user_group_id().await?;
let existing_groups = self.list_user_groups().await?;
let mut groups_by_name = existing_groups
@@ -2138,6 +2170,7 @@ impl<'a> AdminAppState<'a> {
Ok(value) => value,
Err(detail) => return Ok(Err(invalid_request(detail))),
};
let source_user_id = invalid_value!(imported_optional_string(user.get("id")));
let role = invalid_value!(imported_optional_string(user.get("role")))
.unwrap_or_else(|| "user".to_string())
.to_ascii_lowercase();
@@ -2471,6 +2504,9 @@ impl<'a> AdminAppState<'a> {
stats.users.created += 1;
created.id
};
if let Some(source_user_id) = source_user_id {
imported_user_id_map.insert(source_user_id, user_id.clone());
}
let existing_api_keys = self
.list_auth_api_key_export_records_by_user_ids(std::slice::from_ref(&user_id))
@@ -2506,6 +2542,8 @@ impl<'a> AdminAppState<'a> {
));
continue;
};
let source_api_key_id =
invalid_value!(imported_optional_string(key.get("api_key_id")));
let name = invalid_value!(imported_optional_string(key.get("name")));
let allowed_providers = invalid_value!(normalize_imported_user_string_list(
key,
@@ -2538,21 +2576,21 @@ impl<'a> AdminAppState<'a> {
let auto_delete_on_expiry =
invalid_value!(imported_optional_bool(key.get("auto_delete_on_expiry")))
.unwrap_or(false);
let total_requests = invalid_value!(imported_optional_u64(
let imported_total_requests = invalid_value!(imported_optional_u64(
key.get("total_requests"),
"total_requests"
))
.unwrap_or(0);
let total_tokens = invalid_value!(imported_optional_u64(
));
let total_requests = imported_total_requests.unwrap_or(0);
let imported_total_tokens = invalid_value!(imported_optional_u64(
key.get("total_tokens"),
"total_tokens"
))
.unwrap_or(0);
let total_cost_usd = invalid_value!(imported_optional_f64(
));
let total_tokens = imported_total_tokens.unwrap_or(0);
let imported_total_cost_usd = invalid_value!(imported_optional_f64(
key.get("total_cost_usd"),
"total_cost_usd"
))
.unwrap_or(0.0);
));
let total_cost_usd = imported_total_cost_usd.unwrap_or(0.0);
let feature_settings = invalid_value!(imported_optional_json_object(
key.get("feature_settings"),
"feature_settings"
@@ -2624,13 +2662,31 @@ impl<'a> AdminAppState<'a> {
is_active,
)
.await?;
if imported_total_requests.is_some()
|| imported_total_tokens.is_some()
|| imported_total_cost_usd.is_some()
{
let updated_usage = self
.set_api_key_usage_totals(
&existing_key.api_key_id,
imported_total_requests
.unwrap_or(existing_key.total_requests),
imported_total_tokens.unwrap_or(existing_key.total_tokens),
imported_total_cost_usd
.unwrap_or(existing_key.total_cost_usd),
)
.await?;
if updated_usage.is_none() {
return Ok(Err((
http::StatusCode::SERVICE_UNAVAILABLE,
json!({ "detail": "Admin system data unavailable" }),
)));
}
}
if key.contains_key("allowed_api_formats")
|| key.contains_key("allowed_models")
|| key.contains_key("expires_at")
|| key.contains_key("auto_delete_on_expiry")
|| key.contains_key("total_requests")
|| key.contains_key("total_tokens")
|| key.contains_key("total_cost_usd")
{
stats.errors.push(format!(
"用户 '{}' 的现有 API Key 仅覆盖基础字段;高级导入字段保持原值",
@@ -2638,6 +2694,10 @@ impl<'a> AdminAppState<'a> {
));
}
stats.api_keys.updated += 1;
if let Some(source_api_key_id) = source_api_key_id.clone() {
imported_api_key_id_map
.insert(source_api_key_id, existing_key.api_key_id.clone());
}
}
}
continue;
@@ -2680,125 +2740,137 @@ impl<'a> AdminAppState<'a> {
)
.await?;
}
let created_api_key_id = created.api_key_id.clone();
existing_api_keys_by_hash.insert(key_hash, created);
if let Some(source_api_key_id) = source_api_key_id {
imported_api_key_id_map.insert(source_api_key_id, created_api_key_id);
}
stats.api_keys.created += 1;
}
}
if standalone_keys.is_empty() {
return Ok(Ok(json!({
"message": "用户数据导入成功",
"stats": stats,
})));
}
let Some(standalone_owner_id) = standalone_owner_id else {
stats.standalone_keys.skipped += standalone_keys.len() as u64;
stats
.errors
.push("无法导入独立余额 Key: 当前管理员用户记录不存在".to_string());
return Ok(Ok(json!({
"message": "用户数据导入成功",
"stats": stats,
})));
};
let existing_standalone_keys = self
.list_auth_api_key_export_standalone_records()
.await?
.into_iter()
.collect::<Vec<_>>();
let mut existing_standalone_by_hash = existing_standalone_keys
.into_iter()
.map(|record| (record.key_hash.clone(), record))
.collect::<BTreeMap<_, _>>();
for (index, raw_key) in standalone_keys.iter().enumerate() {
let key = match imported_object_field(raw_key, &format!("standalone_keys[{index}]")) {
Ok(value) => value,
Err(detail) => return Ok(Err(invalid_request(detail))),
};
let Some((key_hash, key_encrypted)) =
invalid_value!(self.resolve_imported_system_user_api_key_material(key))
else {
stats.standalone_keys.skipped += 1;
if !standalone_keys.is_empty() {
let Some(standalone_owner_id) = standalone_owner_id else {
stats.standalone_keys.skipped += standalone_keys.len() as u64;
stats
.errors
.push(format!("跳过无效独立余额 Key: standalone_keys[{index}]"));
continue;
.push("无法导入独立余额 Key: 当前管理员用户记录不存在".to_string());
if let Some(summary) = self
.import_admin_system_user_usage_aggregates(
root.get("usage_aggregates"),
&imported_user_id_map,
&imported_api_key_id_map,
merge_mode,
)
.await?
{
stats.usage_aggregates = Some(summary);
}
return Ok(Ok(json!({
"message": "用户数据导入成功",
"stats": stats,
})));
};
let name = invalid_value!(imported_optional_string(key.get("name")));
let allowed_providers = invalid_value!(normalize_imported_user_string_list(
key,
"allowed_providers"
));
let allowed_api_formats = invalid_value!(normalize_imported_user_api_formats(
key,
"allowed_api_formats"
));
let allowed_models =
invalid_value!(normalize_imported_user_string_list(key, "allowed_models"));
let ip_rules = invalid_value!(normalize_imported_user_ip_rules(key));
let rate_limit =
invalid_value!(imported_optional_i32(key.get("rate_limit"), "rate_limit"))
.unwrap_or(0);
let concurrent_limit = invalid_value!(imported_optional_i32(
key.get("concurrent_limit"),
"concurrent_limit"
));
if concurrent_limit.is_some_and(|value| value < 0) {
return Ok(Err(invalid_request("concurrent_limit 必须是非负整数")));
}
let force_capabilities = imported_optional_value(key.get("force_capabilities"));
let is_active =
invalid_value!(imported_optional_bool(key.get("is_active"))).unwrap_or(true);
let expires_at_unix_secs = invalid_value!(imported_rfc3339_to_unix_secs(
key.get("expires_at"),
"expires_at"
));
let auto_delete_on_expiry =
invalid_value!(imported_optional_bool(key.get("auto_delete_on_expiry")))
.unwrap_or(false);
let total_requests = invalid_value!(imported_optional_u64(
key.get("total_requests"),
"total_requests"
))
.unwrap_or(0);
let total_tokens = invalid_value!(imported_optional_u64(
key.get("total_tokens"),
"total_tokens"
))
.unwrap_or(0);
let total_cost_usd = invalid_value!(imported_optional_f64(
key.get("total_cost_usd"),
"total_cost_usd"
))
.unwrap_or(0.0);
let feature_settings = invalid_value!(imported_optional_json_object(
key.get("feature_settings"),
"feature_settings"
)
.and_then(normalize_admin_feature_settings));
let wallet_payload = match key.get("wallet") {
Some(Value::Object(map)) => Some(map),
Some(Value::Null) | None => None,
Some(_) => return Ok(Err(invalid_request("wallet 必须是对象"))),
};
let unlimited =
invalid_value!(imported_optional_bool(key.get("unlimited"))).unwrap_or(false);
let wallet_target =
invalid_value!(normalize_imported_wallet_target(wallet_payload, unlimited));
if let Some(existing_key) = existing_standalone_by_hash.get(&key_hash).cloned() {
match merge_mode {
AdminImportMergeMode::Skip => {
stats.standalone_keys.skipped += 1;
}
AdminImportMergeMode::Error => {
return Ok(Err(invalid_request("独立余额 Key 已存在")));
}
AdminImportMergeMode::Overwrite => {
let updated = self
let existing_standalone_keys = self
.list_auth_api_key_export_standalone_records()
.await?
.into_iter()
.collect::<Vec<_>>();
let mut existing_standalone_by_hash = existing_standalone_keys
.into_iter()
.map(|record| (record.key_hash.clone(), record))
.collect::<BTreeMap<_, _>>();
for (index, raw_key) in standalone_keys.iter().enumerate() {
let key = match imported_object_field(raw_key, &format!("standalone_keys[{index}]"))
{
Ok(value) => value,
Err(detail) => return Ok(Err(invalid_request(detail))),
};
let Some((key_hash, key_encrypted)) =
invalid_value!(self.resolve_imported_system_user_api_key_material(key))
else {
stats.standalone_keys.skipped += 1;
stats
.errors
.push(format!("跳过无效独立余额 Key: standalone_keys[{index}]"));
continue;
};
let source_api_key_id =
invalid_value!(imported_optional_string(key.get("api_key_id")));
let name = invalid_value!(imported_optional_string(key.get("name")));
let allowed_providers = invalid_value!(normalize_imported_user_string_list(
key,
"allowed_providers"
));
let allowed_api_formats = invalid_value!(normalize_imported_user_api_formats(
key,
"allowed_api_formats"
));
let allowed_models =
invalid_value!(normalize_imported_user_string_list(key, "allowed_models"));
let ip_rules = invalid_value!(normalize_imported_user_ip_rules(key));
let rate_limit =
invalid_value!(imported_optional_i32(key.get("rate_limit"), "rate_limit"))
.unwrap_or(0);
let concurrent_limit = invalid_value!(imported_optional_i32(
key.get("concurrent_limit"),
"concurrent_limit"
));
if concurrent_limit.is_some_and(|value| value < 0) {
return Ok(Err(invalid_request("concurrent_limit 必须是非负整数")));
}
let force_capabilities = imported_optional_value(key.get("force_capabilities"));
let is_active =
invalid_value!(imported_optional_bool(key.get("is_active"))).unwrap_or(true);
let expires_at_unix_secs = invalid_value!(imported_rfc3339_to_unix_secs(
key.get("expires_at"),
"expires_at"
));
let auto_delete_on_expiry =
invalid_value!(imported_optional_bool(key.get("auto_delete_on_expiry")))
.unwrap_or(false);
let imported_total_requests = invalid_value!(imported_optional_u64(
key.get("total_requests"),
"total_requests"
));
let total_requests = imported_total_requests.unwrap_or(0);
let imported_total_tokens = invalid_value!(imported_optional_u64(
key.get("total_tokens"),
"total_tokens"
));
let total_tokens = imported_total_tokens.unwrap_or(0);
let imported_total_cost_usd = invalid_value!(imported_optional_f64(
key.get("total_cost_usd"),
"total_cost_usd"
));
let total_cost_usd = imported_total_cost_usd.unwrap_or(0.0);
let feature_settings = invalid_value!(imported_optional_json_object(
key.get("feature_settings"),
"feature_settings"
)
.and_then(normalize_admin_feature_settings));
let wallet_payload = match key.get("wallet") {
Some(Value::Object(map)) => Some(map),
Some(Value::Null) | None => None,
Some(_) => return Ok(Err(invalid_request("wallet 必须是对象"))),
};
let unlimited =
invalid_value!(imported_optional_bool(key.get("unlimited"))).unwrap_or(false);
let wallet_target =
invalid_value!(normalize_imported_wallet_target(wallet_payload, unlimited));
if let Some(existing_key) = existing_standalone_by_hash.get(&key_hash).cloned() {
match merge_mode {
AdminImportMergeMode::Skip => {
stats.standalone_keys.skipped += 1;
}
AdminImportMergeMode::Error => {
return Ok(Err(invalid_request("独立余额 Key 已存在")));
}
AdminImportMergeMode::Overwrite => {
let updated = self
.update_standalone_api_key_basic(
aether_data::repository::auth::UpdateStandaloneApiKeyBasicRecord {
api_key_id: existing_key.api_key_id.clone(),
@@ -2819,94 +2891,134 @@ impl<'a> AdminAppState<'a> {
},
)
.await?;
if updated.is_none() {
return Ok(Err((
http::StatusCode::SERVICE_UNAVAILABLE,
json!({ "detail": "Admin system data unavailable" }),
)));
}
let _ = self
.set_standalone_api_key_active(&existing_key.api_key_id, is_active)
.await?;
if key.contains_key("feature_settings") {
if updated.is_none() {
return Ok(Err((
http::StatusCode::SERVICE_UNAVAILABLE,
json!({ "detail": "Admin system data unavailable" }),
)));
}
let _ = self
.set_standalone_api_key_feature_settings(
&existing_key.api_key_id,
feature_settings.clone(),
)
.set_standalone_api_key_active(&existing_key.api_key_id, is_active)
.await?;
if key.contains_key("feature_settings") {
let _ = self
.set_standalone_api_key_feature_settings(
&existing_key.api_key_id,
feature_settings.clone(),
)
.await?;
}
if imported_total_requests.is_some()
|| imported_total_tokens.is_some()
|| imported_total_cost_usd.is_some()
{
let updated_usage = self
.set_api_key_usage_totals(
&existing_key.api_key_id,
imported_total_requests
.unwrap_or(existing_key.total_requests),
imported_total_tokens.unwrap_or(existing_key.total_tokens),
imported_total_cost_usd
.unwrap_or(existing_key.total_cost_usd),
)
.await?;
if updated_usage.is_none() {
return Ok(Err((
http::StatusCode::SERVICE_UNAVAILABLE,
json!({ "detail": "Admin system data unavailable" }),
)));
}
}
if key.contains_key("expires_at")
|| key.contains_key("auto_delete_on_expiry")
|| key.contains_key("force_capabilities")
{
stats.errors.push(
"现有独立余额 Key 仅覆盖基础字段;高级导入字段保持原值"
.to_string(),
);
}
self.sync_imported_api_key_wallet(
&existing_key.api_key_id,
&wallet_target,
key.get("name")
.and_then(Value::as_str)
.unwrap_or("独立余额 Key"),
)
.await?;
stats.standalone_keys.updated += 1;
if let Some(source_api_key_id) = source_api_key_id.clone() {
imported_api_key_id_map
.insert(source_api_key_id, existing_key.api_key_id.clone());
}
}
if key.contains_key("expires_at")
|| key.contains_key("auto_delete_on_expiry")
|| key.contains_key("force_capabilities")
|| key.contains_key("total_requests")
|| key.contains_key("total_tokens")
|| key.contains_key("total_cost_usd")
{
stats.errors.push(
"现有独立余额 Key 仅覆盖基础字段;高级导入字段保持原值".to_string(),
);
}
self.sync_imported_api_key_wallet(
&existing_key.api_key_id,
&wallet_target,
key.get("name")
.and_then(Value::as_str)
.unwrap_or("独立余额 Key"),
)
.await?;
stats.standalone_keys.updated += 1;
}
continue;
}
continue;
}
let created = self
.create_standalone_api_key(
aether_data::repository::auth::CreateStandaloneApiKeyRecord {
user_id: standalone_owner_id.clone(),
api_key_id: Uuid::new_v4().to_string(),
key_hash: key_hash.clone(),
key_encrypted,
name,
allowed_providers,
allowed_api_formats,
allowed_models,
ip_rules,
rate_limit: Some(rate_limit),
concurrent_limit,
force_capabilities,
is_active,
expires_at_unix_secs,
auto_delete_on_expiry,
total_requests,
total_tokens,
total_cost_usd,
},
)
.await?;
let Some(created) = created else {
return Ok(Err((
http::StatusCode::SERVICE_UNAVAILABLE,
json!({ "detail": "Admin system data unavailable" }),
)));
};
if key.contains_key("feature_settings") {
let _ = self
.set_standalone_api_key_feature_settings(
&created.api_key_id,
feature_settings.clone(),
let created = self
.create_standalone_api_key(
aether_data::repository::auth::CreateStandaloneApiKeyRecord {
user_id: standalone_owner_id.clone(),
api_key_id: Uuid::new_v4().to_string(),
key_hash: key_hash.clone(),
key_encrypted,
name,
allowed_providers,
allowed_api_formats,
allowed_models,
ip_rules,
rate_limit: Some(rate_limit),
concurrent_limit,
force_capabilities,
is_active,
expires_at_unix_secs,
auto_delete_on_expiry,
total_requests,
total_tokens,
total_cost_usd,
},
)
.await?;
let Some(created) = created else {
return Ok(Err((
http::StatusCode::SERVICE_UNAVAILABLE,
json!({ "detail": "Admin system data unavailable" }),
)));
};
if key.contains_key("feature_settings") {
let _ = self
.set_standalone_api_key_feature_settings(
&created.api_key_id,
feature_settings.clone(),
)
.await?;
}
self.sync_imported_api_key_wallet(
&created.api_key_id,
&wallet_target,
created.name.as_deref().unwrap_or("独立余额 Key"),
)
.await?;
let created_api_key_id = created.api_key_id.clone();
existing_standalone_by_hash.insert(key_hash, created);
if let Some(source_api_key_id) = source_api_key_id {
imported_api_key_id_map.insert(source_api_key_id, created_api_key_id);
}
stats.standalone_keys.created += 1;
}
self.sync_imported_api_key_wallet(
&created.api_key_id,
&wallet_target,
created.name.as_deref().unwrap_or("独立余额 Key"),
}
if let Some(summary) = self
.import_admin_system_user_usage_aggregates(
root.get("usage_aggregates"),
&imported_user_id_map,
&imported_api_key_id_map,
merge_mode,
)
.await?;
existing_standalone_by_hash.insert(key_hash, created);
stats.standalone_keys.created += 1;
.await?
{
stats.usage_aggregates = Some(summary);
}
Ok(Ok(json!({
@@ -2915,6 +3027,40 @@ impl<'a> AdminAppState<'a> {
})))
}
async fn import_admin_system_user_usage_aggregates(
&self,
value: Option<&Value>,
user_id_map: &BTreeMap<String, String>,
api_key_id_map: &BTreeMap<String, String>,
merge_mode: AdminImportMergeMode,
) -> Result<Option<AdminSystemUsageAggregateImportSummary>, GatewayError> {
let Some(value) = value else {
return Ok(None);
};
if value.is_null() {
return Ok(None);
}
let snapshot = serde_json::from_value::<AdminSystemUsageAggregateSnapshot>(value.clone())
.map_err(|err| GatewayError::Client {
status: http::StatusCode::BAD_REQUEST,
message: format!("usage_aggregates 格式无效: {err}"),
})?;
if snapshot.stats_daily.is_empty()
&& snapshot.stats_user_daily.is_empty()
&& snapshot.stats_daily_api_key.is_empty()
{
return Ok(None);
}
self.import_admin_system_usage_aggregates(
&snapshot,
user_id_map,
api_key_id_map,
usage_aggregate_import_mode(merge_mode),
)
.await
.map(Some)
}
async fn sync_imported_user_wallet(
&self,
user_id: &str,
@@ -3042,8 +3188,8 @@ mod tests {
imported_optional_bool, imported_optional_f64, imported_optional_i32,
imported_optional_u64, imported_rfc3339_to_unix_secs, imported_string_list_from_value,
normalize_import_endpoint_format, normalize_import_key_formats,
normalize_imported_wallet_target, validate_imported_system_users_export_version,
ImportedProviderKey,
normalize_import_key_raw_payload, normalize_imported_wallet_target,
validate_imported_system_users_export_version, ImportedProviderKey,
};
#[test]
@@ -3117,6 +3263,61 @@ mod tests {
assert!(missing.is_empty());
}
#[test]
fn config_import_filters_key_format_scoped_fields_to_selected_api_formats() {
let raw_key = json!({
"name": "test-key",
"api_key": "sk-test",
"api_formats": ["openai:responses", "openai:video"],
"auth_type_by_format": {
"openai:responses": "api_key",
"openai:video": "bearer"
},
"allow_auth_channel_mismatch_formats": [
"openai:responses",
"openai:video"
]
});
let raw_key = raw_key.as_object().expect("key should be object");
let payload = normalize_import_key_raw_payload(
raw_key,
"api_key",
&["openai:responses".to_string()],
None,
);
assert_eq!(payload["api_formats"], json!(["openai:responses"]));
assert_eq!(
payload["auth_type_by_format"],
json!({ "openai:responses": "api_key" })
);
assert_eq!(
payload["allow_auth_channel_mismatch_formats"],
json!(["openai:responses"])
);
}
#[test]
fn config_import_preserves_explicit_empty_mismatch_scope_after_filtering() {
let raw_key = json!({
"name": "test-key",
"api_key": "sk-test",
"api_formats": ["openai:responses"],
"allow_auth_channel_mismatch_formats": ["openai:video"]
});
let raw_key = raw_key.as_object().expect("key should be object");
let payload = normalize_import_key_raw_payload(
raw_key,
"api_key",
&["openai:responses".to_string()],
None,
);
assert_eq!(payload["allow_auth_channel_mismatch_formats"], json!([]));
}
#[test]
fn import_handles_legacy_string_scalars() {
assert_eq!(
@@ -70,6 +70,26 @@ impl<'a> AdminAppState<'a> {
self.app.purge_admin_system_data(target).await
}
pub(crate) async fn export_admin_system_usage_aggregates(
&self,
) -> Result<aether_data::repository::system::AdminSystemUsageAggregateSnapshot, GatewayError>
{
self.app.export_admin_system_usage_aggregates().await
}
pub(crate) async fn import_admin_system_usage_aggregates(
&self,
snapshot: &aether_data::repository::system::AdminSystemUsageAggregateSnapshot,
user_id_map: &std::collections::BTreeMap<String, String>,
api_key_id_map: &std::collections::BTreeMap<String, String>,
mode: aether_data::repository::system::AdminSystemUsageAggregateImportMode,
) -> Result<aether_data::repository::system::AdminSystemUsageAggregateImportSummary, GatewayError>
{
self.app
.import_admin_system_usage_aggregates(snapshot, user_id_map, api_key_id_map, mode)
.await
}
pub(crate) async fn run_admin_system_cleanup_once(
&self,
) -> Result<crate::maintenance::AdminSystemCleanupSummary, GatewayError> {
@@ -706,6 +706,19 @@ impl<'a> AdminAppState<'a> {
.await
}
pub(crate) async fn set_api_key_usage_totals(
&self,
api_key_id: &str,
total_requests: u64,
total_tokens: u64,
total_cost_usd: f64,
) -> Result<Option<aether_data::repository::auth::StoredAuthApiKeyExportRecord>, GatewayError>
{
self.app
.set_api_key_usage_totals(api_key_id, total_requests, total_tokens, total_cost_usd)
.await
}
pub(crate) async fn delete_user_api_key(
&self,
user_id: &str,
@@ -13,7 +13,8 @@ pub(crate) use crate::handlers::shared::{
effective_catalog_encryption_key, encrypt_catalog_secret_with_fallbacks, json_string_list,
masked_catalog_api_key, normalize_json_array, normalize_json_object, normalize_string_list,
parse_catalog_auth_config_json, provider_catalog_key_supports_format,
provider_key_health_summary, provider_key_status_snapshot_payload, query_param_bool,
query_param_optional_bool, query_param_value, take_secret_prefix, take_secret_suffix,
unix_secs_to_rfc3339, OFFICIAL_EXTERNAL_MODEL_PROVIDERS,
provider_key_health_summary, provider_key_health_summary_at,
provider_key_status_snapshot_payload, query_param_bool, query_param_optional_bool,
query_param_value, take_secret_prefix, take_secret_suffix, unix_secs_to_rfc3339,
OFFICIAL_EXTERNAL_MODEL_PROVIDERS,
};
@@ -2,6 +2,7 @@ use super::enabled_key_capability_short_names;
use crate::handlers::shared::unix_secs_to_rfc3339;
use crate::provider_key_auth::provider_key_effective_api_formats;
use crate::AppState;
use aether_scheduler_core::provider_key_circuit_payload_is_active_open_at;
use serde_json::json;
use std::collections::{BTreeMap, HashMap};
use std::time::{SystemTime, UNIX_EPOCH};
@@ -173,10 +174,10 @@ pub(crate) async fn build_admin_keys_grouped_by_format_payload(
.get("health_score")
.and_then(serde_json::Value::as_f64)
.unwrap_or(1.0),
"circuit_breaker_open": format_circuit
.get("open")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false),
"circuit_breaker_open": provider_key_circuit_payload_is_active_open_at(
&format_circuit,
now_unix_secs,
),
"last_used_at": key.last_used_at_unix_secs.and_then(unix_secs_to_rfc3339),
"created_at": unix_secs_to_rfc3339(key.created_at_unix_ms.unwrap_or(now_unix_secs)),
"updated_at": unix_secs_to_rfc3339(key.updated_at_unix_secs.unwrap_or(now_unix_secs)),
@@ -13,6 +13,7 @@ use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKe
use aether_provider_pool::{
grok_pool_tier_from_quota_bucket, grok_supported_quota_windows_for_tier,
};
use aether_scheduler_core::provider_key_circuit_payload_is_active_open_at;
use serde_json::{json, Map, Value};
use std::borrow::Cow;
use std::time::{SystemTime, UNIX_EPOCH};
@@ -438,27 +439,53 @@ fn chatgpt_web_image_quota_limit(
metadata: &Map<String, Value>,
remaining: Option<f64>,
) -> Option<f64> {
let explicit_limit = metadata
.get("image_quota_total")
.and_then(admin_provider_quota_pure::coerce_json_f64)
.filter(|value| *value > 0.0);
let plan_type = metadata
.get("plan_type")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|value| value.to_ascii_lowercase());
if plan_type.as_deref() == Some("free") {
return Some(25.0);
}
let explicit_limit = metadata
.get("image_quota_total")
.and_then(admin_provider_quota_pure::coerce_json_f64)
.filter(|value| *value > 0.0);
let limit_source = metadata
.get("image_quota_limit_source")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty());
if let Some(limit) = explicit_limit {
return Some(limit);
if !chatgpt_web_image_quota_limit_is_legacy_free_default(
limit,
limit_source,
plan_type.as_deref(),
remaining,
) {
return Some(limit);
}
}
remaining.filter(|value| *value > 0.0)
}
fn chatgpt_web_image_quota_limit_is_legacy_free_default(
limit: f64,
limit_source: Option<&str>,
plan_type: Option<&str>,
remaining: Option<f64>,
) -> bool {
let plan_type_is_free = plan_type
.map(str::trim)
.is_some_and(|value| value.eq_ignore_ascii_case("free"));
if !plan_type_is_free || limit_source.is_some() {
return false;
}
if (limit - 25.0).abs() > f64::EPSILON {
return false;
}
remaining.is_none_or(|value| value < limit)
}
fn model_quota_window_snapshot(
model_name: &str,
item: &Map<String, Value>,
@@ -1778,6 +1805,39 @@ pub(crate) fn provider_key_health_summary(
Option<String>,
bool,
serde_json::Map<String, serde_json::Value>,
) {
provider_key_health_summary_with_circuit_predicate(key, |value| {
value
.get("open")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false)
})
}
pub(crate) fn provider_key_health_summary_at(
key: &StoredProviderCatalogKey,
now_unix_secs: u64,
) -> (
f64,
i64,
Option<String>,
bool,
serde_json::Map<String, serde_json::Value>,
) {
provider_key_health_summary_with_circuit_predicate(key, |value| {
provider_key_circuit_payload_is_active_open_at(value, now_unix_secs)
})
}
fn provider_key_health_summary_with_circuit_predicate(
key: &StoredProviderCatalogKey,
circuit_is_open: impl Fn(&serde_json::Value) -> bool,
) -> (
f64,
i64,
Option<String>,
bool,
serde_json::Map<String, serde_json::Value>,
) {
let health_by_format = key
.health_by_format
@@ -1820,12 +1880,7 @@ pub(crate) fn provider_key_health_summary(
}
}
let any_circuit_open = circuit_by_format.values().any(|value| {
value
.get("open")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false)
});
let any_circuit_open = circuit_by_format.values().any(circuit_is_open);
(
if health_by_format.is_empty() {
@@ -1974,15 +2029,10 @@ pub(crate) fn build_admin_provider_key_response(
last_failure_at,
circuit_breaker_open,
circuit_by_format,
) = provider_key_health_summary(key);
) = provider_key_health_summary_at(key, now_unix_secs);
let circuit_sample = circuit_by_format
.values()
.find(|value| {
value
.get("open")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false)
})
.find(|value| provider_key_circuit_payload_is_active_open_at(value, now_unix_secs))
.or_else(|| circuit_by_format.values().next());
let is_adaptive = key.rpm_limit.is_none();
let effective_limit = if is_adaptive {
@@ -2535,12 +2585,39 @@ mod tests {
assert_eq!(quota.get("code"), Some(&json!("ok")));
assert_eq!(quota.get("plan_type"), Some(&json!("free")));
assert_eq!(quota.get("reset_at"), Some(&json!(1_778_157_172u64)));
assert_eq!(quota.get("usage_ratio"), Some(&json!(0.04)));
assert_eq!(quota.get("usage_ratio"), Some(&json!(0.0)));
assert_eq!(window.get("code"), Some(&json!("image_gen")));
assert_eq!(window.get("remaining_value"), Some(&json!(24.0)));
assert_eq!(window.get("limit_value"), Some(&json!(25.0)));
assert_eq!(window.get("used_value"), Some(&json!(1.0)));
assert_eq!(window.get("remaining_ratio"), Some(&json!(0.96)));
assert_eq!(window.get("limit_value"), Some(&json!(24.0)));
assert_eq!(window.get("used_value"), Some(&json!(0.0)));
assert_eq!(window.get("remaining_ratio"), Some(&json!(1.0)));
}
#[test]
fn provider_key_status_snapshot_payload_ignores_chatgpt_web_legacy_free_25_limit() {
let mut key = sample_catalog_key();
key.upstream_metadata = Some(json!({
"chatgpt_web": {
"updated_at": 1_778_067_246u64,
"plan_type": "free",
"image_quota_remaining": 19.0,
"image_quota_total": 25.0
}
}));
let payload = provider_key_status_snapshot_payload(&key, "chatgpt_web");
let window = payload
.get("quota")
.and_then(Value::as_object)
.and_then(|quota| quota.get("windows"))
.and_then(Value::as_array)
.and_then(|windows| windows.first())
.and_then(Value::as_object)
.expect("image quota window should exist");
assert_eq!(window.get("remaining_value"), Some(&json!(19.0)));
assert_eq!(window.get("limit_value"), Some(&json!(19.0)));
assert_eq!(window.get("used_value"), Some(&json!(0.0)));
}
#[test]
@@ -26,8 +26,9 @@ pub(crate) use self::catalog::{
default_provider_key_status_snapshot, effective_catalog_encryption_key,
encrypt_catalog_secret_with_fallbacks, masked_catalog_api_key, parse_catalog_auth_config_json,
provider_catalog_key_supports_format, provider_key_health_summary,
provider_key_status_snapshot_payload, sync_provider_key_oauth_status_snapshot,
sync_provider_key_quota_status_snapshot, take_secret_prefix, take_secret_suffix,
provider_key_health_summary_at, provider_key_status_snapshot_payload,
sync_provider_key_oauth_status_snapshot, sync_provider_key_quota_status_snapshot,
take_secret_prefix, take_secret_suffix,
};
pub(crate) use self::email_templates::{
admin_email_template_definition, admin_email_template_html_key,
+30
View File
@@ -629,6 +629,36 @@ impl AppState {
Ok(summary)
}
pub(crate) async fn export_admin_system_usage_aggregates(
&self,
) -> Result<aether_data::repository::system::AdminSystemUsageAggregateSnapshot, GatewayError>
{
self.data
.export_admin_system_usage_aggregates()
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn import_admin_system_usage_aggregates(
&self,
snapshot: &aether_data::repository::system::AdminSystemUsageAggregateSnapshot,
user_id_map: &std::collections::BTreeMap<String, String>,
api_key_id_map: &std::collections::BTreeMap<String, String>,
mode: aether_data::repository::system::AdminSystemUsageAggregateImportMode,
) -> Result<aether_data::repository::system::AdminSystemUsageAggregateImportSummary, GatewayError>
{
self.data
.import_admin_system_usage_aggregates(snapshot, user_id_map, api_key_id_map, mode)
.await
.map_err(|err| match err {
aether_data::DataLayerError::InvalidInput(detail) => GatewayError::Client {
status: http::StatusCode::BAD_REQUEST,
message: detail,
},
other => GatewayError::Internal(other.to_string()),
})
}
pub(crate) async fn run_admin_system_cleanup_once(
&self,
) -> Result<crate::maintenance::AdminSystemCleanupSummary, GatewayError> {
@@ -369,6 +369,25 @@ impl AppState {
Ok(api_key)
}
pub(crate) async fn set_api_key_usage_totals(
&self,
api_key_id: &str,
total_requests: u64,
total_tokens: u64,
total_cost_usd: f64,
) -> Result<Option<aether_data::repository::auth::StoredAuthApiKeyExportRecord>, GatewayError>
{
let api_key = self
.data
.set_api_key_usage_totals(api_key_id, total_requests, total_tokens, total_cost_usd)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if api_key.is_some() {
self.invalidate_auth_context_cache();
}
Ok(api_key)
}
pub(crate) async fn set_standalone_api_key_feature_settings(
&self,
api_key_id: &str,
@@ -241,8 +241,8 @@ async fn gateway_provider_keys_expose_circuit_breaker_and_recover_clears_it() {
"open": true,
"open_at": "2026-03-26T12:00:00+00:00",
"reason": "consecutive_failures_8",
"next_probe_at": "2026-03-26T12:01:00+00:00",
"next_probe_at_unix_secs": 1774526460u64,
"next_probe_at": "2099-03-26T12:01:00+00:00",
"next_probe_at_unix_secs": 4078209660u64,
"probe_interval_minutes": 1,
"max_probe_interval_minutes": 32,
"half_open_until": null,
@@ -346,7 +346,7 @@ async fn gateway_handles_admin_key_health_locally_with_trusted_admin_principal()
Some(json!({"openai:chat": {
"open": true,
"open_at": "2026-03-26T12:01:00+00:00",
"next_probe_at": "2026-03-26T12:05:00+00:00",
"next_probe_at": "2099-03-26T12:05:00+00:00",
"half_open_until": null,
"half_open_successes": 1,
"half_open_failures": 0
@@ -399,7 +399,7 @@ async fn gateway_handles_admin_key_health_locally_with_trusted_admin_principal()
payload["circuit_breaker_open_at"],
"2026-03-26T12:01:00+00:00"
);
assert_eq!(payload["next_probe_at"], "2026-03-26T12:05:00+00:00");
assert_eq!(payload["next_probe_at"], "2099-03-26T12:05:00+00:00");
assert_eq!(payload["half_open_successes"], 1);
assert_eq!(payload["half_open_failures"], 0);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
@@ -408,6 +408,85 @@ async fn gateway_handles_admin_key_health_locally_with_trusted_admin_principal()
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_admin_key_health_summary_treats_expired_unix_circuit_as_closed() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/endpoints/health/key/key-openai",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "openai", 10)],
vec![sample_endpoint(
"endpoint-openai",
"provider-openai",
"openai:chat",
"https://api.openai.example",
)],
vec![
sample_key("key-openai", "provider-openai", "openai:chat", "sk-test")
.with_health_fields(
Some(json!({"openai:chat": {
"health_score": 0.7,
"consecutive_failures": 2
}})),
Some(json!({"openai:chat": {
"open": true,
"open_at": "2026-03-26T12:01:00+00:00",
"next_probe_at_unix_secs": 1u64
}})),
),
],
));
let (_upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/endpoints/health/key/key-openai"
))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("request should succeed");
let status = response.status();
let body = response.text().await.expect("body should read");
assert_eq!(status, StatusCode::OK, "body={body}");
let payload: serde_json::Value = serde_json::from_str(&body).expect("json body should parse");
let circuit = &payload["health_by_format"]["openai:chat"]["circuit_breaker"];
assert_eq!(payload["any_circuit_open"], false);
assert_eq!(circuit["open"], false);
assert_eq!(circuit["state"], "closed");
assert_eq!(payload["key_health_score"], 0.7);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_recovers_admin_key_health_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
@@ -442,7 +521,7 @@ async fn gateway_recovers_admin_key_health_locally_with_trusted_admin_principal(
Some(json!({"openai:chat": {
"open": true,
"open_at": "2026-03-26T12:01:00+00:00",
"next_probe_at": "2026-03-26T12:05:00+00:00",
"next_probe_at": "2099-03-26T12:05:00+00:00",
"half_open_until": null,
"half_open_successes": 0,
"half_open_failures": 1
@@ -789,7 +789,7 @@ async fn gateway_handles_admin_global_model_routing_locally_with_trusted_admin_p
"openai:chat": {"health_score": 0.66}
}));
primary_key.circuit_breaker_by_format = Some(json!({
"openai:chat": {"open": true, "next_probe_at": "2026-03-27T15:00:00Z"}
"openai:chat": {"open": true, "next_probe_at": "2099-03-27T15:00:00Z"}
}));
let mut mapped_key = sample_key(
@@ -907,7 +907,7 @@ async fn gateway_handles_admin_global_model_routing_locally_with_trusted_admin_p
openai_keys[0]["circuit_breaker_formats"],
json!(["openai:chat"])
);
assert_eq!(openai_keys[0]["next_probe_at"], "2026-03-27T15:00:00Z");
assert_eq!(openai_keys[0]["next_probe_at"], "2099-03-27T15:00:00Z");
let alt_endpoints = providers[1]["endpoints"]
.as_array()
@@ -2040,7 +2040,8 @@ async fn gateway_handles_admin_monitoring_resilience_circuit_history_locally_wit
"openai:chat": {
"open": true,
"open_at": "2026-03-30T12:00:00+00:00",
"next_probe_at": "2026-03-30T12:05:00+00:00",
"next_probe_at": "2099-03-30T12:05:00+00:00",
"recovery_seconds": 300,
"reason": "错误率过高"
}
})),
@@ -1043,6 +1043,7 @@ async fn gateway_handles_admin_system_users_export_locally_with_trusted_admin_pr
payload["users"][0]["group_names"],
json!(["Restricted GPT"])
);
assert_eq!(payload["users"][0]["id"], json!("user-1"));
assert_eq!(payload["users"][0]["wallet"]["balance"], json!(12.5));
assert_eq!(
payload["users"][0]["wallet"]["recharge_balance"],
@@ -1063,6 +1064,10 @@ async fn gateway_handles_admin_system_users_export_locally_with_trusted_admin_pr
payload["users"][0]["api_keys"][0]["is_standalone"],
json!(false)
);
assert_eq!(
payload["users"][0]["api_keys"][0]["api_key_id"],
json!("key-user-1")
);
assert_eq!(
payload["users"][0]["api_keys"][0]["total_tokens"],
json!(420)
@@ -1071,12 +1076,22 @@ async fn gateway_handles_admin_system_users_export_locally_with_trusted_admin_pr
payload["standalone_keys"][0]["key"],
json!("ak-standalone-live-1")
);
assert_eq!(
payload["standalone_keys"][0]["api_key_id"],
json!("key-standalone-1")
);
assert_eq!(payload["standalone_keys"][0]["total_tokens"], json!(84));
assert_eq!(
payload["standalone_keys"][0]["wallet"]["unlimited"],
json!(true)
);
assert_eq!(payload["standalone_keys"][0].get("is_standalone"), None,);
assert_eq!(payload["usage_aggregates"]["stats_daily"], json!([]));
assert_eq!(payload["usage_aggregates"]["stats_user_daily"], json!([]));
assert_eq!(
payload["usage_aggregates"]["stats_daily_api_key"],
json!([])
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
@@ -4,7 +4,9 @@ use aether_contracts::ExecutionPlan;
use aether_crypto::{
decrypt_python_fernet_ciphertext, encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY,
};
use aether_data::repository::auth::InMemoryAuthApiKeySnapshotRepository;
use aether_data::repository::auth::{
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeyExportRecord, StoredAuthApiKeySnapshot,
};
use aether_data::repository::auth_modules::{
AuthModuleReadRepository, InMemoryAuthModuleReadRepository, StoredOAuthProviderModuleConfig,
};
@@ -27,7 +29,7 @@ use axum::{extract::Request, Json, Router};
use http::StatusCode;
use serde_json::{json, Value};
use super::super::helpers::{sample_endpoint, sample_key, sample_provider};
use super::super::helpers::{hash_api_key, sample_endpoint, sample_key, sample_provider};
use super::super::{
build_router_with_state, build_state_with_execution_runtime_override, start_server, AppState,
};
@@ -122,6 +124,14 @@ fn sample_system_import_payload() -> Value {
"name": "primary",
"api_formats": ["openai:chat"],
"auth_type": "api_key",
"auth_type_by_format": {
"openai:chat": "api_key",
"openai:video": "bearer"
},
"allow_auth_channel_mismatch_formats": [
"openai:chat",
"openai:video"
],
"api_key": "sk-import-123",
"internal_priority": 5,
"is_active": true
@@ -371,6 +381,15 @@ async fn gateway_imports_admin_system_config_locally_and_persists_data() {
.expect("api key should decrypt"),
"sk-import-123"
);
assert_eq!(keys[0].api_formats, Some(json!(["openai:chat"])));
assert_eq!(
keys[0].auth_type_by_format,
Some(json!({ "openai:chat": "api_key" }))
);
assert_eq!(
keys[0].allow_auth_channel_mismatch_formats,
Some(json!(["openai:chat"]))
);
let provider_models = global_model_repository
.list_admin_provider_models(&AdminProviderModelListQuery {
@@ -1052,6 +1071,7 @@ async fn gateway_imports_admin_system_users_locally_and_persists_data() {
.expect("user api keys should load");
assert_eq!(user_api_keys.len(), 1);
assert_eq!(user_api_keys[0].name.as_deref(), Some("Alice CLI"));
assert_eq!(user_api_keys[0].total_requests, 12);
assert_eq!(user_api_keys[0].total_tokens, 3456);
assert_eq!(user_api_keys[0].total_cost_usd, 1.25);
assert_eq!(
@@ -1079,6 +1099,7 @@ async fn gateway_imports_admin_system_users_locally_and_persists_data() {
standalone_keys[0].name.as_deref(),
Some("Imported Standalone")
);
assert_eq!(standalone_keys[0].total_requests, 3);
assert_eq!(standalone_keys[0].total_tokens, 789);
assert_eq!(standalone_keys[0].total_cost_usd, 0.75);
assert_eq!(
@@ -1119,6 +1140,211 @@ async fn gateway_imports_admin_system_users_locally_and_persists_data() {
let _ = upstream_url;
}
#[tokio::test]
async fn gateway_overwrites_existing_admin_system_user_key_usage_totals() {
let user_key_hash = hash_api_key("sk-existing-user-key");
let standalone_key_hash = hash_api_key("sk-existing-standalone-key");
let existing_user = StoredUserAuthRecord::new(
"user-existing".to_string(),
Some("existing@example.com".to_string()),
true,
"existing".to_string(),
Some("existing-hash".to_string()),
"user".to_string(),
"local".to_string(),
None,
None,
None,
true,
false,
Some(chrono::Utc::now()),
Some(chrono::Utc::now()),
)
.expect("existing user should build");
let user_key_snapshot = StoredAuthApiKeySnapshot::new(
"user-existing".to_string(),
"existing".to_string(),
Some("existing@example.com".to_string()),
"user".to_string(),
"local".to_string(),
true,
false,
None,
None,
None,
"key-user-existing".to_string(),
Some("Existing User Key".to_string()),
true,
false,
false,
Some(10),
None,
None,
None,
None,
None,
)
.expect("user key snapshot should build");
let standalone_key_snapshot = StoredAuthApiKeySnapshot::new(
"admin-user-123".to_string(),
"admin".to_string(),
Some("admin@example.com".to_string()),
"admin".to_string(),
"local".to_string(),
true,
false,
None,
None,
None,
"key-standalone-existing".to_string(),
Some("Existing Standalone Key".to_string()),
true,
false,
true,
Some(20),
None,
None,
None,
None,
None,
)
.expect("standalone key snapshot should build");
let auth_repository = Arc::new(
InMemoryAuthApiKeySnapshotRepository::seed(vec![
(Some(user_key_hash.clone()), user_key_snapshot),
(Some(standalone_key_hash.clone()), standalone_key_snapshot),
])
.with_export_records(vec![
StoredAuthApiKeyExportRecord::new(
"user-existing".to_string(),
"key-user-existing".to_string(),
user_key_hash.clone(),
None,
Some("Existing User Key".to_string()),
None,
None,
None,
Some(10),
None,
None,
true,
None,
false,
1,
2,
0.03,
false,
)
.expect("existing user key export should build"),
StoredAuthApiKeyExportRecord::new(
"admin-user-123".to_string(),
"key-standalone-existing".to_string(),
standalone_key_hash.clone(),
None,
Some("Existing Standalone Key".to_string()),
None,
None,
None,
Some(20),
None,
None,
true,
None,
false,
4,
5,
0.06,
true,
)
.expect("existing standalone key export should build"),
]),
);
let user_repository =
Arc::new(aether_data::repository::users::InMemoryUserReadRepository::default());
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_auth_api_key_repository_for_tests(Arc::clone(&auth_repository))
.with_user_reader(user_repository)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_auth_users_for_tests([sample_import_admin_user("admin-user-123"), existing_user])
.with_auth_wallets_for_tests(Vec::<StoredWalletSnapshot>::new());
let gateway = build_router_with_state(state.clone());
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/admin/system/users/import"))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"version": "1.4",
"merge_mode": "overwrite",
"users": [{
"id": "source-user-existing",
"email": "existing@example.com",
"username": "existing",
"password_hash": "existing-hash",
"role": "user",
"is_active": true,
"api_keys": [{
"api_key_id": "source-user-key",
"key_hash": user_key_hash,
"name": "Imported User Key",
"is_active": true,
"total_requests": 222,
"total_tokens": 3333,
"total_cost_usd": 4.56
}]
}],
"standalone_keys": [{
"api_key_id": "source-standalone-key",
"key_hash": standalone_key_hash,
"name": "Imported Standalone Key",
"is_active": true,
"total_requests": 444,
"total_tokens": 5555,
"total_cost_usd": 6.78
}]
}))
.send()
.await
.expect("request should succeed");
let status = response.status();
let payload: Value = response.json().await.expect("json body should parse");
assert_eq!(status, StatusCode::OK, "payload={payload}");
assert_eq!(payload["stats"]["users"]["updated"], json!(1));
assert_eq!(payload["stats"]["api_keys"]["updated"], json!(1));
assert_eq!(payload["stats"]["standalone_keys"]["updated"], json!(1));
let updated_records = state
.list_auth_api_key_export_records_by_ids(&[
"key-user-existing".to_string(),
"key-standalone-existing".to_string(),
])
.await
.expect("api key export records should load");
let user_key = updated_records
.iter()
.find(|record| record.api_key_id == "key-user-existing")
.expect("updated user key should exist");
assert_eq!(user_key.total_requests, 222);
assert_eq!(user_key.total_tokens, 3333);
assert_eq!(user_key.total_cost_usd, 4.56);
let standalone_key = updated_records
.iter()
.find(|record| record.api_key_id == "key-standalone-existing")
.expect("updated standalone key should exist");
assert_eq!(standalone_key.total_requests, 444);
assert_eq!(standalone_key.total_tokens, 5555);
assert_eq!(standalone_key.total_cost_usd, 6.78);
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_imports_admin_system_config_fixture_v22() {
let gateway = build_router_with_state(
@@ -1671,33 +1897,20 @@ async fn gateway_overwrites_oauth_provider_key_credentials_from_admin_system_imp
}
#[tokio::test]
async fn gateway_overwrites_oauth_provider_key_credentials_from_admin_system_import_and_forces_refresh(
async fn gateway_overwrites_oauth_provider_key_credentials_from_admin_system_import_without_refresh(
) {
#[derive(Debug, Clone)]
struct SeenRefreshRequest {
content_type: String,
body: String,
}
let seen_refresh = Arc::new(Mutex::new(None::<SeenRefreshRequest>));
let seen_refresh = Arc::new(Mutex::new(false));
let seen_refresh_clone = Arc::clone(&seen_refresh);
let refresh_hits = Arc::new(Mutex::new(0usize));
let refresh_hits_clone = Arc::clone(&refresh_hits);
let refresh_server = Router::new().route(
"/oauth/token",
post(move |headers: HeaderMap, body: Bytes| {
post(move |_headers: HeaderMap, _body: Bytes| {
let seen_refresh_inner = Arc::clone(&seen_refresh_clone);
let refresh_hits_inner = Arc::clone(&refresh_hits_clone);
async move {
*refresh_hits_inner.lock().expect("mutex should lock") += 1;
*seen_refresh_inner.lock().expect("mutex should lock") = Some(SeenRefreshRequest {
content_type: headers
.get(http::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
body: String::from_utf8(body.to_vec()).unwrap_or_default(),
});
*seen_refresh_inner.lock().expect("mutex should lock") = true;
axum::Json(json!({
"access_token": "oauth-access-token-refreshed",
"refresh_token": "oauth-refresh-token-refreshed",
@@ -1792,21 +2005,8 @@ async fn gateway_overwrites_oauth_provider_key_credentials_from_admin_system_imp
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(*refresh_hits.lock().expect("mutex should lock"), 1);
let seen_refresh = seen_refresh
.lock()
.expect("mutex should lock")
.clone()
.expect("refresh request should be captured");
assert_eq!(
seen_refresh.content_type,
"application/x-www-form-urlencoded"
);
assert!(seen_refresh.body.contains("grant_type=refresh_token"));
assert!(seen_refresh
.body
.contains("refresh_token=oauth-refresh-token-new"));
assert_eq!(*refresh_hits.lock().expect("mutex should lock"), 0);
assert!(!*seen_refresh.lock().expect("mutex should lock"));
let providers = provider_catalog_repository
.list_providers(false)
@@ -1823,7 +2023,7 @@ async fn gateway_overwrites_oauth_provider_key_credentials_from_admin_system_imp
assert_eq!(key.name, "oauth-primary");
assert_eq!(key.oauth_invalid_at_unix_secs, None);
assert_eq!(key.oauth_invalid_reason, None);
assert!(key.expires_at_unix_secs.is_some());
assert_eq!(key.expires_at_unix_secs, None);
assert_eq!(
decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
@@ -1832,7 +2032,7 @@ async fn gateway_overwrites_oauth_provider_key_credentials_from_admin_system_imp
.expect("api key should be present"),
)
.expect("oauth access token should decrypt"),
"oauth-access-token-refreshed"
"oauth-access-token-new"
);
let auth_config = decrypt_python_fernet_ciphertext(
@@ -1845,15 +2045,12 @@ async fn gateway_overwrites_oauth_provider_key_credentials_from_admin_system_imp
let auth_config: Value =
serde_json::from_str(&auth_config).expect("oauth auth config json should parse");
assert_eq!(auth_config["provider_type"], "codex");
assert_eq!(
auth_config["refresh_token"],
"oauth-refresh-token-refreshed"
);
assert_eq!(auth_config["refresh_token"], "oauth-refresh-token-new");
assert_eq!(auth_config["email"], "alice@example.com");
assert_eq!(auth_config["account_id"], "acct-codex-123");
assert_eq!(auth_config["plan_type"], "plus");
assert_eq!(auth_config["token_type"], "Bearer");
assert_eq!(auth_config["expires_at"].as_u64(), key.expires_at_unix_secs);
assert!(auth_config.get("token_type").is_none());
assert!(auth_config.get("expires_at").is_none());
gateway_handle.abort();
refresh_handle.abort();
+23 -2
View File
@@ -157,18 +157,39 @@ fn admin_pool_health_score(key: &StoredProviderCatalogKey) -> f64 {
}
fn admin_pool_circuit_breaker_open(key: &StoredProviderCatalogKey) -> bool {
let now_unix_secs = Utc::now().timestamp().max(0) as u64;
key.circuit_breaker_by_format
.as_ref()
.and_then(Value::as_object)
.map(|formats| {
formats
.values()
.filter_map(Value::as_object)
.any(|item| item.get("open").and_then(Value::as_bool).unwrap_or(false))
.any(|item| admin_pool_circuit_payload_active_open_at(item, now_unix_secs))
})
.unwrap_or(false)
}
fn admin_pool_circuit_payload_active_open_at(value: &Value, now_unix_secs: u64) -> bool {
let Some(item) = value.as_object() else {
return false;
};
if !item.get("open").and_then(Value::as_bool).unwrap_or(false) {
return false;
}
if let Some(next_probe_at) = item.get("next_probe_at_unix_secs").and_then(Value::as_u64) {
return now_unix_secs < next_probe_at;
}
if let Some(next_probe_at) = item
.get("next_probe_at")
.and_then(Value::as_str)
.and_then(|value| chrono::DateTime::parse_from_rfc3339(value).ok())
.and_then(|value| u64::try_from(value.timestamp()).ok())
{
return now_unix_secs < next_probe_at;
}
true
}
fn unix_secs_to_rfc3339(unix_secs: u64) -> Option<String> {
Utc.timestamp_opt(unix_secs as i64, 0)
.single()
@@ -206,6 +206,21 @@ pub fn build_chatgpt_web_image_request_body(
if let Some(user) = request.user.as_ref() {
body.insert("user".to_string(), Value::String(user.clone()));
}
if let Some(quality) = request
.tool
.get("quality")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
body.insert("quality".to_string(), Value::String(quality.to_string()));
}
if let Some(partial_images) = request.tool.get("partial_images").and_then(Value::as_u64) {
body.insert(
"partial_images".to_string(),
Value::Number(Number::from(partial_images)),
);
}
if let Some(output_format) = request
.summary_json
.get("output_format")
@@ -1593,6 +1608,28 @@ mod tests {
assert_eq!(by_size["size"], "1024x1024");
}
#[test]
fn chatgpt_web_preserves_quality_and_partial_images() {
let parts = request_parts("/v1/images/generations", Some("application/json"));
let body = build_chatgpt_web_image_request_body(
&parts,
&json!({
"model": "gpt-image-2",
"prompt": "draw",
"size": "1024x1024",
"quality": "high",
"partial_images": 2,
"output_format": "png"
}),
None,
)
.expect("request should pass");
assert_eq!(body["quality"], "high");
assert_eq!(body["partial_images"], 2);
assert_eq!(body["output_format"], "png");
}
#[test]
fn chatgpt_web_rejects_oversized_resolution_or_size() {
let parts = request_parts("/v1/images/generations", Some("application/json"));
+85 -1
View File
@@ -8,7 +8,9 @@ use crate::maintenance::{
WalletDailyUsageAggregationInput, WalletDailyUsageAggregationResult,
};
use crate::repository::system::{
AdminSystemPurgeSummary, AdminSystemPurgeTarget, AdminSystemStats, StoredSystemConfigEntry,
AdminSystemPurgeSummary, AdminSystemPurgeTarget, AdminSystemStats,
AdminSystemUsageAggregateImportMode, AdminSystemUsageAggregateImportSummary,
AdminSystemUsageAggregateSnapshot, StoredSystemConfigEntry,
};
use crate::DataLayerError;
use sqlx::migrate::MigrateError;
@@ -195,6 +197,37 @@ impl DataBackends {
}
}
pub async fn export_admin_system_usage_aggregates(
&self,
) -> Result<AdminSystemUsageAggregateSnapshot, DataLayerError> {
match self.sql_backend() {
Some(backend) => backend.export_admin_system_usage_aggregates().await,
None => Ok(AdminSystemUsageAggregateSnapshot::default()),
}
}
pub async fn import_admin_system_usage_aggregates(
&self,
snapshot: &AdminSystemUsageAggregateSnapshot,
user_id_map: &std::collections::BTreeMap<String, String>,
api_key_id_map: &std::collections::BTreeMap<String, String>,
mode: AdminSystemUsageAggregateImportMode,
) -> Result<AdminSystemUsageAggregateImportSummary, DataLayerError> {
match self.sql_backend() {
Some(backend) => {
backend
.import_admin_system_usage_aggregates(
snapshot,
user_id_map,
api_key_id_map,
mode,
)
.await
}
None => Ok(AdminSystemUsageAggregateImportSummary::default()),
}
}
pub async fn purge_admin_request_bodies_batch(
&self,
batch_size: usize,
@@ -505,6 +538,57 @@ impl<'a> SqlBackendRef<'a> {
}
}
async fn export_admin_system_usage_aggregates(
self,
) -> Result<AdminSystemUsageAggregateSnapshot, DataLayerError> {
match self {
Self::Postgres(postgres) => postgres.export_admin_system_usage_aggregates().await,
Self::Mysql(mysql) => mysql.export_admin_system_usage_aggregates().await,
Self::Sqlite(sqlite) => sqlite.export_admin_system_usage_aggregates().await,
}
}
async fn import_admin_system_usage_aggregates(
self,
snapshot: &AdminSystemUsageAggregateSnapshot,
user_id_map: &std::collections::BTreeMap<String, String>,
api_key_id_map: &std::collections::BTreeMap<String, String>,
mode: AdminSystemUsageAggregateImportMode,
) -> Result<AdminSystemUsageAggregateImportSummary, DataLayerError> {
match self {
Self::Postgres(postgres) => {
postgres
.import_admin_system_usage_aggregates(
snapshot,
user_id_map,
api_key_id_map,
mode,
)
.await
}
Self::Mysql(mysql) => {
mysql
.import_admin_system_usage_aggregates(
snapshot,
user_id_map,
api_key_id_map,
mode,
)
.await
}
Self::Sqlite(sqlite) => {
sqlite
.import_admin_system_usage_aggregates(
snapshot,
user_id_map,
api_key_id_map,
mode,
)
.await
}
}
}
async fn purge_admin_request_bodies_batch(
self,
batch_size: usize,
+116 -1
View File
@@ -258,9 +258,15 @@ impl SqliteBackend {
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use super::SqliteBackend;
use crate::lifecycle::migrate::run_sqlite_migrations;
use crate::repository::system::AdminSystemPurgeTarget;
use crate::repository::system::{
AdminSystemPurgeTarget, AdminSystemStatsDailyAggregate,
AdminSystemStatsDailyApiKeyAggregate, AdminSystemStatsUserDailyAggregate,
AdminSystemUsageAggregateImportMode, AdminSystemUsageAggregateSnapshot,
};
use crate::{
DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig, StatsDailyAggregationInput,
StatsHourlyAggregationInput, WalletDailyUsageAggregationInput,
@@ -456,6 +462,115 @@ VALUES
assert_eq!(admin_exists, 1);
}
#[tokio::test]
async fn admin_system_usage_aggregates_round_trip_after_sqlite_migrations() {
let config = SqlDatabaseConfig {
driver: DatabaseDriver::Sqlite,
url: "sqlite::memory:".to_string(),
pool: SqlPoolConfig {
max_connections: 1,
..SqlPoolConfig::default()
},
};
let backend = SqliteBackend::from_config(config).expect("backend should build");
run_sqlite_migrations(backend.pool())
.await
.expect("sqlite migrations should run");
sqlx::query(
r#"
INSERT INTO users (id, email, username, role, created_at, updated_at)
VALUES ('target-user-1', 'target@example.com', 'target', 'user', 1, 1)
"#,
)
.execute(backend.pool())
.await
.expect("target user should insert");
sqlx::query(
r#"
INSERT INTO api_keys (id, user_id, key_hash, name, created_at, updated_at)
VALUES ('target-key-1', 'target-user-1', 'hash-target-key', 'target key', 1, 1)
"#,
)
.execute(backend.pool())
.await
.expect("target key should insert");
let snapshot = AdminSystemUsageAggregateSnapshot {
stats_daily: vec![AdminSystemStatsDailyAggregate {
date_unix_secs: 86_400,
total_requests: 9,
success_requests: 8,
error_requests: 1,
input_tokens: 100,
output_tokens: 200,
cache_creation_tokens: 3,
cache_read_tokens: 4,
total_cost: 1.25,
actual_total_cost: 1.0,
is_complete: true,
aggregated_at_unix_secs: Some(90_000),
}],
stats_user_daily: vec![AdminSystemStatsUserDailyAggregate {
user_id: "source-user-1".to_string(),
username: Some("source".to_string()),
date_unix_secs: 86_400,
total_requests: 5,
success_requests: 5,
error_requests: 0,
input_tokens: 50,
output_tokens: 60,
cache_creation_tokens: 1,
cache_read_tokens: 2,
total_cost: 0.5,
}],
stats_daily_api_key: vec![AdminSystemStatsDailyApiKeyAggregate {
api_key_id: "source-key-1".to_string(),
api_key_name: Some("source key".to_string()),
date_unix_secs: 86_400,
total_requests: 4,
success_requests: 3,
error_requests: 1,
input_tokens: 40,
output_tokens: 30,
cache_creation_tokens: 2,
cache_read_tokens: 1,
total_cost: 0.75,
}],
};
let user_id_map =
BTreeMap::from([("source-user-1".to_string(), "target-user-1".to_string())]);
let api_key_id_map =
BTreeMap::from([("source-key-1".to_string(), "target-key-1".to_string())]);
let summary = backend
.import_admin_system_usage_aggregates(
&snapshot,
&user_id_map,
&api_key_id_map,
AdminSystemUsageAggregateImportMode::Overwrite,
)
.await
.expect("usage aggregates should import");
assert_eq!(summary.stats_daily.created, 1);
assert_eq!(summary.stats_user_daily.created, 1);
assert_eq!(summary.stats_daily_api_key.created, 1);
let exported = backend
.export_admin_system_usage_aggregates()
.await
.expect("usage aggregates should export");
assert_eq!(exported.stats_daily.len(), 1);
assert_eq!(exported.stats_daily[0].total_requests, 9);
assert_eq!(exported.stats_daily[0].actual_total_cost, 1.0);
assert_eq!(exported.stats_user_daily.len(), 1);
assert_eq!(exported.stats_user_daily[0].user_id, "target-user-1");
assert_eq!(exported.stats_user_daily[0].total_requests, 5);
assert_eq!(exported.stats_daily_api_key.len(), 1);
assert_eq!(exported.stats_daily_api_key[0].api_key_id, "target-key-1");
assert_eq!(exported.stats_daily_api_key[0].total_requests, 4);
}
#[tokio::test]
async fn admin_system_request_bodies_purge_clears_inline_usage_body_fields() {
let config = SqlDatabaseConfig {
File diff suppressed because it is too large Load Diff
@@ -1017,6 +1017,26 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
Ok(Some(export.clone()))
}
async fn set_api_key_usage_totals(
&self,
api_key_id: &str,
total_requests: u64,
total_tokens: u64,
total_cost_usd: f64,
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
let mut index = self
.index
.write()
.expect("auth api key snapshot repository lock");
let Some(export) = index.export_by_api_key_id.get_mut(api_key_id) else {
return Ok(None);
};
export.total_requests = total_requests;
export.total_tokens = total_tokens;
export.total_cost_usd = total_cost_usd;
Ok(Some(export.clone()))
}
async fn delete_user_api_key(
&self,
user_id: &str,
@@ -692,6 +692,34 @@ WHERE id = ?
self.reload_export_by_id(api_key_id).await
}
async fn set_api_key_usage_totals(
&self,
api_key_id: &str,
total_requests: u64,
total_tokens: u64,
total_cost_usd: f64,
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
sqlx::query(
r#"
UPDATE api_keys
SET total_requests = ?,
total_tokens = ?,
total_cost_usd = ?,
updated_at = ?
WHERE id = ?
"#,
)
.bind(total_requests as i64)
.bind(total_tokens as i64)
.bind(total_cost_usd)
.bind(current_unix_secs() as i64)
.bind(api_key_id)
.execute(&self.pool)
.await
.map_sql_err()?;
self.reload_export_by_id(api_key_id).await
}
async fn delete_user_api_key(
&self,
user_id: &str,
@@ -668,6 +668,40 @@ RETURNING
is_standalone
"#;
const SET_API_KEY_USAGE_TOTALS_SQL: &str = r#"
UPDATE api_keys
SET
total_requests = $2,
total_tokens = $3,
total_cost_usd = $4,
updated_at = NOW()
WHERE id = $1
RETURNING
user_id,
id AS api_key_id,
key_hash,
key_encrypted,
name,
allowed_providers,
allowed_api_formats,
allowed_models,
ip_rules,
rate_limit,
concurrent_limit,
force_capabilities,
feature_settings,
is_active,
CAST(EXTRACT(EPOCH FROM expires_at) AS BIGINT) AS expires_at_unix_secs,
auto_delete_on_expiry,
total_requests,
COALESCE(total_tokens, 0)::BIGINT AS total_tokens,
COALESCE(CAST(total_cost_usd AS DOUBLE PRECISION), 0) AS total_cost_usd,
CAST(EXTRACT(EPOCH FROM last_used_at) AS BIGINT) AS last_used_at_unix_secs,
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs,
is_standalone
"#;
const SET_USER_API_KEY_LOCKED_SQL: &str = r#"
UPDATE api_keys
SET
@@ -1436,6 +1470,24 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
.find(|record| record.user_id == user_id && !record.is_standalone))
}
async fn set_api_key_usage_totals(
&self,
api_key_id: &str,
total_requests: u64,
total_tokens: u64,
total_cost_usd: f64,
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
let row = sqlx::query(SET_API_KEY_USAGE_TOTALS_SQL)
.bind(api_key_id)
.bind(total_requests as i64)
.bind(total_tokens as i64)
.bind(total_cost_usd)
.fetch_optional(&self.pool)
.await
.map_postgres_err()?;
row.as_ref().map(map_auth_api_key_export_row).transpose()
}
async fn delete_user_api_key(
&self,
user_id: &str,
@@ -692,6 +692,34 @@ WHERE id = ?
self.reload_export_by_id(api_key_id).await
}
async fn set_api_key_usage_totals(
&self,
api_key_id: &str,
total_requests: u64,
total_tokens: u64,
total_cost_usd: f64,
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
sqlx::query(
r#"
UPDATE api_keys
SET total_requests = ?,
total_tokens = ?,
total_cost_usd = ?,
updated_at = ?
WHERE id = ?
"#,
)
.bind(total_requests as i64)
.bind(total_tokens as i64)
.bind(total_cost_usd)
.bind(current_unix_secs() as i64)
.bind(api_key_id)
.execute(&self.pool)
.await
.map_sql_err()?;
self.reload_export_by_id(api_key_id).await
}
async fn delete_user_api_key(
&self,
user_id: &str,
@@ -691,6 +691,14 @@ pub trait AuthApiKeyWriteRepository: Send + Sync {
feature_settings: Option<serde_json::Value>,
) -> Result<Option<StoredAuthApiKeyExportRecord>, crate::DataLayerError>;
async fn set_api_key_usage_totals(
&self,
api_key_id: &str,
total_requests: u64,
total_tokens: u64,
total_cost_usd: f64,
) -> Result<Option<StoredAuthApiKeyExportRecord>, crate::DataLayerError>;
async fn delete_user_api_key(
&self,
user_id: &str,
@@ -21,6 +21,85 @@ pub struct AdminSystemStats {
pub total_requests: u64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AdminSystemUsageAggregateImportMode {
Skip,
Overwrite,
Error,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct AdminSystemStatsDailyAggregate {
pub date_unix_secs: u64,
pub total_requests: u64,
pub success_requests: u64,
pub error_requests: u64,
pub input_tokens: u64,
pub output_tokens: u64,
pub cache_creation_tokens: u64,
pub cache_read_tokens: u64,
pub total_cost: f64,
pub actual_total_cost: f64,
pub is_complete: bool,
pub aggregated_at_unix_secs: Option<u64>,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct AdminSystemStatsUserDailyAggregate {
pub user_id: String,
pub username: Option<String>,
pub date_unix_secs: u64,
pub total_requests: u64,
pub success_requests: u64,
pub error_requests: u64,
pub input_tokens: u64,
pub output_tokens: u64,
pub cache_creation_tokens: u64,
pub cache_read_tokens: u64,
pub total_cost: f64,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct AdminSystemStatsDailyApiKeyAggregate {
pub api_key_id: String,
pub api_key_name: Option<String>,
pub date_unix_secs: u64,
pub total_requests: u64,
pub success_requests: u64,
pub error_requests: u64,
pub input_tokens: u64,
pub output_tokens: u64,
pub cache_creation_tokens: u64,
pub cache_read_tokens: u64,
pub total_cost: f64,
}
#[derive(Debug, Clone, Default, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct AdminSystemUsageAggregateSnapshot {
#[serde(default)]
pub stats_daily: Vec<AdminSystemStatsDailyAggregate>,
#[serde(default)]
pub stats_user_daily: Vec<AdminSystemStatsUserDailyAggregate>,
#[serde(default)]
pub stats_daily_api_key: Vec<AdminSystemStatsDailyApiKeyAggregate>,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct AdminSystemUsageAggregateImportCounter {
pub created: u64,
pub updated: u64,
pub skipped: u64,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct AdminSystemUsageAggregateImportSummary {
pub stats_daily: AdminSystemUsageAggregateImportCounter,
pub stats_user_daily: AdminSystemUsageAggregateImportCounter,
pub stats_daily_api_key: AdminSystemUsageAggregateImportCounter,
pub skipped_unmapped_user_daily: u64,
pub skipped_unmapped_api_key_daily: u64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AdminSystemPurgeTarget {
Config,
+42 -17
View File
@@ -108,8 +108,7 @@ macro_rules! impl_materialized_usage_read_repository {
user_ids: &[String],
) -> Result<Vec<$crate::repository::usage::StoredUsageUserTotals>, $crate::DataLayerError>
{
let repository = self.materialize_read_model().await?;
<$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::summarize_usage_totals_by_user_ids(&repository, user_ids).await
<$repository>::summarize_usage_totals_by_user_ids(self, user_ids).await
}
async fn summarize_usage_cache_hit_summary(
@@ -163,8 +162,7 @@ macro_rules! impl_materialized_usage_read_repository {
$crate::repository::usage::StoredUsageDashboardSummary,
$crate::DataLayerError,
> {
let repository = self.materialize_read_model().await?;
<$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::summarize_dashboard_usage(&repository, query).await
<$repository>::summarize_dashboard_usage(self, query).await
}
async fn list_dashboard_daily_breakdown(
@@ -174,8 +172,7 @@ macro_rules! impl_materialized_usage_read_repository {
Vec<$crate::repository::usage::StoredUsageDashboardDailyBreakdownRow>,
$crate::DataLayerError,
> {
let repository = self.materialize_read_model().await?;
<$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::list_dashboard_daily_breakdown(&repository, query).await
<$repository>::list_dashboard_daily_breakdown(self, query).await
}
async fn summarize_dashboard_provider_counts(
@@ -349,8 +346,7 @@ macro_rules! impl_materialized_usage_read_repository {
Vec<$crate::repository::usage::StoredUsageDailySummary>,
$crate::DataLayerError,
> {
let repository = self.materialize_read_model().await?;
<$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::summarize_usage_daily_heatmap(&repository, query).await
<$repository>::summarize_usage_daily_heatmap(self, query).await
}
}
};
@@ -640,15 +636,13 @@ 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()
.map(str::trim)
.filter(|value| !value.is_empty())?
.to_string();
let is_in_flight = matches!(usage.status.as_str(), "pending" | "streaming");
let is_success = provider_api_key_usage_is_success(
usage.status.as_str(),
usage.status_code,
@@ -665,8 +659,14 @@ pub(crate) fn provider_api_key_usage_contribution(
request_count: 1,
success_count: i64::from(is_success),
error_count: i64::from(is_error),
total_tokens: i64::try_from(usage.total_tokens).unwrap_or(i64::MAX),
total_cost_usd: if usage.total_cost_usd.is_finite() {
total_tokens: if is_in_flight {
0
} else {
i64::try_from(usage.total_tokens).unwrap_or(i64::MAX)
},
total_cost_usd: if is_in_flight {
0.0
} else if usage.total_cost_usd.is_finite() {
usage.total_cost_usd.max(0.0)
} else {
0.0
@@ -1029,7 +1029,7 @@ mod tests {
}
#[test]
fn provider_api_key_usage_contribution_tracks_terminal_requests_only() {
fn provider_api_key_usage_contribution_counts_in_flight_requests_once() {
let usage = StoredRequestUsageAudit::new(
"usage-1".to_string(),
"request-1".to_string(),
@@ -1074,11 +1074,36 @@ mod tests {
let mut streaming = usage.clone();
streaming.status = "streaming".to_string();
assert!(provider_api_key_usage_contribution(&streaming).is_none());
let streaming_contribution =
provider_api_key_usage_contribution(&streaming).expect("streaming should count");
assert_eq!(streaming_contribution.request_count, 1);
assert_eq!(streaming_contribution.success_count, 0);
assert_eq!(streaming_contribution.error_count, 0);
assert_eq!(streaming_contribution.total_tokens, 0);
assert_eq!(streaming_contribution.total_cost_usd, 0.0);
assert_eq!(streaming_contribution.total_response_time_ms, 0);
let mut pending = usage;
let mut pending = usage.clone();
pending.status = "pending".to_string();
assert!(provider_api_key_usage_contribution(&pending).is_none());
let pending_contribution =
provider_api_key_usage_contribution(&pending).expect("pending should count");
assert_eq!(pending_contribution.request_count, 1);
assert_eq!(pending_contribution.success_count, 0);
assert_eq!(pending_contribution.error_count, 0);
assert_eq!(pending_contribution.total_tokens, 0);
assert_eq!(pending_contribution.total_cost_usd, 0.0);
assert_eq!(pending_contribution.total_response_time_ms, 0);
let terminal_contribution =
provider_api_key_usage_contribution(&usage).expect("terminal should count");
let delta =
ProviderApiKeyUsageDelta::between(&pending_contribution, &terminal_contribution);
assert_eq!(delta.request_count, 0);
assert_eq!(delta.success_count, 1);
assert_eq!(delta.error_count, 0);
assert_eq!(delta.total_tokens, 20);
assert_eq!(delta.total_cost_usd, 0.25);
assert_eq!(delta.total_response_time_ms, 120);
}
#[test]
+531 -10
View File
@@ -1,14 +1,18 @@
use std::collections::{BTreeMap, HashSet};
use std::time::{SystemTime, UNIX_EPOCH};
use aether_ai_formats::UPSTREAM_IS_STREAM_KEY;
use async_trait::async_trait;
use sqlx::{mysql::MySqlRow, Row};
use sqlx::{mysql::MySqlRow, MySql, QueryBuilder, 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,
usage_request_metadata_client_family, InMemoryUsageReadRepository, PendingUsageCleanupSummary,
StoredRequestUsageAudit, UpsertUsageRecord, UsageWriteRepository,
StoredRequestUsageAudit, StoredUsageDailySummary, StoredUsageDashboardDailyBreakdownRow,
StoredUsageDashboardSummary, StoredUsageUserTotals, UpsertUsageRecord, UsageDailyHeatmapQuery,
UsageDashboardDailyBreakdownQuery, UsageDashboardSummaryQuery, UsageReadRepository,
UsageWriteRepository,
};
use crate::driver::mysql::MysqlPool;
use crate::error::SqlResultExt;
@@ -229,6 +233,467 @@ impl MysqlUsageReadRepository {
.collect::<Result<Vec<_>, _>>()?;
Ok(InMemoryUsageReadRepository::seed(items))
}
async fn summarize_usage_daily_heatmap_raw_from_range(
&self,
created_from_unix_secs: u64,
created_until_unix_secs: u64,
user_id: Option<&str>,
) -> Result<Vec<StoredUsageDailySummary>, DataLayerError> {
let mut sql = String::from(
r#"
SELECT
DATE_FORMAT(FROM_UNIXTIME(created_at_unix_ms), '%Y-%m-%d') AS date,
COUNT(*) AS requests,
COALESCE(SUM(
GREATEST(COALESCE(input_tokens, 0), 0)
+ GREATEST(COALESCE(output_tokens, 0), 0)
+ CASE
WHEN COALESCE(cache_creation_input_tokens, 0) = 0
AND (COALESCE(cache_creation_ephemeral_5m_input_tokens, 0) + COALESCE(cache_creation_ephemeral_1h_input_tokens, 0)) > 0
THEN COALESCE(cache_creation_ephemeral_5m_input_tokens, 0) + COALESCE(cache_creation_ephemeral_1h_input_tokens, 0)
ELSE GREATEST(COALESCE(cache_creation_input_tokens, 0), 0)
END
+ GREATEST(COALESCE(cache_read_input_tokens, 0), 0)
), 0) AS total_tokens,
COALESCE(SUM(COALESCE(total_cost_usd, 0)), 0) AS total_cost_usd,
COALESCE(SUM(COALESCE(actual_total_cost_usd, 0)), 0) AS actual_total_cost_usd
FROM `usage`
WHERE created_at_unix_ms >= ?
AND created_at_unix_ms < ?
AND status NOT IN ('pending', 'streaming')
AND provider_name NOT IN ('unknown', 'pending')
"#,
);
if user_id.is_some() {
sql.push_str(" AND user_id = ?\n");
}
sql.push_str("GROUP BY date ORDER BY date ASC");
let mut query = sqlx::query(&sql)
.bind(to_i64(created_from_unix_secs, "usage.created_at_unix_ms")?)
.bind(to_i64(created_until_unix_secs, "usage.created_at_unix_ms")?);
if let Some(user_id) = user_id {
query = query.bind(user_id.to_string());
}
let rows = query.fetch_all(&self.pool).await.map_sql_err()?;
rows.iter().map(map_mysql_usage_daily_summary).collect()
}
async fn summarize_usage_daily_heatmap_from_daily_aggregates(
&self,
created_from_unix_secs: u64,
created_until_unix_secs: u64,
user_id: Option<&str>,
) -> Result<Vec<StoredUsageDailySummary>, DataLayerError> {
let rows = if let Some(user_id) = user_id {
sqlx::query(
r#"
SELECT
DATE_FORMAT(FROM_UNIXTIME(`date`), '%Y-%m-%d') AS date,
total_requests AS requests,
input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens AS total_tokens,
total_cost AS total_cost_usd,
total_cost AS actual_total_cost_usd
FROM stats_user_daily
WHERE user_id = ?
AND `date` >= ?
AND `date` < ?
AND total_requests > 0
ORDER BY `date` ASC
"#,
)
.bind(user_id)
.bind(to_i64(created_from_unix_secs, "stats_user_daily.date")?)
.bind(to_i64(created_until_unix_secs, "stats_user_daily.date")?)
.fetch_all(&self.pool)
.await
.map_sql_err()?
} else {
sqlx::query(
r#"
SELECT
DATE_FORMAT(FROM_UNIXTIME(`date`), '%Y-%m-%d') AS date,
total_requests AS requests,
input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens AS total_tokens,
total_cost AS total_cost_usd,
actual_total_cost AS actual_total_cost_usd
FROM stats_daily
WHERE `date` >= ?
AND `date` < ?
AND total_requests > 0
ORDER BY `date` ASC
"#,
)
.bind(to_i64(created_from_unix_secs, "stats_daily.date")?)
.bind(to_i64(created_until_unix_secs, "stats_daily.date")?)
.fetch_all(&self.pool)
.await
.map_sql_err()?
};
rows.iter().map(map_mysql_usage_daily_summary).collect()
}
async fn summarize_usage_daily_heatmap(
&self,
query: &UsageDailyHeatmapQuery,
) -> Result<Vec<StoredUsageDailySummary>, DataLayerError> {
let created_until_unix_secs = usage_current_unix_secs().saturating_add(1);
let user_id = query.user_id.as_deref();
let mut summaries = BTreeMap::<String, StoredUsageDailySummary>::new();
for item in self
.summarize_usage_daily_heatmap_from_daily_aggregates(
query.created_from_unix_secs,
created_until_unix_secs,
user_id,
)
.await?
{
summaries.insert(item.date.clone(), item);
}
for item in self
.summarize_usage_daily_heatmap_raw_from_range(
query.created_from_unix_secs,
created_until_unix_secs,
user_id,
)
.await?
{
summaries.entry(item.date.clone()).or_insert(item);
}
Ok(summaries.into_values().collect())
}
async fn summarize_dashboard_usage_from_daily_aggregates(
&self,
query: &UsageDashboardSummaryQuery,
) -> Result<Option<StoredUsageDashboardSummary>, DataLayerError> {
let row = if let Some(user_id) = query.user_id.as_deref() {
sqlx::query(
r#"
SELECT
COALESCE(SUM(total_requests), 0) AS total_requests,
COALESCE(SUM(input_tokens), 0) AS input_tokens,
COALESCE(SUM(input_tokens), 0) AS effective_input_tokens,
COALESCE(SUM(output_tokens), 0) AS output_tokens,
COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS total_tokens,
COALESCE(SUM(cache_creation_tokens), 0) AS cache_creation_tokens,
COALESCE(SUM(cache_read_tokens), 0) AS cache_read_tokens,
COALESCE(SUM(input_tokens + cache_creation_tokens + cache_read_tokens), 0) AS total_input_context,
0.0 AS cache_creation_cost_usd,
0.0 AS cache_read_cost_usd,
COALESCE(SUM(COALESCE(total_cost, 0)), 0) AS total_cost_usd,
COALESCE(SUM(COALESCE(total_cost, 0)), 0) AS actual_total_cost_usd,
COALESCE(SUM(error_requests), 0) AS error_requests,
0.0 AS response_time_sum_ms,
0 AS response_time_samples
FROM stats_user_daily
WHERE user_id = ?
AND `date` >= ?
AND `date` < ?
"#,
)
.bind(user_id)
.bind(to_i64(
query.created_from_unix_secs,
"stats_user_daily.date",
)?)
.bind(to_i64(
query.created_until_unix_secs,
"stats_user_daily.date",
)?)
.fetch_one(&self.pool)
.await
.map_sql_err()?
} else {
sqlx::query(
r#"
SELECT
COALESCE(SUM(total_requests), 0) AS total_requests,
COALESCE(SUM(input_tokens), 0) AS input_tokens,
COALESCE(SUM(input_tokens), 0) AS effective_input_tokens,
COALESCE(SUM(output_tokens), 0) AS output_tokens,
COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS total_tokens,
COALESCE(SUM(cache_creation_tokens), 0) AS cache_creation_tokens,
COALESCE(SUM(cache_read_tokens), 0) AS cache_read_tokens,
COALESCE(SUM(input_tokens + cache_creation_tokens + cache_read_tokens), 0) AS total_input_context,
COALESCE(SUM(COALESCE(cache_creation_cost, 0)), 0) AS cache_creation_cost_usd,
COALESCE(SUM(COALESCE(cache_read_cost, 0)), 0) AS cache_read_cost_usd,
COALESCE(SUM(COALESCE(total_cost, 0)), 0) AS total_cost_usd,
COALESCE(SUM(COALESCE(actual_total_cost, 0)), 0) AS actual_total_cost_usd,
COALESCE(SUM(error_requests), 0) AS error_requests,
0.0 AS response_time_sum_ms,
0 AS response_time_samples
FROM stats_daily
WHERE `date` >= ?
AND `date` < ?
"#,
)
.bind(to_i64(query.created_from_unix_secs, "stats_daily.date")?)
.bind(to_i64(query.created_until_unix_secs, "stats_daily.date")?)
.fetch_one(&self.pool)
.await
.map_sql_err()?
};
let total_requests = row_u64(&row, "total_requests")?;
if total_requests == 0 {
return Ok(None);
}
Ok(Some(StoredUsageDashboardSummary {
total_requests,
input_tokens: row_u64(&row, "input_tokens")?,
effective_input_tokens: row_u64(&row, "effective_input_tokens")?,
output_tokens: row_u64(&row, "output_tokens")?,
total_tokens: row_u64(&row, "total_tokens")?,
cache_creation_tokens: row_u64(&row, "cache_creation_tokens")?,
cache_read_tokens: row_u64(&row, "cache_read_tokens")?,
total_input_context: row_u64(&row, "total_input_context")?,
cache_creation_cost_usd: row.try_get("cache_creation_cost_usd").map_sql_err()?,
cache_read_cost_usd: row.try_get("cache_read_cost_usd").map_sql_err()?,
total_cost_usd: row.try_get("total_cost_usd").map_sql_err()?,
actual_total_cost_usd: row.try_get("actual_total_cost_usd").map_sql_err()?,
error_requests: row_u64(&row, "error_requests")?,
response_time_sum_ms: row.try_get("response_time_sum_ms").map_sql_err()?,
response_time_samples: row_u64(&row, "response_time_samples")?,
}))
}
async fn list_dashboard_daily_breakdown_from_daily_aggregates(
&self,
query: &UsageDashboardDailyBreakdownQuery,
) -> Result<Vec<StoredUsageDashboardDailyBreakdownRow>, DataLayerError> {
let rows = if let Some(user_id) = query.user_id.as_deref() {
sqlx::query(
r#"
SELECT
DATE_FORMAT(FROM_UNIXTIME(`date`), '%Y-%m-%d') AS date,
'aggregate' AS model,
'aggregate' AS provider,
COALESCE(SUM(total_requests), 0) AS requests,
COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS total_tokens,
COALESCE(SUM(COALESCE(total_cost, 0)), 0) AS total_cost_usd,
0.0 AS response_time_sum_ms,
0 AS response_time_samples
FROM stats_user_daily
WHERE user_id = ?
AND `date` >= ?
AND `date` < ?
AND total_requests > 0
GROUP BY `date`
ORDER BY `date` ASC
"#,
)
.bind(user_id)
.bind(to_i64(
query.created_from_unix_secs,
"stats_user_daily.date",
)?)
.bind(to_i64(
query.created_until_unix_secs,
"stats_user_daily.date",
)?)
.fetch_all(&self.pool)
.await
.map_sql_err()?
} else {
sqlx::query(
r#"
SELECT
DATE_FORMAT(FROM_UNIXTIME(`date`), '%Y-%m-%d') AS date,
'aggregate' AS model,
'aggregate' AS provider,
COALESCE(SUM(total_requests), 0) AS requests,
COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS total_tokens,
COALESCE(SUM(COALESCE(total_cost, 0)), 0) AS total_cost_usd,
0.0 AS response_time_sum_ms,
0 AS response_time_samples
FROM stats_daily
WHERE `date` >= ?
AND `date` < ?
AND total_requests > 0
GROUP BY `date`
ORDER BY `date` ASC
"#,
)
.bind(to_i64(query.created_from_unix_secs, "stats_daily.date")?)
.bind(to_i64(query.created_until_unix_secs, "stats_daily.date")?)
.fetch_all(&self.pool)
.await
.map_sql_err()?
};
rows.iter()
.map(|row| {
Ok(StoredUsageDashboardDailyBreakdownRow {
date: row.try_get("date").map_sql_err()?,
model: row.try_get("model").map_sql_err()?,
provider: row.try_get("provider").map_sql_err()?,
requests: row_u64(row, "requests")?,
total_tokens: row_u64(row, "total_tokens")?,
total_cost_usd: row.try_get("total_cost_usd").map_sql_err()?,
response_time_sum_ms: row.try_get("response_time_sum_ms").map_sql_err()?,
response_time_samples: row_u64(row, "response_time_samples")?,
})
})
.collect()
}
async fn summarize_dashboard_usage(
&self,
query: &UsageDashboardSummaryQuery,
) -> Result<StoredUsageDashboardSummary, DataLayerError> {
if let Some(summary) = self
.summarize_dashboard_usage_from_daily_aggregates(query)
.await?
{
return Ok(summary);
}
let repository = self.materialize_read_model().await?;
<InMemoryUsageReadRepository as UsageReadRepository>::summarize_dashboard_usage(
&repository,
query,
)
.await
}
async fn list_dashboard_daily_breakdown(
&self,
query: &UsageDashboardDailyBreakdownQuery,
) -> Result<Vec<StoredUsageDashboardDailyBreakdownRow>, DataLayerError> {
let aggregate_rows = self
.list_dashboard_daily_breakdown_from_daily_aggregates(query)
.await?;
if !aggregate_rows.is_empty() {
return Ok(aggregate_rows);
}
let repository = self.materialize_read_model().await?;
<InMemoryUsageReadRepository as UsageReadRepository>::list_dashboard_daily_breakdown(
&repository,
query,
)
.await
}
async fn summarize_usage_totals_by_user_ids(
&self,
user_ids: &[String],
) -> Result<Vec<StoredUsageUserTotals>, DataLayerError> {
if user_ids.is_empty() {
return Ok(Vec::new());
}
let unique_user_ids = user_ids
.iter()
.cloned()
.collect::<std::collections::BTreeSet<_>>();
let mut totals = BTreeMap::<String, StoredUsageUserTotals>::new();
let mut aggregate_cutoffs = BTreeMap::<String, u64>::new();
let mut aggregate_builder = QueryBuilder::<MySql>::new(
r#"
SELECT
user_id,
COALESCE(SUM(total_requests), 0) AS request_count,
COALESCE(
SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens),
0
) AS total_tokens,
MAX(`date`) AS latest_date
FROM stats_user_daily
WHERE user_id IN (
"#,
);
{
let mut separated = aggregate_builder.separated(", ");
for user_id in &unique_user_ids {
separated.push_bind(user_id.clone());
}
}
aggregate_builder.push(") GROUP BY user_id ORDER BY user_id ASC");
let aggregate_rows = aggregate_builder
.build()
.fetch_all(&self.pool)
.await
.map_sql_err()?;
for row in aggregate_rows {
let user_id: String = row.try_get("user_id").map_sql_err()?;
let latest_date = row.try_get::<i64, _>("latest_date").map_sql_err()?.max(0) as u64;
aggregate_cutoffs.insert(user_id.clone(), latest_date.saturating_add(86_400));
totals.insert(
user_id.clone(),
StoredUsageUserTotals {
user_id,
request_count: row_u64(&row, "request_count")?,
total_tokens: row_u64(&row, "total_tokens")?,
},
);
}
let mut raw_builder = QueryBuilder::<MySql>::new(
r#"
SELECT
`usage`.user_id,
COUNT(*) AS request_count,
COALESCE(SUM(GREATEST(COALESCE(`usage`.total_tokens, 0), 0)), 0) AS total_tokens
FROM `usage`
JOIN (
"#,
);
for (index, user_id) in unique_user_ids.iter().enumerate() {
if index > 0 {
raw_builder.push(" UNION ALL ");
}
let cutoff = aggregate_cutoffs.get(user_id).copied().unwrap_or_default();
raw_builder
.push("SELECT ")
.push_bind(user_id.clone())
.push(" AS user_id, ")
.push_bind(to_i64(cutoff, "usage aggregate cutoff")?)
.push(" AS cutoff_unix_secs");
}
raw_builder.push(
r#"
) AS requested ON requested.user_id = `usage`.user_id
WHERE `usage`.created_at_unix_ms >= requested.cutoff_unix_secs
AND `usage`.status NOT IN ('pending', 'streaming')
AND `usage`.provider_name NOT IN ('unknown', 'pending')
GROUP BY `usage`.user_id
ORDER BY `usage`.user_id ASC
"#,
);
let raw_rows = raw_builder
.build()
.fetch_all(&self.pool)
.await
.map_sql_err()?;
for row in raw_rows {
let user_id: String = row.try_get("user_id").map_sql_err()?;
let entry = totals
.entry(user_id.clone())
.or_insert_with(|| StoredUsageUserTotals {
user_id,
request_count: 0,
total_tokens: 0,
});
entry.request_count = entry
.request_count
.saturating_add(row_u64(&row, "request_count")?);
entry.total_tokens = entry
.total_tokens
.saturating_add(row_u64(&row, "total_tokens")?);
}
Ok(totals.into_values().collect())
}
}
impl_materialized_usage_read_repository!(MysqlUsageReadRepository);
@@ -389,19 +854,28 @@ WHERE provider_api_key_id IS NOT NULL AND provider_api_key_id <> ''
let error_message: Option<String> = row.try_get("error_message").map_sql_err()?;
let entry = stats.entry(key_id).or_default();
entry.request_count += 1;
if provider_api_key_usage_is_success(&status, status_code_u16, error_message.as_deref())
{
let is_success = provider_api_key_usage_is_success(
&status,
status_code_u16,
error_message.as_deref(),
);
let is_in_flight = matches!(status.as_str(), "pending" | "streaming");
if is_success {
entry.success_count += 1;
}
if provider_api_key_usage_is_error(&status, status_code_u16, error_message.as_deref()) {
entry.error_count += 1;
}
entry.total_tokens += row.try_get::<i64, _>("total_tokens").map_sql_err()?;
entry.total_cost_usd += row.try_get::<f64, _>("total_cost_usd").map_sql_err()?;
entry.total_response_time_ms += row
.try_get::<Option<i64>, _>("response_time_ms")
.map_sql_err()?
.unwrap_or_default();
if !is_in_flight {
entry.total_tokens += row.try_get::<i64, _>("total_tokens").map_sql_err()?;
entry.total_cost_usd += row.try_get::<f64, _>("total_cost_usd").map_sql_err()?;
}
if is_success {
entry.total_response_time_ms += row
.try_get::<Option<i64>, _>("response_time_ms")
.map_sql_err()?
.unwrap_or_default();
}
entry.last_used_at = entry.last_used_at.max(
row.try_get::<Option<i64>, _>("updated_at_unix_secs")
.map_sql_err()?,
@@ -990,6 +1464,25 @@ fn row_u64(row: &MySqlRow, field: &str) -> Result<u64, DataLayerError> {
u64::try_from(value).map_err(|_| DataLayerError::UnexpectedValue(format!("{field} negative")))
}
fn map_mysql_usage_daily_summary(
row: &MySqlRow,
) -> Result<StoredUsageDailySummary, DataLayerError> {
Ok(StoredUsageDailySummary {
date: row.try_get("date").map_sql_err()?,
requests: row_u64(row, "requests")?,
total_tokens: row_u64(row, "total_tokens")?,
total_cost_usd: row.try_get("total_cost_usd").map_sql_err()?,
actual_total_cost_usd: row.try_get("actual_total_cost_usd").map_sql_err()?,
})
}
fn usage_current_unix_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_secs())
.unwrap_or_default()
}
#[cfg(test)]
mod tests {
use super::{MysqlUsageReadRepository, MysqlUsageWriteRepository};
@@ -1010,6 +1503,34 @@ mod tests {
let _repository = MysqlUsageWriteRepository::new(pool);
}
#[test]
fn mysql_usage_daily_heatmap_reads_imported_daily_aggregates() {
let source = include_str!("mysql.rs");
assert!(source.contains("summarize_usage_daily_heatmap_from_daily_aggregates"));
assert!(source.contains("FROM stats_daily"));
assert!(source.contains("FROM stats_user_daily"));
assert!(source.contains("summaries.entry(item.date.clone()).or_insert(item)"));
}
#[test]
fn mysql_usage_totals_by_user_ids_reads_imported_user_daily_aggregates() {
let source = include_str!("mysql.rs");
assert!(source.contains("async fn summarize_usage_totals_by_user_ids"));
assert!(source.contains("FROM stats_user_daily"));
assert!(source.contains("MAX(`date`) AS latest_date"));
assert!(source.contains("requested.cutoff_unix_secs"));
}
#[test]
fn mysql_dashboard_reads_imported_daily_aggregates() {
let source = include_str!("mysql.rs");
assert!(source.contains("summarize_dashboard_usage_from_daily_aggregates"));
assert!(source.contains("list_dashboard_daily_breakdown_from_daily_aggregates"));
assert!(source.contains("FROM stats_daily"));
assert!(source.contains("FROM stats_user_daily"));
assert!(source.contains("'aggregate' AS model"));
}
#[tokio::test]
async fn mysql_usage_write_repository_upserts_when_url_is_set() {
let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL")
@@ -1827,11 +1827,39 @@ LIMIT 1
.await
.map_postgres_err()?;
row.map(|row| {
row.try_get::<DateTime<Utc>, _>("cutoff_date")
.map_postgres_err()
})
.transpose()
if let Some(row) = row {
return row
.try_get::<DateTime<Utc>, _>("cutoff_date")
.map(Some)
.map_postgres_err();
}
let row = sqlx::query(
r#"
SELECT MAX(date) AS latest_date
FROM (
SELECT MAX(date) AS date
FROM stats_daily
WHERE total_requests > 0
OR is_complete IS TRUE
UNION ALL
SELECT MAX(date) AS date
FROM stats_user_daily
WHERE total_requests > 0
UNION ALL
SELECT MAX(date) AS date
FROM stats_daily_api_key
WHERE total_requests > 0
) AS imported_daily_aggregates
"#,
)
.fetch_one(&self.pool)
.await
.map_postgres_err()?;
let latest_date = row
.try_get::<Option<DateTime<Utc>>, _>("latest_date")
.map_postgres_err()?;
Ok(latest_date.map(|value| value + chrono::Duration::days(1)))
}
async fn read_stats_hourly_cutoff(&self) -> Result<Option<DateTime<Utc>>, DataLayerError> {
@@ -1867,9 +1895,22 @@ WHERE is_complete IS TRUE
SELECT
COALESCE(SUM(total_requests), 0)::BIGINT AS total_requests,
COALESCE(SUM(input_tokens), 0)::BIGINT AS input_tokens,
COALESCE(SUM(effective_input_tokens), 0)::BIGINT AS effective_input_tokens,
COALESCE(SUM(
CASE
WHEN effective_input_tokens = 0 AND total_input_context = 0 AND input_tokens > 0
THEN input_tokens
ELSE effective_input_tokens
END
), 0)::BIGINT AS effective_input_tokens,
COALESCE(SUM(output_tokens), 0)::BIGINT AS output_tokens,
COALESCE(SUM(effective_input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0)::BIGINT AS total_tokens,
COALESCE(SUM(
CASE
WHEN effective_input_tokens = 0 AND total_input_context = 0 AND input_tokens > 0
THEN input_tokens
ELSE effective_input_tokens
END
+ output_tokens + cache_creation_tokens + cache_read_tokens
), 0)::BIGINT AS total_tokens,
COALESCE(SUM(cache_creation_tokens), 0)::BIGINT AS cache_creation_tokens,
COALESCE(SUM(cache_read_tokens), 0)::BIGINT AS cache_read_tokens,
COALESCE(SUM(total_input_context), 0)::BIGINT AS total_input_context,
@@ -1898,9 +1939,22 @@ WHERE user_id = $1
SELECT
COALESCE(SUM(total_requests), 0)::BIGINT AS total_requests,
COALESCE(SUM(input_tokens), 0)::BIGINT AS input_tokens,
COALESCE(SUM(effective_input_tokens), 0)::BIGINT AS effective_input_tokens,
COALESCE(SUM(
CASE
WHEN effective_input_tokens = 0 AND total_input_context = 0 AND input_tokens > 0
THEN input_tokens
ELSE effective_input_tokens
END
), 0)::BIGINT AS effective_input_tokens,
COALESCE(SUM(output_tokens), 0)::BIGINT AS output_tokens,
COALESCE(SUM(effective_input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0)::BIGINT AS total_tokens,
COALESCE(SUM(
CASE
WHEN effective_input_tokens = 0 AND total_input_context = 0 AND input_tokens > 0
THEN input_tokens
ELSE effective_input_tokens
END
+ output_tokens + cache_creation_tokens + cache_read_tokens
), 0)::BIGINT AS total_tokens,
COALESCE(SUM(cache_creation_tokens), 0)::BIGINT AS cache_creation_tokens,
COALESCE(SUM(cache_read_tokens), 0)::BIGINT AS cache_read_tokens,
COALESCE(SUM(total_input_context), 0)::BIGINT AS total_input_context,
@@ -1926,6 +1980,75 @@ WHERE date >= $1
decode_dashboard_summary_row(&row)
}
async fn list_dashboard_daily_breakdown_from_daily_totals(
&self,
start_day_utc: DateTime<Utc>,
end_day_utc: DateTime<Utc>,
user_id: Option<&str>,
) -> Result<Vec<StoredUsageDashboardDailyBreakdownRow>, DataLayerError> {
if start_day_utc >= end_day_utc {
return Ok(Vec::new());
}
let sql = if user_id.is_some() {
r#"
SELECT
TO_CHAR(date, 'YYYY-MM-DD') AS date,
'aggregate'::TEXT AS model,
'aggregate'::TEXT AS provider,
COALESCE(SUM(total_requests), 0)::BIGINT AS requests,
COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0)::BIGINT AS total_tokens,
COALESCE(SUM(total_cost), 0)::DOUBLE PRECISION AS total_cost_usd,
0::DOUBLE PRECISION AS response_time_sum_ms,
0::BIGINT AS response_time_samples
FROM stats_user_daily
WHERE user_id = $1
AND date >= $2
AND date < $3
AND total_requests > 0
GROUP BY date
ORDER BY date ASC
"#
} else {
r#"
SELECT
TO_CHAR(date, 'YYYY-MM-DD') AS date,
'aggregate'::TEXT AS model,
'aggregate'::TEXT AS provider,
COALESCE(SUM(total_requests), 0)::BIGINT AS requests,
COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0)::BIGINT AS total_tokens,
COALESCE(SUM(total_cost), 0)::DOUBLE PRECISION AS total_cost_usd,
0::DOUBLE PRECISION AS response_time_sum_ms,
0::BIGINT AS response_time_samples
FROM stats_daily
WHERE date >= $1
AND date < $2
AND total_requests > 0
GROUP BY date
ORDER BY date ASC
"#
};
let mut rows = if let Some(user_id) = user_id {
sqlx::query(sql)
.bind(user_id)
.bind(start_day_utc)
.bind(end_day_utc)
.fetch(&self.pool)
} else {
sqlx::query(sql)
.bind(start_day_utc)
.bind(end_day_utc)
.fetch(&self.pool)
};
let mut items = Vec::new();
while let Some(row) = rows.try_next().await.map_postgres_err()? {
items.push(decode_dashboard_daily_breakdown_row(&row)?);
}
Ok(items)
}
async fn summarize_dashboard_usage_raw(
&self,
created_from_unix_secs: u64,
@@ -4241,8 +4364,19 @@ ORDER BY date ASC, total_cost_usd DESC, model ASC, provider_name ASC
};
let mut items = Vec::new();
let mut detailed_dates = std::collections::BTreeSet::<String>::new();
while let Some(row) = rows.try_next().await.map_postgres_err()? {
items.push(decode_dashboard_daily_breakdown_row(&row)?);
let item = decode_dashboard_daily_breakdown_row(&row)?;
detailed_dates.insert(item.date.clone());
items.push(item);
}
for item in self
.list_dashboard_daily_breakdown_from_daily_totals(start_day_utc, end_day_utc, user_id)
.await?
{
if !detailed_dates.contains(&item.date) {
items.push(item);
}
}
Ok(items)
}
@@ -4350,12 +4484,53 @@ ORDER BY date ASC, total_cost_usd DESC, "usage".model ASC, "usage".provider_name
Ok(items)
}
async fn list_dashboard_daily_breakdown_aggregate_segments(
&self,
query: &UsageDashboardDailyBreakdownQuery,
) -> Result<Vec<StoredUsageDashboardDailyBreakdownRow>, DataLayerError> {
let cutoff_utc = match self.read_stats_daily_cutoff_date().await {
Ok(value) => value,
Err(err) if dashboard_should_fallback_to_raw_on_aggregate_error(&err) => {
return Ok(Vec::new());
}
Err(err) => return Err(err),
};
let Some(cutoff_utc) = cutoff_utc else {
return Ok(Vec::new());
};
let start_utc = dashboard_unix_secs_to_utc(query.created_from_unix_secs);
let end_utc = dashboard_unix_secs_to_utc(query.created_until_unix_secs);
let split = split_dashboard_daily_aggregate_range(start_utc, end_utc, cutoff_utc);
let Some((aggregate_start, aggregate_end)) = split.aggregate else {
return Ok(Vec::new());
};
self.list_dashboard_daily_breakdown_from_daily_aggregates(
aggregate_start,
aggregate_end,
query.user_id.as_deref(),
)
.await
}
pub async fn list_dashboard_daily_breakdown(
&self,
query: &UsageDashboardDailyBreakdownQuery,
) -> Result<Vec<StoredUsageDashboardDailyBreakdownRow>, DataLayerError> {
if query.tz_offset_minutes != 0 {
return self.list_dashboard_daily_breakdown_raw(query).await;
let mut items = self
.list_dashboard_daily_breakdown_aggregate_segments(query)
.await?;
let mut aggregate_dates = items
.iter()
.map(|item| item.date.clone())
.collect::<std::collections::BTreeSet<_>>();
for item in self.list_dashboard_daily_breakdown_raw(query).await? {
if aggregate_dates.insert(item.date.clone()) {
items.push(item);
}
}
return Ok(finalize_dashboard_daily_breakdown_rows(items));
}
let cutoff_utc = match self.read_stats_daily_cutoff_date().await {
@@ -7370,6 +7545,7 @@ ORDER BY api_key_id ASC
let mut totals = std::collections::BTreeMap::<String, StoredUsageUserTotals>::new();
if let Some(cutoff_utc) = self.read_stats_daily_cutoff_date().await? {
let mut summary_user_ids = std::collections::BTreeSet::<String>::new();
let mut aggregate_rows = sqlx::query(
r#"
SELECT
@@ -7392,6 +7568,50 @@ ORDER BY user_id ASC
while let Some(row) = aggregate_rows.try_next().await.map_postgres_err()? {
let user_id = row.try_get::<String, _>("user_id").map_postgres_err()?;
let request_count = row
.try_get::<i64, _>("request_count")
.map_postgres_err()?
.max(0) as u64;
let total_tokens = row
.try_get::<i64, _>("total_tokens")
.map_postgres_err()?
.max(0) as u64;
summary_user_ids.insert(user_id.clone());
totals.insert(
user_id.clone(),
StoredUsageUserTotals {
user_id,
request_count,
total_tokens,
},
);
}
let mut daily_rows = sqlx::query(
r#"
SELECT
user_id,
COALESCE(SUM(total_requests), 0)::BIGINT AS request_count,
COALESCE(
SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens),
0
)::BIGINT AS total_tokens
FROM stats_user_daily
WHERE user_id = ANY($1::TEXT[])
AND date < $2
GROUP BY user_id
ORDER BY user_id ASC
"#,
)
.bind(user_ids)
.bind(cutoff_utc)
.fetch(&self.pool);
while let Some(row) = daily_rows.try_next().await.map_postgres_err()? {
let user_id = row.try_get::<String, _>("user_id").map_postgres_err()?;
if summary_user_ids.contains(&user_id) {
continue;
}
let request_count = row
.try_get::<i64, _>("request_count")
.map_postgres_err()?
@@ -24,15 +24,23 @@ WITH aggregated AS (
END
), 0)::BIGINT AS error_count,
COALESCE(SUM(
GREATEST(
COALESCE(
total_tokens,
COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)
),
0
)::BIGINT
CASE
WHEN status IN ('pending', 'streaming') THEN 0
ELSE GREATEST(
COALESCE(
total_tokens,
COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)
),
0
)::BIGINT
END
), 0)::BIGINT AS total_tokens,
COALESCE(SUM(COALESCE(total_cost_usd, 0)), 0)::NUMERIC(20,8) AS total_cost_usd,
COALESCE(SUM(
CASE
WHEN status IN ('pending', 'streaming') THEN 0
ELSE COALESCE(total_cost_usd, 0)
END
), 0)::NUMERIC(20,8) AS total_cost_usd,
COALESCE(SUM(
CASE
WHEN status IN ('completed', 'success', 'ok', 'billed', 'settled')
@@ -47,7 +55,6 @@ WITH aggregated AS (
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
@@ -308,7 +308,9 @@ fn usage_sql_rebuild_matches_online_provider_key_usage_semantics() {
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')"));
.contains("WHEN status NOT IN ('pending', 'streaming')"));
assert!(super::REBUILD_PROVIDER_API_KEY_USAGE_STATS_SQL
.contains("WHEN status IN ('pending', 'streaming') THEN 0"));
}
#[test]
@@ -436,6 +438,30 @@ fn usage_sql_summarize_usage_daily_heatmap_supports_daily_aggregates() {
);
}
#[test]
fn usage_sql_daily_cutoff_falls_back_to_imported_stats_daily() {
let source = include_str!("mod.rs");
assert!(source.contains("FROM stats_summary"));
assert!(source.contains("SELECT MAX(date) AS latest_date"));
assert!(source.contains("FROM stats_daily"));
assert!(source.contains("FROM stats_user_daily"));
assert!(source.contains("FROM stats_daily_api_key"));
assert!(source.contains("value + chrono::Duration::days(1)"));
}
#[test]
fn usage_sql_dashboard_daily_breakdown_falls_back_to_daily_totals() {
let source = include_str!("mod.rs");
assert!(source.contains("list_dashboard_daily_breakdown_aggregate_segments"));
assert!(source.contains("list_dashboard_daily_breakdown_from_daily_totals"));
assert!(source.contains("'aggregate'::TEXT AS model"));
assert!(source.contains("FROM stats_daily"));
assert!(source.contains("FROM stats_user_daily"));
assert!(source.contains("detailed_dates.contains(&item.date)"));
assert!(source.contains("query.tz_offset_minutes != 0"));
assert!(source.contains("aggregate_dates.insert(item.date.clone())"));
}
#[test]
fn usage_sql_summarize_usage_leaderboard_supports_daily_aggregates() {
let source = include_str!("mod.rs");
@@ -533,6 +559,9 @@ fn usage_sql_summarize_usage_totals_by_user_ids_supports_user_summary_aggregates
let source = include_str!("mod.rs");
assert!(source.contains("FROM stats_user_summary"));
assert!(source.contains("all_time_input_tokens"));
assert!(source.contains("FROM stats_user_daily"));
assert!(source.contains("date < $2"));
assert!(source.contains("summary_user_ids.contains(&user_id)"));
}
#[test]
+575 -65
View File
@@ -1,5 +1,6 @@
use std::collections::{BTreeMap, HashSet};
use std::io::Read;
use std::time::{SystemTime, UNIX_EPOCH};
use aether_ai_formats::UPSTREAM_IS_STREAM_KEY;
use aether_data_contracts::repository::usage::{parse_usage_body_ref, UsageBodyField};
@@ -960,6 +961,102 @@ impl SqliteUsageReadRepository {
Self { pool }
}
async fn summarize_usage_daily_heatmap_raw_from_range(
&self,
created_from_unix_secs: u64,
created_until_unix_secs: u64,
user_id: Option<&str>,
) -> Result<Vec<StoredUsageDailySummary>, DataLayerError> {
let mut builder = QueryBuilder::<Sqlite>::new(format!(
r#"
SELECT
date(created_at_unix_ms, 'unixepoch') AS date,
COUNT(*) AS requests,
COALESCE(SUM(
MAX(COALESCE(input_tokens, 0), 0)
+ MAX(COALESCE(output_tokens, 0), 0)
+ {cache_creation_expr}
+ MAX(COALESCE(cache_read_input_tokens, 0), 0)
), 0) AS total_tokens,
COALESCE(SUM(COALESCE(CAST(total_cost_usd AS REAL), 0)), 0) AS total_cost_usd,
COALESCE(SUM(COALESCE(CAST(actual_total_cost_usd AS REAL), 0)), 0)
AS actual_total_cost_usd
FROM "usage"
"#,
cache_creation_expr = SQLITE_USAGE_CACHE_CREATION_TOKENS_EXPR
));
let mut has_where = false;
push_sqlite_usage_where(&mut builder, &mut has_where);
builder
.push("created_at_unix_ms >= ")
.push_bind(created_from_unix_secs as i64);
push_sqlite_usage_where(&mut builder, &mut has_where);
builder
.push("created_at_unix_ms < ")
.push_bind(created_until_unix_secs as i64);
push_sqlite_usage_finalized_filter(&mut builder, &mut has_where);
push_sqlite_usage_optional_text_filter(&mut builder, &mut has_where, "user_id", user_id);
builder.push(" GROUP BY date ORDER BY date ASC");
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
rows.iter().map(map_sqlite_usage_daily_summary).collect()
}
async fn summarize_usage_daily_heatmap_from_daily_aggregates(
&self,
created_from_unix_secs: u64,
created_until_unix_secs: u64,
user_id: Option<&str>,
) -> Result<Vec<StoredUsageDailySummary>, DataLayerError> {
let rows = if let Some(user_id) = user_id {
sqlx::query(
r#"
SELECT
date("date", 'unixepoch') AS date,
total_requests AS requests,
input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens AS total_tokens,
total_cost AS total_cost_usd,
total_cost AS actual_total_cost_usd
FROM stats_user_daily
WHERE user_id = ?
AND "date" >= ?
AND "date" < ?
AND total_requests > 0
ORDER BY "date" ASC
"#,
)
.bind(user_id)
.bind(created_from_unix_secs as i64)
.bind(created_until_unix_secs as i64)
.fetch_all(&self.pool)
.await
.map_sql_err()?
} else {
sqlx::query(
r#"
SELECT
date("date", 'unixepoch') AS date,
total_requests AS requests,
input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens AS total_tokens,
total_cost AS total_cost_usd,
actual_total_cost AS actual_total_cost_usd
FROM stats_daily
WHERE "date" >= ?
AND "date" < ?
AND total_requests > 0
ORDER BY "date" ASC
"#,
)
.bind(created_from_unix_secs as i64)
.bind(created_until_unix_secs as i64)
.fetch_all(&self.pool)
.await
.map_sql_err()?
};
rows.iter().map(map_sqlite_usage_daily_summary).collect()
}
async fn summarize_provider_performance_percentiles(
&self,
query: &UsageProviderPerformanceQuery,
@@ -1463,38 +1560,266 @@ FROM "usage"
return Ok(Vec::new());
}
let mut builder = QueryBuilder::<Sqlite>::new(
let unique_user_ids = user_ids
.iter()
.cloned()
.collect::<std::collections::BTreeSet<_>>();
let mut totals = BTreeMap::<String, StoredUsageUserTotals>::new();
let mut aggregate_cutoffs = BTreeMap::<String, u64>::new();
let mut aggregate_builder = QueryBuilder::<Sqlite>::new(
r#"
SELECT
user_id,
COUNT(*) AS request_count,
COALESCE(SUM(MAX(COALESCE(total_tokens, 0), 0)), 0) AS total_tokens
FROM "usage"
COALESCE(SUM(total_requests), 0) AS request_count,
COALESCE(
SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens),
0
) AS total_tokens,
MAX("date") AS latest_date
FROM stats_user_daily
WHERE user_id IN (
"#,
);
let mut separated = builder.separated(", ");
for user_id in user_ids {
separated.push_bind(user_id.clone());
{
let mut separated = aggregate_builder.separated(", ");
for user_id in &unique_user_ids {
separated.push_bind(user_id.clone());
}
}
separated.push_unseparated(
r#")
AND status NOT IN ('pending', 'streaming')
AND provider_name NOT IN ('unknown', 'pending')
GROUP BY user_id
ORDER BY user_id ASC
aggregate_builder.push(") GROUP BY user_id ORDER BY user_id ASC");
let aggregate_rows = aggregate_builder
.build()
.fetch_all(&self.pool)
.await
.map_sql_err()?;
for row in aggregate_rows {
let user_id: String = row.try_get("user_id").map_sql_err()?;
let latest_date = row.try_get::<i64, _>("latest_date").map_sql_err()?.max(0) as u64;
aggregate_cutoffs.insert(user_id.clone(), latest_date.saturating_add(86_400));
totals.insert(
user_id.clone(),
StoredUsageUserTotals {
user_id,
request_count: row_u64(&row, "request_count")?,
total_tokens: row_u64(&row, "total_tokens")?,
},
);
}
let mut builder = QueryBuilder::<Sqlite>::new(
r#"
SELECT
"usage".user_id,
COUNT(*) AS request_count,
COALESCE(SUM(MAX(COALESCE("usage".total_tokens, 0), 0)), 0) AS total_tokens
FROM "usage"
JOIN (
"#,
);
for (index, user_id) in unique_user_ids.iter().enumerate() {
if index > 0 {
builder.push(" UNION ALL ");
}
let cutoff = aggregate_cutoffs.get(user_id).copied().unwrap_or_default();
builder
.push("SELECT ")
.push_bind(user_id.clone())
.push(" AS user_id, ")
.push_bind(to_i64(cutoff, "usage aggregate cutoff")?)
.push(" AS cutoff_unix_secs");
}
builder.push(
r#"
) AS requested ON requested.user_id = "usage".user_id
WHERE "usage".created_at_unix_ms >= requested.cutoff_unix_secs
AND "usage".status NOT IN ('pending', 'streaming')
AND "usage".provider_name NOT IN ('unknown', 'pending')
GROUP BY "usage".user_id
ORDER BY "usage".user_id ASC
"#,
);
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
for row in rows {
let user_id: String = row.try_get("user_id").map_sql_err()?;
let entry = totals
.entry(user_id.clone())
.or_insert_with(|| StoredUsageUserTotals {
user_id,
request_count: 0,
total_tokens: 0,
});
entry.request_count = entry
.request_count
.saturating_add(row_u64(&row, "request_count")?);
entry.total_tokens = entry
.total_tokens
.saturating_add(row_u64(&row, "total_tokens")?);
}
Ok(totals.into_values().collect())
}
async fn summarize_dashboard_usage_from_daily_aggregates(
&self,
query: &UsageDashboardSummaryQuery,
) -> Result<Option<StoredUsageDashboardSummary>, DataLayerError> {
let row = if let Some(user_id) = query.user_id.as_deref() {
sqlx::query(
r#"
SELECT
COALESCE(SUM(total_requests), 0) AS total_requests,
COALESCE(SUM(input_tokens), 0) AS input_tokens,
COALESCE(SUM(input_tokens), 0) AS effective_input_tokens,
COALESCE(SUM(output_tokens), 0) AS output_tokens,
COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS total_tokens,
COALESCE(SUM(cache_creation_tokens), 0) AS cache_creation_tokens,
COALESCE(SUM(cache_read_tokens), 0) AS cache_read_tokens,
COALESCE(SUM(input_tokens + cache_creation_tokens + cache_read_tokens), 0) AS total_input_context,
0.0 AS cache_creation_cost_usd,
0.0 AS cache_read_cost_usd,
COALESCE(SUM(COALESCE(CAST(total_cost AS REAL), 0)), 0) AS total_cost_usd,
COALESCE(SUM(COALESCE(CAST(total_cost AS REAL), 0)), 0) AS actual_total_cost_usd,
COALESCE(SUM(error_requests), 0) AS error_requests,
0.0 AS response_time_sum_ms,
0 AS response_time_samples
FROM stats_user_daily
WHERE user_id = ?
AND "date" >= ?
AND "date" < ?
"#,
)
.bind(user_id)
.bind(query.created_from_unix_secs as i64)
.bind(query.created_until_unix_secs as i64)
.fetch_one(&self.pool)
.await
.map_sql_err()?
} else {
sqlx::query(
r#"
SELECT
COALESCE(SUM(total_requests), 0) AS total_requests,
COALESCE(SUM(input_tokens), 0) AS input_tokens,
COALESCE(SUM(input_tokens), 0) AS effective_input_tokens,
COALESCE(SUM(output_tokens), 0) AS output_tokens,
COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS total_tokens,
COALESCE(SUM(cache_creation_tokens), 0) AS cache_creation_tokens,
COALESCE(SUM(cache_read_tokens), 0) AS cache_read_tokens,
COALESCE(SUM(input_tokens + cache_creation_tokens + cache_read_tokens), 0) AS total_input_context,
0.0 AS cache_creation_cost_usd,
0.0 AS cache_read_cost_usd,
COALESCE(SUM(COALESCE(CAST(total_cost AS REAL), 0)), 0) AS total_cost_usd,
COALESCE(SUM(COALESCE(CAST(actual_total_cost AS REAL), 0)), 0) AS actual_total_cost_usd,
COALESCE(SUM(error_requests), 0) AS error_requests,
0.0 AS response_time_sum_ms,
0 AS response_time_samples
FROM stats_daily
WHERE "date" >= ?
AND "date" < ?
"#,
)
.bind(query.created_from_unix_secs as i64)
.bind(query.created_until_unix_secs as i64)
.fetch_one(&self.pool)
.await
.map_sql_err()?
};
let total_requests = sqlite_aggregate_u64(&row, "total_requests")?;
if total_requests == 0 {
return Ok(None);
}
Ok(Some(StoredUsageDashboardSummary {
total_requests,
input_tokens: sqlite_aggregate_u64(&row, "input_tokens")?,
effective_input_tokens: sqlite_aggregate_u64(&row, "effective_input_tokens")?,
output_tokens: sqlite_aggregate_u64(&row, "output_tokens")?,
total_tokens: sqlite_aggregate_u64(&row, "total_tokens")?,
cache_creation_tokens: sqlite_aggregate_u64(&row, "cache_creation_tokens")?,
cache_read_tokens: sqlite_aggregate_u64(&row, "cache_read_tokens")?,
total_input_context: sqlite_aggregate_u64(&row, "total_input_context")?,
cache_creation_cost_usd: sqlite_real(&row, "cache_creation_cost_usd")?,
cache_read_cost_usd: sqlite_real(&row, "cache_read_cost_usd")?,
total_cost_usd: sqlite_real(&row, "total_cost_usd")?,
actual_total_cost_usd: sqlite_real(&row, "actual_total_cost_usd")?,
error_requests: sqlite_aggregate_u64(&row, "error_requests")?,
response_time_sum_ms: sqlite_real(&row, "response_time_sum_ms")?,
response_time_samples: sqlite_aggregate_u64(&row, "response_time_samples")?,
}))
}
async fn list_dashboard_daily_breakdown_from_daily_aggregates(
&self,
query: &UsageDashboardDailyBreakdownQuery,
) -> Result<Vec<StoredUsageDashboardDailyBreakdownRow>, DataLayerError> {
let rows = if let Some(user_id) = query.user_id.as_deref() {
sqlx::query(
r#"
SELECT
date("date", 'unixepoch') AS date,
'aggregate' AS model,
'aggregate' AS provider,
COALESCE(SUM(total_requests), 0) AS requests,
COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS total_tokens,
COALESCE(SUM(COALESCE(CAST(total_cost AS REAL), 0)), 0) AS total_cost_usd,
0.0 AS response_time_sum_ms,
0 AS response_time_samples
FROM stats_user_daily
WHERE user_id = ?
AND "date" >= ?
AND "date" < ?
AND total_requests > 0
GROUP BY "date"
ORDER BY "date" ASC
"#,
)
.bind(user_id)
.bind(query.created_from_unix_secs as i64)
.bind(query.created_until_unix_secs as i64)
.fetch_all(&self.pool)
.await
.map_sql_err()?
} else {
sqlx::query(
r#"
SELECT
date("date", 'unixepoch') AS date,
'aggregate' AS model,
'aggregate' AS provider,
COALESCE(SUM(total_requests), 0) AS requests,
COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS total_tokens,
COALESCE(SUM(COALESCE(CAST(total_cost AS REAL), 0)), 0) AS total_cost_usd,
0.0 AS response_time_sum_ms,
0 AS response_time_samples
FROM stats_daily
WHERE "date" >= ?
AND "date" < ?
AND total_requests > 0
GROUP BY "date"
ORDER BY "date" ASC
"#,
)
.bind(query.created_from_unix_secs as i64)
.bind(query.created_until_unix_secs as i64)
.fetch_all(&self.pool)
.await
.map_sql_err()?
};
rows.iter()
.map(|row| {
Ok(StoredUsageUserTotals {
user_id: row.try_get::<String, _>("user_id").map_sql_err()?,
request_count: row.try_get::<i64, _>("request_count").map_sql_err()?.max(0)
as u64,
total_tokens: row.try_get::<i64, _>("total_tokens").map_sql_err()?.max(0)
as u64,
Ok(StoredUsageDashboardDailyBreakdownRow {
date: row.try_get("date").map_sql_err()?,
model: row.try_get("model").map_sql_err()?,
provider: row.try_get("provider").map_sql_err()?,
requests: sqlite_aggregate_u64(row, "requests")?,
total_tokens: sqlite_aggregate_u64(row, "total_tokens")?,
total_cost_usd: sqlite_real(row, "total_cost_usd")?,
response_time_sum_ms: sqlite_real(row, "response_time_sum_ms")?,
response_time_samples: sqlite_aggregate_u64(row, "response_time_samples")?,
})
})
.collect()
@@ -1885,6 +2210,13 @@ ORDER BY created_at_unix_ms ASC, id ASC
return Ok(StoredUsageDashboardSummary::default());
}
if let Some(summary) = self
.summarize_dashboard_usage_from_daily_aggregates(query)
.await?
{
return Ok(summary);
}
let mut builder = QueryBuilder::<Sqlite>::new(format!(
r#"
SELECT
@@ -1959,6 +2291,13 @@ FROM "usage"
return Ok(Vec::new());
}
let aggregate_rows = self
.list_dashboard_daily_breakdown_from_daily_aggregates(query)
.await?;
if !aggregate_rows.is_empty() {
return Ok(aggregate_rows);
}
let date_expr = sqlite_usage_local_date_expr(query.tz_offset_minutes);
let mut builder = QueryBuilder::<Sqlite>::new(format!(
r#"
@@ -3227,53 +3566,54 @@ WHERE provider_id = ?
&self,
query: &UsageDailyHeatmapQuery,
) -> Result<Vec<StoredUsageDailySummary>, DataLayerError> {
let mut builder = QueryBuilder::<Sqlite>::new(format!(
r#"
SELECT
date(created_at_unix_ms, 'unixepoch') AS date,
COUNT(*) AS requests,
COALESCE(SUM(
MAX(COALESCE(input_tokens, 0), 0)
+ MAX(COALESCE(output_tokens, 0), 0)
+ {cache_creation_expr}
+ MAX(COALESCE(cache_read_input_tokens, 0), 0)
), 0) AS total_tokens,
COALESCE(SUM(COALESCE(CAST(total_cost_usd AS REAL), 0)), 0) AS total_cost_usd,
COALESCE(SUM(COALESCE(CAST(actual_total_cost_usd AS REAL), 0)), 0)
AS actual_total_cost_usd
FROM "usage"
"#,
cache_creation_expr = SQLITE_USAGE_CACHE_CREATION_TOKENS_EXPR
));
let mut has_where = false;
push_sqlite_usage_where(&mut builder, &mut has_where);
builder
.push("created_at_unix_ms >= ")
.push_bind(query.created_from_unix_secs as i64);
push_sqlite_usage_finalized_filter(&mut builder, &mut has_where);
push_sqlite_usage_optional_text_filter(
&mut builder,
&mut has_where,
"user_id",
query.user_id.as_deref(),
);
builder.push(" GROUP BY date ORDER BY date ASC");
let created_until_unix_secs = usage_current_unix_secs().saturating_add(1);
let user_id = query.user_id.as_deref();
let mut summaries = BTreeMap::<String, StoredUsageDailySummary>::new();
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
rows.iter()
.map(|row| {
Ok(StoredUsageDailySummary {
date: row.try_get("date").map_sql_err()?,
requests: sqlite_aggregate_u64(row, "requests")?,
total_tokens: sqlite_aggregate_u64(row, "total_tokens")?,
total_cost_usd: sqlite_real(row, "total_cost_usd")?,
actual_total_cost_usd: sqlite_real(row, "actual_total_cost_usd")?,
})
})
.collect()
for item in self
.summarize_usage_daily_heatmap_from_daily_aggregates(
query.created_from_unix_secs,
created_until_unix_secs,
user_id,
)
.await?
{
summaries.insert(item.date.clone(), item);
}
for item in self
.summarize_usage_daily_heatmap_raw_from_range(
query.created_from_unix_secs,
created_until_unix_secs,
user_id,
)
.await?
{
summaries.entry(item.date.clone()).or_insert(item);
}
Ok(summaries.into_values().collect())
}
}
fn map_sqlite_usage_daily_summary(
row: &SqliteRow,
) -> Result<StoredUsageDailySummary, DataLayerError> {
Ok(StoredUsageDailySummary {
date: row.try_get("date").map_sql_err()?,
requests: sqlite_aggregate_u64(row, "requests")?,
total_tokens: sqlite_aggregate_u64(row, "total_tokens")?,
total_cost_usd: sqlite_real(row, "total_cost_usd")?,
actual_total_cost_usd: sqlite_real(row, "actual_total_cost_usd")?,
})
}
fn usage_current_unix_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_secs())
.unwrap_or_default()
}
impl SqliteUsageWriteRepository {
pub fn new(pool: SqlitePool) -> Self {
Self { pool }
@@ -3409,8 +3749,14 @@ SELECT
COUNT(*) AS request_count,
COALESCE(SUM({success_flag_expr}), 0) AS success_count,
COALESCE(SUM({error_flag_expr}), 0) AS error_count,
COALESCE(SUM(MAX(COALESCE(total_tokens, 0), 0)), 0) AS total_tokens,
COALESCE(SUM(COALESCE(CAST(total_cost_usd AS REAL), 0)), 0) AS total_cost_usd,
COALESCE(SUM(CASE
WHEN status IN ('pending', 'streaming') THEN 0
ELSE MAX(COALESCE(total_tokens, 0), 0)
END), 0) AS total_tokens,
COALESCE(SUM(CASE
WHEN status IN ('pending', 'streaming') THEN 0
ELSE COALESCE(CAST(total_cost_usd AS REAL), 0)
END), 0) AS total_cost_usd,
COALESCE(SUM(CASE
WHEN {success_flag_expr} = 1 AND response_time_ms IS NOT NULL
THEN MAX(COALESCE(response_time_ms, 0), 0)
@@ -4012,7 +4358,8 @@ mod tests {
use super::{SqliteUsageReadRepository, SqliteUsageWriteRepository};
use crate::lifecycle::migrate::run_sqlite_migrations;
use crate::repository::usage::{
UpsertUsageRecord, UsageAuditListQuery, UsageDashboardSummaryQuery, UsageReadRepository,
UpsertUsageRecord, UsageAuditListQuery, UsageDailyHeatmapQuery,
UsageDashboardDailyBreakdownQuery, UsageDashboardSummaryQuery, UsageReadRepository,
UsageWriteRepository,
};
@@ -4350,6 +4697,169 @@ INSERT INTO request_candidates (
assert_eq!(summary.total_tokens, 10);
}
#[tokio::test]
async fn sqlite_usage_daily_heatmap_reads_imported_daily_aggregates() {
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");
sqlx::query(
r#"
INSERT INTO stats_daily (
id, "date", total_requests, success_requests, error_requests,
input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens,
total_cost, actual_total_cost, is_complete, created_at, updated_at
) VALUES (
'daily-1', 86400, 9, 8, 1, 10, 20, 3, 4, 1.25, 1.0, 1, 1, 1
);
INSERT INTO stats_user_daily (
id, user_id, username, "date", total_requests, success_requests, error_requests,
input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens,
total_cost, created_at, updated_at
) VALUES (
'user-daily-1', 'user-1', 'user one', 86400, 5, 5, 0, 7, 8, 2, 1, 0.75, 1, 1
);
"#,
)
.execute(&pool)
.await
.expect("daily aggregates should seed");
let reader = SqliteUsageReadRepository::new(pool);
let admin = reader
.summarize_usage_daily_heatmap(&UsageDailyHeatmapQuery {
created_from_unix_secs: 0,
user_id: None,
admin_mode: true,
})
.await
.expect("admin heatmap should load");
assert_eq!(admin.len(), 1);
assert_eq!(admin[0].date, "1970-01-02");
assert_eq!(admin[0].requests, 9);
assert_eq!(admin[0].total_tokens, 37);
assert_eq!(admin[0].actual_total_cost_usd, 1.0);
let user = reader
.summarize_usage_daily_heatmap(&UsageDailyHeatmapQuery {
created_from_unix_secs: 0,
user_id: Some("user-1".to_string()),
admin_mode: false,
})
.await
.expect("user heatmap should load");
assert_eq!(user.len(), 1);
assert_eq!(user[0].date, "1970-01-02");
assert_eq!(user[0].requests, 5);
assert_eq!(user[0].total_tokens, 18);
assert_eq!(user[0].actual_total_cost_usd, 0.75);
}
#[tokio::test]
async fn sqlite_usage_totals_by_user_ids_reads_imported_user_daily_aggregates() {
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");
sqlx::query(
r#"
INSERT INTO stats_user_daily (
id, user_id, username, "date", total_requests, success_requests, error_requests,
input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens,
total_cost, created_at, updated_at
) VALUES (
'user-daily-1', 'user-1', 'user one', 86400, 5, 5, 0, 7, 8, 2, 1, 0.75, 1, 1
);
INSERT INTO "usage" (
request_id, id, user_id, api_key_id, provider_name, model, total_tokens,
status, billing_status, created_at_unix_ms, updated_at_unix_secs
) VALUES
('raw-before-cutoff', 'usage-1', 'user-1', 'api-key-1', 'Provider One', 'model-1', 99,
'completed', 'settled', 90000, 90000),
('raw-after-cutoff', 'usage-2', 'user-1', 'api-key-1', 'Provider One', 'model-1', 7,
'completed', 'settled', 172800, 172800);
"#,
)
.execute(&pool)
.await
.expect("usage totals fixtures should seed");
let reader = SqliteUsageReadRepository::new(pool);
let totals = reader
.summarize_usage_totals_by_user_ids(&["user-1".to_string()])
.await
.expect("user totals should load");
assert_eq!(totals.len(), 1);
assert_eq!(totals[0].user_id, "user-1");
assert_eq!(totals[0].request_count, 6);
assert_eq!(totals[0].total_tokens, 25);
}
#[tokio::test]
async fn sqlite_dashboard_daily_stats_reads_imported_daily_aggregates() {
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");
sqlx::query(
r#"
INSERT INTO stats_daily (
id, "date", total_requests, success_requests, error_requests,
input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens,
total_cost, actual_total_cost, is_complete, created_at, updated_at
) VALUES (
'daily-1', 86400, 9, 8, 1, 10, 20, 3, 4, 1.25, 1.0, 1, 1, 1
);
"#,
)
.execute(&pool)
.await
.expect("daily aggregates should seed");
let reader = SqliteUsageReadRepository::new(pool);
let summary = reader
.summarize_dashboard_usage(&UsageDashboardSummaryQuery {
created_from_unix_secs: 0,
created_until_unix_secs: 172800,
user_id: None,
})
.await
.expect("dashboard summary should load");
assert_eq!(summary.total_requests, 9);
assert_eq!(summary.total_tokens, 37);
let rows = reader
.list_dashboard_daily_breakdown(&UsageDashboardDailyBreakdownQuery {
created_from_unix_secs: 0,
created_until_unix_secs: 172800,
tz_offset_minutes: 480,
user_id: None,
})
.await
.expect("dashboard daily breakdown should load");
assert_eq!(rows.len(), 1);
assert_eq!(rows[0].date, "1970-01-02");
assert_eq!(rows[0].model, "aggregate");
assert_eq!(rows[0].requests, 9);
assert_eq!(rows[0].total_tokens, 37);
}
async fn seed_stats_targets(pool: &sqlx::SqlitePool) {
sqlx::query(
r#"
+69 -3
View File
@@ -212,7 +212,7 @@ mod tests {
}
#[test]
fn chatgpt_web_quota_metadata_enriches_auth_and_normalizes_free_limit() {
fn chatgpt_web_quota_metadata_enriches_auth_and_uses_first_remaining_as_limit() {
let mut metadata = json!({
"image_quota_remaining": 12,
});
@@ -229,8 +229,8 @@ mod tests {
assert_eq!(metadata["plan_type"], json!("free"));
assert_eq!(metadata["email"], json!("user@example.com"));
assert_eq!(metadata["account_id"], json!("acct-1"));
assert_eq!(metadata["image_quota_total"], json!(25.0));
assert_eq!(metadata["image_quota_used"], json!(13.0));
assert_eq!(metadata["image_quota_total"], json!(12.0));
assert_eq!(metadata["image_quota_used"], json!(0.0));
}
#[test]
@@ -252,6 +252,72 @@ mod tests {
assert_eq!(metadata["image_quota_used"], json!(33.0));
}
#[test]
fn chatgpt_web_quota_metadata_does_not_preserve_legacy_free_25_limit() {
let mut metadata = json!({
"plan_type": "free",
"image_quota_remaining": 19,
});
normalize_chatgpt_web_image_quota_limit(
&mut metadata,
Some(&json!({
"chatgpt_web": {
"plan_type": "free",
"image_quota_total": 25
}
})),
);
assert_eq!(metadata["image_quota_total"], json!(19.0));
assert_eq!(metadata["image_quota_used"], json!(0.0));
assert_eq!(
metadata["image_quota_limit_source"],
json!("first_remaining")
);
}
#[test]
fn chatgpt_web_quota_metadata_ignores_upstream_free_25_default() {
let mut metadata = json!({
"plan_type": "free",
"image_quota_remaining": 19,
"image_quota_total": 25,
});
normalize_chatgpt_web_image_quota_limit(&mut metadata, None);
assert_eq!(metadata["image_quota_total"], json!(19.0));
assert_eq!(metadata["image_quota_used"], json!(0.0));
assert_eq!(
metadata["image_quota_limit_source"],
json!("first_remaining")
);
}
#[test]
fn chatgpt_web_quota_metadata_preserves_marked_free_first_limit() {
let mut metadata = json!({
"plan_type": "free",
"image_quota_remaining": 18,
});
normalize_chatgpt_web_image_quota_limit(
&mut metadata,
Some(&json!({
"chatgpt_web": {
"plan_type": "free",
"image_quota_total": 19,
"image_quota_limit_source": "first_remaining"
}
})),
);
assert_eq!(metadata["image_quota_total"], json!(19.0));
assert_eq!(metadata["image_quota_used"], json!(1.0));
assert_eq!(
metadata["image_quota_limit_source"],
json!("first_remaining")
);
}
#[test]
fn windsurf_quota_request_uses_user_status_connect_rpc() {
let spec = build_windsurf_pool_quota_request("key-ws", "session-token-123");
@@ -24,8 +24,6 @@ const CHATGPT_WEB_CLIENT_VERSION: &str = "prod-be885abbfcfe7b1f511e88b3003d9ee44
const CHATGPT_WEB_BUILD_NUMBER: &str = "5955942";
const CHATGPT_WEB_SEC_CH_UA: &str =
r#""Microsoft Edge";v="143", "Chromium";v="143", "Not A(Brand";v="24""#;
const CHATGPT_WEB_FREE_IMAGE_QUOTA_LIMIT: f64 = 25.0;
#[derive(Debug, Clone, Default)]
pub struct ChatGptWebProviderPoolAdapter;
@@ -183,28 +181,47 @@ pub fn normalize_chatgpt_web_image_quota_limit(
};
let remaining = provider_pool_json_f64(object.get("image_quota_remaining"));
let explicit_limit =
let plan_type = chatgpt_web_image_quota_plan_type(object)
.map(ToOwned::to_owned)
.or_else(|| {
existing_limit
.as_ref()
.and_then(|existing| existing.plan_type.clone())
});
let raw_explicit_limit =
provider_pool_json_f64(object.get("image_quota_total")).filter(|value| *value > 0.0);
let plan_type = chatgpt_web_json_string(object.get("plan_type"));
let is_free_plan = plan_type.is_some_and(|value| value.trim().eq_ignore_ascii_case("free"));
let limit = if is_free_plan {
Some(CHATGPT_WEB_FREE_IMAGE_QUOTA_LIMIT)
} else {
explicit_limit
.or_else(|| infer_chatgpt_web_image_quota_limit(plan_type, remaining, existing_limit))
};
let explicit_limit_is_free_default = raw_explicit_limit.is_some_and(|limit| {
is_legacy_chatgpt_web_free_default_limit_value(limit, None, plan_type.as_deref(), remaining)
});
if explicit_limit_is_free_default {
object.remove("image_quota_total");
object.remove("image_quota_limit_source");
}
let explicit_limit = raw_explicit_limit.filter(|_| !explicit_limit_is_free_default);
let limit = explicit_limit
.map(|limit| ChatGptWebImageQuotaLimit {
value: limit,
source: Some("upstream_total".to_string()),
plan_type: plan_type.clone(),
})
.or_else(|| {
infer_chatgpt_web_image_quota_limit(remaining, existing_limit, plan_type.as_deref())
});
if let Some(limit) = limit {
object.insert("image_quota_total".to_string(), json!(limit));
object.insert("image_quota_total".to_string(), json!(limit.value));
if let Some(source) = limit.source.as_deref().filter(|value| !value.is_empty()) {
object.insert("image_quota_limit_source".to_string(), json!(source));
}
if !object.contains_key("image_quota_used") {
if let Some(remaining) = remaining {
object.insert(
"image_quota_used".to_string(),
json!((limit - remaining).max(0.0)),
json!((limit.value - remaining).max(0.0)),
);
} else if object.get("image_quota_blocked").and_then(Value::as_bool) == Some(true) {
object.insert("image_quota_used".to_string(), json!(limit));
object.insert("image_quota_used".to_string(), json!(limit.value));
}
}
}
@@ -222,37 +239,93 @@ fn chatgpt_web_auth_config_string(auth_config: Option<&Value>, fields: &[&str])
})
}
fn chatgpt_web_json_string(value: Option<&Value>) -> Option<&str> {
value
#[derive(Debug, Clone)]
struct ChatGptWebImageQuotaLimit {
value: f64,
source: Option<String>,
plan_type: Option<String>,
}
fn chatgpt_web_image_quota_plan_type(object: &Map<String, Value>) -> Option<&str> {
object
.get("plan_type")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
}
fn existing_chatgpt_web_image_quota_limit(upstream_metadata: Option<&Value>) -> Option<f64> {
upstream_metadata
fn existing_chatgpt_web_image_quota_limit(
upstream_metadata: Option<&Value>,
) -> Option<ChatGptWebImageQuotaLimit> {
let bucket = upstream_metadata
.and_then(Value::as_object)
.and_then(|metadata| metadata.get("chatgpt_web"))
.and_then(Value::as_object)
.and_then(|bucket| provider_pool_json_f64(bucket.get("image_quota_total")))
.filter(|value| *value > 0.0)
.and_then(Value::as_object)?;
let value =
provider_pool_json_f64(bucket.get("image_quota_total")).filter(|value| *value > 0.0)?;
let source = bucket
.get("image_quota_limit_source")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
let plan_type = chatgpt_web_image_quota_plan_type(bucket).map(ToOwned::to_owned);
Some(ChatGptWebImageQuotaLimit {
value,
source,
plan_type,
})
}
fn infer_chatgpt_web_image_quota_limit(
remaining: Option<f64>,
existing_limit: Option<ChatGptWebImageQuotaLimit>,
plan_type: Option<&str>,
) -> Option<ChatGptWebImageQuotaLimit> {
if let Some(existing_limit) = existing_limit {
if !is_legacy_chatgpt_web_free_default_limit(&existing_limit, plan_type, remaining) {
return Some(existing_limit);
}
}
remaining
.filter(|value| *value > 0.0)
.map(|value| ChatGptWebImageQuotaLimit {
value,
source: Some("first_remaining".to_string()),
plan_type: plan_type.map(ToOwned::to_owned),
})
}
fn is_legacy_chatgpt_web_free_default_limit(
existing_limit: &ChatGptWebImageQuotaLimit,
plan_type: Option<&str>,
remaining: Option<f64>,
existing_limit: Option<f64>,
) -> Option<f64> {
let normalized_plan = plan_type.unwrap_or_default().trim().to_ascii_lowercase();
if normalized_plan == "free" {
return Some(CHATGPT_WEB_FREE_IMAGE_QUOTA_LIMIT);
}
) -> bool {
is_legacy_chatgpt_web_free_default_limit_value(
existing_limit.value,
existing_limit.source.as_deref(),
plan_type,
remaining,
)
}
if let Some(existing_limit) = existing_limit.filter(|value| *value > 0.0) {
return Some(existing_limit);
fn is_legacy_chatgpt_web_free_default_limit_value(
value: f64,
source: Option<&str>,
plan_type: Option<&str>,
remaining: Option<f64>,
) -> bool {
let plan_type_is_free = plan_type
.map(str::trim)
.is_some_and(|value| value.eq_ignore_ascii_case("free"));
if !plan_type_is_free || source.is_some() {
return false;
}
remaining.filter(|value| *value > 0.0)
if (value - 25.0).abs() > f64::EPSILON {
return false;
}
remaining.is_none_or(|remaining| remaining < value)
}
pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
+1
View File
@@ -11,6 +11,7 @@ aether-ai-formats.workspace = true
aether-contracts.workspace = true
aether-data-contracts.workspace = true
aether-wallet.workspace = true
chrono.workspace = true
regex.workspace = true
serde.workspace = true
serde_json.workspace = true
+69 -6
View File
@@ -294,10 +294,33 @@ pub fn is_provider_key_circuit_open_at(
.as_ref()
.and_then(serde_json::Value::as_object)
.and_then(|values| values.get(api_format))
.and_then(serde_json::Value::as_object)
else {
return false;
};
provider_key_circuit_payload_is_active_open_at(payload, now_unix_secs)
}
pub fn any_provider_key_circuit_open_at(
key: &StoredProviderCatalogKey,
now_unix_secs: u64,
) -> bool {
key.circuit_breaker_by_format
.as_ref()
.and_then(serde_json::Value::as_object)
.is_some_and(|values| {
values.values().any(|payload| {
provider_key_circuit_payload_is_active_open_at(payload, now_unix_secs)
})
})
}
pub fn provider_key_circuit_payload_is_active_open_at(
payload: &serde_json::Value,
now_unix_secs: u64,
) -> bool {
let Some(payload) = payload.as_object() else {
return false;
};
if !payload
.get("open")
.and_then(serde_json::Value::as_bool)
@@ -305,10 +328,26 @@ pub fn is_provider_key_circuit_open_at(
{
return false;
}
payload
if let Some(next_probe_at) = payload
.get("next_probe_at_unix_secs")
.and_then(serde_json::Value::as_u64)
.is_none_or(|next_probe_at| now_unix_secs < next_probe_at)
{
return now_unix_secs < next_probe_at;
}
if let Some(next_probe_at) = payload
.get("next_probe_at")
.and_then(serde_json::Value::as_str)
.and_then(rfc3339_to_unix_secs)
{
return now_unix_secs < next_probe_at;
}
true
}
fn rfc3339_to_unix_secs(value: &str) -> Option<u64> {
chrono::DateTime::parse_from_rfc3339(value)
.ok()
.and_then(|value| u64::try_from(value.timestamp()).ok())
}
fn available_provider_key_rpm_slots_for_new_user(
@@ -641,9 +680,9 @@ mod tests {
count_recent_rpm_requests_for_provider_key,
count_recent_rpm_requests_for_provider_key_since, effective_provider_key_health_score,
effective_provider_key_rpm_limit, is_candidate_in_recent_failure_cooldown,
is_provider_key_circuit_open, provider_key_health_bucket, provider_key_health_score,
provider_key_rpm_allows_request, provider_key_rpm_allows_request_since,
ProviderKeyHealthBucket,
is_provider_key_circuit_open, is_provider_key_circuit_open_at, provider_key_health_bucket,
provider_key_health_score, provider_key_rpm_allows_request,
provider_key_rpm_allows_request_since, ProviderKeyHealthBucket,
};
fn stored_candidate(
@@ -1334,6 +1373,30 @@ mod tests {
assert!(!is_provider_key_circuit_open(&key, "openai:responses"));
}
#[test]
fn provider_key_circuit_open_at_allows_probe_after_rfc3339_deadline() {
let key = provider_catalog_key("key-a").with_health_fields(
None,
Some(serde_json::json!({
"openai:chat": {
"open": true,
"next_probe_at": "2026-05-24T14:45:27Z"
}
})),
);
assert!(is_provider_key_circuit_open_at(
&key,
"openai:chat",
1_779_633_926
));
assert!(!is_provider_key_circuit_open_at(
&key,
"openai:chat",
1_779_633_927
));
}
#[test]
fn aggregates_provider_key_health_score_with_lower_bound_strategy() {
let key = provider_catalog_key("key-a").with_health_fields(
+9 -8
View File
@@ -27,14 +27,15 @@ pub use candidate::{
SchedulerPriorityMode,
};
pub use health::{
aggregate_provider_key_health_score, count_recent_active_requests_for_api_key,
count_recent_active_requests_for_provider, count_recent_active_requests_for_provider_key,
count_recent_rpm_requests_for_provider_key, count_recent_rpm_requests_for_provider_key_since,
effective_provider_key_health_score, effective_provider_key_rpm_limit,
is_candidate_in_recent_failure_cooldown, is_provider_key_circuit_open,
is_provider_key_circuit_open_at, provider_key_health_bucket, provider_key_health_score,
provider_key_rpm_allows_request, provider_key_rpm_allows_request_since,
ProviderKeyHealthBucket, PROVIDER_KEY_RPM_WINDOW_SECS,
aggregate_provider_key_health_score, any_provider_key_circuit_open_at,
count_recent_active_requests_for_api_key, count_recent_active_requests_for_provider,
count_recent_active_requests_for_provider_key, count_recent_rpm_requests_for_provider_key,
count_recent_rpm_requests_for_provider_key_since, effective_provider_key_health_score,
effective_provider_key_rpm_limit, is_candidate_in_recent_failure_cooldown,
is_provider_key_circuit_open, is_provider_key_circuit_open_at,
provider_key_circuit_payload_is_active_open_at, provider_key_health_bucket,
provider_key_health_score, provider_key_rpm_allows_request,
provider_key_rpm_allows_request_since, ProviderKeyHealthBucket, PROVIDER_KEY_RPM_WINDOW_SECS,
};
pub use model::{
candidate_model_names, extract_global_priority_for_format, matches_model_mapping,
@@ -1,238 +0,0 @@
# Postgres to Aether Single Node Migration
Chinese version: [pg-to-single-node-migration.zh-CN.md](pg-to-single-node-migration.zh-CN.md)
This runbook migrates an existing Docker Compose Postgres deployment to Aether
single-node. In this repository, **single-node** means the default SQLite installer mode:
`install.sh --mode single-node`, a system service backed by SQLite. The Docker Compose
single-node template is `docker-compose.single-node.yml`, exposed through `--mode compose-single-node`.
The migration script is:
```bash
scripts/migrate-pg-to-single-node.sh
```
If the target should stay on Docker Compose instead of becoming a system
service, use the image-based Compose migration script:
```bash
scripts/migrate-pg-compose-to-single-node.sh
```
Both migration scripts pull/install the target single-node version before
downtime, stop only the source `app`, copy Postgres records directly into a
temporary SQLite DB without writing a JSONL file, replace the target
`aether.db`, and start single-node.
You can also use the installer as the unified entrypoint and let `--mode`
select the migration target:
```bash
# In interactive mode, first choose the target deployment mode:
# 1) Docker Compose standard deployment (Postgres + Redis)
# 2) Docker Compose single-node deployment (SQLite)
# 3) System service single-node deployment (SQLite)
# After choosing 2 or 3, choose the data initialization mode:
# 1) Fresh initialization (do not migrate existing data)
# 2) Migrate from an existing Docker Compose PG database
install.sh
# Migrate into a new single-node Docker Compose directory.
install.sh \
--mode compose-single-node \
--migrate-from-compose /root/Aether/docker-compose.yml \
--compose-dir /opt/aether-single \
--replace-existing
# Migrate into the system service + SQLite layout.
sudo install.sh \
--mode single-node \
--migrate-from-compose /root/Aether/docker-compose.yml \
--replace-existing
```
Interactive mode first asks for the target deployment shape. If the target is
`compose-single-node` or `single-node`, the installer then asks for the data
initialization mode: fresh initialization, or migration from an existing Docker
Compose PG database. If you choose migration, it tries to detect the source PG
Compose file from `docker compose ls`, then verifies that the Compose config
contains the default `app` and `postgres` services. If exactly one match is
found, it is used as the default prompt value. If detection is ambiguous or
fails, the installer stops; rerun it with `--migrate-from-compose` to specify
the source compose path.
The installer only normalizes the entrypoint: `compose-single-node` delegates to
`scripts/migrate-pg-compose-to-single-node.sh`, while `single-node` delegates to
`scripts/migrate-pg-to-single-node.sh`.
## What It Does
The script keeps the production cutover window short:
1. Reads the source Compose `.env`.
2. Builds a single-node env file that preserves `JWT_SECRET_KEY`, `ENCRYPTION_KEY` or
`AETHER_GATEWAY_DATA_ENCRYPTION_KEY`, admin settings, port, and app config.
3. Installs the single-node release with `install.sh --mode single-node --skip-start`.
4. Preflights SQLite migrations with the installed single-node binary.
5. Pulls the target single-node image, confirms its `copy` command is available,
and verifies that its Docker image ID matches the currently running source
`app` image ID.
6. Uses the target SQLite schema as the migration plan: same-name source
Postgres tables and columns are copied into the temporary SQLite DB.
7. Applies the compressed body and HTTP body detail policy. The default is full,
and you can opt into an omit mode for large artifacts.
8. Checks that the work directory and target SQLite directory have enough free
disk space for the temporary and final SQLite files.
9. Stops only the source `app` service, leaving Postgres and Redis running.
10. Copies source Postgres records directly into a temporary SQLite database
without generating JSONL files.
11. Replaces the target SQLite DB, including SQLite `-wal`/`-shm` sidecar files
when present, and starts the single-node service.
The image check compares Docker image IDs, not just tag strings. If both source
and target say `latest` but resolve to different image IDs, migration stops.
Upgrade the source PG Compose `app` to the target single-node version first,
verify it is healthy, then run the migration. The scripts also check that the
target image supports direct copy and the request-body omit flag; using a new
script with an old image stops before cutover to avoid missing data.
## Production Cutover
Before production cutover, take a normal server backup or snapshot. Then run:
```bash
sudo scripts/migrate-pg-to-single-node.sh \
--source-compose /root/Aether/docker-compose.yml \
--replace-existing
```
For Docker Compose single-node cutover instead of a system service:
```bash
scripts/migrate-pg-compose-to-single-node.sh \
--source-compose /root/Aether/docker-compose.yml \
--replace-existing
```
The source Postgres compose directory and target single-node compose directory
can be different. For example:
```bash
install.sh \
--mode compose-single-node \
--migrate-from-compose /root/Aether/docker-compose.yml \
--compose-dir /opt/aether-single \
--replace-existing
```
Equivalently, call the lower-level script and pass each target path explicitly:
```bash
scripts/migrate-pg-compose-to-single-node.sh \
--source-compose /root/Aether/docker-compose.yml \
--target-compose /opt/aether-single/docker-compose.single-node.yml \
--target-env /opt/aether-single/.env.single-node \
--target-db /opt/aether-single/data/aether.db \
--replace-existing
```
During cutover, the script stops and removes only the source `app` container to
free the fixed `aether-app` container name. Postgres, Redis, and their volumes
remain in place for rollback.
Defaults:
| Setting | Default |
| --- | --- |
| Source Compose | `docker-compose.yml` |
| Single Node install root | `/opt/aether` |
| Single Node config dir | `/etc/aether` |
| Target SQLite DB | `/opt/aether/data/aether.db` |
| Source app service | `app` |
| Source Postgres service | `postgres` |
| Single Node service | `aether-gateway` |
The script writes migration artifacts under `./data/pg-to-single-node-<timestamp>` next
to the source Compose file unless `--work-dir` is provided.
## Rollback
The script leaves the original Postgres and Redis volumes in place. If cutover
finishes but you need to roll back:
```bash
sudo systemctl stop aether-gateway
cd /root/Aether
docker compose -f docker-compose.yml up -d app
```
For the Compose single-node script, rollback is the same idea: start the app
again from the original Postgres compose file.
If the migration fails before cutover completes, the script attempts to restart
the source `app` service automatically. Pass `--keep-source-stopped-on-error` if
you want to inspect the stopped source deployment manually instead.
## Data Coverage Guard
The migration does not maintain a separate business-domain table list. The
target single-node image first builds a temporary SQLite database with its
normal migrations, then `aether-gateway copy` reads that SQLite schema and copies
matching public Postgres tables and columns.
If the source Postgres database has a non-empty public table that does not exist
in the target SQLite schema, the copy stops instead of silently dropping it. It
ignores lifecycle metadata tables such as `_sqlx_migrations` and
`schema_backfills`. Extra source columns that are absent from the target schema
are not copied.
## Request Body Detail Policy
The production migration migrates all migratable data by default. The only
optional exclusion is request body detail data.
When you choose to skip request bodies, the migration does not copy
`usage_body_blobs`, `usage_http_audits`, or legacy `usage` request body columns
such as `request_body`, `provider_request_body`, `response_body`,
`client_response_body`, and `*_body_compressed`.
Interactive installation lets you choose:
```text
1) Full migration: migrate all migratable data, including request body details
2) Skip request bodies: migrate all other data; skip only request body large fields and HTTP body detail tables; source PG is unchanged
```
For non-interactive full runs:
```bash
scripts/migrate-pg-to-single-node.sh \
--request-body-mode full
```
For non-interactive omit runs:
```bash
scripts/migrate-pg-to-single-node.sh \
--request-body-mode omit
```
`omit` only skips writing those large artifacts and detail tables into the
target SQLite database. It does not delete or clear the source Postgres data.
## Notes
- Single Node requires root or sudo because it writes `/opt/aether`, `/etc/aether`, and
the system service definition.
- The script does not decrypt or re-encrypt provider keys. It preserves the
original encryption key and moves encrypted data as-is.
- Existing target SQLite databases, including `-wal`/`-shm` sidecars, are not
replaced unless `--replace-existing` is provided.
- Disk space checks use `pg_database_size(current_database()) * 2 + 1 GiB` as the
conservative estimate for one SQLite copy. If the work directory and target DB
directory are on the same filesystem, the script requires enough space for both
the temporary and final SQLite files. With `--request-body-mode omit`, the
estimate subtracts `usage_body_blobs` and `usage_http_audits` relation sizes.
- For non-standard source Compose files, set `--app-service` and
`--postgres-service` to match the service names.
@@ -1,221 +0,0 @@
# Postgres 到 Aether Single Node 迁移
英文版:[pg-to-single-node-migration.md](pg-to-single-node-migration.md)
本文档用于把现有 Docker Compose Postgres 部署迁移到 Aether
single-node。当前版本里,**single-node** 指默认 SQLite 安装模式:
`install.sh --mode single-node`,也就是系统服务加 SQLite。Docker Compose
单机模板是 `docker-compose.single-node.yml`,安装脚本入口是
`--mode compose-single-node`
迁移脚本:
```bash
scripts/migrate-pg-to-single-node.sh
```
如果目标形态仍然要保持 Docker Compose,而不是系统服务,使用镜像版迁移脚本:
```bash
scripts/migrate-pg-compose-to-single-node.sh
```
两种迁移脚本都会先拉取/安装目标 single-node 版本,再停止源 `app`,把 Postgres
记录直接写入临时 SQLite DB,不落 JSONL 中间文件;复制成功后替换目标
`aether.db`,最后启动 single-node。
也可以直接用安装脚本作为统一入口,由 `--mode` 选择迁移目标:
```bash
# 交互式执行时,先选择目标部署模式:
# 1) Docker Compose 标准部署(Postgres + Redis
# 2) Docker Compose 单节点部署(SQLite
# 3) 系统服务单节点部署(SQLite)
# 选择 2 或 3 后,再选择数据初始化方式:
# 1) 全新初始化(不迁移现有数据)
# 2) 从现有 Docker Compose PG 数据库迁移
install.sh
# 迁移到新的 single-node Docker Compose 目录
install.sh \
--mode compose-single-node \
--migrate-from-compose /root/Aether/docker-compose.yml \
--compose-dir /opt/aether-single \
--replace-existing
# 迁移到系统服务 + SQLite
sudo install.sh \
--mode single-node \
--migrate-from-compose /root/Aether/docker-compose.yml \
--replace-existing
```
交互模式会先选择目标部署形态。如果目标是 `compose-single-node`
`single-node`,安装脚本会再询问数据初始化方式:全新初始化,或从现有 Docker
Compose PG 数据库迁移。选择迁移后,脚本会通过 `docker compose ls` 自动探测源
PG Compose 文件,并确认该 Compose 配置里存在默认的 `app``postgres` 服务;
如果能唯一识别,会作为默认值带入提示。探测不到或存在多个候选时会直接中止;
此时请用 `--migrate-from-compose` 显式指定源 compose 路径。
安装脚本只是统一参数入口:`compose-single-node` 会委托给
`scripts/migrate-pg-compose-to-single-node.sh``single-node` 会委托给
`scripts/migrate-pg-to-single-node.sh`
## 迁移内容
脚本会尽量缩短生产停机窗口:
1. 读取源 Compose 目录下的 `.env`
2. 生成 single-node 环境文件,保留 `JWT_SECRET_KEY``ENCRYPTION_KEY`
`AETHER_GATEWAY_DATA_ENCRYPTION_KEY`、管理员配置、端口和应用配置。
3. 执行 `install.sh --mode single-node --skip-start`,提前安装 single-node
release,但不启动服务。
4. 使用已安装的 single-node 二进制预检 SQLite schema migration。
5. 拉取目标 single-node 镜像,确认其 `copy` 命令可用,并检查源 `app`
当前运行镜像 ID 与目标镜像 ID 一致。
6. 以目标 SQLite schema 作为迁移计划:把源 Postgres 中同名表、同名字段
复制到临时 SQLite DB。
7. 检查请求体明细迁移策略;默认全部迁移,也可以选择只跳过请求体明细。
8. 检查 work-dir 和目标 SQLite 目录是否有足够空间容纳临时库和正式库。
9. 只停止源 Compose 的 `app` 服务,保留 Postgres 和 Redis 运行,方便回滚。
10. 从源 Postgres 直接复制记录到临时 SQLite 数据库,不生成 JSONL 中间文件。
11. 复制完成后替换目标 SQLite DB,包括 SQLite `-wal``-shm` 边车文件,
然后启动 single-node 系统服务。
镜像一致性检查比较的是 Docker 镜像 ID,不只是 tag 字符串。即使源和目标都写着
`latest`,只要实际镜像 ID 不同,迁移也会中止。请先把源 PG Compose 的 `app`
升级到目标 single-node 相同版本,确认运行正常后再迁移。迁移脚本也会检查目标镜像
是否支持直接 copy 和请求体跳过开关;如果只是换了脚本但镜像还是旧版本,脚本会
直接中止,避免漏迁。
## 生产切换
切换前先做一次常规服务器备份或快照。确认后执行:
```bash
sudo scripts/migrate-pg-to-single-node.sh \
--source-compose /root/Aether/docker-compose.yml \
--replace-existing
```
如果要迁移到 Docker Compose single-node,而不是系统服务:
```bash
scripts/migrate-pg-compose-to-single-node.sh \
--source-compose /root/Aether/docker-compose.yml \
--replace-existing
```
源 Postgres Compose 目录和目标 single-node Compose 目录可以不一样。例如:
```bash
install.sh \
--mode compose-single-node \
--migrate-from-compose /root/Aether/docker-compose.yml \
--compose-dir /opt/aether-single \
--replace-existing
```
等价地,也可以直接调底层脚本并显式传入每个目标路径:
```bash
scripts/migrate-pg-compose-to-single-node.sh \
--source-compose /root/Aether/docker-compose.yml \
--target-compose /opt/aether-single/docker-compose.single-node.yml \
--target-env /opt/aether-single/.env.single-node \
--target-db /opt/aether-single/data/aether.db \
--replace-existing
```
切换时脚本只会停止并移除源 `app` 容器,用来释放固定的 `aether-app`
容器名;Postgres、Redis 和它们的 volume 都会保留,方便回滚。
默认路径和服务名:
| 配置项 | 默认值 |
| --- | --- |
| 源 Compose 文件 | `docker-compose.yml` |
| single-node 安装目录 | `/opt/aether` |
| single-node 配置目录 | `/etc/aether` |
| 目标 SQLite DB | `/opt/aether/data/aether.db` |
| 源 app 服务 | `app` |
| 源 Postgres 服务 | `postgres` |
| single-node 服务 | `aether-gateway` |
除非显式传入 `--work-dir`,脚本会把迁移产物写到源 Compose 文件旁边的
`./data/pg-to-single-node-<timestamp>`
## 回滚
脚本会保留原 Postgres 和 Redis volume。迁移已经完成但需要回滚时:
```bash
sudo systemctl stop aether-gateway
cd /root/Aether
docker compose -f docker-compose.yml up -d app
```
对于 Compose single-node 脚本,回滚思路相同:重新用原 Postgres compose 文件
拉起 `app`
如果迁移在切换完成前失败,脚本默认会尝试自动拉起源 `app` 服务。需要失败后
保持源应用停止以便人工排查时,增加:
```bash
--keep-source-stopped-on-error
```
## 数据覆盖保护
迁移不再维护一份额外的业务表清单。目标 single-node 镜像会先用正常
migrations 建出临时 SQLite 数据库,然后 `aether-gateway copy` 读取这个
SQLite schema,把源 Postgres 里同名表、同名字段复制过去。
如果源 Postgres 里存在非空 public 表,但目标 SQLite schema 中没有同名表,
copy 会直接中止,不会静默丢弃。生命周期元数据表 `_sqlx_migrations`
`schema_backfills` 会被忽略。源表中存在但目标 SQLite 不存在的额外字段不会复制。
## 请求体明细策略
single-node SQLite 生产迁移默认迁移所有可迁移数据,唯一可选的跳过项是请求体明细。
选择“不迁移请求体”时,不会迁移 `usage_body_blobs``usage_http_audits`,也不会迁移 `usage`
表里的 `request_body` / `provider_request_body` / `response_body` /
`client_response_body` / `*_body_compressed` 等请求体大字段。
交互安装时可以选择:
```text
1) 全部迁移:迁移所有可迁移数据,包括请求体明细
2) 不迁移请求体:迁移其他所有数据;仅跳过请求体大字段和 HTTP 请求体明细,源 PG 不清除
```
非交互执行时,全部迁移可以显式指定:
```bash
scripts/migrate-pg-to-single-node.sh \
--request-body-mode full
```
不迁移请求体可以显式指定:
```bash
scripts/migrate-pg-to-single-node.sh \
--request-body-mode omit
```
`omit` 只是不把这些大字段和明细表写进目标 SQLite,不会删除或清空源 Postgres。
## 注意事项
- single-node 安装需要 root 或 sudo 权限,因为会写入 `/opt/aether`
`/etc/aether` 和系统服务定义。
- 脚本不会解密或重新加密供应商密钥;它会沿用源环境的加密密钥,并原样迁移已加密数据。
- 已存在的目标 SQLite DB,包括 `-wal``-shm` 边车文件,只有在传入
`--replace-existing` 时才会被替换。
- 空间检查会用 `pg_database_size(current_database()) * 2 + 1 GiB` 作为单份
SQLite 的保守估算。如果 work-dir 和目标 DB 目录在同一个文件系统,会要求同时
容纳临时 SQLite 和正式 SQLite。选择 `--request-body-mode omit` 时,
估算会扣除 `usage_body_blobs``usage_http_audits` 的表空间。
- 非标准 Compose 服务名需要通过 `--app-service``--postgres-service`
明确指定。
+68
View File
@@ -96,6 +96,7 @@ export interface UsersExportData {
user_groups?: UserGroupExport[]
users: UserExport[]
standalone_keys?: StandaloneKeyExport[]
usage_aggregates?: UsageAggregateSnapshot
}
export interface AggregateExportData {
@@ -120,6 +121,7 @@ export interface UserGroupExport {
}
export interface UserExport {
id?: string
email: string
email_verified?: boolean
username: string
@@ -144,6 +146,7 @@ export interface UserExport {
}
export interface UserApiKeyExport {
api_key_id?: string
key?: string | null
key_hash: string
key_encrypted?: string | null
@@ -161,12 +164,76 @@ export interface UserApiKeyExport {
expires_at?: string | null
auto_delete_on_expiry?: boolean
total_requests?: number
total_tokens?: number
total_cost_usd?: number
}
// 独立余额 Key 导出结构(与 UserApiKeyExport 相同,但不包含 is_standalone
export type StandaloneKeyExport = Omit<UserApiKeyExport, 'is_standalone'>
export interface StatsDailyAggregateExport {
date_unix_secs: number
total_requests: number
success_requests: number
error_requests: number
input_tokens: number
output_tokens: number
cache_creation_tokens: number
cache_read_tokens: number
total_cost: number
actual_total_cost: number
is_complete: boolean
aggregated_at_unix_secs?: number | null
}
export interface StatsUserDailyAggregateExport {
user_id: string
username?: string | null
date_unix_secs: number
total_requests: number
success_requests: number
error_requests: number
input_tokens: number
output_tokens: number
cache_creation_tokens: number
cache_read_tokens: number
total_cost: number
}
export interface StatsDailyApiKeyAggregateExport {
api_key_id: string
api_key_name?: string | null
date_unix_secs: number
total_requests: number
success_requests: number
error_requests: number
input_tokens: number
output_tokens: number
cache_creation_tokens: number
cache_read_tokens: number
total_cost: number
}
export interface UsageAggregateSnapshot {
stats_daily?: StatsDailyAggregateExport[]
stats_user_daily?: StatsUserDailyAggregateExport[]
stats_daily_api_key?: StatsDailyApiKeyAggregateExport[]
}
export interface UsageAggregateImportCounter {
created: number
updated: number
skipped: number
}
export interface UsageAggregateImportSummary {
stats_daily: UsageAggregateImportCounter
stats_user_daily: UsageAggregateImportCounter
stats_daily_api_key: UsageAggregateImportCounter
skipped_unmapped_user_daily: number
skipped_unmapped_api_key_daily: number
}
export interface GlobalModelExport {
name: string
display_name: string
@@ -522,6 +589,7 @@ export interface UsersImportResponse {
users: { created: number; updated: number; skipped: number }
api_keys: { created: number; updated?: number; skipped: number }
standalone_keys?: { created: number; updated?: number; skipped: number }
usage_aggregates?: UsageAggregateImportSummary
errors: string[]
}
}
@@ -1066,7 +1066,7 @@
</div>
<div>
<div class="flex items-center justify-between text-[10px] mb-0.5">
<span class="text-muted-foreground">使用额度</span>
<span class="text-muted-foreground">剩余额度</span>
<span :class="getQuotaRemainingClass(getChatGPTWebQuotaUsedPercent(key))">
{{ getChatGPTWebQuotaRemainingPercent(key).toFixed(1) }}%
</span>
@@ -1080,7 +1080,7 @@
</div>
<div class="flex items-center justify-between text-[9px] text-muted-foreground/70 mt-0.5">
<span>
{{ formatChatGPTWebUsage(getChatGPTWebQuotaDisplay(key)?.image_quota_used) }} /
{{ formatChatGPTWebUsage(getChatGPTWebQuotaDisplay(key)?.image_quota_remaining) }} /
{{ formatChatGPTWebUsage(getChatGPTWebQuotaDisplay(key)?.image_quota_total) }}
</span>
<span v-if="getChatGPTWebQuotaDisplay(key)?.image_quota_reset_at">
@@ -105,6 +105,28 @@ describe('providerKeyQuota', () => {
}, 'grok')).toBe('Auto剩余 40.0% (60/150) | Heavy剩余 0.0% (0/20)')
})
it('formats ChatGPT Web image quota as remaining count', () => {
expect(getQuotaDisplayText({
status_snapshot: {
quota: {
provider_type: 'chatgpt_web',
code: 'ok',
exhausted: false,
windows: [
{
code: 'image_gen',
scope: 'account',
remaining_ratio: 0.96,
used_value: 1,
remaining_value: 24,
limit_value: 25,
},
],
},
},
}, 'chatgpt_web')).toBe('生图剩余 24/25')
})
it('surfaces Windsurf hard account states', () => {
expect(getQuotaDisplayText({
status_snapshot: {
+4 -10
View File
@@ -343,19 +343,13 @@ function getChatGPTWebQuotaText(quota: QuotaStatusSnapshot): string | null {
if (!window) return normalizeText(quota.label)
const remainingPercent = getQuotaWindowRemainingPercent(window)
if (typeof window.remaining_value === 'number' && typeof window.limit_value === 'number' && window.limit_value > 0 && window.remaining_value <= 0) {
return `生图剩余 ${formatQuotaValue(window.remaining_value)}/${formatQuotaValue(window.limit_value)}`
}
if (remainingPercent != null) {
if (typeof window.used_value === 'number' && typeof window.limit_value === 'number' && window.limit_value > 0) {
return `生图剩余 ${formatPercent(remainingPercent)} (${formatQuotaValue(window.used_value)}/${formatQuotaValue(window.limit_value)})`
}
return `生图剩余 ${formatPercent(remainingPercent)}`
}
if (typeof window.remaining_value === 'number' && typeof window.limit_value === 'number' && window.limit_value > 0) {
return `生图剩余 ${formatQuotaValue(window.remaining_value)}/${formatQuotaValue(window.limit_value)}`
}
if (remainingPercent != null) {
return `生图剩余 ${formatPercent(remainingPercent)}`
}
if (typeof window.remaining_value === 'number') {
return `生图剩余 ${formatQuotaValue(window.remaining_value)}`
}
+5 -5
View File
@@ -3700,7 +3700,7 @@ function getQuotaProgressLabel(label: string): string {
}
function getQuotaProgressCountdown(item: QuotaProgressItem) {
if (!['日', '5H', '周', 'Spark5H', 'Spark周', 'Auto', 'Fast', 'Expert', 'Heavy', 'Grok 4.3'].includes(item.label)) return null
if (!['日', '5H', '周', 'Spark5H', 'Spark周', 'Auto', 'Fast', 'Expert', 'Heavy', 'Grok 4.3', '生图'].includes(item.label)) return null
if (item.resetAtSeconds == null && item.resetSeconds == null) return null
return getCodexResetCountdown(
item.resetAtSeconds,
@@ -4109,10 +4109,10 @@ function buildQuotaProgressItemsFromSnapshot(key: PoolKeyDetail): QuotaProgressI
const remainingValue = typeof window?.remaining_value === 'number' ? window.remaining_value : null
const limitValue = typeof window?.limit_value === 'number' ? window.limit_value : null
const usedValue = typeof window?.used_value === 'number' ? window.used_value : null
const detail = usedValue != null && limitValue != null
? `${formatQuotaValue(usedValue)}/${formatQuotaValue(limitValue)}`
: remainingValue != null && limitValue != null
? `${formatQuotaValue(Math.max(limitValue - remainingValue, 0))}/${formatQuotaValue(limitValue)}`
const detail = remainingValue != null && limitValue != null
? `${formatQuotaValue(remainingValue)}/${formatQuotaValue(limitValue)}`
: usedValue != null && limitValue != null
? `${formatQuotaValue(Math.max(limitValue - usedValue, 0))}/${formatQuotaValue(limitValue)}`
: remainingValue != null
? `剩余 ${formatQuotaValue(remainingValue)}`
: undefined
@@ -643,6 +643,47 @@ describe('PoolManagement Codex cycle stats mode', () => {
expect(root.querySelectorAll('button[title="查看评分计算结果"]').length).toBeGreaterThan(0)
})
it('shows ChatGPT Web image quota reset countdown above the quota bar', async () => {
const chatgptWebKey = createPoolKey('chatgpt_web', {
api_formats: ['openai:image'],
status_snapshot: {
oauth: { code: 'valid' },
account: { code: 'ok', blocked: false },
quota: {
code: 'ok',
exhausted: false,
provider_type: 'chatgpt_web',
updated_at: 1_700_000_000,
windows: [
{
code: 'image_gen',
label: '生图',
scope: 'account',
remaining_ratio: 0.96,
remaining_value: 24,
limit_value: 25,
reset_seconds: 3600,
},
],
},
},
})
endpointMocks.getPoolOverview.mockResolvedValue({ items: [createOverview('chatgpt_web')] })
endpointMocks.listPoolKeys.mockResolvedValue(createKeyPage(chatgptWebKey))
endpointMocks.getProvider.mockResolvedValue(createProvider('chatgpt_web', {
api_formats: ['openai:image'],
}))
const root = mountPoolManagement()
await settle()
const resetTexts = Array.from(root.querySelectorAll('[data-testid="pool-quota-reset-text"]'))
.map((element) => element.textContent?.trim())
.filter(Boolean)
expect(resetTexts).toContain('1h')
expect(root.textContent).toContain('生图')
})
it('opens only one score popover across desktop and mobile layouts', async () => {
const scoredKey = createPoolKey('codex', {
pool_score: {
@@ -42,6 +42,9 @@
<li v-if="aggregateImportPreview.user_data.standalone_keys?.length">
独立余额 Keys: {{ aggregateImportPreview.user_data.standalone_keys.length }}
</li>
<li v-if="usageAggregatePreviewCounts.total > 0">
统计聚合: {{ usageAggregatePreviewCounts.total }}
</li>
</ul>
</div>
</div>
@@ -84,7 +87,7 @@
</div>
<p class="text-xs text-muted-foreground">
注意完整备份会先导入配置数据再导入用户数据文件包含用户用户组API Keys 与钱包快照用户 API Keys 需要目标系统使用相同 ENCRYPTION_KEY
注意完整备份会先导入配置数据再导入用户数据文件包含用户用户组API KeysKey 用量钱包快照与统计聚合正常导出的 API Keys 会在导入时使用目标系统密钥重新加密仅当备份中包含 key_encrypted 等未解密密文字段时需要目标系统使用兼容 ENCRYPTION_KEY
</p>
<div
@@ -149,6 +152,9 @@
用户创建 {{ aggregateImportResult.users.stats.users.created }}
API Keys 创建 {{ aggregateImportResult.users.stats.api_keys.created }}
跳过 {{ aggregateImportResult.users.stats.users.skipped }} 个用户
<template v-if="usageAggregateResultText">
统计聚合 {{ usageAggregateResultText }}
</template>
</p>
</div>
</div>
@@ -217,4 +223,36 @@ const warningMessages = computed(() => {
const userErrors = props.aggregateImportResult.users.stats.errors.map((message) => `用户数据: ${message}`)
return [...configErrors, ...userErrors]
})
const usageAggregatePreviewCounts = computed(() => {
const aggregates = props.aggregateImportPreview?.user_data.usage_aggregates
const statsDaily = aggregates?.stats_daily?.length ?? 0
const statsUserDaily = aggregates?.stats_user_daily?.length ?? 0
const statsDailyApiKey = aggregates?.stats_daily_api_key?.length ?? 0
return {
statsDaily,
statsUserDaily,
statsDailyApiKey,
total: statsDaily + statsUserDaily + statsDailyApiKey,
}
})
const usageAggregateResultText = computed(() => {
const aggregates = props.aggregateImportResult?.users.stats.usage_aggregates
if (!aggregates) return ''
const counters = [
aggregates.stats_daily,
aggregates.stats_user_daily,
aggregates.stats_daily_api_key,
]
const created = counters.reduce((sum, item) => sum + item.created, 0)
const updated = counters.reduce((sum, item) => sum + item.updated, 0)
const skipped = counters.reduce((sum, item) => sum + item.skipped, 0)
const total = created + updated + skipped
if (total === 0) return ''
return `创建 ${created},更新 ${updated},跳过 ${skipped}`
})
</script>
-568
View File
@@ -61,19 +61,6 @@ ADMIN_PASSWORD_SOURCE=""
UI_LANG="${AETHER_LANG:-${AETHER_LANGUAGE:-auto}}"
RELEASE_KEEP="${AETHER_RELEASE_KEEP:-3}"
RELEASE_ARCHIVE_URL="${AETHER_RELEASE_ARCHIVE_URL:-${AETHER_DOWNLOAD_URL:-}}"
MIGRATE_FROM_COMPOSE=""
MIGRATE_TARGET_COMPOSE=""
MIGRATE_TARGET_ENV=""
MIGRATE_TARGET_DB=""
MIGRATE_WORK_DIR=""
MIGRATE_APP_SERVICE=""
MIGRATE_POSTGRES_SERVICE=""
MIGRATE_SINGLE_NODE_SERVICE=""
MIGRATE_REPLACE_EXISTING="false"
MIGRATE_REPLACE_TARGET_COMPOSE="false"
MIGRATE_KEEP_SOURCE_STOPPED_ON_ERROR="false"
MIGRATE_INTERACTIVE="false"
MIGRATE_REQUEST_BODY_MODE=""
usage() {
cat <<'EOF'
@@ -104,31 +91,8 @@ Options:
--lang LANG Installer language: zh or en
--skip-start Install files, but do not start Docker Compose or restart the service
--keep-releases N Keep the latest N releases, prune older ones (default: 3, 0=disable)
--migrate-from-compose PATH
Migrate an existing Postgres Compose deployment into the selected single-node mode
--target-compose PATH
Migration target compose file for --mode compose-single-node
--target-env PATH Migration target env file for --mode compose-single-node
--target-db PATH Migration target SQLite DB path
--work-dir PATH Migration working directory
--app-service NAME Source compose app service for migration
--postgres-service NAME
Source compose Postgres service for migration
--single-node-service NAME
Target compose service for --mode compose-single-node migration
--replace-existing Allow replacing an existing target SQLite DB during migration
--replace-target-compose
Overwrite target compose file from the single-node template during migration
--request-body-mode MODE
Request/response body detail handling during migration: full/1 or omit/2
--keep-source-stopped-on-error
Do not auto-restart source app if migration fails after stopping it
-h, --help Show this help
Migration examples:
install.sh --mode compose-single-node --migrate-from-compose /root/Aether/docker-compose.yml --compose-dir /opt/aether-single --replace-existing
sudo install.sh --mode single-node --migrate-from-compose /root/Aether/docker-compose.yml --replace-existing
Environment overrides:
AETHER_REPO, AETHER_SOURCE_REF, AETHER_INSTALL_MODE, AETHER_CHANNEL, AETHER_VERSION
AETHER_LANG or AETHER_LANGUAGE
@@ -312,63 +276,6 @@ parse_args() {
RELEASE_KEEP="$2"
shift 2
;;
--migrate-from-compose)
[[ $# -ge 2 ]] || die "--migrate-from-compose requires a path"
MIGRATE_FROM_COMPOSE="$2"
shift 2
;;
--target-compose)
[[ $# -ge 2 ]] || die "--target-compose requires a path"
MIGRATE_TARGET_COMPOSE="$2"
shift 2
;;
--target-env)
[[ $# -ge 2 ]] || die "--target-env requires a path"
MIGRATE_TARGET_ENV="$2"
shift 2
;;
--target-db)
[[ $# -ge 2 ]] || die "--target-db requires a path"
MIGRATE_TARGET_DB="$2"
shift 2
;;
--work-dir)
[[ $# -ge 2 ]] || die "--work-dir requires a path"
MIGRATE_WORK_DIR="$2"
shift 2
;;
--app-service)
[[ $# -ge 2 ]] || die "--app-service requires a service name"
MIGRATE_APP_SERVICE="$2"
shift 2
;;
--postgres-service)
[[ $# -ge 2 ]] || die "--postgres-service requires a service name"
MIGRATE_POSTGRES_SERVICE="$2"
shift 2
;;
--single-node-service)
[[ $# -ge 2 ]] || die "--single-node-service requires a service name"
MIGRATE_SINGLE_NODE_SERVICE="$2"
shift 2
;;
--replace-existing)
MIGRATE_REPLACE_EXISTING="true"
shift
;;
--replace-target-compose)
MIGRATE_REPLACE_TARGET_COMPOSE="true"
shift
;;
--request-body-mode)
[[ $# -ge 2 ]] || die "--request-body-mode requires a value"
MIGRATE_REQUEST_BODY_MODE="$2"
shift 2
;;
--keep-source-stopped-on-error)
MIGRATE_KEEP_SOURCE_STOPPED_ON_ERROR="true"
shift
;;
-h|--help)
usage
exit 0
@@ -610,44 +517,6 @@ EOF
fi
;;
esac
if [[ -z "${MIGRATE_FROM_COMPOSE}" && "${MODE}" != "compose" ]]; then
if ui_is_zh; then
cat >/dev/tty <<'EOF'
请选择数据初始化方式:
1) 全新初始化(不迁移现有数据)
2) 从现有 Docker Compose PG 数据库迁移
请输入选项 [1]:
EOF
else
cat >/dev/tty <<'EOF'
Choose data initialization mode:
1) Fresh initialization (do not migrate existing data)
2) Migrate from an existing Docker Compose PG database
Enter choice [1]:
EOF
fi
local init_choice
IFS= read -r init_choice </dev/tty || init_choice=""
case "${init_choice:-1}" in
1)
;;
2)
MIGRATE_INTERACTIVE="true"
;;
*)
if ui_is_zh; then
die "无效的数据初始化方式选项: ${init_choice}"
else
die "invalid data initialization choice: ${init_choice}"
fi
;;
esac
fi
else
MODE="single-node"
fi
@@ -1010,434 +879,6 @@ start_compose_deployment() {
run_compose "${compose_args[@]}" up -d
}
migration_options_requested() {
[[ -n "${MIGRATE_TARGET_COMPOSE}" ]] && return 0
[[ -n "${MIGRATE_TARGET_ENV}" ]] && return 0
[[ -n "${MIGRATE_TARGET_DB}" ]] && return 0
[[ -n "${MIGRATE_WORK_DIR}" ]] && return 0
[[ -n "${MIGRATE_APP_SERVICE}" ]] && return 0
[[ -n "${MIGRATE_POSTGRES_SERVICE}" ]] && return 0
[[ -n "${MIGRATE_SINGLE_NODE_SERVICE}" ]] && return 0
[[ "${MIGRATE_REPLACE_EXISTING}" == "true" ]] && return 0
[[ "${MIGRATE_REPLACE_TARGET_COMPOSE}" == "true" ]] && return 0
[[ -n "${MIGRATE_REQUEST_BODY_MODE}" ]] && return 0
[[ "${MIGRATE_KEEP_SOURCE_STOPPED_ON_ERROR}" == "true" ]] && return 0
return 1
}
normalize_migration_request_body_mode() {
case "${MIGRATE_REQUEST_BODY_MODE}" in
""|1|full|all|include)
MIGRATE_REQUEST_BODY_MODE="full"
;;
2|omit|skip)
MIGRATE_REQUEST_BODY_MODE="omit"
;;
*)
die "--request-body-mode must be full/1 or omit/2"
;;
esac
}
prompt_with_default() {
local prompt="$1"
local default_value="$2"
local value
if [[ -n "${default_value}" ]]; then
printf '%s [%s]: ' "${prompt}" "${default_value}" >/dev/tty
else
printf '%s: ' "${prompt}" >/dev/tty
fi
IFS= read -r value </dev/tty || value=""
if [[ -z "${value}" ]]; then
printf '%s\n' "${default_value}"
else
printf '%s\n' "${value}"
fi
}
prompt_yes_no() {
local prompt="$1"
local default_value="$2"
local suffix choice
case "${default_value}" in
yes)
if ui_is_zh; then
suffix="[y/n,默认 y]"
else
suffix="[y/n, default y]"
fi
;;
*)
default_value="no"
if ui_is_zh; then
suffix="[y/n,默认 n]"
else
suffix="[y/n, default n]"
fi
;;
esac
while true; do
printf '%s %s: ' "${prompt}" "${suffix}" >/dev/tty
IFS= read -r choice </dev/tty || choice=""
choice="$(printf '%s' "${choice}" | tr '[:upper:]' '[:lower:]')"
case "${choice:-${default_value}}" in
y|yes)
return 0
;;
n|no)
return 1
;;
*)
if ui_is_zh; then
echo "请输入 y 或 n。" >/dev/tty
else
echo "Enter y or n." >/dev/tty
fi
;;
esac
done
}
docker_compose_ls_config_files() {
local output
output="$(docker compose ls --format json 2>/dev/null || true)"
if [[ -n "${output}" && "${output}" == *ConfigFiles* ]]; then
printf '%s' "${output}" |
tr '{' '\n' |
sed -n 's/.*"ConfigFiles"[[:space:]]*:[[:space:]]*"\([^"]*\)".*/\1/p'
return
fi
docker compose ls 2>/dev/null | awk 'NR > 1 && NF > 0 { print $NF }'
}
compose_file_has_source_services() {
local compose_file="$1"
local app_service="${MIGRATE_APP_SERVICE:-app}"
local postgres_service="${MIGRATE_POSTGRES_SERVICE:-postgres}"
local services
services="$(docker compose -f "${compose_file}" config --services 2>/dev/null || true)"
[[ -n "${services}" ]] || return 1
printf '%s\n' "${services}" | grep -Fxq "${app_service}" || return 1
printf '%s\n' "${services}" | grep -Fxq "${postgres_service}" || return 1
}
append_unique_candidate() {
local candidate="$1"
shift
local existing
for existing in "$@"; do
[[ "${existing}" != "${candidate}" ]] || return 1
done
printf '%s\n' "${candidate}"
}
detect_source_compose_from_docker_compose_ls() {
local config_files compose_file candidate
local -a candidates=()
command -v docker >/dev/null 2>&1 || return 0
docker compose version >/dev/null 2>&1 || return 0
while IFS= read -r config_files || [[ -n "${config_files}" ]]; do
config_files="$(trim_whitespace "${config_files}")"
[[ -n "${config_files}" ]] || continue
# The migration scripts currently accept one source compose file. If the
# source project was launched with multiple compose files, ask explicitly.
[[ "${config_files}" != *,* ]] || continue
compose_file="${config_files}"
[[ -f "${compose_file}" ]] || continue
compose_file="$(absolute_path "${compose_file}")"
compose_file_has_source_services "${compose_file}" || continue
candidate="$(append_unique_candidate "${compose_file}" "${candidates[@]}" || true)"
[[ -z "${candidate}" ]] || candidates+=("${candidate}")
done < <(docker_compose_ls_config_files)
case "${#candidates[@]}" in
0)
return 0
;;
1)
printf '%s\n' "${candidates[0]}"
;;
*)
if interactive_tty_available; then
if ui_is_zh; then
echo "从 docker compose ls 找到多个可能的源 Compose,无法安全自动选择:" >/dev/tty
else
echo "docker compose ls found multiple possible source Compose files and cannot choose safely:" >/dev/tty
fi
printf ' %s\n' "${candidates[@]}" >/dev/tty
fi
return 0
;;
esac
}
collect_interactive_migration_options() {
local detected_source
local prompt
local source_compose_abs
local source_compose_dir
[[ "${MIGRATE_INTERACTIVE}" == "true" ]] || return
interactive_tty_available || die "interactive migration selection requires a terminal"
if ui_is_zh; then
cat >/dev/tty <<'EOF'
迁移会先做预检并拉取/安装目标 single-node,再在切换窗口停止源 app。
源 Postgres 和 Redis 会保留,便于回滚。
EOF
else
cat >/dev/tty <<'EOF'
Migration will preflight and pull/install the target single-node release first.
During cutover it stops only the source app. Source Postgres and Redis remain for rollback.
EOF
fi
if [[ -z "${MIGRATE_FROM_COMPOSE}" ]]; then
detected_source="$(detect_source_compose_from_docker_compose_ls || true)"
if [[ -z "${detected_source}" ]]; then
if ui_is_zh; then
die "未能通过 docker compose ls 唯一识别源 PG Compose;请使用 --migrate-from-compose 显式指定"
else
die "could not uniquely detect source PG Compose from docker compose ls; pass --migrate-from-compose explicitly"
fi
fi
if ui_is_zh; then
printf '已通过 docker compose ls 探测到源 Compose: %s\n' "${detected_source}" >/dev/tty
else
printf 'Detected source Compose from docker compose ls: %s\n' "${detected_source}" >/dev/tty
fi
if ui_is_zh; then
prompt="确认使用该源 Compose 进行迁移"
else
prompt="Use this source Compose for migration"
fi
if prompt_yes_no "${prompt}" "yes"; then
MIGRATE_FROM_COMPOSE="${detected_source}"
else
if ui_is_zh; then
die "已取消迁移;如需指定其他源 Compose,请使用 --migrate-from-compose"
else
die "migration cancelled; pass --migrate-from-compose to use another source Compose"
fi
fi
fi
[[ -n "${MIGRATE_FROM_COMPOSE}" ]] || die "--migrate-from-compose cannot be empty"
source_compose_abs="$(absolute_path "${MIGRATE_FROM_COMPOSE}")"
source_compose_dir="$(dirname "${source_compose_abs}")"
if [[ "${MODE}" == "compose-single-node" ]]; then
if [[ "${COMPOSE_DIR_EXPLICIT}" != "true" ]]; then
COMPOSE_DIR="${source_compose_dir}-single-node"
fi
if ui_is_zh; then
printf '已自动选择目标 single-node Compose 目录: %s\n' "${COMPOSE_DIR}" >/dev/tty
prompt="确认使用该目标目录"
else
printf 'Selected target single-node Compose directory: %s\n' "${COMPOSE_DIR}" >/dev/tty
prompt="Use this target directory"
fi
if ! prompt_yes_no "${prompt}" "yes"; then
if ui_is_zh; then
die "已取消迁移;如需指定其他目标目录,请使用 --compose-dir"
else
die "migration cancelled; pass --compose-dir to use another target directory"
fi
fi
fi
if [[ "${MIGRATE_REPLACE_EXISTING}" != "true" ]]; then
if ui_is_zh; then
prompt="如果目标 SQLite 已存在,是否允许备份后替换"
else
prompt="If the target SQLite DB already exists, allow backup and replacement"
fi
if prompt_yes_no "${prompt}" "no"; then
MIGRATE_REPLACE_EXISTING="true"
fi
fi
if [[ -z "${MIGRATE_REQUEST_BODY_MODE}" ]]; then
if ui_is_zh; then
cat >/dev/tty <<'EOF'
请求体明细迁移策略:
1) 全部迁移:迁移所有可迁移数据,包括请求体明细
2) 不迁移请求体:迁移其他所有数据;仅跳过请求体大字段和 HTTP 请求体明细,源 PG 不清除
请输入选项 [1]:
EOF
else
cat >/dev/tty <<'EOF'
Request/response body detail migration mode:
1) Full migration: migrate all migratable data, including request body details
2) Skip request bodies: migrate all other data; skip only request body large fields and HTTP body detail tables; source PG is unchanged
Enter choice [1]:
EOF
fi
local body_choice
IFS= read -r body_choice </dev/tty || body_choice=""
MIGRATE_REQUEST_BODY_MODE="${body_choice:-1}"
fi
normalize_migration_request_body_mode
}
install_migration_project_file() {
local source_path="$1"
local mode="$2"
local target_path
ensure_tmp_root
target_path="${TMP_ROOT}/$(basename "${source_path}")"
install_project_file "${source_path}" "${target_path}" "${mode}"
printf '%s\n' "${target_path}"
}
run_compose_single_node_migration() {
local migration_script
local target_template
local source_compose_abs
local source_compose_dir
local compose_dir_abs
local target_compose
local target_compose_abs
local target_compose_dir
local target_env
local target_env_abs
local table
local -a migrate_args
source_compose_abs="$(absolute_path "${MIGRATE_FROM_COMPOSE}")"
[[ -f "${source_compose_abs}" ]] || die "source compose file not found: ${MIGRATE_FROM_COMPOSE}"
source_compose_dir="$(dirname "${source_compose_abs}")"
compose_dir_abs="$(absolute_path_maybe_missing "${COMPOSE_DIR}")"
if [[ -n "${MIGRATE_TARGET_COMPOSE}" ]]; then
target_compose="${MIGRATE_TARGET_COMPOSE}"
elif [[ "${compose_dir_abs}" == "${source_compose_dir}" ]]; then
target_compose="${compose_dir_abs}/docker-compose.single-node.yml"
else
target_compose="${compose_dir_abs}/docker-compose.yml"
fi
target_compose_abs="$(absolute_path_maybe_missing "${target_compose}")"
target_compose_dir="$(dirname "${target_compose_abs}")"
if [[ -n "${MIGRATE_TARGET_ENV}" ]]; then
target_env="${MIGRATE_TARGET_ENV}"
elif [[ "$(basename "${target_compose_abs}")" == "docker-compose.yml" ]]; then
target_env="${target_compose_dir}/.env"
else
target_env="${target_compose_dir}/.env.single-node"
fi
target_env_abs="$(absolute_path_maybe_missing "${target_env}")"
[[ "${target_compose_abs}" != "${source_compose_abs}" ]] || die "target compose would overwrite the source compose file; pass --target-compose or --compose-dir"
[[ "${target_env_abs}" != "${source_compose_dir}/.env" ]] || die "target env would overwrite the source .env; pass --target-env or --compose-dir"
migration_script="$(install_migration_project_file "scripts/migrate-pg-compose-to-single-node.sh" "0755")"
target_template="$(install_migration_project_file "docker-compose.single-node.yml" "0644")"
migrate_args=(
"${migration_script}"
--source-compose "${source_compose_abs}"
--target-compose "${target_compose_abs}"
--target-template "${target_template}"
--target-env "${target_env_abs}"
--app-image "$(compose_image)"
)
[[ -z "${MIGRATE_TARGET_DB}" ]] || migrate_args+=(--target-db "${MIGRATE_TARGET_DB}")
[[ -z "${MIGRATE_WORK_DIR}" ]] || migrate_args+=(--work-dir "${MIGRATE_WORK_DIR}")
[[ -z "${MIGRATE_APP_SERVICE}" ]] || migrate_args+=(--app-service "${MIGRATE_APP_SERVICE}")
[[ -z "${MIGRATE_POSTGRES_SERVICE}" ]] || migrate_args+=(--postgres-service "${MIGRATE_POSTGRES_SERVICE}")
[[ -z "${MIGRATE_SINGLE_NODE_SERVICE}" ]] || migrate_args+=(--single-node-service "${MIGRATE_SINGLE_NODE_SERVICE}")
[[ "${MIGRATE_REPLACE_EXISTING}" != "true" ]] || migrate_args+=(--replace-existing)
[[ "${MIGRATE_REPLACE_TARGET_COMPOSE}" != "true" ]] || migrate_args+=(--replace-target-compose)
[[ -z "${MIGRATE_REQUEST_BODY_MODE}" ]] || migrate_args+=(--request-body-mode "${MIGRATE_REQUEST_BODY_MODE}")
[[ "${MIGRATE_KEEP_SOURCE_STOPPED_ON_ERROR}" != "true" ]] || migrate_args+=(--keep-source-stopped-on-error)
bash "${migrate_args[@]}"
}
run_single_node_service_migration() {
local migration_script
local installer
local table
local -a migrate_args
[[ -z "${MIGRATE_TARGET_COMPOSE}" ]] || die "--target-compose is only valid with --mode compose-single-node"
[[ -z "${MIGRATE_TARGET_ENV}" ]] || die "--target-env is only valid with --mode compose-single-node"
[[ -z "${MIGRATE_SINGLE_NODE_SERVICE}" ]] || die "--single-node-service is only valid with --mode compose-single-node"
[[ "${MIGRATE_REPLACE_TARGET_COMPOSE}" != "true" ]] || die "--replace-target-compose is only valid with --mode compose-single-node"
migration_script="$(install_migration_project_file "scripts/migrate-pg-to-single-node.sh" "0755")"
installer="$(install_migration_project_file "install.sh" "0755")"
migrate_args=(
"${migration_script}"
--source-compose "${MIGRATE_FROM_COMPOSE}"
--installer "${installer}"
--install-root "${INSTALL_ROOT}"
--config-dir "${CONFIG_DIR}"
--service-name "${SERVICE_NAME}"
--service-user "${SERVICE_USER}"
--service-group "${SERVICE_GROUP}"
--app-image "$(compose_image)"
--install-channel "${CHANNEL}"
--install-repo "${REPO}"
--install-source-ref "${SOURCE_REF}"
)
[[ -z "${VERSION}" ]] || migrate_args+=(--install-version "${VERSION}")
[[ -z "${ARCHIVE_PATH}" ]] || migrate_args+=(--install-archive "${ARCHIVE_PATH}")
[[ -z "${RELEASE_ARCHIVE_URL}" ]] || migrate_args+=(--install-download-url "${RELEASE_ARCHIVE_URL}")
[[ -z "${MIGRATE_TARGET_DB}" ]] || migrate_args+=(--target-db "${MIGRATE_TARGET_DB}")
[[ -z "${MIGRATE_WORK_DIR}" ]] || migrate_args+=(--work-dir "${MIGRATE_WORK_DIR}")
[[ -z "${MIGRATE_APP_SERVICE}" ]] || migrate_args+=(--app-service "${MIGRATE_APP_SERVICE}")
[[ -z "${MIGRATE_POSTGRES_SERVICE}" ]] || migrate_args+=(--postgres-service "${MIGRATE_POSTGRES_SERVICE}")
[[ "${MIGRATE_REPLACE_EXISTING}" != "true" ]] || migrate_args+=(--replace-existing)
[[ -z "${MIGRATE_REQUEST_BODY_MODE}" ]] || migrate_args+=(--request-body-mode "${MIGRATE_REQUEST_BODY_MODE}")
[[ "${MIGRATE_KEEP_SOURCE_STOPPED_ON_ERROR}" != "true" ]] || migrate_args+=(--keep-source-stopped-on-error)
bash "${migrate_args[@]}"
}
run_migration_from_compose() {
if [[ -n "${MIGRATE_REQUEST_BODY_MODE}" ]]; then
normalize_migration_request_body_mode
fi
case "${MODE}" in
compose-single-node)
run_compose_single_node_migration
;;
single-node)
run_single_node_service_migration
;;
compose)
die "--migrate-from-compose target mode must be compose-single-node or single-node"
;;
*)
die "unsupported migration target mode: ${MODE}"
;;
esac
}
resolve_version() {
if [[ -n "${VERSION}" ]]; then
echo "${VERSION}"
@@ -2725,15 +2166,6 @@ main() {
apply_platform_defaults
select_version
select_mode
collect_interactive_migration_options
if [[ -n "${MIGRATE_FROM_COMPOSE}" ]]; then
run_migration_from_compose
return
fi
if migration_options_requested; then
die "migration options require --migrate-from-compose"
fi
if [[ "${MODE}" == "compose" ]]; then
install_compose_mode
@@ -1,877 +0,0 @@
#!/usr/bin/env bash
set -euo pipefail
SOURCE_COMPOSE="docker-compose.yml"
TARGET_COMPOSE=""
TARGET_COMPOSE_TEMPLATE="./docker-compose.single-node.yml"
TARGET_ENV=""
TARGET_DB=""
WORK_DIR=""
APP_SERVICE="app"
POSTGRES_SERVICE="postgres"
SINGLE_NODE_SERVICE="app"
APP_IMAGE=""
REPLACE_EXISTING="false"
REPLACE_TARGET_COMPOSE="false"
DRY_RUN="false"
KEEP_SOURCE_STOPPED_ON_ERROR="false"
DISK_SPACE_MULTIPLIER="${AETHER_MIGRATION_DISK_SPACE_MULTIPLIER:-2}"
DISK_SPACE_MIN_FREE_BYTES="${AETHER_MIGRATION_MIN_FREE_BYTES:-1073741824}"
REQUEST_BODY_MODE="${AETHER_MIGRATION_REQUEST_BODY_MODE:-full}"
APP_STOPPED="false"
CUTOVER_COMPLETE="false"
SOURCE_COMPOSE_ABS=""
SOURCE_COMPOSE_DIR=""
SOURCE_ENV=""
SOURCE_NETWORK=""
TARGET_COMPOSE_ABS=""
TARGET_COMPOSE_DIR=""
TARGET_ENV_ABS=""
DB_USER=""
DB_NAME=""
DB_PASSWORD=""
NOW=""
usage() {
cat <<'EOF'
Usage: scripts/migrate-pg-compose-to-single-node.sh [options]
Migrate an existing Docker Compose Postgres deployment to Docker Compose single-node.
This script pulls the single-node app image before downtime, stops only the source app,
then copies Postgres data directly into a temporary SQLite DB without writing a JSONL
intermediate file. After the copy succeeds, it starts the single-node Compose app.
Options:
--source-compose PATH Source Postgres docker compose file (default: docker-compose.yml)
--target-compose PATH Target single-node compose file (default: SOURCE_DIR/docker-compose.single-node.yml)
--target-template PATH Template copied when target compose is missing (default: ./docker-compose.single-node.yml)
--target-env PATH Target env file (default: TARGET_COMPOSE_DIR/.env.single-node)
--target-db PATH Final SQLite DB path (default: TARGET_COMPOSE_DIR/data/aether.db)
--work-dir PATH Working directory (default: SOURCE_DIR/data/pg-compose-to-single-node-<timestamp>)
--app-service NAME Source compose app service (default: app)
--postgres-service NAME Source compose Postgres service (default: postgres)
--single-node-service NAME Target single-node compose service (default: app)
--app-image IMAGE Override APP_IMAGE for the target single-node compose env
--replace-existing Allow replacing an existing target SQLite database
--replace-target-compose Overwrite target compose file from the template
--dry-run Pull/preflight/direct-copy without stopping or switching
--request-body-mode MODE Request/response body detail handling: full/1 or omit/2
full: migrate all migratable data, including request body details
omit: migrate all other data; skip only request body large fields and HTTP body detail tables; source PG is unchanged
--keep-source-stopped-on-error
Do not auto-restart source app if migration fails after stopping it
-h, --help Show this help
Default cutover behavior:
1. Copy/prepare docker-compose.single-node.yml and .env.single-node.
2. Pull the target single-node app image while the source app is still running.
3. Verify the running source app image ID matches the target single-node image ID.
4. Preflight SQLite migrations in a temporary DB with the target image.
5. Re-check migration coverage and available disk space.
6. Stop/remove only the source app container; keep Postgres/Redis running.
7. Copy records directly into SQLite, replace TARGET_DB, and start single-node.
EOF
}
log() {
printf '>>> %s\n' "$*"
}
warn() {
printf 'WARN: %s\n' "$*" >&2
}
die() {
printf 'ERROR: %s\n' "$*" >&2
exit 1
}
trim() {
local value="$1"
value="${value#"${value%%[![:space:]]*}"}"
value="${value%"${value##*[![:space:]]}"}"
printf '%s' "$value"
}
strip_optional_quotes() {
local value="$1"
if [[ "${#value}" -ge 2 ]]; then
if [[ "${value:0:1}" == "\"" && "${value: -1}" == "\"" ]]; then
printf '%s' "${value:1:${#value}-2}"
return
fi
if [[ "${value:0:1}" == "'" && "${value: -1}" == "'" ]]; then
printf '%s' "${value:1:${#value}-2}"
return
fi
fi
printf '%s' "$value"
}
absolute_path() {
local path="$1"
local dir
local base
if [[ "$path" == /* ]]; then
printf '%s\n' "$path"
return
fi
dir="$(dirname "$path")"
base="$(basename "$path")"
printf '%s/%s\n' "$(cd "$dir" && pwd -P)" "$base"
}
absolute_path_maybe_missing() {
local path="$1"
if [[ "$path" == /* ]]; then
printf '%s\n' "$path"
else
printf '%s/%s\n' "$(pwd -P)" "$path"
fi
}
parse_args() {
while [[ $# -gt 0 ]]; do
case "$1" in
--source-compose)
[[ $# -ge 2 ]] || die "--source-compose requires a value"
SOURCE_COMPOSE="$2"
shift 2
;;
--target-compose)
[[ $# -ge 2 ]] || die "--target-compose requires a value"
TARGET_COMPOSE="$2"
shift 2
;;
--target-template)
[[ $# -ge 2 ]] || die "--target-template requires a value"
TARGET_COMPOSE_TEMPLATE="$2"
shift 2
;;
--target-env)
[[ $# -ge 2 ]] || die "--target-env requires a value"
TARGET_ENV="$2"
shift 2
;;
--target-db)
[[ $# -ge 2 ]] || die "--target-db requires a value"
TARGET_DB="$2"
shift 2
;;
--work-dir)
[[ $# -ge 2 ]] || die "--work-dir requires a value"
WORK_DIR="$2"
shift 2
;;
--app-service)
[[ $# -ge 2 ]] || die "--app-service requires a value"
APP_SERVICE="$2"
shift 2
;;
--postgres-service)
[[ $# -ge 2 ]] || die "--postgres-service requires a value"
POSTGRES_SERVICE="$2"
shift 2
;;
--single-node-service)
[[ $# -ge 2 ]] || die "--single-node-service requires a value"
SINGLE_NODE_SERVICE="$2"
shift 2
;;
--app-image)
[[ $# -ge 2 ]] || die "--app-image requires a value"
APP_IMAGE="$2"
shift 2
;;
--replace-existing)
REPLACE_EXISTING="true"
shift
;;
--replace-target-compose)
REPLACE_TARGET_COMPOSE="true"
shift
;;
--dry-run)
DRY_RUN="true"
shift
;;
--request-body-mode)
[[ $# -ge 2 ]] || die "--request-body-mode requires a value"
REQUEST_BODY_MODE="$2"
shift 2
;;
--keep-source-stopped-on-error)
KEEP_SOURCE_STOPPED_ON_ERROR="true"
shift
;;
-h|--help)
usage
exit 0
;;
*)
die "unknown argument: $1"
;;
esac
done
}
require_command() {
command -v "$1" >/dev/null 2>&1 || die "required command not found: $1"
}
normalize_request_body_mode() {
case "$REQUEST_BODY_MODE" in
""|1|full|all|include)
REQUEST_BODY_MODE="full"
;;
2|omit|skip)
REQUEST_BODY_MODE="omit"
;;
*)
die "--request-body-mode must be full/1 or omit/2"
;;
esac
}
env_file_get() {
local file="$1"
local wanted="$2"
local line key value
local found=""
while IFS= read -r line || [[ -n "$line" ]]; do
line="${line%$'\r'}"
line="$(trim "$line")"
[[ -z "$line" || "${line:0:1}" == "#" ]] && continue
[[ "$line" == export\ * ]] && line="${line#export }"
key="$(trim "${line%%=*}")"
[[ "$key" == "$wanted" ]] || continue
value="${line#*=}"
found="$(strip_optional_quotes "$(trim "$value")")"
done < "$file"
printf '%s' "$found"
}
validate_env_line_for_copy() {
local line="$1"
local line_no="$2"
[[ "$line" == *'${'* ]] && die "source env line ${line_no} uses variable expansion; write a concrete value before migration"
[[ "$line" == *'$('* ]] && die "source env line ${line_no} uses command substitution; write a concrete value before migration"
[[ "$line" == *'`'* ]] && die "source env line ${line_no} uses command substitution; write a concrete value before migration"
return 0
}
should_skip_single_node_env_key() {
case "$1" in
APP_IMAGE|LOCAL_APP_IMAGE|APP_PORT|DB_HOST|DB_PORT|DB_USER|DB_NAME|DB_PASSWORD|POSTGRES_*|MYSQL_*|REDIS_HOST|REDIS_PORT|REDIS_PASSWORD|REDIS_URL|AETHER_GATEWAY_DATA_REDIS_URL|AETHER_GATEWAY_DATA_REDIS_KEY_PREFIX|DATABASE_URL|AETHER_DATABASE_URL|AETHER_DATABASE_DRIVER|AETHER_GATEWAY_DATA_POSTGRES_URL|AETHER_RUNTIME_BACKEND|AETHER_RUNTIME_REDIS_URL|AETHER_RUNTIME_REDIS_KEY_PREFIX|AETHER_GATEWAY_DEPLOYMENT_TOPOLOGY|AETHER_GATEWAY_NODE_ROLE|AETHER_GATEWAY_STATIC_DIR|AETHER_UPDATE_STRATEGY|AETHER_DOCKER_UPDATE_COMMAND|AETHER_LOG_DIR|AETHER_GATEWAY_AUTO_PREPARE_DATABASE)
return 0
;;
*)
return 1
;;
esac
}
write_single_node_env() {
local output="$1"
local line raw_line key
local line_no=0
local app_port app_image jwt_key encryption_key
app_port="$(env_file_get "$SOURCE_ENV" "APP_PORT")"
app_port="${app_port:-8084}"
app_image="${APP_IMAGE:-$(env_file_get "$SOURCE_ENV" "APP_IMAGE")}"
app_image="${app_image:-ghcr.io/fawney19/aether:latest}"
jwt_key="$(env_file_get "$SOURCE_ENV" "JWT_SECRET_KEY")"
encryption_key="$(env_file_get "$SOURCE_ENV" "ENCRYPTION_KEY")"
: > "$output"
{
printf '# Generated by scripts/migrate-pg-compose-to-single-node.sh from %s\n' "$SOURCE_ENV"
printf '# single-node means Docker Compose app + SQLite for this migration target.\n\n'
} >> "$output"
while IFS= read -r raw_line || [[ -n "$raw_line" ]]; do
line_no=$((line_no + 1))
line="${raw_line%$'\r'}"
line="$(trim "$line")"
[[ -z "$line" || "${line:0:1}" == "#" ]] && continue
[[ "$line" == export\ * ]] && die "source env line ${line_no} uses export; write KEY=VALUE before migration"
[[ "$line" =~ ^[A-Za-z_][A-Za-z0-9_]*= ]] || die "source env line ${line_no} must be KEY=VALUE"
validate_env_line_for_copy "$line" "$line_no"
key="${line%%=*}"
if should_skip_single_node_env_key "$key"; then
continue
fi
printf '%s\n' "$line" >> "$output"
done < "$SOURCE_ENV"
{
printf '\n# Single-node Compose runtime overrides\n'
printf 'APP_IMAGE=%s\n' "$app_image"
printf 'APP_PORT=%s\n' "$app_port"
printf 'AETHER_GATEWAY_STATIC_DIR=/opt/aether/current/frontend\n'
printf 'AETHER_UPDATE_STRATEGY=docker\n'
printf 'AETHER_DOCKER_UPDATE_COMMAND=./update.sh\n'
printf 'AETHER_LOG_DESTINATION=stdout\n'
printf 'AETHER_LOG_FORMAT=pretty\n'
printf 'AETHER_LOG_DIR=/opt/aether/logs\n'
printf 'AETHER_DATABASE_DRIVER=sqlite\n'
printf 'AETHER_DATABASE_URL=sqlite:///opt/aether/data/aether.db\n'
printf 'DATABASE_URL=sqlite:///opt/aether/data/aether.db\n'
printf 'AETHER_RUNTIME_BACKEND=memory\n'
printf 'AETHER_GATEWAY_DEPLOYMENT_TOPOLOGY=single-node\n'
printf 'AETHER_GATEWAY_NODE_ROLE=all\n'
printf 'AETHER_GATEWAY_AUTO_PREPARE_DATABASE=true\n'
printf 'JWT_SECRET_KEY=%s\n' "$jwt_key"
printf 'ENCRYPTION_KEY=%s\n' "$encryption_key"
} >> "$output"
}
source_compose() {
docker compose -f "$SOURCE_COMPOSE_ABS" "$@"
}
target_compose() {
AETHER_ENV_FILE="$TARGET_ENV_ABS" docker compose --env-file "$TARGET_ENV_ABS" -f "$TARGET_COMPOSE_ABS" "$@"
}
target_run_app() {
local database_url="$1"
shift
target_compose run --rm --no-deps \
-v "${WORK_DIR}:/migration" \
-e AETHER_LOG_DESTINATION=stdout \
-e AETHER_DATABASE_DRIVER=sqlite \
-e "AETHER_DATABASE_URL=${database_url}" \
-e "DATABASE_URL=${database_url}" \
"$SINGLE_NODE_SERVICE" "$@"
}
run_psql_stdin() {
local sql_file="$1"
source_compose exec -T \
-e "PGPASSWORD=${DB_PASSWORD}" \
"$POSTGRES_SERVICE" \
psql -h 127.0.0.1 -U "$DB_USER" -d "$DB_NAME" -v ON_ERROR_STOP=1 -At -f - < "$sql_file"
}
source_database_size_bytes() {
local sql_file
local result
sql_file="${WORK_DIR}/source-database-size.sql"
if [[ "$REQUEST_BODY_MODE" == "omit" ]]; then
cat > "$sql_file" <<'SQL'
SELECT GREATEST(
pg_database_size(current_database())
- COALESCE(pg_total_relation_size(to_regclass('public.usage_body_blobs')), 0)
- COALESCE(pg_total_relation_size(to_regclass('public.usage_http_audits')), 0),
0
);
SQL
else
printf 'SELECT pg_database_size(current_database());\n' > "$sql_file"
fi
result="$(run_psql_stdin "$sql_file" | tr -d '[:space:]')"
[[ "$result" =~ ^[0-9]+$ ]] || die "could not determine source Postgres database size"
printf '%s\n' "$result"
}
file_size_bytes() {
local path="$1"
if [[ ! -e "$path" ]]; then
printf '0\n'
return
fi
stat -c '%s' "$path" 2>/dev/null || stat -f '%z' "$path" 2>/dev/null || die "could not stat file: ${path}"
}
target_sqlite_size_bytes() {
local total=0
local suffix
local size
for suffix in "" "-wal" "-shm"; do
size="$(file_size_bytes "${TARGET_DB}${suffix}")"
total=$((total + size))
done
printf '%s\n' "$total"
}
available_bytes_for_path() {
local path="$1"
df -Pk "$path" | awk 'NR == 2 { printf "%.0f\n", $4 * 1024 }'
}
filesystem_key_for_path() {
local path="$1"
df -Pk "$path" | awk 'NR == 2 { print $1 }'
}
format_bytes() {
local bytes="$1"
awk -v bytes="$bytes" 'BEGIN {
if (bytes >= 1073741824) {
printf "%.1f GiB", bytes / 1073741824
} else {
printf "%.1f MiB", bytes / 1048576
}
}'
}
assert_available_space() {
local path="$1"
local required_bytes="$2"
local label="$3"
local available_bytes
available_bytes="$(available_bytes_for_path "$path")"
[[ "$available_bytes" =~ ^[0-9]+$ ]] || die "could not determine free disk space for ${path}"
log "${label} free space: $(format_bytes "$available_bytes"); required: $(format_bytes "$required_bytes")"
if (( available_bytes < required_bytes )); then
die "${label} does not have enough free disk space; required $(format_bytes "$required_bytes"), available $(format_bytes "$available_bytes")"
fi
}
check_disk_space() {
local source_bytes
local estimated_db_bytes
local backup_bytes=0
local target_dir
local work_fs
local target_fs
local required_bytes
[[ "$DISK_SPACE_MULTIPLIER" =~ ^[1-9][0-9]*$ ]] || die "AETHER_MIGRATION_DISK_SPACE_MULTIPLIER must be a positive integer"
[[ "$DISK_SPACE_MIN_FREE_BYTES" =~ ^[0-9]+$ ]] || die "AETHER_MIGRATION_MIN_FREE_BYTES must be a non-negative integer"
source_bytes="$(source_database_size_bytes)"
estimated_db_bytes=$((source_bytes * DISK_SPACE_MULTIPLIER + DISK_SPACE_MIN_FREE_BYTES))
target_dir="$(dirname "$TARGET_DB")"
mkdir -p "$target_dir"
if [[ "$REPLACE_EXISTING" == "true" ]]; then
backup_bytes="$(target_sqlite_size_bytes)"
fi
log "source Postgres size used for disk estimate: $(format_bytes "$source_bytes")"
if [[ "$DRY_RUN" == "true" ]]; then
assert_available_space "$WORK_DIR" "$estimated_db_bytes" "work dir"
return
fi
work_fs="$(filesystem_key_for_path "$WORK_DIR")"
target_fs="$(filesystem_key_for_path "$target_dir")"
if [[ "$work_fs" == "$target_fs" ]]; then
required_bytes=$((estimated_db_bytes * 2 + backup_bytes))
assert_available_space "$WORK_DIR" "$required_bytes" "work/target filesystem"
else
assert_available_space "$WORK_DIR" "$((estimated_db_bytes + backup_bytes))" "work dir"
assert_available_space "$target_dir" "$estimated_db_bytes" "target DB dir"
fi
}
check_request_body_artifacts() {
local sql_file
local result_file
sql_file="${WORK_DIR}/check-request-body-artifacts.sql"
result_file="${WORK_DIR}/request-body-artifacts.txt"
cat > "$sql_file" <<'SQL'
CREATE TEMP TABLE aether_request_body_artifacts (
artifact text PRIMARY KEY
) ON COMMIT PRESERVE ROWS;
DO $$
DECLARE
candidate record;
has_rows boolean;
BEGIN
IF EXISTS (
SELECT 1
FROM information_schema.tables
WHERE table_schema = 'public'
AND table_name = 'usage_body_blobs'
) THEN
EXECUTE 'SELECT EXISTS (SELECT 1 FROM public.usage_body_blobs LIMIT 1)'
INTO has_rows;
IF has_rows THEN
INSERT INTO aether_request_body_artifacts(artifact)
VALUES ('usage_body_blobs')
ON CONFLICT DO NOTHING;
END IF;
END IF;
IF EXISTS (
SELECT 1
FROM information_schema.tables
WHERE table_schema = 'public'
AND table_name = 'usage_http_audits'
) THEN
EXECUTE 'SELECT EXISTS (SELECT 1 FROM public.usage_http_audits LIMIT 1)'
INTO has_rows;
IF has_rows THEN
INSERT INTO aether_request_body_artifacts(artifact)
VALUES ('usage_http_audits')
ON CONFLICT DO NOTHING;
END IF;
END IF;
IF EXISTS (
SELECT 1
FROM information_schema.tables
WHERE table_schema = 'public'
AND table_name = 'usage'
) THEN
FOR candidate IN
SELECT unnest(ARRAY[
'request_body',
'response_body',
'provider_request_body',
'client_response_body',
'request_body_compressed',
'response_body_compressed',
'provider_request_body_compressed',
'client_response_body_compressed'
]) AS column_name
LOOP
IF EXISTS (
SELECT 1
FROM information_schema.columns
WHERE table_schema = 'public'
AND table_name = 'usage'
AND column_name = candidate.column_name
) THEN
EXECUTE format('SELECT EXISTS (SELECT 1 FROM public.usage WHERE %I IS NOT NULL LIMIT 1)', candidate.column_name)
INTO has_rows;
IF has_rows THEN
INSERT INTO aether_request_body_artifacts(artifact)
VALUES ('usage.' || candidate.column_name)
ON CONFLICT DO NOTHING;
END IF;
END IF;
END LOOP;
END IF;
END
$$;
SELECT artifact FROM aether_request_body_artifacts ORDER BY artifact;
SQL
run_psql_stdin "$sql_file" > "$result_file"
if [[ ! -s "$result_file" ]]; then
return
fi
if [[ "$REQUEST_BODY_MODE" == "omit" ]]; then
warn "source has request/response body details that will not be copied into single-node SQLite"
cat "$result_file" >&2
return
fi
warn "source has request/response body details that will be copied into single-node SQLite"
cat "$result_file" >&2
}
source_service_is_running() {
source_compose ps --services --status running | grep -Fxq "$1"
}
source_app_container_id() {
source_compose ps -q "$APP_SERVICE"
}
docker_image_id() {
local image="$1"
docker image inspect -f '{{.Id}}' "$image" 2>/dev/null || true
}
assert_source_and_target_images_match() {
local target_image="$1"
local source_container_id
local source_image_ref
local source_image_id
local target_image_id
source_container_id="$(source_app_container_id)"
[[ -n "$source_container_id" ]] || die "could not resolve running source app container for service: ${APP_SERVICE}"
source_image_ref="$(docker inspect -f '{{.Config.Image}}' "$source_container_id")"
source_image_id="$(docker inspect -f '{{.Image}}' "$source_container_id")"
target_image_id="$(docker_image_id "$target_image")"
[[ -n "$target_image_id" ]] || die "target single-node image is not available locally after pull: ${target_image}"
log "source app image: ${source_image_ref} (${source_image_id})"
log "target single-node image: ${target_image} (${target_image_id})"
if [[ "$source_image_id" != "$target_image_id" ]]; then
die "source app image and target single-node image are different; upgrade the source PG Compose app to ${target_image} before migration"
fi
}
assert_target_copy_command_available() {
local target_image="$1"
local help_output
help_output="$(docker run --rm --entrypoint aether-gateway "$target_image" copy --help 2>&1 || true)"
if [[ "$help_output" != *"--source-driver"* || "$help_output" != *"--target-driver"* || "$help_output" != *"--omit-request-body-details"* ]]; then
die "target single-node image does not support direct PG-to-SQLite copy; use a matching Aether release image"
fi
}
resolve_source_network() {
local container_id
local network
container_id="$(source_compose ps -q "$POSTGRES_SERVICE")"
[[ -n "$container_id" ]] || die "could not resolve container id for source Postgres service: ${POSTGRES_SERVICE}"
SOURCE_NETWORK=""
while IFS= read -r network; do
[[ -n "$network" ]] || continue
SOURCE_NETWORK="$network"
break
done < <(docker inspect -f '{{range $name, $_ := .NetworkSettings.Networks}}{{println $name}}{{end}}' "$container_id")
[[ -n "$SOURCE_NETWORK" ]] || die "could not resolve Docker network for source Postgres service: ${POSTGRES_SERVICE}"
}
prepare_target_compose() {
local template_abs
TARGET_COMPOSE="${TARGET_COMPOSE:-${SOURCE_COMPOSE_DIR}/docker-compose.single-node.yml}"
TARGET_COMPOSE_ABS="$(absolute_path_maybe_missing "$TARGET_COMPOSE")"
TARGET_COMPOSE_DIR="$(dirname "$TARGET_COMPOSE_ABS")"
template_abs="$(absolute_path "$TARGET_COMPOSE_TEMPLATE")"
[[ -f "$template_abs" ]] || die "target compose template not found: ${TARGET_COMPOSE_TEMPLATE}"
mkdir -p "$TARGET_COMPOSE_DIR"
if [[ ! -f "$TARGET_COMPOSE_ABS" || "$REPLACE_TARGET_COMPOSE" == "true" ]]; then
log "installing target single-node compose file at ${TARGET_COMPOSE_ABS}"
install -m 0644 "$template_abs" "$TARGET_COMPOSE_ABS"
else
log "keeping existing target compose file ${TARGET_COMPOSE_ABS}"
fi
TARGET_ENV="${TARGET_ENV:-${TARGET_COMPOSE_DIR}/.env.single-node}"
TARGET_ENV_ABS="$(absolute_path_maybe_missing "$TARGET_ENV")"
}
target_sqlite_file_exists() {
[[ -e "$TARGET_DB" || -e "${TARGET_DB}-wal" || -e "${TARGET_DB}-shm" ]]
}
finalize_target_db() {
local temp_db="$1"
local target_dir
local backup_path
local suffix
target_dir="$(dirname "$TARGET_DB")"
mkdir -p "$target_dir"
if target_sqlite_file_exists; then
[[ "$REPLACE_EXISTING" == "true" ]] || die "target DB already exists: ${TARGET_DB}; pass --replace-existing to replace it"
for suffix in "" "-wal" "-shm"; do
if [[ -e "${TARGET_DB}${suffix}" ]]; then
backup_path="${WORK_DIR}/$(basename "$TARGET_DB")${suffix}.backup.${NOW}"
log "backing up existing target SQLite file to ${backup_path}"
cp -p "${TARGET_DB}${suffix}" "$backup_path"
fi
done
fi
log "installing migrated SQLite DB at ${TARGET_DB}"
install -m 0640 "$temp_db" "$TARGET_DB"
for suffix in "-wal" "-shm"; do
if [[ -e "${temp_db}${suffix}" ]]; then
install -m 0640 "${temp_db}${suffix}" "${TARGET_DB}${suffix}"
else
rm -f "${TARGET_DB}${suffix}"
fi
done
}
cleanup_on_exit() {
local status=$?
if [[ "$status" -eq 0 ]]; then
return
fi
warn "migration failed with exit status ${status}"
if [[ "$APP_STOPPED" == "true" && "$CUTOVER_COMPLETE" != "true" && "$KEEP_SOURCE_STOPPED_ON_ERROR" != "true" ]]; then
warn "attempting to restart source compose app because cutover did not complete"
source_compose up -d "$APP_SERVICE" || warn "source app restart failed; check ${SOURCE_COMPOSE_ABS}"
fi
}
preflight() {
require_command docker
require_command awk
require_command df
docker compose version >/dev/null
SOURCE_COMPOSE_ABS="$(absolute_path "$SOURCE_COMPOSE")"
[[ -f "$SOURCE_COMPOSE_ABS" ]] || die "source compose file not found: ${SOURCE_COMPOSE}"
SOURCE_COMPOSE_DIR="$(dirname "$SOURCE_COMPOSE_ABS")"
SOURCE_ENV="${SOURCE_COMPOSE_DIR}/.env"
[[ -f "$SOURCE_ENV" ]] || die "source env file not found: ${SOURCE_ENV}"
NOW="$(date +%Y%m%d%H%M%S)"
if [[ -z "$WORK_DIR" ]]; then
WORK_DIR="${SOURCE_COMPOSE_DIR}/data/pg-compose-to-single-node-${NOW}"
fi
WORK_DIR="$(absolute_path_maybe_missing "$WORK_DIR")"
mkdir -p "$WORK_DIR"
prepare_target_compose
mkdir -p "${TARGET_COMPOSE_DIR}/logs"
mkdir -p "${TARGET_COMPOSE_DIR}/data"
TARGET_DB="${TARGET_DB:-${TARGET_COMPOSE_DIR}/data/aether.db}"
TARGET_DB="$(absolute_path_maybe_missing "$TARGET_DB")"
DB_USER="$(env_file_get "$SOURCE_ENV" "DB_USER")"
DB_USER="${DB_USER:-postgres}"
DB_NAME="$(env_file_get "$SOURCE_ENV" "DB_NAME")"
DB_NAME="${DB_NAME:-aether}"
DB_PASSWORD="$(env_file_get "$SOURCE_ENV" "DB_PASSWORD")"
DB_PASSWORD="${DB_PASSWORD:-aether}"
[[ -n "$(env_file_get "$SOURCE_ENV" "JWT_SECRET_KEY")" ]] || die "source env must define JWT_SECRET_KEY"
if [[ -z "$(env_file_get "$SOURCE_ENV" "ENCRYPTION_KEY")" && -z "$(env_file_get "$SOURCE_ENV" "AETHER_GATEWAY_DATA_ENCRYPTION_KEY")" ]]; then
die "source env must define ENCRYPTION_KEY or AETHER_GATEWAY_DATA_ENCRYPTION_KEY"
fi
write_single_node_env "$TARGET_ENV_ABS"
chmod 0600 "$TARGET_ENV_ABS"
log "source compose: ${SOURCE_COMPOSE_ABS}"
log "source env: ${SOURCE_ENV}"
log "target compose: ${TARGET_COMPOSE_ABS}"
log "target env: ${TARGET_ENV_ABS}"
log "target SQLite DB: ${TARGET_DB}"
log "work dir: ${WORK_DIR}"
source_service_is_running "$POSTGRES_SERVICE" || die "source Postgres service is not running: ${POSTGRES_SERVICE}"
resolve_source_network
log "source Docker network: ${SOURCE_NETWORK}"
}
source_postgres_url() {
printf 'postgresql://%s:%s@%s:5432/%s' "$DB_USER" "$DB_PASSWORD" "$POSTGRES_SERVICE" "$DB_NAME"
}
copy_source_to_sqlite() {
local target_temp_db="$1"
local target_url
local image
local -a copy_args
image="$(env_file_get "$TARGET_ENV_ABS" "APP_IMAGE")"
[[ -n "$image" ]] || die "target env must define APP_IMAGE"
target_url="sqlite:///migration/$(basename "$target_temp_db")"
rm -f "$target_temp_db" "${target_temp_db}-wal" "${target_temp_db}-shm"
target_run_app "$target_url" --migrate
copy_args=(
copy
--source-driver postgres
--source-url "$(source_postgres_url)"
--target-driver sqlite
--target-url "$target_url"
)
if [[ "$REQUEST_BODY_MODE" == "omit" ]]; then
copy_args+=(--omit-request-body-details)
fi
docker run --rm \
--network "$SOURCE_NETWORK" \
-v "${WORK_DIR}:/migration" \
--env-file "$TARGET_ENV_ABS" \
-e AETHER_LOG_DESTINATION=stdout \
"$image" \
"${copy_args[@]}"
}
main() {
local preflight_db
local dry_run_db
local target_temp_db
local target_image
parse_args "$@"
normalize_request_body_mode
trap cleanup_on_exit EXIT
preflight
preflight_db="${WORK_DIR}/single-node-preflight.db"
dry_run_db="${WORK_DIR}/dry-run-target-aether.db"
target_temp_db="${WORK_DIR}/target-aether.db"
if target_sqlite_file_exists && [[ "$REPLACE_EXISTING" != "true" && "$DRY_RUN" != "true" ]]; then
die "target DB already exists: ${TARGET_DB}; pass --replace-existing to replace it"
fi
log "pulling target single-node image before downtime"
target_compose pull "$SINGLE_NODE_SERVICE"
target_image="$(env_file_get "$TARGET_ENV_ABS" "APP_IMAGE")"
[[ -n "$target_image" ]] || die "target env must define APP_IMAGE"
log "checking target image copy command is available"
assert_target_copy_command_available "$target_image"
log "checking source app image matches target single-node image"
assert_source_and_target_images_match "$target_image"
log "preflighting target SQLite schema migration"
rm -f "$preflight_db" "${preflight_db}-wal" "${preflight_db}-shm"
target_run_app "sqlite:///migration/$(basename "$preflight_db")" --migrate
log "checking request body detail policy"
check_request_body_artifacts
log "checking available disk space before copy"
check_disk_space
if [[ "$DRY_RUN" == "true" ]]; then
warn "dry-run copy happens while the source app may still be writing; use only for rehearsal"
log "copying source Postgres tables directly into dry-run SQLite target"
copy_source_to_sqlite "$dry_run_db"
log "dry run complete; temporary SQLite DB is ${dry_run_db}"
return
fi
log "stopping source app service; Postgres and Redis stay running"
source_compose stop "$APP_SERVICE"
APP_STOPPED="true"
log "checking request body detail policy again after the app has stopped"
check_request_body_artifacts
log "copying source Postgres tables directly into temporary SQLite DB"
copy_source_to_sqlite "$target_temp_db"
finalize_target_db "$target_temp_db"
log "removing stopped source app container to free the app container name"
source_compose rm -f "$APP_SERVICE"
log "starting single-node compose app"
target_compose up -d "$SINGLE_NODE_SERVICE"
CUTOVER_COMPLETE="true"
log "migration complete"
log "source Postgres/Redis volumes were left in place for rollback"
}
main "$@"
File diff suppressed because it is too large Load Diff