mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-10 13:10:21 +08:00
Merge branch 'main' of https://github.com/zhefox/Aether
This commit is contained in:
@@ -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
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
+54
-2
@@ -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,
|
||||
});
|
||||
|
||||
+8
-7
@@ -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,
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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"));
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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()?
|
||||
|
||||
+16
-9
@@ -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]
|
||||
|
||||
@@ -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#"
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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`
|
||||
明确指定。
|
||||
@@ -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: {
|
||||
|
||||
@@ -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)}`
|
||||
}
|
||||
|
||||
@@ -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 Keys、Key 用量、钱包快照与统计聚合。正常导出的 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
@@ -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
Reference in New Issue
Block a user