Merge remote-tracking branch 'origin/pr/383' into codex/pr-376-377-383-384-combined

This commit is contained in:
fawney19
2026-05-06 02:37:09 +08:00
6 changed files with 573 additions and 45 deletions

View File

@@ -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,
};

View File

@@ -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))
}

View File

@@ -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);

View File

@@ -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,