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

feat(stream-bridge): 支持上游 sync 响应转 SSE 流式输出，记录 client/upstream 流模式差异

Browse Source 
- 新增 sync_to_stream 桥接模块，将非 SSE 上游响应转换为 SSE 格式回传给流式客户端
- stream_pump 检测非 SSE 响应头后缓冲整包并通过桥接逻辑重写为 SSE 帧
- proxy handler 同步支持 sync→stream 聚合与转换（覆盖 openai/claude/gemini 四种格式）
- sync_products 补全 openai:cli 的完整流式事件聚合（text delta、reasoning、tool call 等）
- usage runtime 写入 client_requested_stream / upstream_is_stream 到 request_metadata
- SQL 查询层将两个布尔字段从 request_metadata jsonb 中提取并回传给前端
- 前端 status.ts 新增 resolveUsageStreamLabelSegments，优先读取 client_requested_stream
- RequestDetailDrawer 在流式转换场景下显示"客户端→上游"两段 Badge
This commit is contained in:
fawney19
2026-04-23 21:53:50 +08:00
parent 40282c3447
commit 342d4a268c
40 changed files with 3638 additions and 97 deletions
Download Patch File Download Diff File Expand all files Collapse all files

235
apps/aether-gateway/src/handlers/proxy/mod.rs
View File

