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

Merge remote-tracking branch 'origin/pr/538'

Browse Source
This commit is contained in:
fawney19
2026-05-21 23:12:42 +08:00
parent d6eb41aa78 cadc45c5b8
commit 0e6d7539ad
16 changed files with 422 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
apps/aether-gateway/src/execution_runtime/sync/execution.rs
View File

@@ -3,6 +3,7 @@ use std::io::Error as IoError;
use std::sync::Arc;
use std::time::{Duration, Instant};
use aether_ai_serving::UPSTREAM_IS_STREAM_KEY;
use aether_contracts::{
ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionPlan, ExecutionResult,
ExecutionTelemetry,
@@ -1034,7 +1035,7 @@ fn resolve_openai_image_sync_total_timeout_ms(plan: &ExecutionPlan) -> u64 {
fn report_context_upstream_is_stream(report_context: Option<&Value>) -> bool {
report_context
.and_then(|value| value.get("upstream_is_stream"))
.and_then(|value| value.get(UPSTREAM_IS_STREAM_KEY))
.and_then(Value::as_bool)
.unwrap_or(false)
}

75
apps/aether-gateway/src/executor/candidate_loop.rs
View File

@@ -1,5 +1,6 @@
use aether_ai_serving::{
run_ai_attempt_loop, AiAttemptLoopOutcome, AiAttemptLoopPort, AiExecutionAttempt,
UPSTREAM_IS_STREAM_KEY,
};
use aether_data_contracts::repository::candidates::RequestCandidateStatus;
use aether_scheduler_core::{
@@ -471,11 +472,24 @@ fn should_skip_unused_persistence_from_metadata(
metadata.candidate_group_id.is_some() && metadata.pool_key_index.is_some()
}
fn resolve_stream_candidate_watchdog_timeout(plan: &aether_contracts::ExecutionPlan) -> Duration {
fn resolve_stream_candidate_watchdog_timeout(
plan: &aether_contracts::ExecutionPlan,
report_context: Option<&serde_json::Value>,
) -> Duration {
let upstream_is_stream = report_context
.and_then(|context| context.get(UPSTREAM_IS_STREAM_KEY))
.and_then(serde_json::Value::as_bool)
.unwrap_or(true);
let timeout_ms = plan
.timeouts
.as_ref()
.and_then(|timeouts| timeouts.first_byte_ms.or(timeouts.total_ms))
.and_then(|timeouts| {
if upstream_is_stream {
timeouts.first_byte_ms.or(timeouts.total_ms)
} else {
timeouts.total_ms.or(timeouts.first_byte_ms)
}
})
.unwrap_or(DEFAULT_STREAM_CANDIDATE_WATCHDOG_TIMEOUT_MS)
.max(1);
Duration::from_millis(timeout_ms)
@@ -493,7 +507,7 @@ where
Fut:
std::future::Future<Output = Result<Option<Response<Body>>, GatewayError>> + Send + 'static,
{
let timeout_duration = resolve_stream_candidate_watchdog_timeout(plan);
let timeout_duration = resolve_stream_candidate_watchdog_timeout(plan, report_context);
let candidate_started_unix_ms = current_unix_ms();
let mut join_handle = tokio::spawn(execute());
match timeout(timeout_duration, &mut join_handle).await {
@@ -655,19 +669,22 @@ mod tests {
#[test]
fn stream_candidate_watchdog_prefers_first_byte_timeout() {
let timeout =
resolve_stream_candidate_watchdog_timeout(&test_plan(Some(ExecutionTimeouts {
let report_context = json!({"upstream_is_stream": true});
let timeout = resolve_stream_candidate_watchdog_timeout(
&test_plan(Some(ExecutionTimeouts {
first_byte_ms: Some(12_345),
total_ms: Some(90_000),
..ExecutionTimeouts::default()
})));
})),
Some(&report_context),
);
assert_eq!(timeout, Duration::from_millis(12_345));
}
#[test]
fn stream_candidate_watchdog_uses_default_when_timeouts_missing() {
let timeout = resolve_stream_candidate_watchdog_timeout(&test_plan(None));
let timeout = resolve_stream_candidate_watchdog_timeout(&test_plan(None), None);
assert_eq!(
timeout,
@@ -675,6 +692,50 @@ mod tests {
);
}
#[test]
fn stream_candidate_watchdog_prefers_total_timeout_when_upstream_non_stream() {
let report_context = json!({"upstream_is_stream": false});
let timeout = resolve_stream_candidate_watchdog_timeout(
&test_plan(Some(ExecutionTimeouts {
first_byte_ms: Some(300_000),
total_ms: Some(599_000),
..ExecutionTimeouts::default()
})),
Some(&report_context),
);
assert_eq!(timeout, Duration::from_millis(599_000));
}
#[test]
fn stream_candidate_watchdog_falls_back_to_first_byte_when_upstream_non_stream_lacks_total() {
let report_context = json!({"upstream_is_stream": false});
let timeout = resolve_stream_candidate_watchdog_timeout(
&test_plan(Some(ExecutionTimeouts {
first_byte_ms: Some(300_000),
..ExecutionTimeouts::default()
})),
Some(&report_context),
);
assert_eq!(timeout, Duration::from_millis(300_000));
}
#[test]
fn stream_candidate_watchdog_defaults_to_streaming_when_flag_missing() {
let report_context = json!({});
let timeout = resolve_stream_candidate_watchdog_timeout(
&test_plan(Some(ExecutionTimeouts {
first_byte_ms: Some(12_345),
total_ms: Some(90_000),
..ExecutionTimeouts::default()
})),
Some(&report_context),
);
assert_eq!(timeout, Duration::from_millis(12_345));
}
#[test]
fn unused_persistence_skips_pool_internal_candidates() {
assert!(should_skip_unused_persistence(Some(&json!({

3
apps/aether-gateway/src/handlers/public/support/user_me_usage.rs
View File

@@ -1,5 +1,6 @@
use std::collections::{BTreeMap, BTreeSet};
use aether_ai_serving::UPSTREAM_IS_STREAM_KEY;
use aether_billing::{
normalize_input_tokens_for_billing, normalize_total_input_context_for_cache_hit_rate,
};
@@ -314,7 +315,7 @@ fn users_me_usage_upstream_is_stream(item: &StoredRequestUsageAudit) -> bool {
item.request_metadata
.as_ref()
.and_then(serde_json::Value::as_object)
.and_then(|metadata| metadata.get("upstream_is_stream"))
.and_then(|metadata| metadata.get(UPSTREAM_IS_STREAM_KEY))
.and_then(serde_json::Value::as_bool)
.or_else(|| users_me_usage_headers_stream_flag(item.response_headers.as_ref()))
.or_else(|| users_me_usage_infer_upstream_stream_from_captured_bodies(item))

8
crates/aether-admin/src/observability/usage.rs
View File

@@ -1,5 +1,6 @@
use crate::observability::stats::{aggregate_usage_stats, parse_bounded_u32, round_to};
use aether_ai_formats::api::request_path_implies_stream_request;
use aether_ai_formats::UPSTREAM_IS_STREAM_KEY;
use aether_billing::{
normalize_input_tokens_for_billing, normalize_total_input_context_for_cache_hit_rate,
};
@@ -1052,7 +1053,7 @@ fn admin_usage_upstream_is_stream(item: &StoredRequestUsageAudit) -> bool {
item.request_metadata
.as_ref()
.and_then(Value::as_object)
.and_then(|metadata| metadata.get("upstream_is_stream"))
.and_then(|metadata| metadata.get(UPSTREAM_IS_STREAM_KEY))
.and_then(Value::as_bool)
.or_else(|| admin_usage_headers_stream_flag(item.response_headers.as_ref()))
.or_else(|| admin_usage_infer_upstream_stream_from_captured_bodies(item))
@@ -1256,7 +1257,10 @@ pub fn admin_usage_record_json(
.as_object_mut()
.expect("admin usage record payload should be an object");
object.insert("is_stream".to_string(), json!(item.is_stream));
object.insert("upstream_is_stream".to_string(), json!(upstream_is_stream));
object.insert(
UPSTREAM_IS_STREAM_KEY.to_string(),
json!(upstream_is_stream),
);
object.insert(
"client_requested_stream".to_string(),
json!(client_is_stream),

8
crates/aether-ai-formats/src/formats/shared/request.rs
View File

@@ -2,6 +2,14 @@ use base64::Engine as _;
use crate::formats::id::api_format_uses_body_stream_field;
/// JSON key under which `upstream_is_stream` is written into the AI execution
/// report context and propagated into usage metadata. Shared by the producer
/// (`aether-ai-serving::report_context`) and every downstream consumer so that
/// renames cannot silently desync them — a string-literal mismatch here would
/// degrade to default values (e.g. assuming streaming) without any compile-time
/// signal.
pub const UPSTREAM_IS_STREAM_KEY: &str = "upstream_is_stream";
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum UpstreamStreamPolicy {
Auto,

3
crates/aether-ai-formats/src/formats/shared/response.rs
View File

@@ -3,6 +3,7 @@ use std::collections::BTreeMap;
use serde_json::Value;
use crate::contracts::core_success_background_report_kind;
use crate::formats::shared::request::UPSTREAM_IS_STREAM_KEY;
#[derive(Debug, Clone, PartialEq)]
pub struct LocalSyncReportParts {
@@ -66,7 +67,7 @@ fn should_capture_client_sync_success_body(payload: &LocalSyncReportParts) -> bo
.report_context
.as_ref()
.and_then(Value::as_object)
.and_then(|context| context.get("upstream_is_stream"))
.and_then(|context| context.get(UPSTREAM_IS_STREAM_KEY))
.and_then(Value::as_bool)
.unwrap_or(false)
}

2
crates/aether-ai-formats/src/lib.rs
View File

@@ -29,7 +29,7 @@ pub use formats::shared::model_directives::{
};
pub use formats::shared::request::{
endpoint_config_forces_upstream_stream_policy, enforce_request_body_stream_field,
resolve_upstream_is_stream_from_endpoint_config,
resolve_upstream_is_stream_from_endpoint_config, UPSTREAM_IS_STREAM_KEY,
};
pub use protocol::canonical::{
canonical_request_unknown_block_count, canonical_response_unknown_block_count,

1
crates/aether-ai-serving/src/lib.rs
View File

@@ -22,6 +22,7 @@ pub mod request_body_diagnostics;
pub mod runtime_miss;
pub mod surface_spec;
pub use aether_ai_formats::UPSTREAM_IS_STREAM_KEY;
pub use aether_pool_core::{
normalize_enabled_pool_presets, run_pool_scheduler, PoolCandidateFacts, PoolCandidateInput,
PoolCandidateOrchestration, PoolMemberSignals, PoolRuntimeState, PoolScheduledCandidate,

3
crates/aether-ai-serving/src/report_context.rs
View File

@@ -1,6 +1,7 @@
use std::collections::BTreeMap;
use aether_ai_formats::api::ExecutionRuntimeAuthContext;
use aether_ai_formats::UPSTREAM_IS_STREAM_KEY;
use aether_scheduler_core::SchedulerRankingOutcome;
use serde_json::{Map, Value};
@@ -140,7 +141,7 @@ pub fn build_ai_execution_report_context(parts: AiExecutionReportContextParts<'_
Value::Bool(parts.client_requested_stream),
);
object.insert(
"upstream_is_stream".to_string(),
UPSTREAM_IS_STREAM_KEY.to_string(),
Value::Bool(parts.upstream_is_stream),
);
object.insert("has_envelope".to_string(), Value::Bool(parts.has_envelope));

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

@@ -2,6 +2,7 @@ use std::collections::BTreeMap;
use std::sync::Arc;
use std::sync::RwLock;
use aether_ai_formats::UPSTREAM_IS_STREAM_KEY;
use aether_data_contracts::repository::usage::{
parse_usage_body_ref, usage_body_ref, StoredUsageAuditAggregation, StoredUsageAuditSummary,
StoredUsageBreakdownSummaryRow, StoredUsageCacheAffinityHitSummary,
@@ -952,7 +953,7 @@ fn usage_output_tps_uses_generation_time(item: &StoredRequestUsageAudit) -> bool
item.request_metadata
.as_ref()
.and_then(Value::as_object)
.and_then(|metadata| metadata.get("upstream_is_stream"))
.and_then(|metadata| metadata.get(UPSTREAM_IS_STREAM_KEY))
.and_then(Value::as_bool)
.unwrap_or(item.is_stream)
}

81
crates/aether-data/src/repository/usage/mysql.rs
View File

@@ -1,5 +1,6 @@
use std::collections::{BTreeMap, HashSet};
use aether_ai_formats::UPSTREAM_IS_STREAM_KEY;
use async_trait::async_trait;
use sqlx::{mysql::MySqlRow, Row};
@@ -518,13 +519,20 @@ WHERE request_id = ?
continue;
}
let error_message = stale_pending_error_message(&row.status, timeout_minutes);
let candidate_info =
latest_failed_candidate_mysql(&mut tx, &row.request_id).await?;
let (status_code, error_message) = resolve_stale_pending_failure(
candidate_info.as_ref(),
&row.status,
timeout_minutes,
);
let status_code_i64 = i64::from(status_code);
if row.billing_status == "pending" {
sqlx::query(
r#"
UPDATE `usage`
SET status = 'failed',
status_code = 504,
status_code = ?,
error_message = ?,
billing_status = 'void',
finalized_at = ?,
@@ -533,6 +541,7 @@ SET status = 'failed',
WHERE request_id = ?
"#,
)
.bind(status_code_i64)
.bind(&error_message)
.bind(to_i64(now_unix_secs, "usage finalized_at")?)
.bind(&row.request_id)
@@ -550,11 +559,12 @@ WHERE request_id = ?
r#"
UPDATE `usage`
SET status = 'failed',
status_code = 504,
status_code = ?,
error_message = ?
WHERE request_id = ?
"#,
)
.bind(status_code_i64)
.bind(&error_message)
.bind(&row.request_id)
.execute(&mut *tx)
@@ -687,6 +697,67 @@ fn stale_pending_error_message(status: &str, timeout_minutes: u64) -> String {
format!("请求超时: 状态 '{status}' 超过 {timeout_minutes} 分钟未完成")
}
struct FailedCandidateCleanupInfo {
status_code: Option<u16>,
error_message: Option<String>,
}
fn resolve_stale_pending_failure(
candidate: Option<&FailedCandidateCleanupInfo>,
status: &str,
timeout_minutes: u64,
) -> (u16, String) {
match candidate {
Some(info) => (
info.status_code.unwrap_or(502),
info.error_message
.clone()
.unwrap_or_else(|| stale_pending_error_message(status, timeout_minutes)),
),
None => (504, stale_pending_error_message(status, timeout_minutes)),
}
}
async fn latest_failed_candidate_mysql(
tx: &mut sqlx::Transaction<'_, sqlx::MySql>,
request_id: &str,
) -> Result<Option<FailedCandidateCleanupInfo>, DataLayerError> {
let row = sqlx::query(
r#"
SELECT status_code, error_message
FROM request_candidates
WHERE request_id = ?
AND status IN ('failed', 'cancelled')
ORDER BY
COALESCE(finished_at, started_at, created_at) DESC,
retry_index DESC,
candidate_index DESC
LIMIT 1
"#,
)
.bind(request_id)
.fetch_optional(&mut **tx)
.await
.map_sql_err()?;
let Some(row) = row else {
return Ok(None);
};
let status_code = row
.try_get::<Option<i64>, _>("status_code")
.map_sql_err()?
.and_then(|value| u16::try_from(value).ok());
let error_message = row
.try_get::<Option<String>, _>("error_message")
.map_sql_err()?
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty());
Ok(Some(FailedCandidateCleanupInfo {
status_code,
error_message,
}))
}
fn bind_upsert<'q>(
mut query: sqlx::query::Query<'q, sqlx::MySql, sqlx::mysql::MySqlArguments>,
usage: &'q UpsertUsageRecord,
@@ -881,7 +952,7 @@ fn usage_upstream_is_stream(usage: &UpsertUsageRecord) -> bool {
.request_metadata
.as_ref()
.and_then(serde_json::Value::as_object)
.and_then(|metadata| metadata.get("upstream_is_stream"))
.and_then(|metadata| metadata.get(UPSTREAM_IS_STREAM_KEY))
.and_then(serde_json::Value::as_bool)
.unwrap_or_else(|| usage.is_stream.unwrap_or(false))
}
@@ -895,7 +966,7 @@ fn merge_usage_stream_metadata(metadata: &mut Option<serde_json::Value>, upstrea
return;
};
object
.entry("upstream_is_stream")
.entry(UPSTREAM_IS_STREAM_KEY)
.or_insert(serde_json::Value::Bool(upstream));
}

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

@@ -1690,10 +1690,25 @@ SET status = 'completed',
WHERE request_id = $1
"#;
const SELECT_LATEST_FAILED_CANDIDATE_FOR_STALE_REQUESTS_SQL: &str = r#"
SELECT DISTINCT ON (request_id)
request_id,
status_code,
error_message
FROM request_candidates
WHERE request_id = ANY($1)
AND status IN ('failed', 'cancelled')
ORDER BY request_id,
COALESCE(finished_at, started_at, created_at) DESC,
retry_index DESC,
candidate_index DESC,
created_at DESC
"#;
const UPDATE_FAILED_STALE_USAGE_SQL: &str = r#"
UPDATE usage
SET status = 'failed',
status_code = 504,
status_code = $3,
error_message = $2
WHERE request_id = $1
"#;
@@ -1702,7 +1717,7 @@ const UPDATE_FAILED_VOID_STALE_USAGE_SQL: &str = r#"
WITH updated_usage AS (
UPDATE usage
SET status = 'failed',
status_code = 504,
status_code = $4,
error_message = $2,
billing_status = 'void',
finalized_at = $3,
@@ -8292,17 +8307,44 @@ ORDER BY "usage".user_id ASC
.iter()
.map(|row| row.request_id.clone())
.collect::<Vec<_>>();
let completed_request_ids = if request_ids.is_empty() {
Vec::new()
let (completed_request_ids, failed_candidate_info) = if request_ids.is_empty() {
(Vec::new(), std::collections::HashMap::new())
} else {
sqlx::query(SELECT_COMPLETED_PENDING_REQUEST_IDS_SQL)
.bind(request_ids)
let completed = sqlx::query(SELECT_COMPLETED_PENDING_REQUEST_IDS_SQL)
.bind(&request_ids)
.fetch_all(&mut *tx)
.await
.map_postgres_err()?
.iter()
.map(|row| row.try_get("request_id").map_postgres_err())
.collect::<Result<Vec<String>, DataLayerError>>()?
.collect::<Result<Vec<String>, DataLayerError>>()?;
let failed_rows =
sqlx::query(SELECT_LATEST_FAILED_CANDIDATE_FOR_STALE_REQUESTS_SQL)
.bind(&request_ids)
.fetch_all(&mut *tx)
.await
.map_postgres_err()?;
let mut failed_map = std::collections::HashMap::new();
for row in failed_rows {
let request_id: String = row.try_get("request_id").map_postgres_err()?;
let status_code = row
.try_get::<Option<i32>, _>("status_code")
.map_postgres_err()?
.and_then(|value| u16::try_from(value).ok());
let error_message = row
.try_get::<Option<String>, _>("error_message")
.map_postgres_err()?
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty());
failed_map.insert(
request_id,
FailedCandidateCleanupInfo {
status_code,
error_message,
},
);
}
(completed, failed_map)
};
for row in stale_rows {
@@ -8322,12 +8364,16 @@ ORDER BY "usage".user_id ASC
continue;
}
let error_message = stale_pending_error_message(&row.status, timeout_minutes);
let candidate_info = failed_candidate_info.get(&row.request_id);
let (status_code, error_message) =
resolve_stale_pending_failure(candidate_info, &row.status, timeout_minutes);
let status_code_i32 = i32::from(status_code);
if row.billing_status == "pending" {
sqlx::query(UPDATE_FAILED_VOID_STALE_USAGE_SQL)
.bind(&row.request_id)
.bind(&error_message)
.bind(now)
.bind(status_code_i32)
.execute(&mut *tx)
.await
.map_postgres_err()?;
@@ -8335,6 +8381,7 @@ ORDER BY "usage".user_id ASC
sqlx::query(UPDATE_FAILED_STALE_USAGE_SQL)
.bind(&row.request_id)
.bind(&error_message)
.bind(status_code_i32)
.execute(&mut *tx)
.await
.map_postgres_err()?;
@@ -8785,10 +8832,31 @@ struct StalePendingUsageRow {
billing_status: String,
}
struct FailedCandidateCleanupInfo {
status_code: Option<u16>,
error_message: Option<String>,
}
fn stale_pending_error_message(status: &str, timeout_minutes: u64) -> String {
format!("请求超时: 状态 '{status}' 超过 {timeout_minutes} 分钟未完成")
}
fn resolve_stale_pending_failure(
candidate: Option<&FailedCandidateCleanupInfo>,
status: &str,
timeout_minutes: u64,
) -> (u16, String) {
match candidate {
Some(info) => (
info.status_code.unwrap_or(502),
info.error_message
.clone()
.unwrap_or_else(|| stale_pending_error_message(status, timeout_minutes)),
),
None => (504, stale_pending_error_message(status, timeout_minutes)),
}
}
async fn find_usage_by_request_id_in_tx(
tx: &mut sqlx::Transaction<'_, Postgres>,
request_id: &str,

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

@@ -825,6 +825,14 @@ fn usage_sql_clears_stale_failure_fields_for_non_failed_status_updates() {
));
}
#[test]
fn stale_cleanup_failed_candidate_sql_orders_by_effective_timestamp() {
let sql = super::SELECT_LATEST_FAILED_CANDIDATE_FOR_STALE_REQUESTS_SQL;
assert!(sql.contains("COALESCE(finished_at, started_at, created_at) DESC"));
assert!(!sql.contains("finished_at DESC NULLS LAST"));
assert!(!sql.contains("started_at DESC NULLS LAST"));
}
#[test]
fn usage_sql_does_not_allow_streaming_to_regress_back_to_pending() {
assert!(super::UPSERT_SQL.contains(

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

@@ -1,6 +1,7 @@
use std::collections::{BTreeMap, HashSet};
use std::io::Read;
use aether_ai_formats::UPSTREAM_IS_STREAM_KEY;
use aether_data_contracts::repository::usage::{parse_usage_body_ref, UsageBodyField};
use async_trait::async_trait;
use flate2::read::GzDecoder;
@@ -3548,13 +3549,20 @@ WHERE request_id = ?
continue;
}
let error_message = stale_pending_error_message(&row.status, timeout_minutes);
let candidate_info =
latest_failed_candidate_sqlite(&mut tx, &row.request_id).await?;
let (status_code, error_message) = resolve_stale_pending_failure(
candidate_info.as_ref(),
&row.status,
timeout_minutes,
);
let status_code_i64 = i64::from(status_code);
if row.billing_status == "pending" {
sqlx::query(
r#"
UPDATE "usage"
SET status = 'failed',
status_code = 504,
status_code = ?,
error_message = ?,
billing_status = 'void',
finalized_at = ?,
@@ -3563,6 +3571,7 @@ SET status = 'failed',
WHERE request_id = ?
"#,
)
.bind(status_code_i64)
.bind(&error_message)
.bind(to_i64(now_unix_secs, "usage finalized_at")?)
.bind(&row.request_id)
@@ -3580,11 +3589,12 @@ WHERE request_id = ?
r#"
UPDATE "usage"
SET status = 'failed',
status_code = 504,
status_code = ?,
error_message = ?
WHERE request_id = ?
"#,
)
.bind(status_code_i64)
.bind(&error_message)
.bind(&row.request_id)
.execute(&mut *tx)
@@ -3707,6 +3717,67 @@ fn stale_pending_error_message(status: &str, timeout_minutes: u64) -> String {
format!("请求超时: 状态 '{status}' 超过 {timeout_minutes} 分钟未完成")
}
struct FailedCandidateCleanupInfo {
status_code: Option<u16>,
error_message: Option<String>,
}
fn resolve_stale_pending_failure(
candidate: Option<&FailedCandidateCleanupInfo>,
status: &str,
timeout_minutes: u64,
) -> (u16, String) {
match candidate {
Some(info) => (
info.status_code.unwrap_or(502),
info.error_message
.clone()
.unwrap_or_else(|| stale_pending_error_message(status, timeout_minutes)),
),
None => (504, stale_pending_error_message(status, timeout_minutes)),
}
}
async fn latest_failed_candidate_sqlite(
tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>,
request_id: &str,
) -> Result<Option<FailedCandidateCleanupInfo>, DataLayerError> {
let row = sqlx::query(
r#"
SELECT status_code, error_message
FROM request_candidates
WHERE request_id = ?
AND status IN ('failed', 'cancelled')
ORDER BY
COALESCE(finished_at, started_at, created_at) DESC,
retry_index DESC,
candidate_index DESC
LIMIT 1
"#,
)
.bind(request_id)
.fetch_optional(&mut **tx)
.await
.map_sql_err()?;
let Some(row) = row else {
return Ok(None);
};
let status_code = row
.try_get::<Option<i64>, _>("status_code")
.map_sql_err()?
.and_then(|value| u16::try_from(value).ok());
let error_message = row
.try_get::<Option<String>, _>("error_message")
.map_sql_err()?
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty());
Ok(Some(FailedCandidateCleanupInfo {
status_code,
error_message,
}))
}
fn bind_upsert<'q>(
mut query: sqlx::query::Query<'q, sqlx::Sqlite, sqlx::sqlite::SqliteArguments<'q>>,
usage: &'q UpsertUsageRecord,
@@ -3898,7 +3969,7 @@ fn usage_upstream_is_stream(usage: &UpsertUsageRecord) -> bool {
.request_metadata
.as_ref()
.and_then(serde_json::Value::as_object)
.and_then(|metadata| metadata.get("upstream_is_stream"))
.and_then(|metadata| metadata.get(UPSTREAM_IS_STREAM_KEY))
.and_then(serde_json::Value::as_bool)
.unwrap_or_else(|| usage.is_stream.unwrap_or(false))
}
@@ -3912,7 +3983,7 @@ fn merge_usage_stream_metadata(metadata: &mut Option<serde_json::Value>, upstrea
return;
};
object
.entry("upstream_is_stream")
.entry(UPSTREAM_IS_STREAM_KEY)
.or_insert(serde_json::Value::Bool(upstream));
}
@@ -4139,6 +4210,91 @@ ORDER BY request_id
assert_eq!(snapshot, ("void".to_string(), Some(10)));
}
#[tokio::test]
async fn sqlite_usage_write_repository_cleanup_uses_failed_candidate_status_when_present() {
let pool = sqlx::sqlite::SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.expect("sqlite pool should connect");
run_sqlite_migrations(&pool)
.await
.expect("sqlite migrations should run");
seed_stats_targets(&pool).await;
let repository = SqliteUsageWriteRepository::new(pool.clone());
repository
.upsert(sample_usage(
"request-upstream-reset",
"pending",
"pending",
1,
))
.await
.expect("pending usage should upsert");
repository
.upsert(sample_usage("request-stuck", "pending", "pending", 1))
.await
.expect("pending usage should upsert");
// request-upstream-reset has a failed candidate carrying a concrete 502 status
// and a connection-reset message — cleanup should use them instead of 504.
// request-stuck has only a still-pending candidate, so cleanup should fall back to 504.
sqlx::query(
r#"
INSERT INTO request_candidates (
id,
request_id,
candidate_index,
retry_index,
status,
status_code,
error_message,
is_cached,
created_at,
started_at,
finished_at
) VALUES
('candidate-reset', 'request-upstream-reset', 0, 0, 'failed', 502, 'upstream connection reset by peer', 0, 1, 2, 3),
('candidate-stuck', 'request-stuck', 0, 0, 'pending', NULL, NULL, 0, 1, NULL, NULL)
"#,
)
.execute(&pool)
.await
.expect("request candidates should seed");
let summary = repository
.cleanup_stale_pending_requests(2, 10, 5, 5)
.await
.expect("cleanup should run");
assert_eq!(summary.recovered, 0);
assert_eq!(summary.failed, 2);
let reset = repository
.find_by_request_id("request-upstream-reset")
.await
.expect("upstream-reset usage should load")
.expect("upstream-reset usage should exist");
assert_eq!(reset.status, "failed");
assert_eq!(reset.status_code, Some(502));
assert_eq!(
reset.error_message.as_deref(),
Some("upstream connection reset by peer")
);
let stuck = repository
.find_by_request_id("request-stuck")
.await
.expect("stuck usage should load")
.expect("stuck usage should exist");
assert_eq!(stuck.status, "failed");
assert_eq!(stuck.status_code, Some(504));
assert!(stuck
.error_message
.as_deref()
.is_some_and(|message| message.contains("超过 5 分钟未完成")));
}
#[tokio::test]
async fn sqlite_usage_read_repository_reads_usage_contract_views() {
let pool = sqlx::sqlite::SqlitePoolOptions::new()

5
crates/aether-usage-runtime/src/request_metadata.rs
View File

@@ -1,6 +1,7 @@
use aether_ai_formats::api::{
sanitize_request_path, sanitize_request_path_and_query, sanitize_request_query_string,
};
use aether_ai_formats::UPSTREAM_IS_STREAM_KEY;
use aether_contracts::ExecutionPlan;
use serde_json::{json, Map, Value};
@@ -74,7 +75,7 @@ fn copy_allowed_metadata_fields(source: &Map<String, Value>, target: &mut Map<St
copy_non_empty_string(source, target, "user_agent");
copy_non_empty_string(source, target, "client_family");
copy_bool(source, target, "client_requested_stream");
copy_bool(source, target, "upstream_is_stream");
copy_bool(source, target, UPSTREAM_IS_STREAM_KEY);
copy_non_null_value(source, target, "client_session_affinity");
copy_bool(source, target, "api_key_is_standalone");
copy_non_empty_string(source, target, "request_path");
@@ -114,7 +115,7 @@ fn move_allowed_metadata_fields(mut source: Map<String, Value>, target: &mut Map
remove_non_empty_string(&mut source, target, "user_agent");
remove_non_empty_string(&mut source, target, "client_family");
remove_bool(&mut source, target, "client_requested_stream");
remove_bool(&mut source, target, "upstream_is_stream");
remove_bool(&mut source, target, UPSTREAM_IS_STREAM_KEY);
remove_non_null_value(&mut source, target, "client_session_affinity");
remove_bool(&mut source, target, "api_key_is_standalone");
remove_non_empty_string(&mut source, target, "request_path");

7
crates/aether-usage-runtime/src/write.rs
View File

@@ -1,5 +1,6 @@
use std::collections::BTreeMap;
use aether_ai_formats::UPSTREAM_IS_STREAM_KEY;
use aether_contracts::{ExecutionPlan, ExecutionTelemetry};
use aether_data_contracts::repository::usage::{UpsertUsageRecord, UsageBodyCaptureState};
use aether_data_contracts::DataLayerError;
@@ -719,7 +720,7 @@ pub fn build_sync_terminal_usage_payload_seed(
.report_context
.as_ref()
.and_then(Value::as_object)
.and_then(|context| context.get("upstream_is_stream"))
.and_then(|context| context.get(UPSTREAM_IS_STREAM_KEY))
.and_then(Value::as_bool)
.unwrap_or(false);
let provider_response_full = if upstream_is_stream && payload.body_base64.is_some() {
@@ -1632,9 +1633,9 @@ fn build_runtime_request_metadata_seed_from_parts(
Value::Bool(client_requested_stream),
);
}
if let Some(upstream_is_stream) = context_bool(context, "upstream_is_stream") {
if let Some(upstream_is_stream) = context_bool(context, UPSTREAM_IS_STREAM_KEY) {
metadata.insert(
"upstream_is_stream".to_string(),
UPSTREAM_IS_STREAM_KEY.to_string(),
Value::Bool(upstream_is_stream),
);
}