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

refactor(gateway): 重构 ai pipeline 规划链路

Browse Source
This commit is contained in:
fawney19
2026-04-14 01:27:04 +08:00
parent 37bb120d20
commit 53acfbabf6
74 changed files with 7357 additions and 4380 deletions
Download Patch File Download Diff File Expand all files Collapse all files

36
crates/aether-data/src/repository/proxy_nodes/sql.rs
View File

@@ -124,7 +124,7 @@ SET
heartbeat_interval = COALESCE($2, heartbeat_interval),
active_connections = COALESCE($3, active_connections),
avg_latency_ms = COALESCE($4, avg_latency_ms),
proxy_metadata = COALESCE($5, proxy_metadata),
proxy_metadata = COALESCE($5::json, proxy_metadata),
total_requests = total_requests + GREATEST(COALESCE($6, 0), 0),
failed_requests = failed_requests + GREATEST(COALESCE($7, 0), 0),
dns_failures = dns_failures + GREATEST(COALESCE($8, 0), 0),
@@ -205,11 +205,11 @@ VALUES (
COALESCE($8, 0),
COALESCE($9, 0),
$10,
$11,
$11::json,
$12,
$13,
FALSE,
$14
$14::json
)
"#;
@@ -226,10 +226,10 @@ SET
active_connections = COALESCE($8, active_connections),
total_requests = COALESCE($9, total_requests),
avg_latency_ms = COALESCE($10, avg_latency_ms),
hardware_info = COALESCE($11, hardware_info),
hardware_info = COALESCE($11::json, hardware_info),
estimated_max_concurrency = COALESCE($12, estimated_max_concurrency),
tunnel_mode = $13,
proxy_metadata = COALESCE($14, proxy_metadata),
proxy_metadata = COALESCE($14::json, proxy_metadata),
updated_at = NOW()
WHERE id = $1
"#;
@@ -248,7 +248,7 @@ const UPDATE_PROXY_NODE_REMOTE_CONFIG_SQL: &str = r#"
UPDATE proxy_nodes
SET
name = COALESCE($2, name),
remote_config = $3,
remote_config = $3::json,
config_version = config_version + 1,
updated_at = NOW()
WHERE id = $1
@@ -748,3 +748,27 @@ VALUES (
self.find_proxy_node(&mutation.node_id).await
}
}
#[cfg(test)]
mod tests {
#[test]
fn proxy_node_sql_uses_json_casts_for_json_columns() {
assert!(super::APPLY_HEARTBEAT_SQL
.contains("proxy_metadata = COALESCE($5::json, proxy_metadata)"));
assert!(super::INSERT_PROXY_NODE_SQL
.contains("\n $11::json,\n $12,\n $13,\n FALSE,\n $14::json\n"));
assert!(super::UPDATE_PROXY_NODE_REGISTRATION_SQL
.contains("hardware_info = COALESCE($11::json, hardware_info)"));
assert!(super::UPDATE_PROXY_NODE_REGISTRATION_SQL
.contains("proxy_metadata = COALESCE($14::json, proxy_metadata)"));
assert!(super::UPDATE_PROXY_NODE_REMOTE_CONFIG_SQL.contains("remote_config = $3::json"));
}
#[test]
fn proxy_node_sql_does_not_use_jsonb_casts() {
assert!(!super::APPLY_HEARTBEAT_SQL.contains("::jsonb"));
assert!(!super::INSERT_PROXY_NODE_SQL.contains("::jsonb"));
assert!(!super::UPDATE_PROXY_NODE_REGISTRATION_SQL.contains("::jsonb"));
assert!(!super::UPDATE_PROXY_NODE_REMOTE_CONFIG_SQL.contains("::jsonb"));
}
}

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