@@ -5,10 +5,15 @@ use self::local::{
};
use super::internal::resolve_local_proxy_execution_path;
pub(crate) use super::public::matches_model_mapping_for_models;
use crate::ai_pipeline_api;
use crate::ai_pipeline_api::{
aggregate_claude_stream_sync_response, aggregate_gemini_stream_sync_response,
aggregate_openai_chat_stream_sync_response, aggregate_openai_cli_stream_sync_response,
maybe_bridge_standard_sync_json_to_stream,
};
use crate::api::response::{
build_client_response, build_local_auth_rejection_response, build_local_http_error_response,
build_local_overloaded_response, build_local_user_rpm_limited_response,
build_client_response, build_client_response_from_parts, build_local_auth_rejection_response,
build_local_http_error_response, build_local_overloaded_response,
build_local_user_rpm_limited_response,
};
use crate::constants::{
DEPENDENCY_REASON_HEADER, EXECUTION_PATH_CONTROL_EXECUTE_STREAM,
@@ -53,7 +58,7 @@ use axum::body::{to_bytes, Body, Bytes};
use axum::extract::{ConnectInfo, Request, State};
use axum::http::{self, header::HeaderName, header::HeaderValue, Response};
use sha2::{Digest, Sha256};
use std::time::Instant;
use std::{collections::BTreeMap, time::Instant};
use tracing::{debug, info, warn};
const OPENAI_CHAT_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL: &str =
@@ -406,8 +411,13 @@ async fn maybe_forward_public_request_to_tunnel_owner(
message: format!("owner gateway affinity forward failed: {err}"),
})?;
let mut response =
build_client_response(upstream_response, &request_context.trace_id, Some(decision))?;
let mut response = build_sync_aware_affinity_forward_response(
request_context,
buffered_body,
decision,
upstream_response,
)
.await?;
response.headers_mut().insert(
HeaderName::from_static(TUNNEL_AFFINITY_OWNER_INSTANCE_HEADER),
HeaderValue::from_str(owner.gateway_instance_id.as_str())
@@ -416,6 +426,219 @@ async fn maybe_forward_public_request_to_tunnel_owner(
Ok(Some(response))
}
fn upstream_response_is_sse(headers: &reqwest::header::HeaderMap) -> bool {
headers
.get(http::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.is_some_and(|value| value.to_ascii_lowercase().contains("text/event-stream"))
}
fn collect_upstream_response_headers(
headers: &reqwest::header::HeaderMap,
) -> BTreeMap<String, String> {
headers
.iter()
.map(|(name, value)| {
(
name.as_str().to_string(),
value.to_str().unwrap_or_default().to_string(),
)
})
.collect()
}
fn aggregate_sync_sse_response_for_client(
decision: &GatewayControlDecision,
public_path: &str,
body: &[u8],
) -> Option<serde_json::Value> {
let api_format = decision
.auth_endpoint_signature
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty());
match api_format {
Some(value) if value.eq_ignore_ascii_case("openai:chat") => {
aggregate_openai_chat_stream_sync_response(body)
}
Some(value)
if value.eq_ignore_ascii_case("openai:cli")
|| value.eq_ignore_ascii_case("openai:compact") =>
{
aggregate_openai_cli_stream_sync_response(body)
}
Some(value)
if value.eq_ignore_ascii_case("claude:chat")
|| value.eq_ignore_ascii_case("claude:cli") =>
{
aggregate_claude_stream_sync_response(body)
}
Some(value)
if value.eq_ignore_ascii_case("gemini:chat")
|| value.eq_ignore_ascii_case("gemini:cli") =>
{
aggregate_gemini_stream_sync_response(body)
}
_ if public_path == "/v1/chat/completions" => {
aggregate_openai_chat_stream_sync_response(body)
}
_ if public_path == "/v1/responses" || public_path == "/v1/responses/compact" => {
aggregate_openai_cli_stream_sync_response(body)
}
_ if public_path == "/v1/messages" => aggregate_claude_stream_sync_response(body),
_ if decision.route_family.as_deref() == Some("gemini")
&& (public_path.contains(":generateContent")
|| public_path.contains(":streamGenerateContent")) =>
{
aggregate_gemini_stream_sync_response(body)
}
_ => None,
}
}
fn build_sync_json_proxy_response(
status_code: u16,
upstream_headers: &BTreeMap<String, String>,
body_json: &serde_json::Value,
trace_id: &str,
decision: &GatewayControlDecision,
) -> Result<Response<Body>, GatewayError> {
let mut headers = upstream_headers.clone();
headers.remove("content-encoding");
headers.remove("content-length");
headers.insert("content-type".to_string(), "application/json".to_string());
let body_bytes =
serde_json::to_vec(body_json).map_err(|err| GatewayError::Internal(err.to_string()))?;
headers.insert("content-length".to_string(), body_bytes.len().to_string());
build_client_response_from_parts(
status_code,
&headers,
Body::from(body_bytes),
trace_id,
Some(decision),
)
}
fn resolve_affinity_forward_client_api_format(
decision: &GatewayControlDecision,
public_path: &str,
) -> Option<&'static str> {
let api_format = decision
.auth_endpoint_signature
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty());
match api_format {
Some(value) if value.eq_ignore_ascii_case("openai:chat") => Some("openai:chat"),
Some(value) if value.eq_ignore_ascii_case("openai:cli") => Some("openai:cli"),
Some(value) if value.eq_ignore_ascii_case("openai:compact") => Some("openai:compact"),
Some(value) if value.eq_ignore_ascii_case("claude:chat") => Some("claude:chat"),
Some(value) if value.eq_ignore_ascii_case("claude:cli") => Some("claude:cli"),
Some(value) if value.eq_ignore_ascii_case("gemini:chat") => Some("gemini:chat"),
Some(value) if value.eq_ignore_ascii_case("gemini:cli") => Some("gemini:cli"),
_ if public_path == "/v1/chat/completions" => Some("openai:chat"),
_ if public_path == "/v1/responses" => Some("openai:cli"),
_ if public_path == "/v1/responses/compact" => Some("openai:compact"),
_ if public_path == "/v1/messages" => Some("claude:chat"),
_ if decision.route_family.as_deref() == Some("gemini")
&& (public_path.contains(":generateContent")
|| public_path.contains(":streamGenerateContent")) =>
{
Some("gemini:chat")
}
_ => None,
}
}
fn build_stream_sse_proxy_response(
status_code: u16,
upstream_headers: &BTreeMap<String, String>,
sse_body: &[u8],
trace_id: &str,
decision: &GatewayControlDecision,
) -> Result<Response<Body>, GatewayError> {
let mut headers = upstream_headers.clone();
headers.remove("content-encoding");
headers.remove("content-length");
headers.insert("content-type".to_string(), "text/event-stream".to_string());
headers.insert("content-length".to_string(), sse_body.len().to_string());
build_client_response_from_parts(
status_code,
&headers,
Body::from(sse_body.to_vec()),
trace_id,
Some(decision),
)
}
async fn build_sync_aware_affinity_forward_response(
request_context: &GatewayPublicRequestContext,
buffered_body: Option<&Bytes>,
decision: &GatewayControlDecision,
upstream_response: reqwest::Response,
) -> Result<Response<Body>, GatewayError> {
let Some(buffered_body) = buffered_body else {
return build_client_response(upstream_response, &request_context.trace_id, Some(decision));
};
let stream_request = request_wants_stream(request_context, buffered_body);
let upstream_is_sse = upstream_response_is_sse(upstream_response.headers());
if (!stream_request && !upstream_is_sse) || (stream_request && upstream_is_sse) {
return build_client_response(upstream_response, &request_context.trace_id, Some(decision));
}
let status_code = upstream_response.status().as_u16();
let headers = collect_upstream_response_headers(upstream_response.headers());
let body_bytes = upstream_response
.bytes()
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if stream_request {
if (200..300).contains(&status_code) {
if let Some(client_api_format) = resolve_affinity_forward_client_api_format(
decision,
request_context.request_path.as_str(),
) {
if let Ok(body_json) = serde_json::from_slice::<serde_json::Value>(&body_bytes) {
if let Some(outcome) = maybe_bridge_standard_sync_json_to_stream(
&body_json,
client_api_format,
client_api_format,
None,
)? {
return build_stream_sse_proxy_response(
status_code,
&headers,
&outcome.sse_body,
&request_context.trace_id,
decision,
);
}
}
}
}
} else if let Some(body_json) = aggregate_sync_sse_response_for_client(
decision,
request_context.request_path.as_str(),
&body_bytes,
) {
return build_sync_json_proxy_response(
status_code,
&headers,
&body_json,
&request_context.trace_id,
decision,
);
}
build_client_response_from_parts(
status_code,
&headers,
Body::from(body_bytes),
&request_context.trace_id,
Some(decision),
)
}
pub(crate) async fn proxy_request(
State(state): State<AppState>,
ConnectInfo(remote_addr): ConnectInfo<std::net::SocketAddr>,

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

@@ -232,6 +232,73 @@ fn users_me_usage_request_body_implies_default_non_stream(item: &StoredRequestUs
!body.contains_key("stream") && users_me_usage_api_format_defaults_to_non_stream(item)
}
fn users_me_usage_headers_stream_flag(headers: Option<&serde_json::Value>) -> Option<bool> {
let object = headers.and_then(serde_json::Value::as_object)?;
let raw = object
.iter()
.find(|(key, _)| key.eq_ignore_ascii_case("content-type"))
.and_then(|(_, value)| match value {
serde_json::Value::String(text) => Some(text.as_str()),
serde_json::Value::Array(values) => values.iter().find_map(serde_json::Value::as_str),
_ => None,
})?
.trim();
if raw.is_empty() {
return None;
}
let normalized = raw.to_ascii_lowercase();
Some(
normalized.contains("event-stream")
|| normalized.contains("eventstream")
|| normalized.contains("x-ndjson"),
)
}
fn users_me_usage_body_is_sse_capture(value: Option<&serde_json::Value>) -> bool {
let Some(object) = value.and_then(serde_json::Value::as_object) else {
return false;
};
object
.get("chunks")
.and_then(serde_json::Value::as_array)
.is_some()
&& object
.get("metadata")
.and_then(serde_json::Value::as_object)
.and_then(|metadata| metadata.get("stream"))
.and_then(serde_json::Value::as_bool)
.unwrap_or(false)
}
fn users_me_usage_infer_client_stream_from_captured_bodies(
item: &StoredRequestUsageAudit,
) -> Option<bool> {
let provider_stream = users_me_usage_body_is_sse_capture(item.response_body.as_ref());
let client_stream = users_me_usage_body_is_sse_capture(item.client_response_body.as_ref());
if client_stream {
Some(true)
} else if provider_stream && item.client_response_body.is_some() {
Some(false)
} else {
None
}
}
fn users_me_usage_infer_upstream_stream_from_captured_bodies(
item: &StoredRequestUsageAudit,
) -> Option<bool> {
let provider_stream = users_me_usage_body_is_sse_capture(item.response_body.as_ref());
let client_stream = users_me_usage_body_is_sse_capture(item.client_response_body.as_ref());
if provider_stream {
Some(true)
} else if client_stream && item.response_body.is_some() {
Some(false)
} else {
None
}
}
fn users_me_usage_client_is_stream(item: &StoredRequestUsageAudit) -> bool {
item.request_metadata
.as_ref()
@@ -239,7 +306,20 @@ fn users_me_usage_client_is_stream(item: &StoredRequestUsageAudit) -> bool {
.and_then(|metadata| metadata.get("client_requested_stream"))
.and_then(serde_json::Value::as_bool)
.or_else(|| users_me_usage_request_body_stream_flag(item))
.or_else(|| users_me_usage_headers_stream_flag(item.client_response_headers.as_ref()))
.or_else(|| users_me_usage_request_body_implies_default_non_stream(item).then_some(false))
.or_else(|| users_me_usage_infer_client_stream_from_captured_bodies(item))
.unwrap_or(item.is_stream)
}
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(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))
.unwrap_or(item.is_stream)
}
@@ -256,6 +336,7 @@ fn build_users_me_usage_record_payload(
let cache_creation_input_tokens = users_me_usage_cache_creation_tokens(item);
let rate_multiplier = item.settlement_rate_multiplier();
let client_is_stream = users_me_usage_client_is_stream(item);
let upstream_is_stream = users_me_usage_upstream_is_stream(item);
let mut payload = json!({
"id": item.id,
"model": item.model,
@@ -271,7 +352,7 @@ fn build_users_me_usage_record_payload(
"response_time_ms": item.response_time_ms,
"first_byte_time_ms": item.first_byte_time_ms,
"is_stream": item.is_stream,
"upstream_is_stream": item.is_stream,
"upstream_is_stream": upstream_is_stream,
"client_requested_stream": client_is_stream,
"client_is_stream": client_is_stream,
"status": item.status,
@@ -307,6 +388,7 @@ fn build_users_me_usage_record_payload(
fn build_users_me_usage_active_payload(item: &StoredRequestUsageAudit) -> serde_json::Value {
let cache_creation_input_tokens = users_me_usage_cache_creation_tokens(item);
let client_is_stream = users_me_usage_client_is_stream(item);
let upstream_is_stream = users_me_usage_upstream_is_stream(item);
let mut payload = json!({
"id": item.id,
"status": item.status,
@@ -325,7 +407,7 @@ fn build_users_me_usage_active_payload(item: &StoredRequestUsageAudit) -> serde_
"api_format": item.api_format,
"endpoint_api_format": item.endpoint_api_format,
"is_stream": item.is_stream,
"upstream_is_stream": item.is_stream,
"upstream_is_stream": upstream_is_stream,
"client_requested_stream": client_is_stream,
"client_is_stream": client_is_stream,
"has_format_conversion": item.has_format_conversion,
@@ -1151,7 +1233,7 @@ mod tests {
use super::{
build_users_me_usage_active_payload, build_users_me_usage_record_payload,
users_me_usage_client_is_stream,
users_me_usage_client_is_stream, users_me_usage_upstream_is_stream,
};
fn sample_usage(status: &str) -> StoredRequestUsageAudit {
@@ -1303,4 +1385,105 @@ mod tests {
assert_eq!(active_payload["client_requested_stream"], false);
assert_eq!(active_payload["client_is_stream"], false);
}
#[test]
fn user_usage_upstream_stream_prefers_request_metadata_flag() {
let item = StoredRequestUsageAudit {
is_stream: false,
request_metadata: Some(json!({
"client_requested_stream": false,
"upstream_is_stream": true
})),
..sample_usage("completed")
};
let record_payload =
build_users_me_usage_record_payload(&item, false, &BTreeMap::new(), false);
assert_eq!(record_payload["is_stream"], false);
assert_eq!(record_payload["upstream_is_stream"], true);
assert_eq!(record_payload["client_requested_stream"], false);
assert_eq!(record_payload["client_is_stream"], false);
let active_payload = build_users_me_usage_active_payload(&item);
assert_eq!(active_payload["is_stream"], false);
assert_eq!(active_payload["upstream_is_stream"], true);
assert_eq!(active_payload["client_requested_stream"], false);
assert_eq!(active_payload["client_is_stream"], false);
}
#[test]
fn user_usage_stream_modes_fall_back_to_captured_response_bodies_when_request_metadata_is_missing(
) {
let item = StoredRequestUsageAudit {
is_stream: true,
response_body: Some(json!({
"chunks": [
{"type": "response.created"},
{"type": "response.output_text.delta", "delta": "Hello"}
],
"metadata": {
"stream": true,
"stored_chunks": 2,
"total_chunks": 2
}
})),
client_response_body: Some(json!({
"id": "resp-1",
"object": "response",
"status": "completed",
"output": []
})),
..sample_usage("completed")
};
assert!(!users_me_usage_client_is_stream(&item));
assert!(users_me_usage_upstream_is_stream(&item));
let record_payload =
build_users_me_usage_record_payload(&item, false, &BTreeMap::new(), false);
assert_eq!(record_payload["is_stream"], true);
assert_eq!(record_payload["upstream_is_stream"], true);
assert_eq!(record_payload["client_requested_stream"], false);
assert_eq!(record_payload["client_is_stream"], false);
let active_payload = build_users_me_usage_active_payload(&item);
assert_eq!(active_payload["is_stream"], true);
assert_eq!(active_payload["upstream_is_stream"], true);
assert_eq!(active_payload["client_requested_stream"], false);
assert_eq!(active_payload["client_is_stream"], false);
}
#[test]
fn user_usage_stream_modes_fall_back_to_captured_response_headers_when_bodies_are_detached() {
let item = StoredRequestUsageAudit {
is_stream: true,
response_headers: Some(json!({
"content-type": "text/event-stream; charset=utf-8"
})),
client_response_headers: Some(json!({
"content-type": "application/json"
})),
response_body_ref: Some("usage://request/req-1/response_body".to_string()),
client_response_body_ref: Some(
"usage://request/req-1/client_response_body".to_string(),
),
..sample_usage("completed")
};
assert!(!users_me_usage_client_is_stream(&item));
assert!(users_me_usage_upstream_is_stream(&item));
let record_payload =
build_users_me_usage_record_payload(&item, false, &BTreeMap::new(), false);
assert_eq!(record_payload["is_stream"], true);
assert_eq!(record_payload["upstream_is_stream"], true);
assert_eq!(record_payload["client_requested_stream"], false);
assert_eq!(record_payload["client_is_stream"], false);
let active_payload = build_users_me_usage_active_payload(&item);
assert_eq!(active_payload["is_stream"], true);
assert_eq!(active_payload["upstream_is_stream"], true);
assert_eq!(active_payload["client_requested_stream"], false);
assert_eq!(active_payload["client_is_stream"], false);
}
}