fix: harden frontdoor and usage ingestion

This commit is contained in:
fawney19
2026-05-22 23:57:38 +08:00
parent 6ef6cbade2
commit d18b13a91a
25 changed files with 1363 additions and 252 deletions
@@ -317,12 +317,7 @@ pub(super) async fn execute_provider_quota_plan(
match state.execute_execution_runtime_sync_plan(None, &plan).await {
Ok(result) => Ok(ProviderQuotaExecutionOutcome::Response(result)),
Err(err) => {
let error = match err {
GatewayError::UpstreamUnavailable { message, .. }
| GatewayError::ControlUnavailable { message, .. }
| GatewayError::Client { message, .. }
| GatewayError::Internal(message) => message,
};
let error = err.into_message();
let proxy_node_id = plan
.proxy
.as_ref()
@@ -301,12 +301,7 @@ fn admin_provider_ops_decode_response_bytes(
}
fn admin_provider_ops_gateway_error_message(error: GatewayError) -> String {
match error {
GatewayError::UpstreamUnavailable { message, .. }
| GatewayError::ControlUnavailable { message, .. }
| GatewayError::Client { message, .. }
| GatewayError::Internal(message) => message,
}
error.into_message()
}
pub(super) fn admin_provider_ops_verify_execution_error_message(error: &str) -> String {
@@ -16,7 +16,12 @@ use aether_runtime_state::{DataLayerError, RuntimeState};
use futures_util::future::join_all;
use std::collections::{BTreeMap, BTreeSet};
use std::time::{SystemTime, UNIX_EPOCH};
use tracing::warn;
use tracing::{info, warn};
const DEFAULT_POOL_RUNTIME_WINDOW_METRIC_KEY_LIMIT: usize = 512;
const MAX_POOL_RUNTIME_WINDOW_METRIC_KEY_LIMIT: usize = 10_000;
const POOL_RUNTIME_WINDOW_METRIC_KEY_LIMIT_ENV: &str =
"AETHER_GATEWAY_ADMIN_POOL_RUNTIME_WINDOW_METRIC_KEY_LIMIT";
fn current_unix_secs() -> u64 {
SystemTime::now()
@@ -29,6 +34,20 @@ fn should_load_active_probe_members(pool_config: &AdminProviderPoolConfig) -> bo
pool_config.probing_enabled
}
fn pool_runtime_window_metric_key_limit() -> usize {
std::env::var(POOL_RUNTIME_WINDOW_METRIC_KEY_LIMIT_ENV)
.ok()
.and_then(|value| value.trim().parse::<usize>().ok())
.filter(|value| *value > 0)
.unwrap_or(DEFAULT_POOL_RUNTIME_WINDOW_METRIC_KEY_LIMIT)
.clamp(1, MAX_POOL_RUNTIME_WINDOW_METRIC_KEY_LIMIT)
}
fn bounded_runtime_window_metric_key_ids(key_ids: &[String], limit: usize) -> &[String] {
let end = key_ids.len().min(limit.max(1));
&key_ids[..end]
}
pub(crate) async fn read_admin_provider_pool_cooldown_counts(
runtime: &RuntimeState,
provider_ids: &[String],
@@ -54,8 +73,21 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
) -> AdminProviderPoolRuntimeState {
let mut state = AdminProviderPoolRuntimeState::default();
let cooldown_keys = pool_cooldown_keys(provider_id, key_ids);
let cost_keys = pool_cost_keys(provider_id, key_ids);
let latency_keys = pool_latency_keys(provider_id, key_ids);
let metric_key_limit = pool_runtime_window_metric_key_limit();
let metric_key_ids = bounded_runtime_window_metric_key_ids(key_ids, metric_key_limit);
if metric_key_ids.len() < key_ids.len() {
info!(
event_name = "admin_pool_runtime_window_metrics_truncated",
log_type = "event",
provider_id,
total_key_count = key_ids.len(),
scanned_key_count = metric_key_ids.len(),
metric_key_limit,
"gateway limited admin pool runtime cost/latency window reads"
);
}
let cost_keys = pool_cost_keys(provider_id, metric_key_ids);
let latency_keys = pool_latency_keys(provider_id, metric_key_ids);
let sticky_sessions_enabled = pool_config.sticky_session_ttl_seconds > 0
&& admin_provider_pool_cache_affinity_enabled(pool_config);
@@ -179,7 +211,7 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
.map(|cost_key| runtime.score_range_by_min(cost_key, cost_window_start)),
)
.await;
for (key_id, members) in key_ids.iter().zip(cost_results) {
for (key_id, members) in metric_key_ids.iter().zip(cost_results) {
let total = members
.unwrap_or_default()
.iter()
@@ -197,7 +229,7 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
.map(|latency_key| runtime.score_range_by_min(latency_key, latency_window_start)),
)
.await;
for (key_id, members) in key_ids.iter().zip(latency_results) {
for (key_id, members) in metric_key_ids.iter().zip(latency_results) {
let samples = members
.unwrap_or_default()
.iter()
@@ -265,3 +297,30 @@ pub(crate) async fn read_admin_provider_pool_key_cooldown_reason(
.kv_get(&pool_cooldown_key(provider_id, key_id))
.await
}
#[cfg(test)]
mod tests {
use super::bounded_runtime_window_metric_key_ids;
#[test]
fn runtime_window_metric_key_ids_are_bounded() {
let key_ids = vec![
"key-1".to_string(),
"key-2".to_string(),
"key-3".to_string(),
];
let bounded = bounded_runtime_window_metric_key_ids(&key_ids, 2);
assert_eq!(bounded, &key_ids[..2]);
}
#[test]
fn runtime_window_metric_key_ids_keep_at_least_one_key() {
let key_ids = vec!["key-1".to_string(), "key-2".to_string()];
let bounded = bounded_runtime_window_metric_key_ids(&key_ids, 0);
assert_eq!(bounded, &key_ids[..1]);
}
}
@@ -662,10 +662,5 @@ fn admin_provider_oauth_decode_response_bytes(
}
fn admin_provider_oauth_gateway_error_message(error: GatewayError) -> String {
match error {
GatewayError::UpstreamUnavailable { message, .. }
| GatewayError::ControlUnavailable { message, .. }
| GatewayError::Client { message, .. }
| GatewayError::Internal(message) => message,
}
error.into_message()
}
@@ -354,12 +354,7 @@ pub(crate) async fn maybe_build_internal_finalize_video_response(
}
pub(crate) fn gateway_error_message(error: GatewayError) -> String {
match error {
GatewayError::UpstreamUnavailable { message, .. }
| GatewayError::ControlUnavailable { message, .. }
| GatewayError::Client { message, .. }
| GatewayError::Internal(message) => message,
}
error.into_message()
}
pub(crate) fn build_internal_tunnel_heartbeat_ack(
+390 -52
View File
@@ -21,13 +21,13 @@ use crate::constants::{
EXECUTION_PATH_EXECUTION_RUNTIME_STREAM, EXECUTION_PATH_EXECUTION_RUNTIME_SYNC,
EXECUTION_PATH_LOCAL_AI_PUBLIC, EXECUTION_PATH_LOCAL_API_KEY_CONCURRENCY_LIMITED,
EXECUTION_PATH_LOCAL_AUTH_DENIED, EXECUTION_PATH_LOCAL_EXECUTION_LOOP_DETECTED,
EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS, EXECUTION_PATH_LOCAL_INVALID_REQUEST,
EXECUTION_PATH_LOCAL_OVERLOADED, EXECUTION_PATH_LOCAL_PROXY_PASSTHROUGH_REMOVED,
EXECUTION_PATH_LOCAL_RATE_LIMITED, EXECUTION_PATH_LOCAL_ROUTE_NOT_FOUND,
EXECUTION_PATH_PUBLIC_PROXY_PASSTHROUGH, EXECUTION_RUNTIME_LOOP_GUARD_HEADER,
FORWARDED_FOR_HEADER, FORWARDED_HOST_HEADER, FORWARDED_PROTO_HEADER, GATEWAY_HEADER,
LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER, TRACE_ID_HEADER,
TRUSTED_AUTH_ACCESS_ALLOWED_HEADER, TRUSTED_AUTH_API_KEY_ID_HEADER,
EXECUTION_PATH_LOCAL_EXECUTION_PLANNING_TIMEOUT, EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS,
EXECUTION_PATH_LOCAL_INVALID_REQUEST, EXECUTION_PATH_LOCAL_OVERLOADED,
EXECUTION_PATH_LOCAL_PROXY_PASSTHROUGH_REMOVED, EXECUTION_PATH_LOCAL_RATE_LIMITED,
EXECUTION_PATH_LOCAL_ROUTE_NOT_FOUND, EXECUTION_PATH_PUBLIC_PROXY_PASSTHROUGH,
EXECUTION_RUNTIME_LOOP_GUARD_HEADER, FORWARDED_FOR_HEADER, FORWARDED_HOST_HEADER,
FORWARDED_PROTO_HEADER, GATEWAY_HEADER, LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER,
TRACE_ID_HEADER, TRUSTED_AUTH_ACCESS_ALLOWED_HEADER, TRUSTED_AUTH_API_KEY_ID_HEADER,
TRUSTED_AUTH_BALANCE_HEADER, TRUSTED_AUTH_USER_ID_HEADER, TUNNEL_AFFINITY_FORWARDED_BY_HEADER,
TUNNEL_AFFINITY_OWNER_INSTANCE_HEADER,
};
@@ -65,7 +65,11 @@ use axum::extract::{ConnectInfo, Request, State};
use axum::http::{self, header::HeaderName, header::HeaderValue, Response};
use futures_util::StreamExt;
use sha2::{Digest, Sha256};
use std::{collections::BTreeMap, time::Instant};
use std::{
collections::BTreeMap,
error::Error as StdError,
time::{Duration, Instant},
};
use tracing::{debug, info, warn};
const OPENAI_CHAT_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL: &str =
@@ -89,23 +93,192 @@ const LOCAL_EXECUTION_LOOP_DETECTED_DETAIL: &str =
"Gateway detected an execution runtime request loop back into the local frontdoor";
const AUTH_API_KEY_CONCURRENCY_LIMIT_REACHED_DETAIL: &str =
"当前 API Key 并发请求数已达上限,请稍后重试";
const REQUEST_BODY_READ_TIMEOUT_DETAIL: &str =
"Request body read timed out before the gateway could route the request";
const REQUEST_BODY_READ_FAILED_DETAIL: &str = "Failed to read request body";
const LOCAL_EXECUTION_PLANNING_TIMEOUT_DETAIL: &str =
"当前 AI 请求在本地执行规划阶段超时,请稍后重试";
const EXECUTION_PATH_TUNNEL_AFFINITY_FORWARD: &str = "tunnel_affinity_forward";
const MANAGEMENT_TOKEN_PREFIX: &str = "ae-";
const LEGACY_MANAGEMENT_TOKEN_PREFIX: &str = "ae_";
fn build_request_body_normalization_error_response(
#[derive(Debug, Clone, Copy)]
struct RequestBodyBufferPolicy {
max_bytes: u64,
read_timeout: Duration,
}
impl RequestBodyBufferPolicy {
fn from_state(state: &AppState) -> Self {
Self {
max_bytes: crate::headers::max_request_body_bytes(),
read_timeout: state.frontdoor_runtime_guards.request_body_read_timeout,
}
}
#[cfg(test)]
fn for_tests(max_bytes: u64, read_timeout: Duration) -> Self {
Self {
max_bytes,
read_timeout,
}
}
}
#[derive(Debug)]
enum RequestBodyBufferError {
Normalization(RequestBodyNormalizationError),
TooLarge { limit_bytes: u64 },
Timeout { timeout_ms: u64 },
ReadFailed { message: String },
}
impl RequestBodyBufferError {
fn http_status(&self) -> http::StatusCode {
match self {
Self::Normalization(error) => error.http_status(),
Self::TooLarge { .. } => http::StatusCode::PAYLOAD_TOO_LARGE,
Self::Timeout { .. } => http::StatusCode::REQUEST_TIMEOUT,
Self::ReadFailed { .. } => http::StatusCode::BAD_REQUEST,
}
}
fn client_message(&self) -> String {
match self {
Self::Normalization(error) => error.client_message(),
Self::TooLarge { limit_bytes } => format!("Request body exceeds {limit_bytes} bytes"),
Self::Timeout { .. } => REQUEST_BODY_READ_TIMEOUT_DETAIL.to_string(),
Self::ReadFailed { .. } => REQUEST_BODY_READ_FAILED_DETAIL.to_string(),
}
}
fn reason(&self) -> &'static str {
match self {
Self::Normalization(error) => match error {
RequestBodyNormalizationError::UnsupportedContentEncoding(_) => {
"unsupported_content_encoding"
}
RequestBodyNormalizationError::DecodeFailed { .. } => "decode_failed",
RequestBodyNormalizationError::DecompressedBodyTooLarge { .. } => {
"decompressed_body_too_large"
}
RequestBodyNormalizationError::RequestBodyTooLarge { .. } => {
"request_body_too_large"
}
},
Self::TooLarge { .. } => "request_body_too_large",
Self::Timeout { .. } => "request_body_read_timeout",
Self::ReadFailed { .. } => "request_body_read_failed",
}
}
}
async fn buffer_and_normalize_request_body(
request_body: &mut Option<Body>,
headers: &mut http::HeaderMap,
body_owner_expectation: &'static str,
trace_id: &str,
method: &http::Method,
path_and_query: &str,
phase: &'static str,
policy: RequestBodyBufferPolicy,
) -> Result<Bytes, RequestBodyBufferError> {
if let Err(err) =
crate::headers::check_request_content_length_with_limit(headers, policy.max_bytes)
{
return Err(RequestBodyBufferError::Normalization(err));
}
let read_started_at = Instant::now();
let timeout_ms = policy.read_timeout.as_millis() as u64;
info!(
event_name = "frontdoor_request_body_buffer_started",
log_type = "event",
trace_id,
method = %method,
path = %path_and_query,
phase,
max_body_bytes = policy.max_bytes,
timeout_ms,
"gateway started buffering request body"
);
let body_limit = usize::try_from(policy.max_bytes).unwrap_or(usize::MAX);
let body = match tokio::time::timeout(
policy.read_timeout,
to_bytes(
request_body.take().expect(body_owner_expectation),
body_limit,
),
)
.await
{
Ok(Ok(body)) => body,
Ok(Err(err)) if request_body_collection_exceeded_limit(&err) => {
return Err(RequestBodyBufferError::TooLarge {
limit_bytes: policy.max_bytes,
});
}
Ok(Err(err)) => {
return Err(RequestBodyBufferError::ReadFailed {
message: err.to_string(),
});
}
Err(_) => {
return Err(RequestBodyBufferError::Timeout { timeout_ms });
}
};
let normalized = crate::headers::normalize_request_body_headers_and_bytes_with_limit(
headers,
body,
policy.max_bytes,
)
.map_err(RequestBodyBufferError::Normalization)?;
info!(
event_name = "frontdoor_request_body_buffer_completed",
log_type = "event",
trace_id,
method = %method,
path = %path_and_query,
phase,
body_bytes = normalized.len(),
elapsed_ms = read_started_at.elapsed().as_millis() as u64,
"gateway completed request body buffering"
);
Ok(normalized)
}
fn request_body_collection_exceeded_limit(error: &(dyn StdError + 'static)) -> bool {
let mut current = Some(error);
while let Some(error) = current {
if error.to_string().contains("length limit exceeded") {
return true;
}
current = error.source();
}
false
}
fn build_request_body_buffer_error_response(
trace_id: &str,
request_context: &GatewayPublicRequestContext,
error: &RequestBodyNormalizationError,
error: &RequestBodyBufferError,
) -> Result<Response<Body>, GatewayError> {
warn!(
event_name = "frontdoor_request_body_normalization_failed",
event_name = "frontdoor_request_body_buffer_failed",
log_type = "ops",
trace_id,
method = %request_context.request_method,
path = %request_context.request_path_and_query(),
error = %error,
"gateway rejected request with invalid encoded body"
status_code = error.http_status().as_u16(),
reason = error.reason(),
detail = %error.client_message(),
read_error = match error {
RequestBodyBufferError::ReadFailed { message } => message.as_str(),
_ => "",
},
"gateway rejected request body before local execution planning"
);
build_local_http_error_response(
trace_id,
@@ -115,36 +288,16 @@ fn build_request_body_normalization_error_response(
)
}
async fn buffer_and_normalize_request_body(
request_body: &mut Option<Body>,
headers: &mut http::HeaderMap,
body_owner_expectation: &'static str,
) -> Result<Result<Bytes, RequestBodyNormalizationError>, GatewayError> {
if let Err(err) = crate::headers::check_request_content_length(headers) {
return Ok(Err(err));
}
let body = to_bytes(
request_body.take().expect(body_owner_expectation),
usize::MAX,
)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
Ok(crate::headers::normalize_request_body_headers_and_bytes(
headers, body,
))
}
fn finalize_request_body_normalization_rejection(
fn finalize_request_body_buffer_rejection(
state: &AppState,
request_context: &GatewayPublicRequestContext,
remote_addr: &std::net::SocketAddr,
started_at: &std::time::Instant,
trace_id: &str,
request_permit: Option<aether_runtime::AdmissionPermit>,
error: &RequestBodyNormalizationError,
error: &RequestBodyBufferError,
) -> Result<Response<Body>, GatewayError> {
let response =
build_request_body_normalization_error_response(trace_id, request_context, error)?;
let response = build_request_body_buffer_error_response(trace_id, request_context, error)?;
Ok(finalize_gateway_response_with_context(
state,
response,
@@ -156,6 +309,59 @@ fn finalize_request_body_normalization_rejection(
))
}
fn local_execution_planning_timeout_parts(error: &GatewayError) -> Option<(&'static str, u64)> {
match error {
GatewayError::LocalExecutionPlanningTimeout {
phase, timeout_ms, ..
} => Some((*phase, *timeout_ms)),
_ => None,
}
}
fn finalize_local_execution_planning_timeout(
state: &AppState,
request_context: &GatewayPublicRequestContext,
remote_addr: &std::net::SocketAddr,
started_at: &std::time::Instant,
trace_id: &str,
request_permit: Option<aether_runtime::AdmissionPermit>,
control_decision: Option<&GatewayControlDecision>,
phase: &'static str,
timeout_ms: u64,
) -> Result<Response<Body>, GatewayError> {
warn!(
event_name = "frontdoor_local_execution_planning_timeout",
log_type = "ops",
trace_id,
method = %request_context.request_method,
path = %request_context.request_path_and_query(),
route_family = control_decision
.and_then(|decision| decision.route_family.as_deref())
.unwrap_or("-"),
route_kind = control_decision
.and_then(|decision| decision.route_kind.as_deref())
.unwrap_or("-"),
phase,
timeout_ms,
"gateway failed local execution before a candidate could be selected"
);
let response = build_local_http_error_response(
trace_id,
control_decision,
http::StatusCode::GATEWAY_TIMEOUT,
LOCAL_EXECUTION_PLANNING_TIMEOUT_DETAIL,
)?;
Ok(finalize_gateway_response_with_context(
state,
response,
remote_addr,
request_context,
EXECUTION_PATH_LOCAL_EXECUTION_PLANNING_TIMEOUT,
started_at,
request_permit,
))
}
fn local_execution_outcome_label(outcome: &LocalExecutionRequestOutcome) -> &'static str {
match outcome {
LocalExecutionRequestOutcome::Responded(_) => "responded",
@@ -991,16 +1197,22 @@ pub(crate) async fn proxy_request(
}
let mut request_body = Some(body);
let local_proxy_body = if local_proxy_route_requires_buffered_body(&request_context) {
let body_buffer_policy = RequestBodyBufferPolicy::from_state(&state);
let body = buffer_and_normalize_request_body(
&mut request_body,
&mut parts.headers,
"local proxy body buffering should own request body",
&trace_id,
&parts.method,
&request_context.request_path_and_query(),
"local_proxy",
body_buffer_policy,
)
.await?;
.await;
match body {
Ok(body) => Some(body),
Err(err) => {
return finalize_request_body_normalization_rejection(
return finalize_request_body_buffer_rejection(
&state,
&request_context,
&remote_addr,
@@ -1171,16 +1383,22 @@ pub(crate) async fn proxy_request(
&& request_enables_control_execute(&parts.headers);
let buffered_body = if should_buffer_body {
let body_buffer_policy = RequestBodyBufferPolicy::from_state(&state);
let body = buffer_and_normalize_request_body(
&mut request_body,
&mut parts.headers,
"buffered auth/execution runtime path should own request body",
&trace_id,
&parts.method,
&request_context.request_path_and_query(),
"auth_execution",
body_buffer_policy,
)
.await?;
.await;
match body {
Ok(body) => Some(body),
Err(err) => {
return finalize_request_body_normalization_rejection(
return finalize_request_body_buffer_rejection(
&state,
&request_context,
&remote_addr,
@@ -1335,14 +1553,34 @@ pub(crate) async fn proxy_request(
let stream_request = request_wants_stream(&request_context, &parts.headers, buffered_body);
let mut local_execution_exhaustion = None;
if stream_request {
let stream_outcome = maybe_execute_stream_request(
let stream_outcome = match maybe_execute_stream_request(
&state,
&parts,
buffered_body,
&trace_id,
control_decision,
)
.await?;
.await
{
Ok(outcome) => outcome,
Err(err) => {
if let Some((phase, timeout_ms)) = local_execution_planning_timeout_parts(&err)
{
return finalize_local_execution_planning_timeout(
&state,
&request_context,
&remote_addr,
&started_at,
&trace_id,
request_permit.take(),
control_decision,
phase,
timeout_ms,
);
}
return Err(err);
}
};
debug!(
event_name = "proxy_stream_local_execute_outcome",
log_type = "debug",
@@ -1380,9 +1618,34 @@ pub(crate) async fn proxy_request(
LocalExecutionRequestOutcome::NoPath => {}
}
}
match maybe_execute_sync_request(&state, &parts, buffered_body, &trace_id, control_decision)
.await?
let sync_outcome = match maybe_execute_sync_request(
&state,
&parts,
buffered_body,
&trace_id,
control_decision,
)
.await
{
Ok(outcome) => outcome,
Err(err) => {
if let Some((phase, timeout_ms)) = local_execution_planning_timeout_parts(&err) {
return finalize_local_execution_planning_timeout(
&state,
&request_context,
&remote_addr,
&started_at,
&trace_id,
request_permit.take(),
control_decision,
phase,
timeout_ms,
);
}
return Err(err);
}
};
match sync_outcome {
LocalExecutionRequestOutcome::Responded(execution_runtime_response) => {
let execution_runtime_response = restore_redacted_sync_execution_response(
execution_runtime_response,
@@ -1406,15 +1669,35 @@ pub(crate) async fn proxy_request(
LocalExecutionRequestOutcome::NoPath => {}
}
if parts.method != http::Method::POST {
match maybe_execute_stream_request(
let stream_outcome = match maybe_execute_stream_request(
&state,
&parts,
buffered_body,
&trace_id,
control_decision,
)
.await?
.await
{
Ok(outcome) => outcome,
Err(err) => {
if let Some((phase, timeout_ms)) = local_execution_planning_timeout_parts(&err)
{
return finalize_local_execution_planning_timeout(
&state,
&request_context,
&remote_addr,
&started_at,
&trace_id,
request_permit.take(),
control_decision,
phase,
timeout_ms,
);
}
return Err(err);
}
};
match stream_outcome {
LocalExecutionRequestOutcome::Responded(execution_runtime_response) => {
let execution_runtime_response = restore_redacted_stream_execution_response(
execution_runtime_response,
@@ -1977,14 +2260,17 @@ fn local_execution_runtime_miss_route_detail(
#[cfg(test)]
mod tests {
use std::time::Duration;
use super::{
api_key_remote_ip_allowed, diagnostic_is_auth_api_key_concurrency_limited,
local_execution_runtime_miss_detail, restore_redacted_stream_execution_response,
restore_redacted_sync_execution_response, GatewayControlDecision,
LocalExecutionRuntimeMissDiagnostic,
api_key_remote_ip_allowed, buffer_and_normalize_request_body,
diagnostic_is_auth_api_key_concurrency_limited, local_execution_runtime_miss_detail,
restore_redacted_stream_execution_response, restore_redacted_sync_execution_response,
GatewayControlDecision, LocalExecutionRuntimeMissDiagnostic, RequestBodyBufferError,
RequestBodyBufferPolicy,
};
use axum::body::{to_bytes, Body};
use axum::http::{header, Response};
use axum::body::{to_bytes, Body, Bytes};
use axum::http::{header, HeaderMap, Method, Response};
use serde_json::json;
#[test]
@@ -2112,6 +2398,58 @@ mod tests {
assert!(!message.contains(&sentinel));
}
#[tokio::test]
async fn request_body_buffer_rejects_chunked_body_when_limit_is_exceeded() {
let mut body = Some(Body::from(Bytes::from_static(b"abcdef")));
let mut headers = HeaderMap::new();
let err = buffer_and_normalize_request_body(
&mut body,
&mut headers,
"test owns body",
"trace-body-large",
&Method::POST,
"/v1/responses",
"test",
RequestBodyBufferPolicy::for_tests(5, Duration::from_secs(1)),
)
.await
.expect_err("body exceeding the ingress limit should fail");
assert!(matches!(
err,
RequestBodyBufferError::TooLarge { limit_bytes: 5 }
));
}
#[tokio::test]
async fn request_body_buffer_times_out_instead_of_waiting_forever() {
let stream = async_stream::stream! {
yield Ok::<Bytes, std::io::Error>(Bytes::from_static(b"{"));
std::future::pending::<()>().await;
};
let mut body = Some(Body::from_stream(stream));
let mut headers = HeaderMap::new();
let err = buffer_and_normalize_request_body(
&mut body,
&mut headers,
"test owns body",
"trace-body-timeout",
&Method::POST,
"/v1/responses",
"test",
RequestBodyBufferPolicy::for_tests(1024, Duration::from_millis(5)),
)
.await
.expect_err("body buffering should time out");
assert!(matches!(
err,
RequestBodyBufferError::Timeout { timeout_ms: 5 }
));
}
#[test]
fn runtime_miss_detail_returns_model_specific_stream_message_when_candidates_are_unavailable() {
let decision = GatewayControlDecision::synthetic(
@@ -136,12 +136,7 @@ pub(super) fn announcements_internal_error_response(detail: impl Into<String>) -
}
pub(super) fn announcements_internal_detail(err: GatewayError) -> String {
match err {
GatewayError::UpstreamUnavailable { message, .. }
| GatewayError::ControlUnavailable { message, .. }
| GatewayError::Client { message, .. }
| GatewayError::Internal(message) => message,
}
err.into_message()
}
pub(super) fn parse_optional_rfc3339_unix_secs(