@@ -8,9 +8,10 @@ use async_trait::async_trait;
use serde_json::Value;
use super::{
strip_deprecated_usage_display_fields, StoredProviderApiKeyUsageSummary,
StoredProviderUsageSummary, StoredProviderUsageWindow, StoredRequestUsageAudit,
UpsertUsageRecord, UsageAuditListQuery, UsageReadRepository, UsageWriteRepository,
strip_deprecated_usage_display_fields, usage_can_recover_terminal_failure,
StoredProviderApiKeyUsageSummary, StoredProviderUsageSummary, StoredProviderUsageWindow,
StoredRequestUsageAudit, UpsertUsageRecord, UsageAuditListQuery, UsageReadRepository,
UsageWriteRepository,
};
use crate::DataLayerError;
@@ -450,6 +451,12 @@ impl UsageWriteRepository for InMemoryUsageReadRepository {
if existing.is_some_and(|existing| {
usage_status_is_finalized(existing.status.as_str())
&& usage_status_is_lifecycle(usage.status.as_str())
&& !usage_can_recover_terminal_failure(
existing.status.as_str(),
existing.billing_status.as_str(),
usage.status.as_str(),
usage.billing_status.as_str(),
)
}) {
return Ok(existing.expect("existing usage should be present").clone());
}
@@ -856,6 +863,166 @@ mod tests {
assert_eq!(stored.finalized_at_unix_secs, Some(101));
}
#[tokio::test]
async fn upsert_allows_streaming_recovery_after_void_failure() {
let repository = InMemoryUsageReadRepository::default();
repository
.upsert(UpsertUsageRecord {
request_id: "req-recover-1".to_string(),
user_id: Some("user-1".to_string()),
api_key_id: Some("api-key-1".to_string()),
username: None,
api_key_name: None,
provider_name: "OpenAI".to_string(),
model: "gpt-5".to_string(),
target_model: None,
provider_id: Some("provider-1".to_string()),
provider_endpoint_id: Some("endpoint-1".to_string()),
provider_api_key_id: Some("provider-key-1".to_string()),
request_type: Some("chat".to_string()),
api_format: Some("openai:chat".to_string()),
api_family: Some("openai".to_string()),
endpoint_kind: Some("chat".to_string()),
endpoint_api_format: Some("openai:chat".to_string()),
provider_api_family: Some("openai".to_string()),
provider_endpoint_kind: Some("chat".to_string()),
has_format_conversion: Some(false),
is_stream: Some(false),
input_tokens: None,
output_tokens: None,
total_tokens: None,
cache_creation_input_tokens: None,
cache_creation_ephemeral_5m_input_tokens: None,
cache_creation_ephemeral_1h_input_tokens: None,
cache_read_input_tokens: None,
cache_creation_cost_usd: None,
cache_read_cost_usd: None,
output_price_per_1m: None,
total_cost_usd: Some(0.0),
actual_total_cost_usd: Some(0.0),
status_code: Some(503),
error_message: Some("provider timeout".to_string()),
error_category: Some("provider_error".to_string()),
response_time_ms: Some(90),
first_byte_time_ms: None,
status: "failed".to_string(),
billing_status: "void".to_string(),
request_headers: None,
request_body: None,
request_body_ref: None,
provider_request_headers: None,
provider_request_body: None,
provider_request_body_ref: None,
response_headers: None,
response_body: None,
response_body_ref: None,
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
candidate_id: None,
candidate_index: None,
key_name: None,
planner_kind: None,
route_family: None,
route_kind: None,
execution_path: None,
local_execution_runtime_miss_reason: None,
request_metadata: None,
finalized_at_unix_secs: Some(101),
created_at_unix_ms: Some(100),
updated_at_unix_secs: 101,
})
.await
.expect("failed usage should upsert");
repository
.upsert(UpsertUsageRecord {
request_id: "req-recover-1".to_string(),
user_id: Some("user-1".to_string()),
api_key_id: Some("api-key-1".to_string()),
username: None,
api_key_name: None,
provider_name: "OpenAI".to_string(),
model: "gpt-5".to_string(),
target_model: Some("gpt-5-mini".to_string()),
provider_id: Some("provider-1".to_string()),
provider_endpoint_id: Some("endpoint-1".to_string()),
provider_api_key_id: Some("provider-key-1".to_string()),
request_type: Some("chat".to_string()),
api_format: Some("openai:chat".to_string()),
api_family: Some("openai".to_string()),
endpoint_kind: Some("chat".to_string()),
endpoint_api_format: Some("openai:chat".to_string()),
provider_api_family: Some("openai".to_string()),
provider_endpoint_kind: Some("chat".to_string()),
has_format_conversion: Some(true),
is_stream: Some(true),
input_tokens: Some(10),
output_tokens: None,
total_tokens: None,
cache_creation_input_tokens: None,
cache_creation_ephemeral_5m_input_tokens: None,
cache_creation_ephemeral_1h_input_tokens: None,
cache_read_input_tokens: None,
cache_creation_cost_usd: None,
cache_read_cost_usd: None,
output_price_per_1m: None,
total_cost_usd: None,
actual_total_cost_usd: None,
status_code: None,
error_message: None,
error_category: None,
response_time_ms: Some(45),
first_byte_time_ms: Some(12),
status: "streaming".to_string(),
billing_status: "pending".to_string(),
request_headers: None,
request_body: None,
request_body_ref: None,
provider_request_headers: None,
provider_request_body: None,
provider_request_body_ref: None,
response_headers: None,
response_body: None,
response_body_ref: None,
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
candidate_id: Some("cand-1".to_string()),
candidate_index: Some(1),
key_name: Some("primary".to_string()),
planner_kind: Some("claude_cli_sync".to_string()),
route_family: Some("claude".to_string()),
route_kind: Some("cli".to_string()),
execution_path: Some("remote".to_string()),
local_execution_runtime_miss_reason: None,
request_metadata: Some(json!({
"trace_id": "trace-recovered"
})),
finalized_at_unix_secs: None,
created_at_unix_ms: Some(100),
updated_at_unix_secs: 102,
})
.await
.expect("recovery usage should upsert");
let stored = repository
.find_by_request_id("req-recover-1")
.await
.expect("usage lookup should succeed")
.expect("usage should exist");
assert_eq!(stored.status, "streaming");
assert_eq!(stored.billing_status, "pending");
assert_eq!(stored.status_code, None);
assert_eq!(stored.error_message, None);
assert_eq!(stored.finalized_at_unix_secs, None);
assert_eq!(
stored.request_metadata,
Some(json!({ "trace_id": "trace-recovered" }))
);
assert_eq!(stored.total_tokens, 10);
}
#[tokio::test]
async fn seed_hydrates_legacy_body_ref_metadata_into_typed_fields() {
let repository = InMemoryUsageReadRepository::seed(vec![StoredRequestUsageAudit {

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

@@ -10,6 +10,25 @@ pub(crate) use aether_data_contracts::repository::usage::{
pub use memory::InMemoryUsageReadRepository;
pub use sql::SqlxUsageReadRepository;
pub(crate) fn incoming_usage_can_recover_terminal_failure(
incoming_status: &str,
incoming_billing_status: &str,
) -> bool {
incoming_billing_status == "pending"
&& matches!(incoming_status, "pending" | "streaming" | "completed")
}
pub(crate) fn usage_can_recover_terminal_failure(
existing_status: &str,
existing_billing_status: &str,
incoming_status: &str,
incoming_billing_status: &str,
) -> bool {
existing_billing_status == "void"
&& matches!(existing_status, "failed" | "cancelled")
&& incoming_usage_can_recover_terminal_failure(incoming_status, incoming_billing_status)
}
pub(crate) fn strip_deprecated_usage_display_fields(
mut usage: UpsertUsageRecord,
) -> UpsertUsageRecord {
@@ -20,7 +39,10 @@ pub(crate) fn strip_deprecated_usage_display_fields(
#[cfg(test)]
mod tests {
use super::{strip_deprecated_usage_display_fields, UpsertUsageRecord};
use super::{
incoming_usage_can_recover_terminal_failure, strip_deprecated_usage_display_fields,
usage_can_recover_terminal_failure, UpsertUsageRecord,
};
#[test]
fn strip_deprecated_usage_display_fields_clears_legacy_display_columns() {
@@ -97,4 +119,48 @@ mod tests {
assert_eq!(usage.provider_name, "OpenAI");
assert_eq!(usage.model, "gpt-5");
}
#[test]
fn incoming_usage_recovery_only_applies_to_pending_lifecycle_states() {
assert!(incoming_usage_can_recover_terminal_failure(
"completed",
"pending"
));
assert!(incoming_usage_can_recover_terminal_failure(
"streaming",
"pending"
));
assert!(!incoming_usage_can_recover_terminal_failure(
"failed", "void"
));
assert!(!incoming_usage_can_recover_terminal_failure(
"completed",
"settled"
));
}
#[test]
fn usage_recovery_requires_void_failure_to_be_followed_by_pending_lifecycle_state() {
assert!(usage_can_recover_terminal_failure(
"failed",
"void",
"completed",
"pending"
));
assert!(usage_can_recover_terminal_failure(
"cancelled",
"void",
"streaming",
"pending"
));
assert!(!usage_can_recover_terminal_failure(
"completed",
"pending",
"completed",
"pending"
));
assert!(!usage_can_recover_terminal_failure(
"failed", "void", "failed", "void"
));
}
}

55
crates/aether-data/src/repository/usage/sql.rs
View File

@@ -12,9 +12,9 @@ use std::io::{Read, Write};
use uuid::Uuid;
use super::{
strip_deprecated_usage_display_fields, StoredProviderApiKeyUsageSummary,
StoredProviderUsageSummary, StoredRequestUsageAudit, UpsertUsageRecord, UsageAuditListQuery,
UsageReadRepository, UsageWriteRepository,
incoming_usage_can_recover_terminal_failure, strip_deprecated_usage_display_fields,
StoredProviderApiKeyUsageSummary, StoredProviderUsageSummary, StoredRequestUsageAudit,
UpsertUsageRecord, UsageAuditListQuery, UsageReadRepository, UsageWriteRepository,
};
use crate::postgres::PostgresTransactionRunner;
use crate::{error::SqlxResultExt, DataLayerError};
@@ -45,6 +45,24 @@ const DELETE_USAGE_BODY_BLOB_SQL: &str = r#"
DELETE FROM usage_body_blobs
WHERE body_ref = $1
"#;
const RESET_STALE_VOID_USAGE_SQL: &str = r#"
UPDATE "usage"
SET
billing_status = 'pending',
finalized_at = NULL
WHERE request_id = $1
AND billing_status = 'void'
AND status IN ('failed', 'cancelled')
"#;
const RESET_STALE_VOID_USAGE_SETTLEMENT_SNAPSHOT_SQL: &str = r#"
UPDATE usage_settlement_snapshots
SET
billing_status = 'pending',
finalized_at = NULL,
updated_at = NOW()
WHERE request_id = $1
AND billing_status = 'void'
"#;
const UPSERT_USAGE_HTTP_AUDIT_SQL: &str = r#"
INSERT INTO usage_http_audits (
request_id,
@@ -1317,6 +1335,22 @@ impl SqlxUsageReadRepository {
self.tx_runner
.run_read_write(|tx| {
Box::pin(async move {
if incoming_usage_can_recover_terminal_failure(
usage.status.as_str(),
usage.billing_status.as_str(),
) {
sqlx::query(RESET_STALE_VOID_USAGE_SQL)
.bind(&usage.request_id)
.execute(&mut **tx)
.await
.map_postgres_err()?;
sqlx::query(RESET_STALE_VOID_USAGE_SETTLEMENT_SNAPSHOT_SQL)
.bind(&usage.request_id)
.execute(&mut **tx)
.await
.map_postgres_err()?;
}
let request_headers_json = json_bind_text(usage.request_headers.as_ref())?;
let request_body_storage =
prepare_usage_body_storage(usage.request_body.as_ref())?;
@@ -3095,6 +3129,21 @@ mod tests {
));
}
#[test]
fn usage_sql_recovers_void_failures_before_upsert_and_settlement() {
assert!(super::RESET_STALE_VOID_USAGE_SQL.contains("UPDATE \"usage\""));
assert!(super::RESET_STALE_VOID_USAGE_SQL.contains("billing_status = 'pending'"));
assert!(super::RESET_STALE_VOID_USAGE_SQL.contains("finalized_at = NULL"));
assert!(super::RESET_STALE_VOID_USAGE_SQL.contains("status IN ('failed', 'cancelled')"));
assert!(super::RESET_STALE_VOID_USAGE_SETTLEMENT_SNAPSHOT_SQL
.contains("UPDATE usage_settlement_snapshots"));
assert!(super::RESET_STALE_VOID_USAGE_SETTLEMENT_SNAPSHOT_SQL
.contains("billing_status = 'pending'"));
assert!(
super::RESET_STALE_VOID_USAGE_SETTLEMENT_SNAPSHOT_SQL.contains("finalized_at = NULL")
);
}
#[test]
fn prepare_usage_body_storage_detaches_small_payloads_into_blob_storage() {
let payload = json!({"message": "hello"});