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

fix(gateway): SSE 流跳过成功探测与 prefetch，统计聚合 SQL 显式 CAST BIGINT

Browse Source 
- stream: event-stream 响应跳过 direct finalize prefetch 与成功失败转移探测，避免消费流首帧
- stream_pump: 仅在已知 content-length 时缓冲非 SSE 响应
- maintenance: usage_billing_facts 聚合列显式 CAST 为 BIGINT，避免类型不匹配
This commit is contained in:
fawney19
2026-04-24 18:14:59 +08:00
parent f695238e8a
commit 3cc54deb9c
3 changed files with 127 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files

57
apps/aether-gateway/src/execution_runtime/stream/execution.rs
View File

@@ -619,14 +619,7 @@ fn should_skip_direct_finalize_prefetch(
has_private_stream_normalizer: bool, has_private_stream_normalizer: bool,
has_local_stream_rewriter: bool, has_local_stream_rewriter: bool,
) -> bool { ) -> bool {
if direct_stream_finalize_kind.is_none() if direct_stream_finalize_kind.is_none() {
|| has_private_stream_normalizer
|| has_local_stream_rewriter
{
return false;
}
if !provider_api_format.eq_ignore_ascii_case(client_api_format) {
return false; return false;
} }
@@ -635,6 +628,18 @@ fn should_skip_direct_finalize_prefetch(
.filter(|value| !value.is_empty()) .filter(|value| !value.is_empty())
.unwrap_or_default() .unwrap_or_default()
.to_ascii_lowercase(); .to_ascii_lowercase();
if content_type.contains("text/event-stream") {
return true;
}
if has_private_stream_normalizer || has_local_stream_rewriter {
return false;
}
if !provider_api_format.eq_ignore_ascii_case(client_api_format) {
return false;
}
if content_type.is_empty() { if content_type.is_empty() {
return true; return true;
} }
@@ -642,6 +647,18 @@ fn should_skip_direct_finalize_prefetch(
!(content_type.contains("json") || content_type.ends_with("+json")) !(content_type.contains("json") || content_type.ends_with("+json"))
} }
fn should_probe_success_failover_before_stream(headers: &BTreeMap<String, String>) -> bool {
let content_type = headers
.get("content-type")
.map(String::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or_default()
.to_ascii_lowercase();
content_type.contains("json") || content_type.ends_with("+json")
}
async fn probe_local_stream_success_failover_text<R>( async fn probe_local_stream_success_failover_text<R>(
buffered_frames: &mut VecDeque<StreamFrame>, buffered_frames: &mut VecDeque<StreamFrame>,
lines: &mut FramedRead<R, LinesCodec>, lines: &mut FramedRead<R, LinesCodec>,
@@ -712,7 +729,7 @@ async fn execute_stream_from_frame_stream(
}; };
let mut buffered_frames = VecDeque::new(); let mut buffered_frames = VecDeque::new();
let mut stream_terminal_summary: Option<ExecutionStreamTerminalSummary> = None; let mut stream_terminal_summary: Option<ExecutionStreamTerminalSummary> = None;
if status_code == 200 { if status_code == 200 && should_probe_success_failover_before_stream(&headers) {
let success_probe_text = let success_probe_text =
probe_local_stream_success_failover_text(&mut buffered_frames, &mut lines).await?; probe_local_stream_success_failover_text(&mut buffered_frames, &mut lines).await?;
if should_retry_next_local_candidate_stream( if should_retry_next_local_candidate_stream(
@@ -2131,7 +2148,10 @@ mod tests {
use serde_json::{json, Value}; use serde_json::{json, Value};
use tokio::sync::{watch, Notify}; use tokio::sync::{watch, Notify};
use super::{execute_execution_runtime_stream, should_skip_direct_finalize_prefetch}; use super::{
execute_execution_runtime_stream, should_probe_success_failover_before_stream,
should_skip_direct_finalize_prefetch,
};
use crate::control::GatewayControlDecision; use crate::control::GatewayControlDecision;
use crate::tunnel::{tunnel_protocol, TunnelProxyConn}; use crate::tunnel::{tunnel_protocol, TunnelProxyConn};
use crate::AppState; use crate::AppState;
@@ -2195,8 +2215,8 @@ mod tests {
} }
#[test] #[test]
fn keeps_prefetch_for_cross_format_or_rewritten_streams() { fn skips_prefetch_for_event_streams_even_when_cross_format_or_rewritten() {
assert!(!should_skip_direct_finalize_prefetch( assert!(should_skip_direct_finalize_prefetch(
Some("claude_cli_sync_finalize"), Some("claude_cli_sync_finalize"),
Some("text/event-stream"), Some("text/event-stream"),
"openai:chat", "openai:chat",
@@ -2206,6 +2226,19 @@ mod tests {
)); ));
} }
#[test]
fn skips_success_failover_probe_for_event_streams() {
assert!(!should_probe_success_failover_before_stream(
&BTreeMap::from([(
"content-type".to_string(),
"text/event-stream; charset=utf-8".to_string(),
)])
));
assert!(should_probe_success_failover_before_stream(
&BTreeMap::from([("content-type".to_string(), "application/json".to_string(),)])
));
}
#[tokio::test] #[tokio::test]
async fn execute_execution_runtime_stream_records_first_data_as_streaming_before_terminal_telemetry( async fn execute_execution_runtime_stream_records_first_data_as_streaming_before_terminal_telemetry(
) { ) {

41
apps/aether-gateway/src/execution_runtime/stream_pump.rs
View File

@@ -60,7 +60,7 @@ pub(crate) fn build_direct_execution_frame_stream(
let mut stream_terminal_observer = StreamingStandardTerminalObserver::default(); let mut stream_terminal_observer = StreamingStandardTerminalObserver::default();
let mut observer_buffered = Vec::new(); let mut observer_buffered = Vec::new();
if !should_treat_upstream_response_as_stream(&headers, &observer_context) { if should_buffer_non_stream_response(&headers, &observer_context) {
let original_headers = headers.clone(); let original_headers = headers.clone();
match buffer_non_sse_upstream_body(response, started_at).await { match buffer_non_sse_upstream_body(response, started_at).await {
Ok(buffered) => { Ok(buffered) => {
@@ -402,6 +402,20 @@ fn should_treat_upstream_response_as_stream(
.is_some_and(|value| value.eq_ignore_ascii_case(crate::ai_pipeline::KIRO_ENVELOPE_NAME)) .is_some_and(|value| value.eq_ignore_ascii_case(crate::ai_pipeline::KIRO_ENVELOPE_NAME))
} }
fn should_buffer_non_stream_response(
headers: &BTreeMap<String, String>,
report_context: &Value,
) -> bool {
if should_treat_upstream_response_as_stream(headers, report_context) {
return false;
}
headers
.get("content-length")
.and_then(|value| value.trim().parse::<u64>().ok())
.is_some()
}
async fn buffer_non_sse_upstream_body( async fn buffer_non_sse_upstream_body(
response: DirectUpstreamResponse, response: DirectUpstreamResponse,
started_at: Instant, started_at: Instant,
@@ -662,7 +676,10 @@ mod tests {
use serde_json::Value; use serde_json::Value;
use tokio::sync::watch; use tokio::sync::watch;
use super::{build_direct_execution_frame_stream, should_treat_upstream_response_as_stream}; use super::{
build_direct_execution_frame_stream, should_buffer_non_stream_response,
should_treat_upstream_response_as_stream,
};
use crate::execution_runtime::transport::{ use crate::execution_runtime::transport::{
execute_stream_plan_via_local_tunnel, DirectSyncExecutionRuntime, DirectUpstreamResponse, execute_stream_plan_via_local_tunnel, DirectSyncExecutionRuntime, DirectUpstreamResponse,
}; };
@@ -693,6 +710,26 @@ mod tests {
)); ));
} }
#[test]
fn buffers_non_sse_response_only_when_content_length_is_known() {
let report_context = serde_json::json!({
"provider_api_format": "openai:chat",
"client_api_format": "openai:chat",
});
assert!(!should_buffer_non_stream_response(
&BTreeMap::from([("content-type".into(), "application/json".into())]),
&report_context
));
assert!(should_buffer_non_stream_response(
&BTreeMap::from([
("content-type".into(), "application/json".into()),
("content-length".into(), "128".into()),
]),
&report_context
));
}
#[tokio::test] #[tokio::test]
async fn direct_execution_frame_stream_reports_ttfb_after_first_upstream_chunk() { async fn direct_execution_frame_stream_reports_ttfb_after_first_upstream_chunk() {
let listener = crate::test_support::bind_loopback_listener() let listener = crate::test_support::bind_loopback_listener()

86
apps/aether-gateway/src/maintenance/runtime.rs
View File

@@ -4049,8 +4049,8 @@ SELECT
AND settled_usage.billing_status = 'settled' AND settled_usage.billing_status = 'settled'
AND COALESCE(CAST(settled_usage.total_cost_usd AS DOUBLE PRECISION), 0) > 0 AND COALESCE(CAST(settled_usage.total_cost_usd AS DOUBLE PRECISION), 0) > 0
) AS settled_last_finalized_at_unix_secs, ) AS settled_last_finalized_at_unix_secs,
COUNT(id) AS total_requests, CAST(COUNT(id) AS BIGINT) AS total_requests,
COALESCE( CAST(COALESCE(
SUM( SUM(
CASE CASE
WHEN status_code >= 400 WHEN status_code >= 400
@@ -4060,10 +4060,10 @@ SELECT
END END
), ),
0 0
) AS error_requests, ) AS BIGINT) AS error_requests,
COALESCE(SUM(input_tokens), 0) AS input_tokens, CAST(COALESCE(SUM(input_tokens), 0) AS BIGINT) AS input_tokens,
COALESCE(SUM(output_tokens), 0) AS output_tokens, CAST(COALESCE(SUM(output_tokens), 0) AS BIGINT) AS output_tokens,
COALESCE( CAST(COALESCE(
SUM( SUM(
CASE CASE
WHEN COALESCE(cache_creation_input_tokens, 0) = 0 WHEN COALESCE(cache_creation_input_tokens, 0) = 0
@@ -4077,8 +4077,8 @@ SELECT
END END
), ),
0 0
) AS cache_creation_tokens, ) AS BIGINT) AS cache_creation_tokens,
COALESCE(SUM(cache_read_input_tokens), 0) AS cache_read_tokens, CAST(COALESCE(SUM(cache_read_input_tokens), 0) AS BIGINT) AS cache_read_tokens,
CAST(COALESCE(SUM(total_cost_usd), 0) AS DOUBLE PRECISION) AS total_cost, CAST(COALESCE(SUM(total_cost_usd), 0) AS DOUBLE PRECISION) AS total_cost,
CAST(COALESCE(SUM(actual_total_cost_usd), 0) AS DOUBLE PRECISION) AS actual_total_cost, CAST(COALESCE(SUM(actual_total_cost_usd), 0) AS DOUBLE PRECISION) AS actual_total_cost,
COALESCE( COALESCE(
@@ -4091,7 +4091,7 @@ SELECT
), ),
0 0
) AS response_time_sum_ms, ) AS response_time_sum_ms,
COALESCE( CAST(COALESCE(
SUM( SUM(
CASE CASE
WHEN response_time_ms IS NOT NULL THEN 1 WHEN response_time_ms IS NOT NULL THEN 1
@@ -4099,7 +4099,7 @@ SELECT
END END
), ),
0 0
) AS response_time_samples, ) AS BIGINT) AS response_time_samples,
CAST(COALESCE(AVG(response_time_ms), 0) AS DOUBLE PRECISION) AS avg_response_time_ms CAST(COALESCE(AVG(response_time_ms), 0) AS DOUBLE PRECISION) AS avg_response_time_ms
FROM usage_billing_facts AS usage FROM usage_billing_facts AS usage
WHERE created_at >= $1 WHERE created_at >= $1
@@ -4193,8 +4193,8 @@ const UPSERT_STATS_HOURLY_USER_SQL: &str = r#"
WITH aggregated AS ( WITH aggregated AS (
SELECT SELECT
user_id, user_id,
COUNT(id) AS total_requests, CAST(COUNT(id) AS BIGINT) AS total_requests,
COALESCE( CAST(COALESCE(
SUM( SUM(
CASE CASE
WHEN status_code >= 400 WHEN status_code >= 400
@@ -4204,10 +4204,10 @@ WITH aggregated AS (
END END
), ),
0 0
) AS error_requests, ) AS BIGINT) AS error_requests,
COALESCE(SUM(input_tokens), 0) AS input_tokens, CAST(COALESCE(SUM(input_tokens), 0) AS BIGINT) AS input_tokens,
COALESCE(SUM(output_tokens), 0) AS output_tokens, CAST(COALESCE(SUM(output_tokens), 0) AS BIGINT) AS output_tokens,
COALESCE( CAST(COALESCE(
SUM( SUM(
CASE CASE
WHEN COALESCE(cache_creation_input_tokens, 0) = 0 WHEN COALESCE(cache_creation_input_tokens, 0) = 0
@@ -4221,8 +4221,8 @@ WITH aggregated AS (
END END
), ),
0 0
) AS cache_creation_tokens, ) AS BIGINT) AS cache_creation_tokens,
COALESCE(SUM(cache_read_input_tokens), 0) AS cache_read_tokens, CAST(COALESCE(SUM(cache_read_input_tokens), 0) AS BIGINT) AS cache_read_tokens,
CAST(COALESCE(SUM(total_cost_usd), 0) AS DOUBLE PRECISION) AS total_cost, CAST(COALESCE(SUM(total_cost_usd), 0) AS DOUBLE PRECISION) AS total_cost,
CAST(COALESCE(SUM(actual_total_cost_usd), 0) AS DOUBLE PRECISION) AS actual_total_cost, CAST(COALESCE(SUM(actual_total_cost_usd), 0) AS DOUBLE PRECISION) AS actual_total_cost,
CAST( CAST(
@@ -4238,7 +4238,7 @@ WITH aggregated AS (
0 0
) AS DOUBLE PRECISION ) AS DOUBLE PRECISION
) AS settled_total_cost, ) AS settled_total_cost,
COALESCE( CAST(COALESCE(
SUM( SUM(
CASE CASE
WHEN billing_status = 'settled' WHEN billing_status = 'settled'
@@ -4248,8 +4248,8 @@ WITH aggregated AS (
END END
), ),
0 0
) AS settled_total_requests, ) AS BIGINT) AS settled_total_requests,
COALESCE( CAST(COALESCE(
SUM( SUM(
CASE CASE
WHEN billing_status = 'settled' WHEN billing_status = 'settled'
@@ -4259,8 +4259,8 @@ WITH aggregated AS (
END END
), ),
0 0
) AS settled_input_tokens, ) AS BIGINT) AS settled_input_tokens,
COALESCE( CAST(COALESCE(
SUM( SUM(
CASE CASE
WHEN billing_status = 'settled' WHEN billing_status = 'settled'
@@ -4270,8 +4270,8 @@ WITH aggregated AS (
END END
), ),
0 0
) AS settled_output_tokens, ) AS BIGINT) AS settled_output_tokens,
COALESCE( CAST(COALESCE(
SUM( SUM(
CASE CASE
WHEN billing_status = 'settled' WHEN billing_status = 'settled'
@@ -4281,8 +4281,8 @@ WITH aggregated AS (
END END
), ),
0 0
) AS settled_cache_creation_tokens, ) AS BIGINT) AS settled_cache_creation_tokens,
COALESCE( CAST(COALESCE(
SUM( SUM(
CASE CASE
WHEN billing_status = 'settled' WHEN billing_status = 'settled'
@@ -4292,7 +4292,7 @@ WITH aggregated AS (
END END
), ),
0 0
) AS settled_cache_read_tokens, ) AS BIGINT) AS settled_cache_read_tokens,
MIN( MIN(
CASE CASE
WHEN billing_status = 'settled' WHEN billing_status = 'settled'
@@ -4321,7 +4321,7 @@ WITH aggregated AS (
), ),
0 0
) AS response_time_sum_ms, ) AS response_time_sum_ms,
COALESCE( CAST(COALESCE(
SUM( SUM(
CASE CASE
WHEN response_time_ms IS NOT NULL THEN 1 WHEN response_time_ms IS NOT NULL THEN 1
@@ -4329,7 +4329,7 @@ WITH aggregated AS (
END END
), ),
0 0
) AS response_time_samples ) AS BIGINT) AS response_time_samples
FROM usage_billing_facts AS usage FROM usage_billing_facts AS usage
WHERE created_at >= $1 WHERE created_at >= $1
AND created_at < $2 AND created_at < $2
@@ -4417,9 +4417,9 @@ const UPSERT_STATS_HOURLY_MODEL_SQL: &str = r#"
WITH aggregated AS ( WITH aggregated AS (
SELECT SELECT
model, model,
COUNT(id) AS total_requests, CAST(COUNT(id) AS BIGINT) AS total_requests,
COALESCE(SUM(input_tokens), 0) AS input_tokens, CAST(COALESCE(SUM(input_tokens), 0) AS BIGINT) AS input_tokens,
COALESCE(SUM(output_tokens), 0) AS output_tokens, CAST(COALESCE(SUM(output_tokens), 0) AS BIGINT) AS output_tokens,
CAST(COALESCE(SUM(total_cost_usd), 0) AS DOUBLE PRECISION) AS total_cost, CAST(COALESCE(SUM(total_cost_usd), 0) AS DOUBLE PRECISION) AS total_cost,
COALESCE( COALESCE(
SUM( SUM(
@@ -4431,7 +4431,7 @@ WITH aggregated AS (
), ),
0 0
) AS response_time_sum_ms, ) AS response_time_sum_ms,
COALESCE( CAST(COALESCE(
SUM( SUM(
CASE CASE
WHEN response_time_ms IS NOT NULL THEN 1 WHEN response_time_ms IS NOT NULL THEN 1
@@ -4439,7 +4439,7 @@ WITH aggregated AS (
END END
), ),
0 0
) AS response_time_samples, ) AS BIGINT) AS response_time_samples,
CAST(COALESCE(AVG(response_time_ms), 0) AS DOUBLE PRECISION) AS avg_response_time_ms CAST(COALESCE(AVG(response_time_ms), 0) AS DOUBLE PRECISION) AS avg_response_time_ms
FROM usage_billing_facts AS usage FROM usage_billing_facts AS usage
WHERE created_at >= $1 WHERE created_at >= $1
@@ -4494,9 +4494,9 @@ WITH aggregated AS (
SELECT SELECT
user_id, user_id,
model, model,
COUNT(id) AS total_requests, CAST(COUNT(id) AS BIGINT) AS total_requests,
COALESCE(SUM(input_tokens), 0) AS input_tokens, CAST(COALESCE(SUM(input_tokens), 0) AS BIGINT) AS input_tokens,
COALESCE(SUM(output_tokens), 0) AS output_tokens, CAST(COALESCE(SUM(output_tokens), 0) AS BIGINT) AS output_tokens,
CAST(COALESCE(SUM(total_cost_usd), 0) AS DOUBLE PRECISION) AS total_cost, CAST(COALESCE(SUM(total_cost_usd), 0) AS DOUBLE PRECISION) AS total_cost,
COALESCE( COALESCE(
SUM( SUM(
@@ -4508,7 +4508,7 @@ WITH aggregated AS (
), ),
0 0
) AS response_time_sum_ms, ) AS response_time_sum_ms,
COALESCE( CAST(COALESCE(
SUM( SUM(
CASE CASE
WHEN response_time_ms IS NOT NULL THEN 1 WHEN response_time_ms IS NOT NULL THEN 1
@@ -4516,7 +4516,7 @@ WITH aggregated AS (
END END
), ),
0 0
) AS response_time_samples ) AS BIGINT) AS response_time_samples
FROM usage_billing_facts AS usage FROM usage_billing_facts AS usage
WHERE created_at >= $1 WHERE created_at >= $1
AND created_at < $2 AND created_at < $2
@@ -4569,9 +4569,9 @@ const UPSERT_STATS_HOURLY_PROVIDER_SQL: &str = r#"
WITH aggregated AS ( WITH aggregated AS (
SELECT SELECT
provider_name, provider_name,
COUNT(id) AS total_requests, CAST(COUNT(id) AS BIGINT) AS total_requests,
COALESCE(SUM(input_tokens), 0) AS input_tokens, CAST(COALESCE(SUM(input_tokens), 0) AS BIGINT) AS input_tokens,
COALESCE(SUM(output_tokens), 0) AS output_tokens, CAST(COALESCE(SUM(output_tokens), 0) AS BIGINT) AS output_tokens,
CAST(COALESCE(SUM(total_cost_usd), 0) AS DOUBLE PRECISION) AS total_cost CAST(COALESCE(SUM(total_cost_usd), 0) AS DOUBLE PRECISION) AS total_cost
FROM usage_billing_facts AS usage FROM usage_billing_facts AS usage
WHERE created_at >= $1 WHERE created_at >= $1