mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +08:00
fix(gateway): SSE 流跳过成功探测与 prefetch,统计聚合 SQL 显式 CAST BIGINT
- stream: event-stream 响应跳过 direct finalize prefetch 与成功失败转移探测,避免消费流首帧 - stream_pump: 仅在已知 content-length 时缓冲非 SSE 响应 - maintenance: usage_billing_facts 聚合列显式 CAST 为 BIGINT,避免类型不匹配
This commit is contained in:
@@ -619,14 +619,7 @@ fn should_skip_direct_finalize_prefetch(
|
||||
has_private_stream_normalizer: bool,
|
||||
has_local_stream_rewriter: bool,
|
||||
) -> bool {
|
||||
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) {
|
||||
if direct_stream_finalize_kind.is_none() {
|
||||
return false;
|
||||
}
|
||||
|
||||
@@ -635,6 +628,18 @@ fn should_skip_direct_finalize_prefetch(
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or_default()
|
||||
.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() {
|
||||
return true;
|
||||
}
|
||||
@@ -642,6 +647,18 @@ fn should_skip_direct_finalize_prefetch(
|
||||
!(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>(
|
||||
buffered_frames: &mut VecDeque<StreamFrame>,
|
||||
lines: &mut FramedRead<R, LinesCodec>,
|
||||
@@ -712,7 +729,7 @@ async fn execute_stream_from_frame_stream(
|
||||
};
|
||||
let mut buffered_frames = VecDeque::new();
|
||||
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 =
|
||||
probe_local_stream_success_failover_text(&mut buffered_frames, &mut lines).await?;
|
||||
if should_retry_next_local_candidate_stream(
|
||||
@@ -2131,7 +2148,10 @@ mod tests {
|
||||
use serde_json::{json, Value};
|
||||
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::tunnel::{tunnel_protocol, TunnelProxyConn};
|
||||
use crate::AppState;
|
||||
@@ -2195,8 +2215,8 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keeps_prefetch_for_cross_format_or_rewritten_streams() {
|
||||
assert!(!should_skip_direct_finalize_prefetch(
|
||||
fn skips_prefetch_for_event_streams_even_when_cross_format_or_rewritten() {
|
||||
assert!(should_skip_direct_finalize_prefetch(
|
||||
Some("claude_cli_sync_finalize"),
|
||||
Some("text/event-stream"),
|
||||
"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]
|
||||
async fn execute_execution_runtime_stream_records_first_data_as_streaming_before_terminal_telemetry(
|
||||
) {
|
||||
|
||||
@@ -60,7 +60,7 @@ pub(crate) fn build_direct_execution_frame_stream(
|
||||
let mut stream_terminal_observer = StreamingStandardTerminalObserver::default();
|
||||
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();
|
||||
match buffer_non_sse_upstream_body(response, started_at).await {
|
||||
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))
|
||||
}
|
||||
|
||||
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(
|
||||
response: DirectUpstreamResponse,
|
||||
started_at: Instant,
|
||||
@@ -662,7 +676,10 @@ mod tests {
|
||||
use serde_json::Value;
|
||||
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::{
|
||||
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]
|
||||
async fn direct_execution_frame_stream_reports_ttfb_after_first_upstream_chunk() {
|
||||
let listener = crate::test_support::bind_loopback_listener()
|
||||
|
||||
Reference in New Issue
Block a user