mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
Merge remote-tracking branch 'origin/pr/383' into codex/pr-376-377-383-384-combined
This commit is contained in:
@@ -10,6 +10,7 @@ pub(crate) mod ndjson;
|
||||
mod oauth_retry;
|
||||
#[cfg(test)]
|
||||
pub(crate) mod remote_compat;
|
||||
mod response_header_rules;
|
||||
mod server;
|
||||
pub(crate) mod stream;
|
||||
mod stream_pump;
|
||||
@@ -30,6 +31,9 @@ pub(crate) use self::fallback::{
|
||||
should_retry_next_local_candidate_stream, should_retry_next_local_candidate_sync,
|
||||
should_stop_local_candidate_failover_stream, should_stop_local_candidate_failover_sync,
|
||||
};
|
||||
pub(crate) use self::response_header_rules::{
|
||||
apply_endpoint_response_header_rules, attach_provider_response_headers_to_report_context,
|
||||
};
|
||||
pub(crate) use crate::orchestration::{
|
||||
append_local_failover_policy_to_value, LocalFailoverAnalysis, LocalFailoverDecision,
|
||||
};
|
||||
|
||||
@@ -0,0 +1,101 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use aether_contracts::ExecutionPlan;
|
||||
use serde_json::{Map, Value};
|
||||
use tracing::warn;
|
||||
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
const RESPONSE_HEADER_RULES_KEY: &str = "response_header_rules";
|
||||
const RESPONSE_HEADER_RULES_CAMEL_KEY: &str = "responseHeaderRules";
|
||||
const PROVIDER_RESPONSE_HEADERS_CONTEXT_KEY: &str = "provider_response_headers";
|
||||
const RESPONSE_HEADER_RULE_PROTECTED_KEYS: &[&str] = &["content-length"];
|
||||
|
||||
fn endpoint_response_header_rules_from_config(config: Option<&Value>) -> Option<&Value> {
|
||||
let config = config?.as_object()?;
|
||||
config
|
||||
.get(RESPONSE_HEADER_RULES_KEY)
|
||||
.or_else(|| config.get(RESPONSE_HEADER_RULES_CAMEL_KEY))
|
||||
.filter(|value| !value.is_null())
|
||||
}
|
||||
|
||||
async fn read_endpoint_response_header_rules(state: &AppState, endpoint_id: &str) -> Option<Value> {
|
||||
let endpoint_id = endpoint_id.trim();
|
||||
if endpoint_id.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let endpoint_id = endpoint_id.to_string();
|
||||
|
||||
match state
|
||||
.read_provider_catalog_endpoints_by_ids(std::slice::from_ref(&endpoint_id))
|
||||
.await
|
||||
{
|
||||
Ok(endpoints) => endpoints.into_iter().next().and_then(|endpoint| {
|
||||
endpoint_response_header_rules_from_config(endpoint.config.as_ref()).cloned()
|
||||
}),
|
||||
Err(err) => {
|
||||
warn!(
|
||||
event_name = "response_header_rules_endpoint_read_failed",
|
||||
log_type = "ops",
|
||||
endpoint_id = %endpoint_id,
|
||||
error = ?err,
|
||||
"gateway failed to read endpoint response header rules; skipping response header edits"
|
||||
);
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn apply_endpoint_response_header_rules(
|
||||
state: &AppState,
|
||||
plan: &ExecutionPlan,
|
||||
headers: &mut BTreeMap<String, String>,
|
||||
response_body: Option<&Value>,
|
||||
) -> Result<(), GatewayError> {
|
||||
let Some(rules) = read_endpoint_response_header_rules(state, plan.endpoint_id.as_str()).await
|
||||
else {
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
if !rules.is_array() {
|
||||
warn!(
|
||||
event_name = "response_header_rules_invalid_shape",
|
||||
log_type = "ops",
|
||||
endpoint_id = %plan.endpoint_id,
|
||||
"gateway skipped endpoint response header rules because response_header_rules is not an array"
|
||||
);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let empty_body = Value::Null;
|
||||
let body = response_body.unwrap_or(&empty_body);
|
||||
if !crate::provider_transport::apply_local_header_rules(
|
||||
headers,
|
||||
Some(&rules),
|
||||
RESPONSE_HEADER_RULE_PROTECTED_KEYS,
|
||||
body,
|
||||
response_body,
|
||||
) {
|
||||
return Err(GatewayError::Internal(
|
||||
"response_header_rules 应用失败".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) fn attach_provider_response_headers_to_report_context(
|
||||
report_context: Option<Value>,
|
||||
provider_headers: &BTreeMap<String, String>,
|
||||
) -> Option<Value> {
|
||||
let provider_headers = serde_json::to_value(provider_headers).ok()?;
|
||||
let mut object = match report_context {
|
||||
Some(Value::Object(object)) => object,
|
||||
Some(other) => Map::from_iter([("seed".to_string(), other)]),
|
||||
None => Map::new(),
|
||||
};
|
||||
object.insert(
|
||||
PROVIDER_RESPONSE_HEADERS_CONTEXT_KEY.to_string(),
|
||||
provider_headers,
|
||||
);
|
||||
Some(Value::Object(object))
|
||||
}
|
||||
@@ -65,6 +65,7 @@ use crate::execution_runtime::transport::{
|
||||
DirectSyncExecutionRuntime, DirectUpstreamStreamExecution, ExecutionRuntimeTransportError,
|
||||
};
|
||||
use crate::execution_runtime::{
|
||||
apply_endpoint_response_header_rules, attach_provider_response_headers_to_report_context,
|
||||
local_failover_response_text, resolve_core_stream_direct_finalize_report_kind,
|
||||
resolve_core_stream_error_finalize_report_kind,
|
||||
resolve_local_candidate_failover_analysis_stream, should_fallback_to_control_stream,
|
||||
@@ -839,6 +840,8 @@ async fn execute_stream_from_frame_stream(
|
||||
"execution runtime stream must start with headers frame".to_string(),
|
||||
));
|
||||
};
|
||||
let report_context =
|
||||
attach_provider_response_headers_to_report_context(report_context, &headers);
|
||||
let mut buffered_frames = VecDeque::new();
|
||||
let mut stream_terminal_summary: Option<ExecutionStreamTerminalSummary> = None;
|
||||
if status_code == 200 && should_probe_success_failover_before_stream(&headers) {
|
||||
@@ -1069,6 +1072,10 @@ async fn execute_stream_from_frame_stream(
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let mut client_headers = headers.clone();
|
||||
apply_endpoint_response_header_rules(state, &plan, &mut client_headers, body_json.as_ref())
|
||||
.await?;
|
||||
|
||||
let payload = build_stream_sync_payload(
|
||||
trace_id,
|
||||
stream_error_finalize_kind
|
||||
@@ -1078,7 +1085,7 @@ async fn execute_stream_from_frame_stream(
|
||||
.to_string(),
|
||||
report_context,
|
||||
status_code,
|
||||
headers,
|
||||
client_headers,
|
||||
body_json,
|
||||
body_base64,
|
||||
None,
|
||||
@@ -1510,6 +1517,8 @@ async fn execute_stream_from_frame_stream(
|
||||
});
|
||||
}
|
||||
|
||||
apply_endpoint_response_header_rules(state, &plan, &mut headers, None).await?;
|
||||
|
||||
let request_id = request_id.to_string();
|
||||
let candidate_id = candidate_id.map(ToOwned::to_owned);
|
||||
let (tx, mut rx) = mpsc::channel::<Result<Bytes, IoError>>(16);
|
||||
|
||||
@@ -30,7 +30,8 @@ use crate::execution_runtime::remote_compat::post_sync_plan_to_remote_execution_
|
||||
use crate::execution_runtime::submission::submit_local_core_error_or_sync_finalize;
|
||||
use crate::execution_runtime::transport::DirectSyncExecutionRuntime;
|
||||
use crate::execution_runtime::{
|
||||
analyze_local_candidate_failover_sync, local_failover_response_text,
|
||||
analyze_local_candidate_failover_sync, apply_endpoint_response_header_rules,
|
||||
attach_provider_response_headers_to_report_context, local_failover_response_text,
|
||||
resolve_core_sync_error_finalize_report_kind, should_fallback_to_control_sync,
|
||||
should_finalize_sync_response, LocalFailoverDecision,
|
||||
};
|
||||
@@ -539,6 +540,11 @@ pub(crate) async fn execute_execution_runtime_sync(
|
||||
}
|
||||
let status_code = result.status_code;
|
||||
let has_body_bytes = body_base64.is_some();
|
||||
let report_context =
|
||||
attach_provider_response_headers_to_report_context(report_context, &headers);
|
||||
let mut client_headers = headers.clone();
|
||||
apply_endpoint_response_header_rules(state, &plan, &mut client_headers, body_json.as_ref())
|
||||
.await?;
|
||||
let explicit_finalize = should_finalize_sync_response(report_kind.as_deref());
|
||||
let mapped_error_finalize_kind =
|
||||
resolve_core_sync_error_finalize_report_kind(plan_kind, &result, body_json.as_ref());
|
||||
@@ -549,7 +555,7 @@ pub(crate) async fn execute_execution_runtime_sync(
|
||||
plan_kind,
|
||||
&report_context,
|
||||
status_code,
|
||||
&headers,
|
||||
&client_headers,
|
||||
&body_json,
|
||||
&body_base64,
|
||||
&result.telemetry,
|
||||
@@ -694,7 +700,7 @@ pub(crate) async fn execute_execution_runtime_sync(
|
||||
finalize_report_kind,
|
||||
report_context,
|
||||
status_code,
|
||||
headers,
|
||||
client_headers,
|
||||
body_json,
|
||||
body_base64,
|
||||
telemetry,
|
||||
@@ -863,7 +869,7 @@ pub(crate) async fn execute_execution_runtime_sync(
|
||||
report_kind.unwrap_or_default(),
|
||||
report_context,
|
||||
status_code,
|
||||
headers,
|
||||
client_headers,
|
||||
body_json,
|
||||
body_base64,
|
||||
telemetry,
|
||||
|
||||
Reference in New Issue
Block a user