mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-12 06:00:20 +08:00
feat(gateway): harden provider request execution
Preserve exact request payloads and model client surface and API operation explicitly. Add Anthropic compatibility profiles, bounded stream commitment, and scoped OAuth retry behavior across provider transports.
This commit is contained in:
@@ -223,8 +223,8 @@ async fn execute_grok_app_chat(
|
||||
|
||||
let elapsed_ms = started_at.elapsed().as_millis() as u64;
|
||||
if !(200..300).contains(&status_code) {
|
||||
let decoded = decode_response_body_bytes(&headers, &raw_body).unwrap_or(raw_body);
|
||||
let text = String::from_utf8_lossy(&decoded).to_string();
|
||||
let decoded = decode_response_body_bytes(&headers, &raw_body)?;
|
||||
let text = String::from_utf8_lossy(decoded.as_ref()).to_string();
|
||||
return Ok(GrokCollected {
|
||||
status_code,
|
||||
headers,
|
||||
@@ -274,8 +274,8 @@ async fn execute_grok_app_chat_stream(
|
||||
&mut adapter,
|
||||
)
|
||||
.await?;
|
||||
let decoded = decode_response_body_bytes(&headers, &raw_body).unwrap_or(raw_body);
|
||||
let text = String::from_utf8_lossy(&decoded).to_string();
|
||||
let decoded = decode_response_body_bytes(&headers, &raw_body)?;
|
||||
let text = String::from_utf8_lossy(decoded.as_ref()).to_string();
|
||||
let elapsed_ms = started_at.elapsed().as_millis() as u64;
|
||||
let collected = GrokCollected {
|
||||
status_code,
|
||||
|
||||
@@ -42,7 +42,67 @@ pub(crate) use self::response_header_rules::{
|
||||
pub(crate) use crate::orchestration::{
|
||||
append_local_failover_policy_to_value, LocalFailoverAnalysis, LocalFailoverDecision,
|
||||
};
|
||||
pub(crate) use aether_ai_serving::AdaptationMode;
|
||||
pub(crate) use aether_ai_serving::{ConversionMode, ExecutionStrategy};
|
||||
|
||||
pub(crate) fn ai_attempt_retry_scope_from_failure_disposition(
|
||||
disposition: crate::orchestration::FailureDisposition,
|
||||
) -> aether_ai_serving::AiAttemptRetryScope {
|
||||
use crate::orchestration::{FailureRetryAction, FailureScope};
|
||||
use aether_ai_serving::AiAttemptRetryScope;
|
||||
|
||||
match disposition.failure_scope {
|
||||
FailureScope::Credential | FailureScope::CredentialModel => AiAttemptRetryScope::Credential,
|
||||
FailureScope::Endpoint => AiAttemptRetryScope::Endpoint,
|
||||
FailureScope::Provider => AiAttemptRetryScope::Provider,
|
||||
FailureScope::None => match disposition.retry_action {
|
||||
FailureRetryAction::NextCredential => AiAttemptRetryScope::Credential,
|
||||
FailureRetryAction::NextEndpoint => AiAttemptRetryScope::Endpoint,
|
||||
FailureRetryAction::Stop
|
||||
| FailureRetryAction::SameCredential
|
||||
| FailureRetryAction::NextCandidate => AiAttemptRetryScope::Candidate,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod retry_scope_tests {
|
||||
use aether_ai_serving::AiAttemptRetryScope;
|
||||
|
||||
use super::ai_attempt_retry_scope_from_failure_disposition;
|
||||
use crate::orchestration::{classify_failure_disposition, LocalFailoverClassification};
|
||||
|
||||
#[test]
|
||||
fn anthropic_failure_scope_survives_runtime_mapping() {
|
||||
let retry_scope = |status_code| {
|
||||
ai_attempt_retry_scope_from_failure_disposition(classify_failure_disposition(
|
||||
"claude:messages",
|
||||
LocalFailoverClassification::RetryUpstreamFailure,
|
||||
status_code,
|
||||
))
|
||||
};
|
||||
|
||||
assert_eq!(retry_scope(429), AiAttemptRetryScope::Credential);
|
||||
assert_eq!(retry_scope(500), AiAttemptRetryScope::Endpoint);
|
||||
assert_eq!(retry_scope(529), AiAttemptRetryScope::Provider);
|
||||
assert_eq!(retry_scope(400), AiAttemptRetryScope::Candidate);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_anthropic_retry_keeps_existing_candidate_order() {
|
||||
let disposition = classify_failure_disposition(
|
||||
"openai:chat",
|
||||
LocalFailoverClassification::RetryUpstreamFailure,
|
||||
429,
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
ai_attempt_retry_scope_from_failure_disposition(disposition),
|
||||
AiAttemptRetryScope::Candidate
|
||||
);
|
||||
assert!(!disposition.preserve_upstream_error);
|
||||
}
|
||||
}
|
||||
pub use server::{
|
||||
build_execution_runtime_router, build_execution_runtime_router_with_request_concurrency_limit,
|
||||
build_execution_runtime_router_with_request_gates, serve_execution_runtime_tcp,
|
||||
@@ -57,12 +117,14 @@ pub async fn prewarm_direct_h2c_sender_cache_from_env_for_startup(
|
||||
.map_err(|err| err.to_string())
|
||||
}
|
||||
|
||||
pub(crate) use stream::execute_execution_runtime_stream;
|
||||
pub(crate) use stream::{
|
||||
execute_execution_runtime_stream, execute_execution_runtime_stream_with_retry_scope,
|
||||
};
|
||||
pub(crate) use stream_pump::build_direct_execution_frame_stream;
|
||||
pub(crate) use sync::{
|
||||
execute_execution_runtime_sync, maybe_build_local_sync_finalize_response,
|
||||
maybe_build_local_video_error_response, maybe_build_local_video_success_outcome,
|
||||
resolve_local_sync_error_background_report_kind,
|
||||
execute_execution_runtime_sync, execute_execution_runtime_sync_with_retry_scope,
|
||||
maybe_build_local_sync_finalize_response, maybe_build_local_video_error_response,
|
||||
maybe_build_local_video_success_outcome, resolve_local_sync_error_background_report_kind,
|
||||
resolve_local_sync_success_background_report_kind, LocalVideoSyncSuccessBuild,
|
||||
LocalVideoSyncSuccessOutcome,
|
||||
};
|
||||
@@ -220,6 +282,16 @@ pub(crate) fn append_execution_contract_fields(
|
||||
"provider_contract".to_string(),
|
||||
Value::String(provider_contract.to_string()),
|
||||
);
|
||||
let default_adaptation_mode = if execution_strategy == ExecutionStrategy::LocalCrossFormat
|
||||
|| conversion_mode != ConversionMode::None
|
||||
{
|
||||
AdaptationMode::CrossFormat
|
||||
} else {
|
||||
AdaptationMode::NativeTransparent
|
||||
};
|
||||
object
|
||||
.entry("adaptation_mode".to_string())
|
||||
.or_insert_with(|| Value::String(default_adaptation_mode.as_str().to_string()));
|
||||
}
|
||||
|
||||
pub(crate) fn append_execution_contract_fields_to_value(
|
||||
@@ -263,6 +335,7 @@ mod tests {
|
||||
assert_eq!(value["conversion_mode"], "bidirectional");
|
||||
assert_eq!(value["client_contract"], "openai:chat");
|
||||
assert_eq!(value["provider_contract"], "gemini:generate_content");
|
||||
assert_eq!(value["adaptation_mode"], "cross_format");
|
||||
assert_eq!(value["provider_api_format"], "gemini:generate_content");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,6 +1,10 @@
|
||||
use aether_contracts::ExecutionPlan;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::orchestration::{
|
||||
oauth_status_may_be_invalid as status_may_be_oauth_invalid,
|
||||
oauth_status_proves_access_token_invalid as status_proves_access_token_invalid,
|
||||
};
|
||||
use crate::state::AgentIdentityAuthConfigFence;
|
||||
use crate::{provider_transport::LocalOAuthRefreshError, AppState};
|
||||
|
||||
@@ -70,17 +74,17 @@ pub(crate) async fn refresh_oauth_plan_auth_for_retry(
|
||||
// A bearer-token response cannot authorize refreshing an Agent Identity
|
||||
// installed under the same key id while the request was in flight.
|
||||
return false;
|
||||
} else if transport
|
||||
.provider
|
||||
.provider_type
|
||||
.trim()
|
||||
.eq_ignore_ascii_case("codex")
|
||||
&& transport.key.auth_type.trim().eq_ignore_ascii_case("oauth")
|
||||
&& !request_authorization.is_some_and(|authorization| {
|
||||
bearer_authorization_matches_transport(authorization, &transport)
|
||||
})
|
||||
{
|
||||
return false;
|
||||
} else if aether_provider_transport::supports_local_generic_oauth_request_auth_resolution(
|
||||
&transport,
|
||||
) {
|
||||
if let Some(current_authorization) = generic_oauth_transport_authorization(&transport) {
|
||||
if !request_authorization.is_some_and(|authorization| {
|
||||
authorizations_use_same_access_token(authorization, ¤t_authorization)
|
||||
}) {
|
||||
replace_execution_plan_authorization(plan, current_authorization);
|
||||
return true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if transport.key.decrypted_auth_config.is_none()
|
||||
@@ -164,63 +168,31 @@ fn execution_plan_authorization(plan: &ExecutionPlan) -> Option<&str> {
|
||||
.map(|(_, value)| value.as_str())
|
||||
}
|
||||
|
||||
fn bearer_authorization_matches_transport(
|
||||
authorization: &str,
|
||||
fn generic_oauth_transport_authorization(
|
||||
transport: &aether_provider_transport::GatewayProviderTransportSnapshot,
|
||||
) -> bool {
|
||||
let current_token = transport.key.decrypted_api_key.trim();
|
||||
!current_token.is_empty()
|
||||
&& authorization
|
||||
.trim()
|
||||
.strip_prefix("Bearer ")
|
||||
.map(str::trim)
|
||||
.is_some_and(|token| token == current_token)
|
||||
) -> Option<String> {
|
||||
aether_provider_transport::resolve_local_generic_oauth_transport_authorization(transport)
|
||||
}
|
||||
|
||||
fn status_may_be_oauth_invalid(status_code: u16, response_text: Option<&str>) -> bool {
|
||||
if status_code == 401 {
|
||||
return true;
|
||||
fn authorizations_use_same_access_token(left: &str, right: &str) -> bool {
|
||||
match (bearer_access_token(left), bearer_access_token(right)) {
|
||||
(Some(left), Some(right)) => left == right,
|
||||
_ => left.trim() == right.trim(),
|
||||
}
|
||||
if status_code != 403 {
|
||||
return false;
|
||||
}
|
||||
|
||||
let Some(response_text) = response_text else {
|
||||
return true;
|
||||
};
|
||||
let response_text = response_text.to_ascii_lowercase();
|
||||
["oauth", "token", "auth", "credential", "expired"]
|
||||
.iter()
|
||||
.any(|needle| response_text.contains(needle))
|
||||
}
|
||||
|
||||
fn status_proves_access_token_invalid(status_code: u16, response_text: Option<&str>) -> bool {
|
||||
if status_code == 401 {
|
||||
return true;
|
||||
}
|
||||
if status_code != 403 {
|
||||
return false;
|
||||
}
|
||||
fn bearer_access_token(authorization: &str) -> Option<&str> {
|
||||
let mut parts = authorization.split_ascii_whitespace();
|
||||
let scheme = parts.next()?;
|
||||
let token = parts.next()?;
|
||||
(scheme.eq_ignore_ascii_case("bearer") && parts.next().is_none()).then_some(token)
|
||||
}
|
||||
|
||||
let Some(response_text) = response_text else {
|
||||
return false;
|
||||
};
|
||||
let response_text = response_text.to_ascii_lowercase();
|
||||
[
|
||||
"oauth_token_invalid",
|
||||
"invalid_token",
|
||||
"invalid access token",
|
||||
"access token invalid",
|
||||
"access token expired",
|
||||
"expired access token",
|
||||
"authentication token has been invalidated",
|
||||
"token has been invalidated",
|
||||
"personal access token owner is inactive",
|
||||
"biscuit_baker_service_auth_credential_error_status",
|
||||
"security token included in the request is expired",
|
||||
]
|
||||
.iter()
|
||||
.any(|needle| response_text.contains(needle))
|
||||
fn replace_execution_plan_authorization(plan: &mut ExecutionPlan, authorization: String) {
|
||||
plan.headers
|
||||
.retain(|name, _| !name.eq_ignore_ascii_case("authorization"));
|
||||
plan.headers
|
||||
.insert("authorization".to_string(), authorization);
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -230,14 +202,15 @@ mod tests {
|
||||
status_proves_access_token_invalid,
|
||||
};
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use aether_contracts::{ExecutionPlan, RequestBody};
|
||||
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
|
||||
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
ProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
StoredProviderCatalogProvider,
|
||||
ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
use axum::routing::post;
|
||||
use axum::{extract::Request, Json, Router};
|
||||
@@ -252,8 +225,65 @@ mod tests {
|
||||
403,
|
||||
Some("The security token included in the request is expired")
|
||||
));
|
||||
assert!(status_may_be_oauth_invalid(403, None));
|
||||
assert!(status_may_be_oauth_invalid(
|
||||
403,
|
||||
Some("oauth_token_invalid")
|
||||
));
|
||||
assert!(!status_may_be_oauth_invalid(403, None));
|
||||
assert!(!status_may_be_oauth_invalid(
|
||||
403,
|
||||
Some(
|
||||
r#"{"type":"error","error":{"type":"permission_error","message":"this token is not authorized for the workspace"}}"#
|
||||
)
|
||||
));
|
||||
assert!(!status_may_be_oauth_invalid(
|
||||
403,
|
||||
Some(
|
||||
r#"{"error":{"type":"permission_error","message":"the authentication token has been invalidated for this workspace"}}"#
|
||||
)
|
||||
));
|
||||
assert!(status_may_be_oauth_invalid(
|
||||
403,
|
||||
Some(
|
||||
r#"{"type":"error","error":{"type":"authentication_error","message":"credential expired"}}"#
|
||||
)
|
||||
));
|
||||
assert!(status_may_be_oauth_invalid(
|
||||
403,
|
||||
Some(r#"{"error":{"type":"oauth_token_invalid","message":"sign in again"}}"#)
|
||||
));
|
||||
assert!(status_may_be_oauth_invalid(
|
||||
403,
|
||||
Some(
|
||||
r#"{"error":{"code":"biscuit_baker_service_auth_credential_error_status","message":"Personal access token owner is inactive."}}"#
|
||||
)
|
||||
));
|
||||
assert!(!status_may_be_oauth_invalid(
|
||||
403,
|
||||
Some(
|
||||
r#"{"error":{"type":"invalid_request_error","message":"Your authentication token has been invalidated. Please try signing in again."}}"#
|
||||
)
|
||||
));
|
||||
assert!(!status_may_be_oauth_invalid(
|
||||
403,
|
||||
Some(
|
||||
r#"{"error":{"type":"invalid_request_error","message":"invalid request: token budget is invalid"}}"#
|
||||
)
|
||||
));
|
||||
assert!(!status_may_be_oauth_invalid(403, Some("quota exceeded")));
|
||||
assert!(!status_may_be_oauth_invalid(
|
||||
403,
|
||||
Some("invalid request: max token budget is invalid")
|
||||
));
|
||||
assert!(!status_may_be_oauth_invalid(
|
||||
403,
|
||||
Some("invalid_token_budget")
|
||||
));
|
||||
assert!(!status_may_be_oauth_invalid(403, Some("not authorized")));
|
||||
assert!(!status_may_be_oauth_invalid(
|
||||
403,
|
||||
Some("authorization denied")
|
||||
));
|
||||
assert!(!status_may_be_oauth_invalid(429, Some("token bucket")));
|
||||
}
|
||||
|
||||
@@ -282,7 +312,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn auto_removes_request_proven_oauth_failure_after_terminal_refresh_failure() {
|
||||
async fn retains_codex_key_after_request_proven_terminal_refresh_failure() {
|
||||
let token_hits = Arc::new(Mutex::new(0usize));
|
||||
let token_hits_clone = Arc::clone(&token_hits);
|
||||
let token_server = Router::new().route(
|
||||
@@ -437,7 +467,233 @@ mod tests {
|
||||
.list_keys_by_ids(&["key-codex-oauth-retry".to_string()])
|
||||
.await
|
||||
.expect("keys should read");
|
||||
assert!(keys.is_empty());
|
||||
assert_eq!(keys.len(), 1);
|
||||
assert!(keys[0].oauth_invalid_at_unix_secs.is_some());
|
||||
assert!(keys[0]
|
||||
.oauth_invalid_reason
|
||||
.as_deref()
|
||||
.is_some_and(|reason| reason.contains("[REFRESH_FAILED]")
|
||||
&& reason.contains("Token 续期失败 (401)")));
|
||||
|
||||
token_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn stale_claude_code_request_reuses_rotated_access_token_without_second_refresh() {
|
||||
let refresh_hits = Arc::new(AtomicUsize::new(0));
|
||||
let refresh_hits_for_server = Arc::clone(&refresh_hits);
|
||||
let token_server = Router::new().route(
|
||||
"/oauth/token",
|
||||
post(move |_request: Request| {
|
||||
let hits = Arc::clone(&refresh_hits_for_server);
|
||||
async move {
|
||||
hits.fetch_add(1, Ordering::SeqCst);
|
||||
Json(json!({
|
||||
"access_token": "fresh-claude-access-token",
|
||||
"refresh_token": "fresh-claude-refresh-token",
|
||||
"expires_in": 3600,
|
||||
"token_type": "Bearer"
|
||||
}))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let provider = StoredProviderCatalogProvider::new(
|
||||
"provider-claude-code".to_string(),
|
||||
"Claude Code".to_string(),
|
||||
Some("https://api.anthropic.com".to_string()),
|
||||
"claude_code".to_string(),
|
||||
)
|
||||
.expect("provider should build");
|
||||
let endpoint = StoredProviderCatalogEndpoint::new(
|
||||
"endpoint-claude-code".to_string(),
|
||||
"provider-claude-code".to_string(),
|
||||
"claude:messages".to_string(),
|
||||
None,
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("endpoint should build")
|
||||
.with_transport_fields(
|
||||
"https://api.anthropic.com".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("endpoint transport should build");
|
||||
let encrypted_api_key = encrypt_python_fernet_plaintext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
"stale-claude-access-token",
|
||||
)
|
||||
.expect("api key ciphertext should build");
|
||||
let encrypted_auth_config = encrypt_python_fernet_plaintext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
r#"{"provider_type":"claude_code","access_token":"stale-claude-access-token","refresh_token":"stale-claude-refresh-token","expires_at":4102444800}"#,
|
||||
)
|
||||
.expect("auth config ciphertext should build");
|
||||
let mut key = StoredProviderCatalogKey::new(
|
||||
"key-claude-code".to_string(),
|
||||
"provider-claude-code".to_string(),
|
||||
"Claude OAuth".to_string(),
|
||||
"oauth".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build")
|
||||
.with_transport_fields(
|
||||
Some(json!(["claude:messages"])),
|
||||
encrypted_api_key,
|
||||
Some(encrypted_auth_config),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.expect("key transport should build");
|
||||
key.expires_at_unix_secs = Some(4_102_444_800);
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![provider],
|
||||
vec![endpoint],
|
||||
vec![key],
|
||||
));
|
||||
let (token_url, token_handle) = start_test_server(token_server).await;
|
||||
let oauth_refresh =
|
||||
crate::provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![
|
||||
Arc::new(
|
||||
crate::provider_transport::oauth_refresh::GenericOAuthRefreshAdapter::default()
|
||||
.with_token_url_for_tests(
|
||||
"claude_code",
|
||||
format!("{token_url}/oauth/token"),
|
||||
),
|
||||
),
|
||||
]);
|
||||
let state = crate::AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(
|
||||
crate::data::GatewayDataState::with_provider_catalog_repository_for_tests(
|
||||
provider_catalog_repository.clone(),
|
||||
)
|
||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
|
||||
)
|
||||
.with_oauth_refresh_coordinator_for_tests(oauth_refresh);
|
||||
let stale_transport = state
|
||||
.read_provider_transport_snapshot(
|
||||
"provider-claude-code",
|
||||
"endpoint-claude-code",
|
||||
"key-claude-code",
|
||||
)
|
||||
.await
|
||||
.expect("stale transport should load")
|
||||
.expect("stale transport should exist");
|
||||
let stale_plan = ExecutionPlan {
|
||||
request_id: "req-claude-oauth-fence".to_string(),
|
||||
candidate_id: None,
|
||||
provider_name: Some("claude_code".to_string()),
|
||||
provider_id: "provider-claude-code".to_string(),
|
||||
endpoint_id: "endpoint-claude-code".to_string(),
|
||||
key_id: "key-claude-code".to_string(),
|
||||
method: "POST".to_string(),
|
||||
url: "https://api.anthropic.com/v1/messages".to_string(),
|
||||
headers: BTreeMap::from([(
|
||||
"authorization".to_string(),
|
||||
"Bearer stale-claude-access-token".to_string(),
|
||||
)]),
|
||||
content_type: Some("application/json".to_string()),
|
||||
content_encoding: None,
|
||||
body: RequestBody::from_json(json!({"model": "claude-sonnet-4-5"})),
|
||||
stream: false,
|
||||
client_api_format: "claude:messages".to_string(),
|
||||
provider_api_format: "claude:messages".to_string(),
|
||||
model_name: Some("claude-sonnet-4-5".to_string()),
|
||||
proxy: None,
|
||||
transport_profile: None,
|
||||
timeouts: None,
|
||||
};
|
||||
|
||||
let mut first_plan = stale_plan.clone();
|
||||
assert!(
|
||||
refresh_oauth_plan_auth_for_retry(
|
||||
&state,
|
||||
&mut first_plan,
|
||||
401,
|
||||
Some(r#"{"error":"invalid_token"}"#),
|
||||
"trace-claude-oauth-fence-first",
|
||||
)
|
||||
.await
|
||||
);
|
||||
assert_eq!(
|
||||
first_plan.headers.get("authorization").map(String::as_str),
|
||||
Some("Bearer fresh-claude-access-token")
|
||||
);
|
||||
assert_eq!(refresh_hits.load(Ordering::SeqCst), 1);
|
||||
|
||||
let stale_force_result = state
|
||||
.force_local_oauth_refresh_entry(&stale_transport)
|
||||
.await
|
||||
.expect("stale force should reuse the persisted winner")
|
||||
.expect("stale force should return the winner entry");
|
||||
assert_eq!(
|
||||
stale_force_result.auth_header_value,
|
||||
"Bearer fresh-claude-access-token"
|
||||
);
|
||||
assert_eq!(refresh_hits.load(Ordering::SeqCst), 1);
|
||||
|
||||
let mut stale_in_flight_plan = stale_plan;
|
||||
assert!(
|
||||
refresh_oauth_plan_auth_for_retry(
|
||||
&state,
|
||||
&mut stale_in_flight_plan,
|
||||
401,
|
||||
Some(r#"{"error":"invalid_token"}"#),
|
||||
"trace-claude-oauth-fence-stale",
|
||||
)
|
||||
.await
|
||||
);
|
||||
assert_eq!(
|
||||
stale_in_flight_plan
|
||||
.headers
|
||||
.get("authorization")
|
||||
.map(String::as_str),
|
||||
Some("Bearer fresh-claude-access-token")
|
||||
);
|
||||
assert_eq!(refresh_hits.load(Ordering::SeqCst), 1);
|
||||
|
||||
let mut admin_replacement = provider_catalog_repository
|
||||
.list_keys_by_ids(&["key-claude-code".to_string()])
|
||||
.await
|
||||
.expect("Claude key should load")
|
||||
.pop()
|
||||
.expect("Claude key should exist");
|
||||
admin_replacement.encrypted_api_key = Some(
|
||||
encrypt_python_fernet_plaintext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
"admin-claude-access-token",
|
||||
)
|
||||
.expect("admin access token should encrypt"),
|
||||
);
|
||||
admin_replacement.expires_at_unix_secs = Some(4_102_444_800);
|
||||
provider_catalog_repository
|
||||
.update_key(&admin_replacement)
|
||||
.await
|
||||
.expect("admin replacement should persist");
|
||||
|
||||
let admin_result = state
|
||||
.force_local_oauth_refresh_entry(&stale_transport)
|
||||
.await
|
||||
.expect("stale force should reuse the admin replacement")
|
||||
.expect("admin replacement should resolve");
|
||||
assert_eq!(
|
||||
admin_result.auth_header_value,
|
||||
"Bearer admin-claude-access-token"
|
||||
);
|
||||
assert_eq!(refresh_hits.load(Ordering::SeqCst), 1);
|
||||
|
||||
token_handle.abort();
|
||||
}
|
||||
|
||||
@@ -375,6 +375,8 @@ impl IntoResponse for ExecutionRuntimeAppError {
|
||||
| ExecutionRuntimeTransportError::BrowserClientBuild(_)
|
||||
| ExecutionRuntimeTransportError::BrowserBody(_)
|
||||
| ExecutionRuntimeTransportError::UpstreamRequest(_)
|
||||
| ExecutionRuntimeTransportError::UpstreamResponseTooLarge { .. }
|
||||
| ExecutionRuntimeTransportError::UpstreamResponseDecode { .. }
|
||||
| ExecutionRuntimeTransportError::RelayError(_)
|
||||
| ExecutionRuntimeTransportError::InvalidJson(_),
|
||||
) => StatusCode::BAD_GATEWAY,
|
||||
|
||||
@@ -0,0 +1,527 @@
|
||||
use std::time::Duration;
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::execution_runtime::MAX_STREAM_PREFETCH_BYTES;
|
||||
|
||||
const ANTHROPIC_PRECOMMIT_MAX_WAIT: Duration = Duration::from_millis(750);
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(super) enum StreamCommitPolicy {
|
||||
OnResponseHeaders,
|
||||
OnFirstClassifiedBody,
|
||||
OnFirstAnthropicSemanticEvent {
|
||||
max_bytes: usize,
|
||||
max_wait: Duration,
|
||||
},
|
||||
}
|
||||
|
||||
impl StreamCommitPolicy {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(super) fn for_response(
|
||||
has_direct_finalize: bool,
|
||||
content_type: Option<&str>,
|
||||
provider_api_format: &str,
|
||||
client_api_format: &str,
|
||||
has_private_stream_normalizer: bool,
|
||||
has_local_stream_rewriter: bool,
|
||||
force_prefetch: bool,
|
||||
) -> Self {
|
||||
if !has_direct_finalize {
|
||||
return Self::OnFirstClassifiedBody;
|
||||
}
|
||||
|
||||
if force_prefetch {
|
||||
return Self::OnFirstClassifiedBody;
|
||||
}
|
||||
|
||||
let content_type = content_type
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or_default()
|
||||
.to_ascii_lowercase();
|
||||
if content_type.contains("text/event-stream") {
|
||||
if provider_api_format.eq_ignore_ascii_case("claude:messages")
|
||||
&& provider_api_format.eq_ignore_ascii_case(client_api_format)
|
||||
&& !has_private_stream_normalizer
|
||||
&& !has_local_stream_rewriter
|
||||
{
|
||||
return Self::OnFirstAnthropicSemanticEvent {
|
||||
max_bytes: MAX_STREAM_PREFETCH_BYTES,
|
||||
max_wait: ANTHROPIC_PRECOMMIT_MAX_WAIT,
|
||||
};
|
||||
}
|
||||
return Self::OnResponseHeaders;
|
||||
}
|
||||
|
||||
if has_private_stream_normalizer || has_local_stream_rewriter {
|
||||
return Self::OnFirstClassifiedBody;
|
||||
}
|
||||
|
||||
if !provider_api_format.eq_ignore_ascii_case(client_api_format) {
|
||||
return Self::OnFirstClassifiedBody;
|
||||
}
|
||||
|
||||
if content_type.is_empty() {
|
||||
return Self::OnResponseHeaders;
|
||||
}
|
||||
|
||||
if content_type.contains("json") || content_type.ends_with("+json") {
|
||||
Self::OnFirstClassifiedBody
|
||||
} else {
|
||||
Self::OnResponseHeaders
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) const fn commits_on_response_headers(self) -> bool {
|
||||
matches!(self, Self::OnResponseHeaders)
|
||||
}
|
||||
|
||||
pub(super) const fn requires_bounded_frame_wait(self) -> bool {
|
||||
matches!(self, Self::OnFirstAnthropicSemanticEvent { .. })
|
||||
}
|
||||
|
||||
pub(super) const fn max_precommit_wait(self) -> Option<Duration> {
|
||||
match self {
|
||||
Self::OnFirstAnthropicSemanticEvent { max_wait, .. } => Some(max_wait),
|
||||
Self::OnResponseHeaders | Self::OnFirstClassifiedBody => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) const fn is_native_anthropic(self) -> bool {
|
||||
matches!(self, Self::OnFirstAnthropicSemanticEvent { .. })
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(super) enum StreamCommitState {
|
||||
Uncommitted,
|
||||
Committed,
|
||||
Terminal,
|
||||
}
|
||||
|
||||
#[derive(Debug, PartialEq)]
|
||||
pub(super) enum StreamPrecommitObservation {
|
||||
Pending,
|
||||
Commit,
|
||||
UpstreamError { status_code: u16, body_json: Value },
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub(super) struct StreamCommitGate {
|
||||
policy: StreamCommitPolicy,
|
||||
state: StreamCommitState,
|
||||
observed_bytes: usize,
|
||||
anthropic: AnthropicSsePrecommitInspector,
|
||||
}
|
||||
|
||||
impl StreamCommitGate {
|
||||
pub(super) fn new(policy: StreamCommitPolicy) -> Self {
|
||||
let state = if policy.commits_on_response_headers() {
|
||||
StreamCommitState::Committed
|
||||
} else {
|
||||
StreamCommitState::Uncommitted
|
||||
};
|
||||
Self {
|
||||
policy,
|
||||
state,
|
||||
observed_bytes: 0,
|
||||
anthropic: AnthropicSsePrecommitInspector::default(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) const fn state(&self) -> StreamCommitState {
|
||||
self.state
|
||||
}
|
||||
|
||||
pub(super) const fn is_uncommitted(&self) -> bool {
|
||||
matches!(self.state, StreamCommitState::Uncommitted)
|
||||
}
|
||||
|
||||
pub(super) fn observe_provider_bytes(&mut self, chunk: &[u8]) -> StreamPrecommitObservation {
|
||||
if self.state != StreamCommitState::Uncommitted {
|
||||
return StreamPrecommitObservation::Commit;
|
||||
}
|
||||
|
||||
let StreamCommitPolicy::OnFirstAnthropicSemanticEvent { max_bytes, .. } = self.policy
|
||||
else {
|
||||
return StreamPrecommitObservation::Pending;
|
||||
};
|
||||
|
||||
self.observed_bytes = self.observed_bytes.saturating_add(chunk.len());
|
||||
match self.anthropic.observe(chunk, max_bytes) {
|
||||
AnthropicSseObservation::Pending => {}
|
||||
AnthropicSseObservation::SemanticEvent => {
|
||||
self.state = StreamCommitState::Committed;
|
||||
return StreamPrecommitObservation::Commit;
|
||||
}
|
||||
AnthropicSseObservation::Error(body_json) => {
|
||||
self.state = StreamCommitState::Terminal;
|
||||
return StreamPrecommitObservation::UpstreamError {
|
||||
status_code: anthropic_error_status_code(&body_json),
|
||||
body_json,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
if self.observed_bytes >= max_bytes {
|
||||
self.commit();
|
||||
StreamPrecommitObservation::Commit
|
||||
} else {
|
||||
StreamPrecommitObservation::Pending
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn commit(&mut self) {
|
||||
if self.state == StreamCommitState::Uncommitted {
|
||||
self.state = StreamCommitState::Committed;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
enum AnthropicSseObservation {
|
||||
Pending,
|
||||
SemanticEvent,
|
||||
Error(Value),
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct AnthropicSsePrecommitInspector {
|
||||
buffered: Vec<u8>,
|
||||
}
|
||||
|
||||
impl AnthropicSsePrecommitInspector {
|
||||
fn observe(&mut self, chunk: &[u8], max_bytes: usize) -> AnthropicSseObservation {
|
||||
let remaining = max_bytes.saturating_sub(self.buffered.len());
|
||||
let truncated = chunk.len() > remaining;
|
||||
self.buffered
|
||||
.extend_from_slice(&chunk[..chunk.len().min(remaining)]);
|
||||
|
||||
while let Some((record_end, separator_len)) = find_sse_record_boundary(&self.buffered) {
|
||||
let record = self.buffered[..record_end].to_vec();
|
||||
self.buffered.drain(..record_end + separator_len);
|
||||
match classify_anthropic_sse_record(&record) {
|
||||
AnthropicSseObservation::Pending => {}
|
||||
decision => return decision,
|
||||
}
|
||||
}
|
||||
|
||||
if truncated {
|
||||
AnthropicSseObservation::SemanticEvent
|
||||
} else {
|
||||
AnthropicSseObservation::Pending
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn find_sse_record_boundary(buffer: &[u8]) -> Option<(usize, usize)> {
|
||||
let mut cursor = 0;
|
||||
while cursor < buffer.len() {
|
||||
let (line_end, line_ending_len) = next_sse_line_ending(buffer, cursor)?;
|
||||
let next_line_start = line_end + line_ending_len;
|
||||
let Some((next_line_end, next_line_ending_len)) =
|
||||
next_sse_line_ending(buffer, next_line_start)
|
||||
else {
|
||||
return None;
|
||||
};
|
||||
if next_line_end == next_line_start {
|
||||
return Some((
|
||||
line_end,
|
||||
line_ending_len.saturating_add(next_line_ending_len),
|
||||
));
|
||||
}
|
||||
cursor = next_line_start;
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn next_sse_line_ending(buffer: &[u8], start: usize) -> Option<(usize, usize)> {
|
||||
let relative = buffer
|
||||
.get(start..)?
|
||||
.iter()
|
||||
.position(|byte| matches!(byte, b'\r' | b'\n'))?;
|
||||
let index = start + relative;
|
||||
let ending_len = if buffer[index] == b'\r' && buffer.get(index + 1) == Some(&b'\n') {
|
||||
2
|
||||
} else {
|
||||
1
|
||||
};
|
||||
Some((index, ending_len))
|
||||
}
|
||||
|
||||
fn classify_anthropic_sse_record(record: &[u8]) -> AnthropicSseObservation {
|
||||
let Ok(record) = std::str::from_utf8(record) else {
|
||||
return AnthropicSseObservation::Pending;
|
||||
};
|
||||
let normalized_record = record.replace("\r\n", "\n").replace('\r', "\n");
|
||||
let mut event_type = None;
|
||||
let mut data = String::new();
|
||||
for line in normalized_record.lines() {
|
||||
if line.starts_with(':') {
|
||||
continue;
|
||||
}
|
||||
if let Some(value) = line.strip_prefix("event:") {
|
||||
let value = value.trim();
|
||||
if !value.is_empty() {
|
||||
event_type = Some(value);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
if let Some(value) = line.strip_prefix("data:") {
|
||||
if !data.is_empty() {
|
||||
data.push('\n');
|
||||
}
|
||||
data.push_str(value.trim_start());
|
||||
}
|
||||
}
|
||||
if data.trim().is_empty() {
|
||||
return AnthropicSseObservation::Pending;
|
||||
}
|
||||
|
||||
let Ok(body_json) = serde_json::from_str::<Value>(data.trim()) else {
|
||||
return AnthropicSseObservation::Pending;
|
||||
};
|
||||
let payload_type = body_json.get("type").and_then(Value::as_str).map(str::trim);
|
||||
if event_type == Some("error") || payload_type == Some("error") {
|
||||
return AnthropicSseObservation::Error(body_json);
|
||||
}
|
||||
|
||||
let semantic_type = match (event_type, payload_type) {
|
||||
(Some(event_type), Some(payload_type)) if event_type == payload_type => Some(event_type),
|
||||
(None, Some(payload_type)) => Some(payload_type),
|
||||
_ => None,
|
||||
};
|
||||
if semantic_type.is_some_and(is_anthropic_semantic_event_type) {
|
||||
AnthropicSseObservation::SemanticEvent
|
||||
} else {
|
||||
AnthropicSseObservation::Pending
|
||||
}
|
||||
}
|
||||
|
||||
fn is_anthropic_semantic_event_type(event_type: &str) -> bool {
|
||||
matches!(
|
||||
event_type,
|
||||
"message_start"
|
||||
| "content_block_start"
|
||||
| "content_block_delta"
|
||||
| "content_block_stop"
|
||||
| "message_delta"
|
||||
| "message_stop"
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn anthropic_error_status_code(body_json: &Value) -> u16 {
|
||||
let error_type = body_json
|
||||
.get("error")
|
||||
.and_then(|error| error.get("type"))
|
||||
.or_else(|| body_json.get("type"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.unwrap_or_default();
|
||||
match error_type {
|
||||
"invalid_request_error" => 400,
|
||||
"authentication_error" => 401,
|
||||
"permission_error" => 403,
|
||||
"not_found_error" => 404,
|
||||
"request_too_large" => 413,
|
||||
"rate_limit_error" => 429,
|
||||
"overloaded_error" => 529,
|
||||
_ => 500,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::time::Duration;
|
||||
|
||||
use super::{
|
||||
anthropic_error_status_code, StreamCommitGate, StreamCommitPolicy, StreamCommitState,
|
||||
StreamPrecommitObservation,
|
||||
};
|
||||
|
||||
fn native_anthropic_policy() -> StreamCommitPolicy {
|
||||
StreamCommitPolicy::OnFirstAnthropicSemanticEvent {
|
||||
max_bytes: 16_384,
|
||||
max_wait: Duration::from_millis(750),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn policy_selects_bounded_anthropic_gate_only_for_native_same_format_sse() {
|
||||
let native = StreamCommitPolicy::for_response(
|
||||
true,
|
||||
Some("text/event-stream; charset=utf-8"),
|
||||
"claude:messages",
|
||||
"claude:messages",
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
);
|
||||
assert!(native.is_native_anthropic());
|
||||
assert_eq!(
|
||||
native.max_precommit_wait(),
|
||||
Some(Duration::from_millis(750))
|
||||
);
|
||||
assert!(StreamCommitPolicy::for_response(
|
||||
true,
|
||||
Some("text/event-stream"),
|
||||
"openai:chat",
|
||||
"claude:messages",
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
)
|
||||
.commits_on_response_headers());
|
||||
assert!(StreamCommitPolicy::for_response(
|
||||
true,
|
||||
Some("text/event-stream"),
|
||||
"claude:messages",
|
||||
"claude:messages",
|
||||
false,
|
||||
true,
|
||||
false,
|
||||
)
|
||||
.commits_on_response_headers());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gate_detects_anthropic_error_across_every_chunk_boundary() {
|
||||
let event = b"event: error\r\ndata: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",\"message\":\"busy\"}}\r\n\r\n";
|
||||
for split in 1..event.len() {
|
||||
let mut gate = StreamCommitGate::new(native_anthropic_policy());
|
||||
let first_observation = gate.observe_provider_bytes(&event[..split]);
|
||||
if matches!(
|
||||
first_observation,
|
||||
StreamPrecommitObservation::UpstreamError {
|
||||
status_code: 529,
|
||||
..
|
||||
}
|
||||
) {
|
||||
assert_eq!(event[split - 1], b'\r');
|
||||
} else {
|
||||
assert_eq!(first_observation, StreamPrecommitObservation::Pending);
|
||||
assert!(matches!(
|
||||
gate.observe_provider_bytes(&event[split..]),
|
||||
StreamPrecommitObservation::UpstreamError {
|
||||
status_code: 529,
|
||||
..
|
||||
}
|
||||
));
|
||||
}
|
||||
assert_eq!(gate.state(), StreamCommitState::Terminal);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gate_detects_cr_only_and_mixed_line_ending_errors() {
|
||||
for event in [
|
||||
"event: error\rdata: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\"}}\r\r",
|
||||
"event: error\r\ndata: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\"}}\n\r",
|
||||
] {
|
||||
for split in 1..event.len() {
|
||||
let mut gate = StreamCommitGate::new(native_anthropic_policy());
|
||||
assert_eq!(
|
||||
gate.observe_provider_bytes(&event.as_bytes()[..split]),
|
||||
StreamPrecommitObservation::Pending,
|
||||
"gate committed before complete mixed-line event at split {split}",
|
||||
);
|
||||
assert!(matches!(
|
||||
gate.observe_provider_bytes(&event.as_bytes()[split..]),
|
||||
StreamPrecommitObservation::UpstreamError {
|
||||
status_code: 529,
|
||||
..
|
||||
}
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unknown_and_ping_events_do_not_commit_before_anthropic_error() {
|
||||
let mut gate = StreamCommitGate::new(native_anthropic_policy());
|
||||
assert_eq!(
|
||||
gate.observe_provider_bytes(
|
||||
b"event: future_event\ndata: {\"type\":\"future_event\",\"value\":1}\n\n"
|
||||
),
|
||||
StreamPrecommitObservation::Pending
|
||||
);
|
||||
assert_eq!(
|
||||
gate.observe_provider_bytes(b"event: ping\ndata: {\"type\":\"ping\"}\n\n"),
|
||||
StreamPrecommitObservation::Pending
|
||||
);
|
||||
assert!(matches!(
|
||||
gate.observe_provider_bytes(
|
||||
b"event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"rate_limit_error\"}}\n\n"
|
||||
),
|
||||
StreamPrecommitObservation::UpstreamError {
|
||||
status_code: 429,
|
||||
..
|
||||
}
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn first_semantic_event_commits_before_later_error_in_same_chunk() {
|
||||
let mut gate = StreamCommitGate::new(native_anthropic_policy());
|
||||
let observation = gate.observe_provider_bytes(
|
||||
concat!(
|
||||
"event: message_start\n",
|
||||
"data: {\"type\":\"message_start\",\"message\":{}}\n\n",
|
||||
"event: error\n",
|
||||
"data: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\"}}\n\n",
|
||||
)
|
||||
.as_bytes(),
|
||||
);
|
||||
|
||||
assert_eq!(observation, StreamPrecommitObservation::Commit);
|
||||
assert_eq!(gate.state(), StreamCommitState::Committed);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transport_fragment_count_does_not_commit_an_incomplete_anthropic_error() {
|
||||
let policy = StreamCommitPolicy::OnFirstAnthropicSemanticEvent {
|
||||
max_bytes: 1024,
|
||||
max_wait: Duration::from_millis(750),
|
||||
};
|
||||
let mut gate = StreamCommitGate::new(policy);
|
||||
let event = b"event: error\ndata: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\"}}\n\n";
|
||||
for byte in &event[..event.len() - 1] {
|
||||
assert_eq!(
|
||||
gate.observe_provider_bytes(std::slice::from_ref(byte)),
|
||||
StreamPrecommitObservation::Pending,
|
||||
);
|
||||
}
|
||||
assert!(matches!(
|
||||
gate.observe_provider_bytes(&event[event.len() - 1..]),
|
||||
StreamPrecommitObservation::UpstreamError {
|
||||
status_code: 529,
|
||||
..
|
||||
}
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn anthropic_error_status_mapping_matches_messages_api_taxonomy() {
|
||||
for (error_type, status_code) in [
|
||||
("invalid_request_error", 400),
|
||||
("authentication_error", 401),
|
||||
("permission_error", 403),
|
||||
("not_found_error", 404),
|
||||
("request_too_large", 413),
|
||||
("rate_limit_error", 429),
|
||||
("overloaded_error", 529),
|
||||
("api_error", 500),
|
||||
] {
|
||||
let body = serde_json::json!({
|
||||
"type": "error",
|
||||
"error": { "type": error_type, "message": "upstream failure" }
|
||||
});
|
||||
assert_eq!(
|
||||
anthropic_error_status_code(&body),
|
||||
status_code,
|
||||
"unexpected status for {error_type}"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,3 +1,4 @@
|
||||
use aether_ai_serving::AiAttemptRetryScope;
|
||||
use aether_contracts::{ExecutionError, ExecutionPlan, ExecutionTelemetry};
|
||||
use aether_data_contracts::repository::candidates::RequestCandidateStatus;
|
||||
use aether_scheduler_core::SchedulerRequestCandidateStatusUpdate;
|
||||
@@ -14,16 +15,17 @@ use tracing::warn;
|
||||
use crate::api::response::attach_control_metadata_headers;
|
||||
use crate::clock::current_unix_ms as current_request_candidate_unix_ms;
|
||||
use crate::control::GatewayControlDecision;
|
||||
use crate::execution_runtime::ai_attempt_retry_scope_from_failure_disposition;
|
||||
use crate::execution_runtime::submission::{
|
||||
resolve_core_error_background_report_kind, submit_local_core_error_or_sync_finalize,
|
||||
};
|
||||
use crate::log_ids::short_request_id;
|
||||
use crate::orchestration::{
|
||||
apply_local_execution_effect, resolve_local_failover_analysis_for_attempt,
|
||||
with_upstream_response_report_context, LocalAdaptiveRateLimitEffect, LocalAttemptFailureEffect,
|
||||
LocalExecutionEffect, LocalExecutionEffectContext, LocalFailoverAnalysis,
|
||||
LocalFailoverDecision, LocalHealthFailureEffect, LocalOAuthInvalidationEffect,
|
||||
LocalPoolErrorEffect,
|
||||
apply_local_execution_effect, classify_failure_disposition,
|
||||
resolve_local_failover_analysis_for_attempt, with_upstream_response_report_context,
|
||||
LocalAdaptiveRateLimitEffect, LocalAttemptFailureEffect, LocalExecutionEffect,
|
||||
LocalExecutionEffectContext, LocalFailoverAnalysis, LocalFailoverDecision,
|
||||
LocalHealthFailureEffect, LocalOAuthInvalidationEffect, LocalPoolErrorEffect,
|
||||
};
|
||||
use crate::request_candidate_runtime::record_report_request_candidate_status;
|
||||
use crate::request_diagnostics::attach_current_request_diagnostics_to_report_context;
|
||||
@@ -409,9 +411,13 @@ pub(super) async fn handle_prefetch_provider_private_stream_error(
|
||||
mut headers: std::collections::BTreeMap<String, String>,
|
||||
telemetry: Option<ExecutionTelemetry>,
|
||||
buffered_body: &[u8],
|
||||
upstream_status_code: u16,
|
||||
status_code: u16,
|
||||
body_json: Value,
|
||||
retry_scope_out: Option<&mut AiAttemptRetryScope>,
|
||||
retry_fallback_out: Option<&mut Option<Response<Body>>>,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
let upstream_headers = headers.clone();
|
||||
headers.remove("content-encoding");
|
||||
headers.remove("content-length");
|
||||
headers.insert("content-type".to_string(), "application/json".to_string());
|
||||
@@ -441,6 +447,29 @@ pub(super) async fn handle_prefetch_provider_private_stream_error(
|
||||
failure_analysis.decision,
|
||||
LocalFailoverDecision::RetryNextCandidate
|
||||
) {
|
||||
let failure_disposition = classify_failure_disposition(
|
||||
&plan.provider_api_format,
|
||||
failure_analysis.classification,
|
||||
status_code,
|
||||
);
|
||||
if let Some(retry_scope) = retry_scope_out {
|
||||
*retry_scope = ai_attempt_retry_scope_from_failure_disposition(failure_disposition);
|
||||
}
|
||||
if failure_disposition.preserve_upstream_error {
|
||||
if let Some(retry_fallback) = retry_fallback_out {
|
||||
*retry_fallback = Some(attach_control_metadata_headers(
|
||||
crate::api::response::build_client_response_from_parts(
|
||||
upstream_status_code,
|
||||
&upstream_headers,
|
||||
Body::from(buffered_body.to_vec()),
|
||||
trace_id,
|
||||
Some(decision),
|
||||
)?,
|
||||
Some(request_id),
|
||||
candidate_id,
|
||||
)?);
|
||||
}
|
||||
}
|
||||
warn!(
|
||||
event_name = "local_stream_candidate_retry_scheduled",
|
||||
log_type = "event",
|
||||
|
||||
@@ -1,4 +1,7 @@
|
||||
mod commit_policy;
|
||||
mod error;
|
||||
mod execution;
|
||||
|
||||
pub(crate) use execution::execute_execution_runtime_stream;
|
||||
pub(crate) use execution::{
|
||||
execute_execution_runtime_stream, execute_execution_runtime_stream_with_retry_scope,
|
||||
};
|
||||
|
||||
@@ -21,12 +21,14 @@ use crate::ai_serving::api::{
|
||||
};
|
||||
use crate::execution_runtime::ndjson::encode_stream_frame_ndjson;
|
||||
use crate::execution_runtime::transport::{
|
||||
format_hyper_error_chain, format_wreq_upstream_request_error,
|
||||
stream_first_byte_timeout_message, DirectUpstreamResponse,
|
||||
append_upstream_response_body_chunk, decode_response_body_bytes, format_hyper_error_chain,
|
||||
format_wreq_upstream_request_error, stream_first_byte_timeout_message, DirectUpstreamResponse,
|
||||
};
|
||||
use crate::execution_runtime::DirectUpstreamStreamExecution;
|
||||
use crate::GatewayError;
|
||||
|
||||
const STREAM_USAGE_OBSERVER_MAX_LINE_BYTES: usize = 1024 * 1024;
|
||||
|
||||
pub(crate) fn build_direct_execution_frame_stream(
|
||||
execution: DirectUpstreamStreamExecution,
|
||||
) -> impl Stream<Item = Result<Bytes, IoError>> + Send + 'static {
|
||||
@@ -39,6 +41,7 @@ pub(crate) fn build_direct_execution_frame_stream(
|
||||
provider_api_format,
|
||||
stream_summary_report_context,
|
||||
prefetched_body,
|
||||
stream_precommit_committed: _,
|
||||
response,
|
||||
started_at,
|
||||
stream_first_byte_timeout,
|
||||
@@ -712,6 +715,23 @@ struct BufferedUpstreamBodyError {
|
||||
first_byte_timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
fn append_buffered_upstream_body_chunk(
|
||||
body_bytes: &mut Vec<u8>,
|
||||
chunk: &[u8],
|
||||
ttfb_ms: Option<u64>,
|
||||
upstream_bytes: &mut u64,
|
||||
) -> Result<(), BufferedUpstreamBodyError> {
|
||||
*upstream_bytes = upstream_bytes.saturating_add(chunk.len() as u64);
|
||||
append_upstream_response_body_chunk(body_bytes, chunk).map_err(|error| {
|
||||
BufferedUpstreamBodyError {
|
||||
message: error.to_string(),
|
||||
ttfb_ms,
|
||||
upstream_bytes: *upstream_bytes,
|
||||
first_byte_timeout: None,
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn response_headers_indicate_sse(headers: &BTreeMap<String, String>) -> bool {
|
||||
headers
|
||||
.get("content-type")
|
||||
@@ -770,8 +790,12 @@ async fn buffer_non_sse_upstream_body(
|
||||
if ttfb_ms.is_none() {
|
||||
ttfb_ms = Some(started_at.elapsed().as_millis() as u64);
|
||||
}
|
||||
upstream_bytes += chunk.len() as u64;
|
||||
body_bytes.extend_from_slice(&chunk);
|
||||
append_buffered_upstream_body_chunk(
|
||||
&mut body_bytes,
|
||||
&chunk,
|
||||
ttfb_ms,
|
||||
&mut upstream_bytes,
|
||||
)?;
|
||||
}
|
||||
Err(message) => {
|
||||
return Err(BufferedUpstreamBodyError {
|
||||
@@ -817,8 +841,12 @@ async fn buffer_non_sse_upstream_body(
|
||||
if ttfb_ms.is_none() {
|
||||
ttfb_ms = Some(started_at.elapsed().as_millis() as u64);
|
||||
}
|
||||
upstream_bytes += chunk.len() as u64;
|
||||
body_bytes.extend_from_slice(&chunk);
|
||||
append_buffered_upstream_body_chunk(
|
||||
&mut body_bytes,
|
||||
&chunk,
|
||||
ttfb_ms,
|
||||
&mut upstream_bytes,
|
||||
)?;
|
||||
}
|
||||
Err(err) => {
|
||||
let message = format_error_chain(&err);
|
||||
@@ -871,8 +899,12 @@ async fn buffer_non_sse_upstream_body(
|
||||
if ttfb_ms.is_none() {
|
||||
ttfb_ms = Some(started_at.elapsed().as_millis() as u64);
|
||||
}
|
||||
upstream_bytes += chunk.len() as u64;
|
||||
body_bytes.extend_from_slice(&chunk);
|
||||
append_buffered_upstream_body_chunk(
|
||||
&mut body_bytes,
|
||||
&chunk,
|
||||
ttfb_ms,
|
||||
&mut upstream_bytes,
|
||||
)?;
|
||||
}
|
||||
Err(err) => {
|
||||
let message = format_hyper_error_chain(&err);
|
||||
@@ -925,8 +957,12 @@ async fn buffer_non_sse_upstream_body(
|
||||
if ttfb_ms.is_none() {
|
||||
ttfb_ms = Some(started_at.elapsed().as_millis() as u64);
|
||||
}
|
||||
upstream_bytes += chunk.len() as u64;
|
||||
body_bytes.extend_from_slice(&chunk);
|
||||
append_buffered_upstream_body_chunk(
|
||||
&mut body_bytes,
|
||||
&chunk,
|
||||
ttfb_ms,
|
||||
&mut upstream_bytes,
|
||||
)?;
|
||||
}
|
||||
Err(err) => {
|
||||
let message = format_wreq_upstream_request_error(&err);
|
||||
@@ -974,8 +1010,12 @@ async fn buffer_non_sse_upstream_body(
|
||||
if ttfb_ms.is_none() {
|
||||
ttfb_ms = Some(started_at.elapsed().as_millis() as u64);
|
||||
}
|
||||
upstream_bytes += chunk.len() as u64;
|
||||
body_bytes.extend_from_slice(&chunk);
|
||||
append_buffered_upstream_body_chunk(
|
||||
&mut body_bytes,
|
||||
&chunk,
|
||||
ttfb_ms,
|
||||
&mut upstream_bytes,
|
||||
)?;
|
||||
}
|
||||
Ok(None) => break,
|
||||
Err(message) => {
|
||||
@@ -1015,13 +1055,13 @@ fn maybe_bridge_non_sse_sync_json_to_stream(
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let decoded_body_bytes = decode_non_sse_response_body_bytes(headers, body_bytes)
|
||||
.unwrap_or_else(|| body_bytes.to_vec());
|
||||
if !response_body_is_json(headers, &decoded_body_bytes) {
|
||||
let decoded_body_bytes = decode_response_body_bytes(headers, body_bytes)
|
||||
.map_err(|error| GatewayError::Internal(error.to_string()))?;
|
||||
if !response_body_is_json(headers, decoded_body_bytes.as_ref()) {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let body_json: Value = serde_json::from_slice(&decoded_body_bytes)
|
||||
let body_json: Value = serde_json::from_slice(decoded_body_bytes.as_ref())
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
let client_api_format = report_context
|
||||
.get("client_api_format")
|
||||
@@ -1046,33 +1086,6 @@ fn rewrite_headers_for_bridged_sse_response(
|
||||
rewritten
|
||||
}
|
||||
|
||||
fn decode_non_sse_response_body_bytes(
|
||||
headers: &BTreeMap<String, String>,
|
||||
body_bytes: &[u8],
|
||||
) -> Option<Vec<u8>> {
|
||||
let encoding = headers
|
||||
.get("content-encoding")
|
||||
.map(String::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(|value| value.to_ascii_lowercase());
|
||||
match encoding.as_deref() {
|
||||
Some("gzip") => {
|
||||
let mut decoder = flate2::read::GzDecoder::new(body_bytes);
|
||||
let mut out = Vec::new();
|
||||
std::io::Read::read_to_end(&mut decoder, &mut out).ok()?;
|
||||
Some(out)
|
||||
}
|
||||
Some("deflate") => {
|
||||
let mut decoder = flate2::read::DeflateDecoder::new(body_bytes);
|
||||
let mut out = Vec::new();
|
||||
std::io::Read::read_to_end(&mut decoder, &mut out).ok()?;
|
||||
Some(out)
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn response_body_is_json(headers: &BTreeMap<String, String>, body_bytes: &[u8]) -> bool {
|
||||
if headers
|
||||
.get("content-type")
|
||||
@@ -1159,16 +1172,39 @@ fn observe_normalized_bytes(
|
||||
observer_buffered: &mut Vec<u8>,
|
||||
normalized: &[u8],
|
||||
) {
|
||||
if normalized.is_empty() {
|
||||
if normalized.is_empty()
|
||||
|| observer
|
||||
.latest_summary()
|
||||
.and_then(|summary| summary.parser_error.as_deref())
|
||||
.is_some()
|
||||
{
|
||||
return;
|
||||
}
|
||||
observer_buffered.extend_from_slice(normalized);
|
||||
while let Some(line_end) = observer_buffered.iter().position(|byte| *byte == b'\n') {
|
||||
let line = observer_buffered.drain(..=line_end).collect::<Vec<_>>();
|
||||
if let Err(err) = observer.push_line(report_context, line) {
|
||||
observer.disable_with_error(err.to_string());
|
||||
|
||||
let mut remaining = normalized;
|
||||
while !remaining.is_empty() {
|
||||
let line_part_len = remaining
|
||||
.iter()
|
||||
.position(|byte| *byte == b'\n')
|
||||
.map_or(remaining.len(), |index| index + 1);
|
||||
if observer_buffered.len().saturating_add(line_part_len)
|
||||
> STREAM_USAGE_OBSERVER_MAX_LINE_BYTES
|
||||
{
|
||||
observer.disable_with_error(format!(
|
||||
"stream usage event exceeded {STREAM_USAGE_OBSERVER_MAX_LINE_BYTES} bytes"
|
||||
));
|
||||
observer_buffered.clear();
|
||||
break;
|
||||
return;
|
||||
}
|
||||
observer_buffered.extend_from_slice(&remaining[..line_part_len]);
|
||||
remaining = &remaining[line_part_len..];
|
||||
if observer_buffered.last() == Some(&b'\n') {
|
||||
let line = std::mem::take(observer_buffered);
|
||||
if let Err(err) = observer.push_line(report_context, line) {
|
||||
observer.disable_with_error(err.to_string());
|
||||
observer_buffered.clear();
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1193,9 +1229,11 @@ mod tests {
|
||||
use tokio::sync::watch;
|
||||
|
||||
use super::{
|
||||
build_direct_execution_frame_stream, should_buffer_non_stream_response,
|
||||
should_treat_upstream_response_as_stream,
|
||||
build_direct_execution_frame_stream, observe_normalized_bytes,
|
||||
should_buffer_non_stream_response, should_treat_upstream_response_as_stream,
|
||||
STREAM_USAGE_OBSERVER_MAX_LINE_BYTES,
|
||||
};
|
||||
use crate::ai_serving::api::StreamingStandardTerminalObserver;
|
||||
use crate::execution_runtime::transport::{
|
||||
execute_stream_plan_via_local_tunnel, DirectSyncExecutionRuntime, DirectUpstreamResponse,
|
||||
};
|
||||
@@ -1260,6 +1298,27 @@ mod tests {
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn oversized_usage_line_disables_observation_without_retaining_the_line() {
|
||||
let mut observer = StreamingStandardTerminalObserver::default();
|
||||
let report_context = serde_json::json!({
|
||||
"provider_api_format": "claude:messages",
|
||||
"client_api_format": "claude:messages",
|
||||
});
|
||||
let mut buffered = Vec::new();
|
||||
let oversized = vec![b'x'; STREAM_USAGE_OBSERVER_MAX_LINE_BYTES + 1];
|
||||
|
||||
observe_normalized_bytes(&mut observer, &report_context, &mut buffered, &oversized);
|
||||
|
||||
assert!(buffered.is_empty());
|
||||
assert!(observer
|
||||
.latest_summary()
|
||||
.and_then(|summary| summary.parser_error.as_deref())
|
||||
.is_some_and(|error| error.contains("stream usage event exceeded")));
|
||||
observe_normalized_bytes(&mut observer, &report_context, &mut buffered, b"ignored");
|
||||
assert!(buffered.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn direct_execution_frame_stream_reports_ttfb_after_first_upstream_chunk() {
|
||||
let listener = crate::test_support::bind_loopback_listener()
|
||||
|
||||
@@ -496,6 +496,15 @@ fn classify_local_sync_error_kind(
|
||||
{
|
||||
return LocalCoreSyncErrorKind::RateLimit;
|
||||
}
|
||||
if status_code == 413
|
||||
|| fingerprint.contains("request_too_large")
|
||||
|| fingerprint.contains("request too large")
|
||||
|| fingerprint.contains("payload_too_large")
|
||||
|| fingerprint.contains("payload too large")
|
||||
|| fingerprint.contains("request entity too large")
|
||||
{
|
||||
return LocalCoreSyncErrorKind::RequestTooLarge;
|
||||
}
|
||||
if fingerprint.contains("contextlength")
|
||||
|| fingerprint.contains("contentlengthexceeded")
|
||||
|| fingerprint.contains("context window")
|
||||
@@ -534,6 +543,7 @@ fn default_status_code_for_local_sync_error_kind(kind: LocalCoreSyncErrorKind) -
|
||||
LocalCoreSyncErrorKind::InvalidRequest | LocalCoreSyncErrorKind::ContextLengthExceeded => {
|
||||
400
|
||||
}
|
||||
LocalCoreSyncErrorKind::RequestTooLarge => 413,
|
||||
LocalCoreSyncErrorKind::Authentication => 401,
|
||||
LocalCoreSyncErrorKind::PermissionDenied => 403,
|
||||
LocalCoreSyncErrorKind::NotFound => 404,
|
||||
@@ -792,6 +802,76 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn local_core_error_maps_request_too_large_without_changing_openai_shape() {
|
||||
let claude_payload = core_finalize_payload(
|
||||
"claude_chat_sync_finalize",
|
||||
"claude:messages",
|
||||
"openai:chat",
|
||||
413,
|
||||
json!({
|
||||
"error": {
|
||||
"type": "invalid_request_error",
|
||||
"message": "request body is too large"
|
||||
}
|
||||
}),
|
||||
);
|
||||
let claude_response = maybe_build_local_core_error_response(
|
||||
"trace-sync-claude-too-large",
|
||||
&test_decision(),
|
||||
&claude_payload,
|
||||
)
|
||||
.expect("response build should not error")
|
||||
.expect("response should exist");
|
||||
assert_eq!(
|
||||
claude_response.status(),
|
||||
http::StatusCode::PAYLOAD_TOO_LARGE
|
||||
);
|
||||
let claude_body: serde_json::Value = serde_json::from_slice(
|
||||
&to_bytes(claude_response.into_body(), usize::MAX)
|
||||
.await
|
||||
.expect("body should read"),
|
||||
)
|
||||
.expect("body should decode");
|
||||
assert_eq!(claude_body["type"], "error");
|
||||
assert_eq!(claude_body["error"]["type"], "request_too_large");
|
||||
|
||||
let openai_payload = core_finalize_payload(
|
||||
"openai_chat_sync_finalize",
|
||||
"openai:chat",
|
||||
"claude:messages",
|
||||
200,
|
||||
json!({
|
||||
"type": "error",
|
||||
"error": {
|
||||
"type": "request_too_large",
|
||||
"message": "request body is too large"
|
||||
}
|
||||
}),
|
||||
);
|
||||
let openai_response = maybe_build_local_core_error_response(
|
||||
"trace-sync-openai-too-large",
|
||||
&test_decision(),
|
||||
&openai_payload,
|
||||
)
|
||||
.expect("response build should not error")
|
||||
.expect("response should exist");
|
||||
assert_eq!(
|
||||
openai_response.status(),
|
||||
http::StatusCode::PAYLOAD_TOO_LARGE
|
||||
);
|
||||
let openai_body: serde_json::Value = serde_json::from_slice(
|
||||
&to_bytes(openai_response.into_body(), usize::MAX)
|
||||
.await
|
||||
.expect("body should read"),
|
||||
)
|
||||
.expect("body should decode");
|
||||
assert_eq!(
|
||||
openai_body["error"]["type"], "context_length_exceeded",
|
||||
"OpenAI compatibility shape should remain unchanged"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn local_core_sync_finalize_rejects_gemini_http_200_without_visible_output() {
|
||||
let mut payload = core_finalize_payload(
|
||||
|
||||
@@ -3,7 +3,7 @@ use std::io::Error as IoError;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use aether_ai_serving::UPSTREAM_IS_STREAM_KEY;
|
||||
use aether_ai_serving::{AiAttemptExecutionOutcome, AiAttemptRetryScope, UPSTREAM_IS_STREAM_KEY};
|
||||
use aether_contracts::{
|
||||
ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionPlan, ExecutionResult,
|
||||
ExecutionTelemetry,
|
||||
@@ -55,17 +55,18 @@ use crate::execution_runtime::submission::{
|
||||
resolve_local_sync_error_status_code, submit_local_core_error_or_sync_finalize,
|
||||
};
|
||||
use crate::execution_runtime::transport::{
|
||||
build_execution_response_body, build_request_body, collect_response_headers,
|
||||
decode_response_body_bytes, format_hyper_error_chain, format_upstream_request_error,
|
||||
format_wreq_upstream_request_error, response_body_is_json, send_request, DirectHttpResponse,
|
||||
DirectSyncExecutionRuntime, ExecutionRuntimeTransportError,
|
||||
append_upstream_response_body_chunk, build_execution_response_body, build_request_body,
|
||||
collect_response_headers, decode_response_body_bytes, execution_response_body_mode,
|
||||
format_hyper_error_chain, format_upstream_request_error, format_wreq_upstream_request_error,
|
||||
response_body_is_json, send_request, DirectHttpResponse, DirectSyncExecutionRuntime,
|
||||
ExecutionRuntimeTransportError,
|
||||
};
|
||||
use crate::execution_runtime::windsurf::maybe_execute_windsurf_sync;
|
||||
use crate::execution_runtime::{
|
||||
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,
|
||||
ai_attempt_retry_scope_from_failure_disposition, 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,
|
||||
};
|
||||
use crate::log_ids::short_request_id;
|
||||
use crate::orchestration::{
|
||||
@@ -113,6 +114,29 @@ struct SyncExecutionFailure {
|
||||
message: String,
|
||||
status_code: Option<u16>,
|
||||
latency_ms: Option<u64>,
|
||||
fallback_kind: Option<SyncExecutionFailureFallbackKind>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
enum SyncExecutionFailureFallbackKind {
|
||||
UpstreamResponseTooLarge,
|
||||
UpstreamResponseDecode,
|
||||
}
|
||||
|
||||
impl SyncExecutionFailureFallbackKind {
|
||||
fn error_type(self) -> &'static str {
|
||||
match self {
|
||||
Self::UpstreamResponseTooLarge => "upstream_response_too_large",
|
||||
Self::UpstreamResponseDecode => "upstream_response_decode_failed",
|
||||
}
|
||||
}
|
||||
|
||||
fn client_message(self) -> &'static str {
|
||||
match self {
|
||||
Self::UpstreamResponseTooLarge => "Upstream response too large",
|
||||
Self::UpstreamResponseDecode => "Failed to decode upstream response",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct SyncAttemptTerminalGuard {
|
||||
@@ -269,11 +293,23 @@ async fn record_sync_attempt_forced_terminal_state(
|
||||
|
||||
impl SyncExecutionFailure {
|
||||
fn from_transport(err: ExecutionRuntimeTransportError) -> Self {
|
||||
let fallback_kind = match &err {
|
||||
ExecutionRuntimeTransportError::UpstreamResponseTooLarge { .. } => {
|
||||
Some(SyncExecutionFailureFallbackKind::UpstreamResponseTooLarge)
|
||||
}
|
||||
ExecutionRuntimeTransportError::UpstreamResponseDecode { .. } => {
|
||||
Some(SyncExecutionFailureFallbackKind::UpstreamResponseDecode)
|
||||
}
|
||||
_ => None,
|
||||
};
|
||||
Self {
|
||||
error_type: "execution_runtime_unavailable",
|
||||
error_type: fallback_kind
|
||||
.map(SyncExecutionFailureFallbackKind::error_type)
|
||||
.unwrap_or("execution_runtime_unavailable"),
|
||||
message: err.to_string(),
|
||||
status_code: None,
|
||||
status_code: fallback_kind.map(|_| StatusCode::BAD_GATEWAY.as_u16()),
|
||||
latency_ms: None,
|
||||
fallback_kind,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -285,10 +321,94 @@ impl SyncExecutionFailure {
|
||||
),
|
||||
status_code: Some(StatusCode::GATEWAY_TIMEOUT.as_u16()),
|
||||
latency_ms: Some(elapsed_ms),
|
||||
fallback_kind: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn build_sync_execution_failure_fallback_body(
|
||||
client_api_format: &str,
|
||||
kind: SyncExecutionFailureFallbackKind,
|
||||
) -> Value {
|
||||
let message = kind.client_message();
|
||||
let error_type = kind.error_type();
|
||||
match crate::ai_serving::normalize_api_format_alias(client_api_format).as_str() {
|
||||
"claude:messages" => json!({
|
||||
"type": "error",
|
||||
"error": {
|
||||
"type": "upstream_error",
|
||||
"message": message,
|
||||
}
|
||||
}),
|
||||
"gemini:generate_content" => json!({
|
||||
"error": {
|
||||
"code": StatusCode::BAD_GATEWAY.as_u16(),
|
||||
"message": message,
|
||||
"status": "BAD_GATEWAY",
|
||||
}
|
||||
}),
|
||||
_ => json!({
|
||||
"error": {
|
||||
"type": "upstream_error",
|
||||
"message": message,
|
||||
"code": error_type,
|
||||
}
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
fn build_sync_execution_failure_fallback_response(
|
||||
failure: &SyncExecutionFailure,
|
||||
plan: &ExecutionPlan,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
let Some(kind) = failure.fallback_kind else {
|
||||
return Ok(None);
|
||||
};
|
||||
let body_json = build_sync_execution_failure_fallback_body(&plan.client_api_format, kind);
|
||||
let body_bytes = serde_json::to_vec(&body_json)
|
||||
.map_err(|error| GatewayError::Internal(error.to_string()))?;
|
||||
let headers = BTreeMap::from([
|
||||
("content-type".to_string(), "application/json".to_string()),
|
||||
("content-length".to_string(), body_bytes.len().to_string()),
|
||||
]);
|
||||
let response = build_client_response_from_parts(
|
||||
StatusCode::BAD_GATEWAY.as_u16(),
|
||||
&headers,
|
||||
Body::from(body_bytes),
|
||||
trace_id,
|
||||
Some(decision),
|
||||
)?;
|
||||
attach_control_metadata_headers(
|
||||
response,
|
||||
Some(plan.request_id.as_str()),
|
||||
plan.candidate_id.as_deref(),
|
||||
)
|
||||
.map(Some)
|
||||
}
|
||||
|
||||
fn maybe_store_sync_execution_failure_fallback(
|
||||
failure: &SyncExecutionFailure,
|
||||
plan: &ExecutionPlan,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
retry_scope_out: &mut Option<&mut AiAttemptRetryScope>,
|
||||
retry_fallback_out: &mut Option<&mut Option<Response<Body>>>,
|
||||
) -> Result<(), GatewayError> {
|
||||
if failure.fallback_kind.is_none() {
|
||||
return Ok(());
|
||||
}
|
||||
if let Some(retry_scope) = retry_scope_out.as_deref_mut() {
|
||||
*retry_scope = AiAttemptRetryScope::Candidate;
|
||||
}
|
||||
if let Some(retry_fallback) = retry_fallback_out.as_deref_mut() {
|
||||
*retry_fallback =
|
||||
build_sync_execution_failure_fallback_response(failure, plan, trace_id, decision)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
struct ImplicitSyncFinalizeOutcome {
|
||||
payload: GatewaySyncReportRequest,
|
||||
outcome: LocalCoreSyncFinalizeOutcome,
|
||||
@@ -1297,11 +1417,12 @@ async fn execute_openai_image_sync_upstream_sse_candidate(
|
||||
),
|
||||
)
|
||||
})?;
|
||||
append_upstream_response_body_chunk(&mut body_bytes, &chunk)
|
||||
.map_err(SyncExecutionFailure::from_transport)?;
|
||||
let elapsed_ms = started_at.elapsed().as_millis() as u64;
|
||||
progress
|
||||
.observe_chunk(&chunk, status_code, elapsed_ms)
|
||||
.await;
|
||||
body_bytes.extend_from_slice(&chunk);
|
||||
}
|
||||
}
|
||||
DirectHttpResponse::HyperH2c(response) => {
|
||||
@@ -1314,11 +1435,12 @@ async fn execute_openai_image_sync_upstream_sse_candidate(
|
||||
)),
|
||||
)
|
||||
})?;
|
||||
append_upstream_response_body_chunk(&mut body_bytes, &chunk)
|
||||
.map_err(SyncExecutionFailure::from_transport)?;
|
||||
let elapsed_ms = started_at.elapsed().as_millis() as u64;
|
||||
progress
|
||||
.observe_chunk(&chunk, status_code, elapsed_ms)
|
||||
.await;
|
||||
body_bytes.extend_from_slice(&chunk);
|
||||
}
|
||||
}
|
||||
DirectHttpResponse::BrowserWreq(response) => {
|
||||
@@ -1331,24 +1453,30 @@ async fn execute_openai_image_sync_upstream_sse_candidate(
|
||||
),
|
||||
)
|
||||
})?;
|
||||
append_upstream_response_body_chunk(&mut body_bytes, &chunk)
|
||||
.map_err(SyncExecutionFailure::from_transport)?;
|
||||
let elapsed_ms = started_at.elapsed().as_millis() as u64;
|
||||
progress
|
||||
.observe_chunk(&chunk, status_code, elapsed_ms)
|
||||
.await;
|
||||
body_bytes.extend_from_slice(&chunk);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let decoded_body_bytes =
|
||||
decode_response_body_bytes(&headers, &body_bytes).unwrap_or_else(|| body_bytes.clone());
|
||||
let decoded_body_bytes = decode_response_body_bytes(&headers, &body_bytes)
|
||||
.map_err(SyncExecutionFailure::from_transport)?;
|
||||
let elapsed_ms = started_at.elapsed().as_millis() as u64;
|
||||
let upstream_bytes = body_bytes.len() as u64;
|
||||
progress.finish(status_code, elapsed_ms).await;
|
||||
|
||||
let body =
|
||||
build_execution_response_body(&headers, &body_bytes, &decoded_body_bytes, plan.stream)
|
||||
.map_err(SyncExecutionFailure::from_transport)?;
|
||||
let body = build_execution_response_body(
|
||||
&headers,
|
||||
&body_bytes,
|
||||
decoded_body_bytes.as_ref(),
|
||||
plan.stream,
|
||||
execution_response_body_mode(plan),
|
||||
)
|
||||
.map_err(SyncExecutionFailure::from_transport)?;
|
||||
|
||||
Ok(ExecutionResult {
|
||||
request_id: plan.request_id.clone(),
|
||||
@@ -1451,6 +1579,8 @@ fn build_openai_image_sync_json_heartbeat_response(
|
||||
report_context,
|
||||
false,
|
||||
Some(progress_snapshot),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.await,
|
||||
)
|
||||
@@ -1631,10 +1761,49 @@ pub(crate) async fn execute_execution_runtime_sync(
|
||||
report_context,
|
||||
true,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn execute_execution_runtime_sync_with_retry_scope(
|
||||
state: &AppState,
|
||||
request_path: &str,
|
||||
plan: ExecutionPlan,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
plan_kind: &str,
|
||||
report_kind: Option<String>,
|
||||
report_context: Option<serde_json::Value>,
|
||||
) -> Result<AiAttemptExecutionOutcome<Response<Body>>, GatewayError> {
|
||||
let mut retry_scope = AiAttemptRetryScope::Candidate;
|
||||
let mut fallback_response = None;
|
||||
let response = execute_execution_runtime_sync_impl(
|
||||
state,
|
||||
request_path,
|
||||
plan,
|
||||
trace_id,
|
||||
decision,
|
||||
plan_kind,
|
||||
report_kind,
|
||||
report_context,
|
||||
true,
|
||||
None,
|
||||
Some(&mut retry_scope),
|
||||
Some(&mut fallback_response),
|
||||
)
|
||||
.await?;
|
||||
Ok(match response {
|
||||
Some(response) => AiAttemptExecutionOutcome::Responded(response),
|
||||
None => AiAttemptExecutionOutcome::Retry {
|
||||
scope: retry_scope,
|
||||
fallback_response,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)] // internal function, grouping would add unnecessary indirection
|
||||
async fn execute_execution_runtime_sync_impl(
|
||||
state: &AppState,
|
||||
@@ -1647,6 +1816,8 @@ async fn execute_execution_runtime_sync_impl(
|
||||
mut report_context: Option<serde_json::Value>,
|
||||
allow_json_heartbeat: bool,
|
||||
progress_snapshot: Option<Arc<Mutex<OpenAiImageSyncProgressSnapshot>>>,
|
||||
mut retry_scope_out: Option<&mut AiAttemptRetryScope>,
|
||||
mut retry_fallback_out: Option<&mut Option<Response<Body>>>,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
if allow_json_heartbeat
|
||||
&& should_enable_openai_image_sync_json_heartbeat(plan_kind, &plan, report_context.as_ref())
|
||||
@@ -1751,6 +1922,14 @@ async fn execute_execution_runtime_sync_impl(
|
||||
{
|
||||
Ok(result) => result,
|
||||
Err(err) => {
|
||||
maybe_store_sync_execution_failure_fallback(
|
||||
&err,
|
||||
&plan,
|
||||
trace_id,
|
||||
decision,
|
||||
&mut retry_scope_out,
|
||||
&mut retry_fallback_out,
|
||||
)?;
|
||||
warn!(
|
||||
event_name = "sync_execution_runtime_unavailable",
|
||||
log_type = "ops",
|
||||
@@ -1932,6 +2111,14 @@ async fn execute_execution_runtime_sync_impl(
|
||||
{
|
||||
Ok(result) => result,
|
||||
Err(err) => {
|
||||
maybe_store_sync_execution_failure_fallback(
|
||||
&err,
|
||||
&plan,
|
||||
trace_id,
|
||||
decision,
|
||||
&mut retry_scope_out,
|
||||
&mut retry_fallback_out,
|
||||
)?;
|
||||
warn!(
|
||||
event_name = "sync_execution_runtime_unavailable",
|
||||
log_type = "ops",
|
||||
@@ -2256,6 +2443,38 @@ async fn execute_execution_runtime_sync_impl(
|
||||
local_failover_analysis.decision,
|
||||
LocalFailoverDecision::RetryNextCandidate
|
||||
) {
|
||||
let failure_disposition = crate::orchestration::classify_failure_disposition(
|
||||
&plan.provider_api_format,
|
||||
local_failover_analysis.classification,
|
||||
result.status_code,
|
||||
);
|
||||
if let Some(retry_scope) = retry_scope_out.as_deref_mut() {
|
||||
*retry_scope =
|
||||
ai_attempt_retry_scope_from_failure_disposition(failure_disposition);
|
||||
}
|
||||
if failure_disposition.preserve_upstream_error {
|
||||
if let Some(retry_fallback) = retry_fallback_out.as_deref_mut() {
|
||||
let mut fallback_headers = headers.clone();
|
||||
apply_endpoint_response_header_rules(
|
||||
state,
|
||||
&plan,
|
||||
&mut fallback_headers,
|
||||
body_json.as_ref(),
|
||||
)
|
||||
.await?;
|
||||
*retry_fallback = Some(attach_control_metadata_headers(
|
||||
build_client_response_from_parts(
|
||||
result.status_code,
|
||||
&fallback_headers,
|
||||
Body::from(body_bytes.clone()),
|
||||
trace_id,
|
||||
Some(decision),
|
||||
)?,
|
||||
Some(plan.request_id.as_str()),
|
||||
plan.candidate_id.as_deref(),
|
||||
)?);
|
||||
}
|
||||
}
|
||||
let terminal_unix_secs = current_request_candidate_unix_ms();
|
||||
let error_trace_report_context = with_sync_error_trace_context(
|
||||
report_context.as_ref(),
|
||||
@@ -2909,6 +3128,60 @@ mod tests {
|
||||
.with_execution_runtime_candidate(true)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn oversized_upstream_response_builds_claude_502_retry_fallback() {
|
||||
let mut plan = test_openai_image_plan(false);
|
||||
plan.client_api_format = "claude:messages".to_string();
|
||||
plan.provider_api_format = "claude:messages".to_string();
|
||||
let decision = GatewayControlDecision::synthetic(
|
||||
"/v1/messages",
|
||||
Some("ai_public".to_string()),
|
||||
Some("claude".to_string()),
|
||||
Some("messages".to_string()),
|
||||
Some("claude:messages".to_string()),
|
||||
)
|
||||
.with_execution_runtime_candidate(true);
|
||||
let failure = SyncExecutionFailure::from_transport(
|
||||
ExecutionRuntimeTransportError::UpstreamResponseTooLarge {
|
||||
phase: crate::execution_runtime::transport::UpstreamResponseBodyPhase::Wire,
|
||||
limit_bytes: 8,
|
||||
},
|
||||
);
|
||||
|
||||
assert_eq!(failure.status_code, Some(StatusCode::BAD_GATEWAY.as_u16()));
|
||||
assert_eq!(
|
||||
failure.fallback_kind,
|
||||
Some(SyncExecutionFailureFallbackKind::UpstreamResponseTooLarge)
|
||||
);
|
||||
let mut retry_scope = AiAttemptRetryScope::Provider;
|
||||
let mut retry_fallback = None;
|
||||
{
|
||||
let mut retry_scope_out = Some(&mut retry_scope);
|
||||
let mut retry_fallback_out = Some(&mut retry_fallback);
|
||||
maybe_store_sync_execution_failure_fallback(
|
||||
&failure,
|
||||
&plan,
|
||||
"trace-too-large",
|
||||
&decision,
|
||||
&mut retry_scope_out,
|
||||
&mut retry_fallback_out,
|
||||
)
|
||||
.expect("fallback response should build");
|
||||
}
|
||||
|
||||
assert_eq!(retry_scope, AiAttemptRetryScope::Candidate);
|
||||
let response = retry_fallback.expect("oversized response should provide a fallback");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::BAD_GATEWAY);
|
||||
let body = to_bytes(response.into_body(), 1024)
|
||||
.await
|
||||
.expect("fallback body should read");
|
||||
let body: Value = serde_json::from_slice(&body).expect("fallback body should be json");
|
||||
assert_eq!(body["type"], "error");
|
||||
assert_eq!(body["error"]["type"], "upstream_error");
|
||||
assert_eq!(body["error"]["message"], "Upstream response too large");
|
||||
}
|
||||
|
||||
fn test_kiro_sync_plan() -> ExecutionPlan {
|
||||
ExecutionPlan {
|
||||
request_id: "req-kiro-sync-cache-1".to_string(),
|
||||
|
||||
@@ -14,8 +14,19 @@ pub(super) fn decode_execution_result_body(
|
||||
let Some(body) = body else {
|
||||
return Ok((Vec::new(), None, None));
|
||||
};
|
||||
let ResponseBody {
|
||||
json_body,
|
||||
body_bytes_b64,
|
||||
} = body;
|
||||
|
||||
if let Some(json_body) = body.json_body {
|
||||
if let Some(body_bytes_b64) = body_bytes_b64 {
|
||||
let bytes = base64::engine::general_purpose::STANDARD
|
||||
.decode(&body_bytes_b64)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
return Ok((bytes, json_body, Some(body_bytes_b64)));
|
||||
}
|
||||
|
||||
if let Some(json_body) = json_body {
|
||||
remove_header_case_insensitive(headers, "content-encoding");
|
||||
remove_header_case_insensitive(headers, "content-length");
|
||||
headers
|
||||
@@ -27,13 +38,6 @@ pub(super) fn decode_execution_result_body(
|
||||
return Ok((bytes, Some(json_body), None));
|
||||
}
|
||||
|
||||
if let Some(body_bytes_b64) = body.body_bytes_b64 {
|
||||
let bytes = base64::engine::general_purpose::STANDARD
|
||||
.decode(&body_bytes_b64)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
return Ok((bytes, None, Some(body_bytes_b64)));
|
||||
}
|
||||
|
||||
Ok((Vec::new(), None, None))
|
||||
}
|
||||
|
||||
@@ -52,6 +56,7 @@ mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use aether_contracts::ResponseBody;
|
||||
use base64::Engine as _;
|
||||
use serde_json::json;
|
||||
|
||||
use super::decode_execution_result_body;
|
||||
@@ -81,4 +86,36 @@ mod tests {
|
||||
Some(body_bytes.len().to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dual_body_prefers_wire_bytes_and_retains_parsed_json() {
|
||||
let raw = br#"{ "unknown": true, "ok": true }"#;
|
||||
let encoded = base64::engine::general_purpose::STANDARD.encode(raw);
|
||||
let raw_len = raw.len().to_string();
|
||||
let mut headers = BTreeMap::from([
|
||||
("content-encoding".to_string(), "gzip".to_string()),
|
||||
("content-length".to_string(), raw_len.clone()),
|
||||
]);
|
||||
|
||||
let (body_bytes, body_json, body_base64) = decode_execution_result_body(
|
||||
Some(ResponseBody {
|
||||
json_body: Some(json!({"unknown": true, "ok": true})),
|
||||
body_bytes_b64: Some(encoded.clone()),
|
||||
}),
|
||||
&mut headers,
|
||||
)
|
||||
.expect("body should decode");
|
||||
|
||||
assert_eq!(body_bytes, raw);
|
||||
assert_eq!(body_json, Some(json!({"unknown": true, "ok": true})));
|
||||
assert_eq!(body_base64.as_deref(), Some(encoded.as_str()));
|
||||
assert_eq!(
|
||||
headers.get("content-encoding").map(String::as_str),
|
||||
Some("gzip")
|
||||
);
|
||||
assert_eq!(
|
||||
headers.get("content-length").map(String::as_str),
|
||||
Some(raw_len.as_str())
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3,6 +3,7 @@ mod execution;
|
||||
pub(crate) use execution::{
|
||||
build_openai_image_sync_json_whitespace_heartbeat_stream,
|
||||
build_sync_json_whitespace_heartbeat_stream, execute_execution_runtime_sync,
|
||||
execute_execution_runtime_sync_with_retry_scope,
|
||||
};
|
||||
|
||||
#[allow(unused_imports)]
|
||||
|
||||
@@ -390,7 +390,7 @@ fn build_best_effort_local_core_error_body_converts_sync_errors_across_standard_
|
||||
"type": "error",
|
||||
"error": {
|
||||
"message": "backend busy",
|
||||
"type": "api_error",
|
||||
"type": "overloaded_error",
|
||||
"code": "UNAVAILABLE"
|
||||
}
|
||||
}),
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
use std::borrow::Cow;
|
||||
use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
|
||||
use std::error::Error as _;
|
||||
use std::future::Future;
|
||||
@@ -8,11 +9,12 @@ use std::sync::{Arc, LazyLock, Mutex as StdMutex, OnceLock, RwLock as StdRwLock}
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use aether_contracts::{
|
||||
ExecutionPlan, ExecutionResult, ExecutionTelemetry, ProxySnapshot, ResolvedTransportProfile,
|
||||
ResponseBody, EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER,
|
||||
ExecutionPlan, ExecutionResponseBodyMode, ExecutionResult, ExecutionTelemetry, ProxySnapshot,
|
||||
ResolvedTransportProfile, ResponseBody, EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER,
|
||||
EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER, EXECUTION_REQUEST_HTTP1_ONLY_HEADER,
|
||||
TRANSPORT_BACKEND_BROWSER_WREQ, TRANSPORT_BACKEND_REQWEST_RUSTLS,
|
||||
TRANSPORT_HTTP_MODE_H2C_PRIOR_KNOWLEDGE, TRANSPORT_HTTP_MODE_HTTP1_ONLY,
|
||||
EXECUTION_RESPONSE_BODY_MODE_HEADER, TRANSPORT_BACKEND_BROWSER_WREQ,
|
||||
TRANSPORT_BACKEND_REQWEST_RUSTLS, TRANSPORT_HTTP_MODE_H2C_PRIOR_KNOWLEDGE,
|
||||
TRANSPORT_HTTP_MODE_HTTP1_ONLY,
|
||||
};
|
||||
use aether_data::repository::proxy_nodes::ProxyNodeTrafficMutation;
|
||||
use aether_http::{apply_http_client_config, HttpClientConfig};
|
||||
@@ -35,6 +37,7 @@ use reqwest::redirect::Policy;
|
||||
use serde::Serialize;
|
||||
use serde_json::json;
|
||||
use serde_json::Value;
|
||||
use sha2::Digest as _;
|
||||
use thiserror::Error;
|
||||
use tokio::net::TcpStream;
|
||||
use tokio::sync::OnceCell as TokioOnceCell;
|
||||
@@ -107,6 +110,7 @@ type DirectHyperH2cSenderCacheCell = TokioOnceCell<Arc<DirectHyperH2cSenderCache
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
|
||||
struct DirectReqwestClientCacheKey {
|
||||
upstream_origin: Option<String>,
|
||||
pool_partition: Option<String>,
|
||||
connect_timeout_ms: Option<u64>,
|
||||
proxy_url: Option<String>,
|
||||
follow_redirects: bool,
|
||||
@@ -444,8 +448,11 @@ pub(crate) fn format_upstream_request_error(err: &reqwest::Error) -> String {
|
||||
}
|
||||
|
||||
if let Some(url) = err.url() {
|
||||
let (sanitized_detail, sanitized_url) =
|
||||
sanitize_upstream_request_error_detail(&detail, url.as_str());
|
||||
detail = sanitized_detail;
|
||||
detail.push_str(" [url=");
|
||||
detail.push_str(url.as_str());
|
||||
detail.push_str(&sanitized_url);
|
||||
detail.push(']');
|
||||
}
|
||||
if !kinds.is_empty() {
|
||||
@@ -457,6 +464,25 @@ pub(crate) fn format_upstream_request_error(err: &reqwest::Error) -> String {
|
||||
detail
|
||||
}
|
||||
|
||||
fn sanitize_upstream_request_error_detail(detail: &str, upstream_url: &str) -> (String, String) {
|
||||
let sanitized_url = sanitize_upstream_url_text(upstream_url);
|
||||
(detail.replace(upstream_url, &sanitized_url), sanitized_url)
|
||||
}
|
||||
|
||||
fn sanitize_upstream_url_text(upstream_url: &str) -> String {
|
||||
if let Ok(mut parsed_url) = reqwest::Url::parse(upstream_url) {
|
||||
parsed_url.set_query(None);
|
||||
parsed_url.set_fragment(None);
|
||||
return parsed_url.to_string();
|
||||
}
|
||||
|
||||
let suffix_offset = upstream_url
|
||||
.char_indices()
|
||||
.find_map(|(offset, character)| matches!(character, '?' | '#').then_some(offset))
|
||||
.unwrap_or(upstream_url.len());
|
||||
upstream_url[..suffix_offset].to_string()
|
||||
}
|
||||
|
||||
pub(crate) fn format_wreq_upstream_request_error(err: &wreq::Error) -> String {
|
||||
let mut kinds = Vec::new();
|
||||
if err.is_connect() {
|
||||
@@ -490,8 +516,12 @@ pub(crate) fn format_wreq_upstream_request_error(err: &wreq::Error) -> String {
|
||||
}
|
||||
|
||||
if let Some(uri) = err.uri() {
|
||||
let uri = uri.to_string();
|
||||
let (sanitized_detail, sanitized_uri) =
|
||||
sanitize_upstream_request_error_detail(&detail, &uri);
|
||||
detail = sanitized_detail;
|
||||
detail.push_str(" [uri=");
|
||||
detail.push_str(&uri.to_string());
|
||||
detail.push_str(&sanitized_uri);
|
||||
detail.push(']');
|
||||
}
|
||||
if !kinds.is_empty() {
|
||||
@@ -547,12 +577,60 @@ pub(crate) enum ExecutionRuntimeTransportError {
|
||||
BrowserBody(String),
|
||||
#[error("failed to execute upstream request: {0}")]
|
||||
UpstreamRequest(String),
|
||||
#[error("upstream response {phase} body exceeds {limit_bytes} bytes")]
|
||||
UpstreamResponseTooLarge {
|
||||
phase: UpstreamResponseBodyPhase,
|
||||
limit_bytes: usize,
|
||||
},
|
||||
#[error("failed to decode upstream response body with content-encoding {encoding}: {message}")]
|
||||
UpstreamResponseDecode { encoding: String, message: String },
|
||||
#[error("hub relay request failed: {0}")]
|
||||
RelayError(String),
|
||||
#[error("upstream response is not valid JSON: {0}")]
|
||||
InvalidJson(serde_json::Error),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) enum UpstreamResponseBodyPhase {
|
||||
Wire,
|
||||
Decoded,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for UpstreamResponseBodyPhase {
|
||||
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
formatter.write_str(match self {
|
||||
Self::Wire => "wire",
|
||||
Self::Decoded => "decoded",
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn append_upstream_response_body_chunk(
|
||||
body: &mut Vec<u8>,
|
||||
chunk: &[u8],
|
||||
) -> Result<(), ExecutionRuntimeTransportError> {
|
||||
append_upstream_response_body_chunk_with_limit(
|
||||
body,
|
||||
chunk,
|
||||
crate::headers::max_internal_buffered_body_bytes(),
|
||||
)
|
||||
}
|
||||
|
||||
fn append_upstream_response_body_chunk_with_limit(
|
||||
body: &mut Vec<u8>,
|
||||
chunk: &[u8],
|
||||
limit_bytes: usize,
|
||||
) -> Result<(), ExecutionRuntimeTransportError> {
|
||||
if body.len() > limit_bytes || chunk.len() > limit_bytes.saturating_sub(body.len()) {
|
||||
return Err(ExecutionRuntimeTransportError::UpstreamResponseTooLarge {
|
||||
phase: UpstreamResponseBodyPhase::Wire,
|
||||
limit_bytes,
|
||||
});
|
||||
}
|
||||
body.extend_from_slice(chunk);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct RelayRequestMeta {
|
||||
provider_id: String,
|
||||
@@ -608,6 +686,7 @@ pub(crate) struct DirectUpstreamStreamExecution {
|
||||
pub(crate) provider_api_format: String,
|
||||
pub(crate) stream_summary_report_context: Value,
|
||||
pub(crate) prefetched_body: VecDeque<Result<Bytes, String>>,
|
||||
pub(crate) stream_precommit_committed: bool,
|
||||
pub(crate) response: DirectUpstreamResponse,
|
||||
pub(crate) started_at: Instant,
|
||||
pub(crate) stream_first_byte_timeout: Option<Duration>,
|
||||
@@ -654,16 +733,16 @@ impl DirectSyncExecutionRuntime {
|
||||
});
|
||||
let (body_bytes, stream_ttfb_ms) =
|
||||
response.bytes_with_stream_timeout(plan, started_at).await?;
|
||||
let decoded_body_bytes = decode_response_body_bytes(&headers, &body_bytes)
|
||||
.unwrap_or_else(|| body_bytes.to_vec());
|
||||
let decoded_body_bytes = decode_response_body_bytes(&headers, &body_bytes)?;
|
||||
let elapsed_ms = started_at.elapsed().as_millis() as u64;
|
||||
let upstream_bytes = body_bytes.len() as u64;
|
||||
|
||||
let body = build_execution_response_body(
|
||||
&headers,
|
||||
&body_bytes,
|
||||
&decoded_body_bytes,
|
||||
decoded_body_bytes.as_ref(),
|
||||
plan.stream,
|
||||
execution_response_body_mode(plan),
|
||||
)?;
|
||||
|
||||
Ok(ExecutionResult {
|
||||
@@ -713,6 +792,7 @@ impl DirectSyncExecutionRuntime {
|
||||
provider_api_format: plan.provider_api_format.clone(),
|
||||
stream_summary_report_context,
|
||||
prefetched_body: VecDeque::new(),
|
||||
stream_precommit_committed: false,
|
||||
response: response.into_direct_upstream_response(),
|
||||
started_at,
|
||||
stream_first_byte_timeout: resolve_stream_first_byte_timeout(plan),
|
||||
@@ -827,6 +907,7 @@ pub(crate) async fn execute_stream_plan_via_local_tunnel(
|
||||
provider_api_format: plan.provider_api_format.clone(),
|
||||
stream_summary_report_context: build_stream_summary_report_context(plan),
|
||||
prefetched_body: VecDeque::new(),
|
||||
stream_precommit_committed: false,
|
||||
response: DirectUpstreamResponse::LocalTunnel(response),
|
||||
started_at,
|
||||
stream_first_byte_timeout: resolve_stream_first_byte_timeout(plan),
|
||||
@@ -962,8 +1043,7 @@ async fn execute_sync_plan_via_local_tunnel_inner(
|
||||
let proxy_timing = execution_header_for_log(&headers, "x-proxy-timing").unwrap_or("-");
|
||||
let (body_bytes, stream_ttfb_ms) =
|
||||
collect_local_tunnel_response_body(response, plan, started_at).await?;
|
||||
let decoded_body_bytes =
|
||||
decode_response_body_bytes(&headers, &body_bytes).unwrap_or_else(|| body_bytes.clone());
|
||||
let decoded_body_bytes = decode_response_body_bytes(&headers, &body_bytes)?;
|
||||
let elapsed_ms = started_at.elapsed().as_millis() as u64;
|
||||
let upstream_bytes = body_bytes.len() as u64;
|
||||
if status_code >= 400 {
|
||||
@@ -1000,8 +1080,13 @@ async fn execute_sync_plan_via_local_tunnel_inner(
|
||||
);
|
||||
}
|
||||
|
||||
let body =
|
||||
build_execution_response_body(&headers, &body_bytes, &decoded_body_bytes, plan.stream)?;
|
||||
let body = build_execution_response_body(
|
||||
&headers,
|
||||
&body_bytes,
|
||||
decoded_body_bytes.as_ref(),
|
||||
plan.stream,
|
||||
execution_response_body_mode(plan),
|
||||
)?;
|
||||
|
||||
Ok(ExecutionResult {
|
||||
request_id: plan.request_id.clone(),
|
||||
@@ -1044,7 +1129,7 @@ async fn collect_local_tunnel_response_body(
|
||||
if plan.stream && first_byte_ms.is_none() && !chunk.is_empty() {
|
||||
first_byte_ms = Some(started_at.elapsed().as_millis() as u64);
|
||||
}
|
||||
body_bytes.extend_from_slice(&chunk);
|
||||
append_upstream_response_body_chunk(&mut body_bytes, &chunk)?;
|
||||
}
|
||||
|
||||
Ok((body_bytes, first_byte_ms))
|
||||
@@ -1154,6 +1239,7 @@ async fn send_request_inner(
|
||||
let client_select_started_at = Instant::now();
|
||||
let client = build_client(
|
||||
&plan.url,
|
||||
&plan.key_id,
|
||||
plan.timeouts.as_ref(),
|
||||
plan.proxy.as_ref(),
|
||||
plan.transport_profile.as_ref(),
|
||||
@@ -1204,23 +1290,23 @@ impl DirectHttpResponse {
|
||||
}
|
||||
|
||||
pub(crate) async fn bytes(self) -> Result<Bytes, ExecutionRuntimeTransportError> {
|
||||
let started_at = Instant::now();
|
||||
match self {
|
||||
DirectHttpResponse::Reqwest(response) => response.bytes().await.map_err(|err| {
|
||||
ExecutionRuntimeTransportError::UpstreamRequest(format_upstream_request_error(&err))
|
||||
}),
|
||||
DirectHttpResponse::HyperH2c(response) => response
|
||||
.into_body()
|
||||
.collect()
|
||||
.await
|
||||
.map(|collected| collected.to_bytes())
|
||||
.map_err(|err| {
|
||||
ExecutionRuntimeTransportError::UpstreamRequest(format_hyper_error_chain(&err))
|
||||
}),
|
||||
DirectHttpResponse::BrowserWreq(response) => response.bytes().await.map_err(|err| {
|
||||
ExecutionRuntimeTransportError::BrowserBody(format_wreq_upstream_request_error(
|
||||
&err,
|
||||
))
|
||||
}),
|
||||
DirectHttpResponse::Reqwest(response) => {
|
||||
collect_reqwest_stream_body(response, started_at, None)
|
||||
.await
|
||||
.map(|(body, _)| body)
|
||||
}
|
||||
DirectHttpResponse::HyperH2c(response) => {
|
||||
collect_hyper_stream_body(response, started_at, None)
|
||||
.await
|
||||
.map(|(body, _)| body)
|
||||
}
|
||||
DirectHttpResponse::BrowserWreq(response) => {
|
||||
collect_wreq_stream_body(response, started_at, None)
|
||||
.await
|
||||
.map(|(body, _)| body)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1308,7 +1394,7 @@ async fn collect_reqwest_stream_body(
|
||||
if first_byte_ms.is_none() && !chunk.is_empty() {
|
||||
first_byte_ms = Some(started_at.elapsed().as_millis() as u64);
|
||||
}
|
||||
body_bytes.extend_from_slice(&chunk);
|
||||
append_upstream_response_body_chunk(&mut body_bytes, &chunk)?;
|
||||
}
|
||||
|
||||
Ok((Bytes::from(body_bytes), first_byte_ms))
|
||||
@@ -1338,7 +1424,7 @@ async fn collect_hyper_stream_body(
|
||||
if first_byte_ms.is_none() && !chunk.is_empty() {
|
||||
first_byte_ms = Some(started_at.elapsed().as_millis() as u64);
|
||||
}
|
||||
body_bytes.extend_from_slice(&chunk);
|
||||
append_upstream_response_body_chunk(&mut body_bytes, &chunk)?;
|
||||
}
|
||||
|
||||
Ok((Bytes::from(body_bytes), first_byte_ms))
|
||||
@@ -1368,7 +1454,7 @@ async fn collect_wreq_stream_body(
|
||||
if first_byte_ms.is_none() && !chunk.is_empty() {
|
||||
first_byte_ms = Some(started_at.elapsed().as_millis() as u64);
|
||||
}
|
||||
body_bytes.extend_from_slice(&chunk);
|
||||
append_upstream_response_body_chunk(&mut body_bytes, &chunk)?;
|
||||
}
|
||||
|
||||
Ok((Bytes::from(body_bytes), first_byte_ms))
|
||||
@@ -2555,6 +2641,7 @@ fn resolve_local_tunnel_node_id(state: &AppState, proxy: Option<&ProxySnapshot>)
|
||||
|
||||
fn build_client(
|
||||
request_url: &str,
|
||||
key_id: &str,
|
||||
timeouts: Option<&aether_contracts::ExecutionTimeouts>,
|
||||
proxy: Option<&ProxySnapshot>,
|
||||
transport_profile: Option<&ResolvedTransportProfile>,
|
||||
@@ -2564,6 +2651,7 @@ fn build_client(
|
||||
let resolved_proxy_url = resolve_proxy_url(proxy)?;
|
||||
let cache_key = direct_reqwest_client_cache_key(
|
||||
request_url,
|
||||
key_id,
|
||||
timeouts,
|
||||
resolved_proxy_url,
|
||||
transport_profile,
|
||||
@@ -2610,7 +2698,10 @@ pub(crate) fn prewarm_direct_reqwest_client_cache_for_plan(plan: &ExecutionPlan)
|
||||
candidate_id = ?plan.candidate_id,
|
||||
provider_id = %plan.provider_id,
|
||||
endpoint_id = %plan.endpoint_id,
|
||||
key_id = %plan.key_id,
|
||||
key_partition = ?direct_reqwest_pool_partition(
|
||||
plan.transport_profile.as_ref(),
|
||||
&plan.key_id,
|
||||
),
|
||||
"gateway direct reqwest client prewarm skipped"
|
||||
);
|
||||
}
|
||||
@@ -2638,6 +2729,7 @@ fn try_prewarm_direct_reqwest_client_cache_for_plan(
|
||||
let resolved_proxy_url = resolve_proxy_url(plan.proxy.as_ref())?;
|
||||
let cache_key = direct_reqwest_client_cache_key(
|
||||
&plan.url,
|
||||
&plan.key_id,
|
||||
plan.timeouts.as_ref(),
|
||||
resolved_proxy_url,
|
||||
plan.transport_profile.as_ref(),
|
||||
@@ -2877,6 +2969,7 @@ fn mark_direct_reqwest_client_cache_not_warming(cache_key: &DirectReqwestClientC
|
||||
|
||||
fn direct_reqwest_client_cache_key(
|
||||
request_url: &str,
|
||||
key_id: &str,
|
||||
timeouts: Option<&aether_contracts::ExecutionTimeouts>,
|
||||
proxy_url: Option<String>,
|
||||
transport_profile: Option<&ResolvedTransportProfile>,
|
||||
@@ -2886,6 +2979,7 @@ fn direct_reqwest_client_cache_key(
|
||||
upstream_origin: direct_reqwest_cache_per_origin()
|
||||
.then(|| direct_reqwest_upstream_origin(request_url))
|
||||
.flatten(),
|
||||
pool_partition: direct_reqwest_pool_partition(transport_profile, key_id),
|
||||
connect_timeout_ms: timeouts.and_then(|timeouts| timeouts.connect_ms),
|
||||
proxy_url,
|
||||
follow_redirects: transport_controls.follow_redirects == Some(true),
|
||||
@@ -2895,6 +2989,17 @@ fn direct_reqwest_client_cache_key(
|
||||
}
|
||||
}
|
||||
|
||||
fn direct_reqwest_pool_partition(
|
||||
transport_profile: Option<&ResolvedTransportProfile>,
|
||||
key_id: &str,
|
||||
) -> Option<String> {
|
||||
let key_id = key_id.trim();
|
||||
transport_profile
|
||||
.filter(|profile| profile.pool_scope.trim().eq_ignore_ascii_case("key"))
|
||||
.filter(|_| !key_id.is_empty())
|
||||
.map(|_| format!("{:x}", sha2::Sha256::digest(key_id.as_bytes())))
|
||||
}
|
||||
|
||||
fn direct_reqwest_cache_per_origin() -> bool {
|
||||
std::env::var(DIRECT_REQWEST_CACHE_PER_ORIGIN_ENV)
|
||||
.ok()
|
||||
@@ -3719,11 +3824,13 @@ pub(crate) fn build_request_headers(
|
||||
}
|
||||
for (key, value) in headers {
|
||||
let normalized_key = key.trim().to_ascii_lowercase();
|
||||
if is_hop_by_hop_header(&normalized_key)
|
||||
if crate::headers::should_skip_request_header(&normalized_key)
|
||||
|| is_hop_by_hop_header(&normalized_key)
|
||||
|| normalized_key == "content-encoding"
|
||||
|| normalized_key == EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER
|
||||
|| normalized_key == EXECUTION_REQUEST_HTTP1_ONLY_HEADER
|
||||
|| normalized_key == EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER
|
||||
|| normalized_key == EXECUTION_RESPONSE_BODY_MODE_HEADER
|
||||
{
|
||||
continue;
|
||||
}
|
||||
@@ -3766,6 +3873,23 @@ fn resolve_execution_transport_controls(
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn execution_response_body_mode(plan: &ExecutionPlan) -> ExecutionResponseBodyMode {
|
||||
if plan.stream
|
||||
|| plan.body.body_bytes_b64.is_none()
|
||||
|| !plan
|
||||
.client_api_format
|
||||
.trim()
|
||||
.eq_ignore_ascii_case(plan.provider_api_format.trim())
|
||||
{
|
||||
return ExecutionResponseBodyMode::StructuredJson;
|
||||
}
|
||||
|
||||
ExecutionResponseBodyMode::from_header_value(execution_transport_header_value(
|
||||
&plan.headers,
|
||||
EXECUTION_RESPONSE_BODY_MODE_HEADER,
|
||||
))
|
||||
}
|
||||
|
||||
fn execution_transport_header_value<'a>(
|
||||
headers: &'a BTreeMap<String, String>,
|
||||
target: &str,
|
||||
@@ -3844,10 +3968,22 @@ fn execution_log_url_host(url: &str) -> String {
|
||||
.unwrap_or_else(|| "-".to_string())
|
||||
}
|
||||
|
||||
pub(crate) fn decode_response_body_bytes(
|
||||
pub(crate) fn decode_response_body_bytes<'a>(
|
||||
headers: &BTreeMap<String, String>,
|
||||
body_bytes: &[u8],
|
||||
) -> Option<Vec<u8>> {
|
||||
body_bytes: &'a [u8],
|
||||
) -> Result<Cow<'a, [u8]>, ExecutionRuntimeTransportError> {
|
||||
decode_response_body_bytes_with_limit(
|
||||
headers,
|
||||
body_bytes,
|
||||
crate::headers::max_internal_buffered_body_bytes(),
|
||||
)
|
||||
}
|
||||
|
||||
fn decode_response_body_bytes_with_limit<'a>(
|
||||
headers: &BTreeMap<String, String>,
|
||||
body_bytes: &'a [u8],
|
||||
limit_bytes: usize,
|
||||
) -> Result<Cow<'a, [u8]>, ExecutionRuntimeTransportError> {
|
||||
let encoding = headers
|
||||
.get("content-encoding")
|
||||
.map(String::as_str)
|
||||
@@ -3857,20 +3993,43 @@ pub(crate) fn decode_response_body_bytes(
|
||||
match encoding.as_deref() {
|
||||
Some("gzip") => {
|
||||
let mut decoder = GzDecoder::new(body_bytes);
|
||||
let mut out = Vec::new();
|
||||
decoder.read_to_end(&mut out).ok()?;
|
||||
Some(out)
|
||||
read_upstream_response_decoder_with_limit("gzip", &mut decoder, limit_bytes)
|
||||
.map(Cow::Owned)
|
||||
}
|
||||
Some("deflate") => {
|
||||
let mut decoder = DeflateDecoder::new(body_bytes);
|
||||
let mut out = Vec::new();
|
||||
decoder.read_to_end(&mut out).ok()?;
|
||||
Some(out)
|
||||
read_upstream_response_decoder_with_limit("deflate", &mut decoder, limit_bytes)
|
||||
.map(Cow::Owned)
|
||||
}
|
||||
_ => None,
|
||||
_ => Ok(Cow::Borrowed(body_bytes)),
|
||||
}
|
||||
}
|
||||
|
||||
fn read_upstream_response_decoder_with_limit(
|
||||
encoding: &str,
|
||||
decoder: &mut impl Read,
|
||||
limit_bytes: usize,
|
||||
) -> Result<Vec<u8>, ExecutionRuntimeTransportError> {
|
||||
let read_limit = u64::try_from(limit_bytes)
|
||||
.unwrap_or(u64::MAX)
|
||||
.saturating_add(1);
|
||||
let mut limited = decoder.take(read_limit);
|
||||
let mut out = Vec::new();
|
||||
limited.read_to_end(&mut out).map_err(|error| {
|
||||
ExecutionRuntimeTransportError::UpstreamResponseDecode {
|
||||
encoding: encoding.to_string(),
|
||||
message: error.to_string(),
|
||||
}
|
||||
})?;
|
||||
if out.len() > limit_bytes {
|
||||
return Err(ExecutionRuntimeTransportError::UpstreamResponseTooLarge {
|
||||
phase: UpstreamResponseBodyPhase::Decoded,
|
||||
limit_bytes,
|
||||
});
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
pub(crate) fn response_body_is_json(headers: &BTreeMap<String, String>, body_bytes: &[u8]) -> bool {
|
||||
let content_type = headers
|
||||
.get("content-type")
|
||||
@@ -3893,6 +4052,7 @@ pub(crate) fn build_execution_response_body(
|
||||
body_bytes: &[u8],
|
||||
decoded_body_bytes: &[u8],
|
||||
stream: bool,
|
||||
response_body_mode: ExecutionResponseBodyMode,
|
||||
) -> Result<Option<ResponseBody>, ExecutionRuntimeTransportError> {
|
||||
if body_bytes.is_empty() {
|
||||
return Ok(None);
|
||||
@@ -3903,7 +4063,8 @@ pub(crate) fn build_execution_response_body(
|
||||
.map_err(ExecutionRuntimeTransportError::InvalidJson)?;
|
||||
return Ok(Some(ResponseBody {
|
||||
json_body: Some(body_json),
|
||||
body_bytes_b64: None,
|
||||
body_bytes_b64: (response_body_mode == ExecutionResponseBodyMode::PreserveBytes)
|
||||
.then(|| base64::engine::general_purpose::STANDARD.encode(body_bytes)),
|
||||
}));
|
||||
}
|
||||
|
||||
@@ -3932,12 +4093,13 @@ pub(crate) fn build_execution_response_body(
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
use std::io::Read;
|
||||
use std::io::{Read, Write};
|
||||
use std::sync::{Arc, Mutex, MutexGuard, OnceLock};
|
||||
|
||||
use aether_contracts::{
|
||||
ExecutionPlan, ExecutionTimeouts, ProxySnapshot, RequestBody, ResolvedTransportProfile,
|
||||
EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER, EXECUTION_REQUEST_HTTP1_ONLY_HEADER,
|
||||
ExecutionPlan, ExecutionResponseBodyMode, ExecutionTimeouts, ProxySnapshot, RequestBody,
|
||||
ResolvedTransportProfile, EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER,
|
||||
EXECUTION_REQUEST_HTTP1_ONLY_HEADER, EXECUTION_RESPONSE_BODY_MODE_HEADER,
|
||||
TRANSPORT_BACKEND_BROWSER_WREQ, TRANSPORT_BACKEND_REQWEST_RUSTLS, TRANSPORT_HTTP_MODE_AUTO,
|
||||
TRANSPORT_HTTP_MODE_H2C_PRIOR_KNOWLEDGE, TRANSPORT_HTTP_MODE_HTTP1_ONLY,
|
||||
};
|
||||
@@ -3956,13 +4118,14 @@ mod tests {
|
||||
use tokio::sync::watch;
|
||||
|
||||
use super::{
|
||||
build_browser_wreq_client, build_client, build_direct_tunnel_request_meta,
|
||||
build_execution_response_body, build_request_headers, execute_sync_plan,
|
||||
append_upstream_response_body_chunk_with_limit, build_browser_wreq_client, build_client,
|
||||
build_direct_tunnel_request_meta, build_execution_response_body, build_request_headers,
|
||||
decode_response_body_bytes_with_limit, execute_sync_plan, execution_response_body_mode,
|
||||
record_manual_proxy_request_failure, record_manual_proxy_request_outcome,
|
||||
record_manual_proxy_request_success, record_manual_proxy_stream_error,
|
||||
resolve_execution_transport_controls, resolve_non_stream_total_timeout,
|
||||
resolve_stream_first_byte_timeout, response_body_is_json, DirectSyncExecutionRuntime,
|
||||
ExecutionRuntimeTransportError, ExecutionTransportControls,
|
||||
ExecutionRuntimeTransportError, ExecutionTransportControls, UpstreamResponseBodyPhase,
|
||||
};
|
||||
use crate::constants::{
|
||||
EXECUTION_RUNTIME_LOOP_GUARD_HEADER, EXECUTION_RUNTIME_LOOP_GUARD_VIA_TOKEN,
|
||||
@@ -3976,6 +4139,106 @@ mod tests {
|
||||
|
||||
const LOCAL_HTTP_SUCCESS_TIMEOUT_MS: u64 = 15_000;
|
||||
|
||||
#[test]
|
||||
fn upstream_error_url_sanitization_removes_secrets_everywhere() {
|
||||
let upstream_url =
|
||||
"https://api.example.test/v1/messages?key=query-secret&alt=sse#fragment-secret";
|
||||
let detail = format!(
|
||||
"error sending request for url ({upstream_url}); source repeated {upstream_url}"
|
||||
);
|
||||
|
||||
let (sanitized_detail, sanitized_url) =
|
||||
super::sanitize_upstream_request_error_detail(&detail, upstream_url);
|
||||
|
||||
assert_eq!(sanitized_url, "https://api.example.test/v1/messages");
|
||||
assert_eq!(
|
||||
sanitized_detail,
|
||||
"error sending request for url (https://api.example.test/v1/messages); source repeated https://api.example.test/v1/messages"
|
||||
);
|
||||
assert!(!sanitized_detail.contains("query-secret"));
|
||||
assert!(!sanitized_detail.contains("fragment-secret"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_header_materialization_strips_all_aether_internal_headers() {
|
||||
let headers = BTreeMap::from([
|
||||
("authorization".to_string(), "Bearer upstream".to_string()),
|
||||
("x-aether-grok-runtime".to_string(), "1".to_string()),
|
||||
("x-aether-future-control".to_string(), "private".to_string()),
|
||||
]);
|
||||
|
||||
let materialized = build_request_headers(&headers, None, false)
|
||||
.expect("provider request headers should materialize");
|
||||
|
||||
assert_eq!(
|
||||
materialized
|
||||
.get("authorization")
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("Bearer upstream")
|
||||
);
|
||||
assert!(!materialized.contains_key("x-aether-grok-runtime"));
|
||||
assert!(!materialized.contains_key("x-aether-future-control"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn upstream_response_wire_limit_allows_exact_body_and_rejects_next_byte() {
|
||||
let mut body = Vec::new();
|
||||
append_upstream_response_body_chunk_with_limit(&mut body, b"1234", 5)
|
||||
.expect("chunk below limit should append");
|
||||
append_upstream_response_body_chunk_with_limit(&mut body, b"5", 5)
|
||||
.expect("body exactly at limit should append");
|
||||
|
||||
let error = append_upstream_response_body_chunk_with_limit(&mut body, b"6", 5)
|
||||
.expect_err("body above limit should fail");
|
||||
|
||||
assert_eq!(body, b"12345");
|
||||
assert!(matches!(
|
||||
error,
|
||||
ExecutionRuntimeTransportError::UpstreamResponseTooLarge {
|
||||
phase: UpstreamResponseBodyPhase::Wire,
|
||||
limit_bytes: 5,
|
||||
}
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn upstream_response_gzip_decode_limit_rejects_decompression_bomb() {
|
||||
let payload = vec![b'x'; 9];
|
||||
let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default());
|
||||
encoder
|
||||
.write_all(&payload)
|
||||
.expect("gzip payload should encode");
|
||||
let encoded = encoder.finish().expect("gzip payload should finish");
|
||||
let headers = BTreeMap::from([("content-encoding".to_string(), "gzip".to_string())]);
|
||||
|
||||
let error = decode_response_body_bytes_with_limit(&headers, &encoded, 8)
|
||||
.expect_err("decoded body above limit should fail");
|
||||
|
||||
assert!(matches!(
|
||||
error,
|
||||
ExecutionRuntimeTransportError::UpstreamResponseTooLarge {
|
||||
phase: UpstreamResponseBodyPhase::Decoded,
|
||||
limit_bytes: 8,
|
||||
}
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn upstream_response_gzip_decode_limit_allows_exact_body() {
|
||||
let payload = b"12345678";
|
||||
let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default());
|
||||
encoder
|
||||
.write_all(payload)
|
||||
.expect("gzip payload should encode");
|
||||
let encoded = encoder.finish().expect("gzip payload should finish");
|
||||
let headers = BTreeMap::from([("content-encoding".to_string(), "gzip".to_string())]);
|
||||
|
||||
let decoded = decode_response_body_bytes_with_limit(&headers, &encoded, payload.len())
|
||||
.expect("decoded body exactly at limit should pass");
|
||||
|
||||
assert_eq!(decoded.as_ref(), payload);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gateway_frontdoor_self_loop_guard_matches_loopback_public_ai_route() {
|
||||
assert!(gateway_frontdoor_self_loop_guard_matches_with_port(
|
||||
@@ -4034,6 +4297,7 @@ mod tests {
|
||||
for proxy_url in ["socks5://127.0.0.1:1080", "socks5h://127.0.0.1:1080"] {
|
||||
build_client(
|
||||
"https://api.example.test/v1/chat/completions",
|
||||
"key-test",
|
||||
Some(&timeouts),
|
||||
Some(&aether_contracts::ProxySnapshot {
|
||||
enabled: Some(true),
|
||||
@@ -4103,6 +4367,7 @@ mod tests {
|
||||
|
||||
let left = super::direct_reqwest_client_cache_key(
|
||||
"http://127.0.0.1:18184/v1/chat/completions",
|
||||
"key-1",
|
||||
Some(&timeouts),
|
||||
None,
|
||||
Some(&h2c_profile),
|
||||
@@ -4110,6 +4375,7 @@ mod tests {
|
||||
);
|
||||
let right = super::direct_reqwest_client_cache_key(
|
||||
"http://127.0.0.1:18184/v1/responses",
|
||||
"key-1",
|
||||
Some(&timeouts),
|
||||
None,
|
||||
Some(&same_h2c_profile),
|
||||
@@ -4117,6 +4383,7 @@ mod tests {
|
||||
);
|
||||
let different_mode = super::direct_reqwest_client_cache_key(
|
||||
"http://127.0.0.1:18184/v1/chat/completions",
|
||||
"key-1",
|
||||
Some(&timeouts),
|
||||
None,
|
||||
Some(&http1_profile),
|
||||
@@ -4124,6 +4391,7 @@ mod tests {
|
||||
);
|
||||
let different_proxy = super::direct_reqwest_client_cache_key(
|
||||
"http://127.0.0.1:18184/v1/chat/completions",
|
||||
"key-1",
|
||||
Some(&timeouts),
|
||||
Some("http://127.0.0.1:8080".into()),
|
||||
Some(&h2c_profile),
|
||||
@@ -4138,6 +4406,72 @@ mod tests {
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn direct_reqwest_client_cache_key_partitions_key_scoped_pools_by_hashed_key_id() {
|
||||
let profile = ResolvedTransportProfile {
|
||||
profile_id: "key-scoped-profile".into(),
|
||||
backend: TRANSPORT_BACKEND_REQWEST_RUSTLS.into(),
|
||||
http_mode: TRANSPORT_HTTP_MODE_AUTO.into(),
|
||||
pool_scope: " key ".into(),
|
||||
header_fingerprint: None,
|
||||
extra: None,
|
||||
};
|
||||
let first_key_id = "plain-key-identity-alpha";
|
||||
let second_key_id = "plain-key-identity-beta";
|
||||
let cache_key = |key_id| {
|
||||
super::direct_reqwest_client_cache_key(
|
||||
"https://api.example.test/v1/messages",
|
||||
key_id,
|
||||
None,
|
||||
None,
|
||||
Some(&profile),
|
||||
ExecutionTransportControls::default(),
|
||||
)
|
||||
};
|
||||
|
||||
let first = cache_key(first_key_id);
|
||||
let first_key_id_with_whitespace = format!(" {first_key_id} ");
|
||||
let first_with_whitespace = cache_key(&first_key_id_with_whitespace);
|
||||
let second = cache_key(second_key_id);
|
||||
let empty = cache_key(" ");
|
||||
|
||||
assert_eq!(first, first_with_whitespace);
|
||||
assert_ne!(first, second);
|
||||
assert_eq!(first.pool_partition.as_deref().map(str::len), Some(64));
|
||||
assert!(empty.pool_partition.is_none());
|
||||
let debug = format!("{first:?} {second:?}");
|
||||
assert!(!debug.contains(first_key_id));
|
||||
assert!(!debug.contains(second_key_id));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn direct_reqwest_client_cache_key_shares_non_key_scoped_pools() {
|
||||
let profile = ResolvedTransportProfile {
|
||||
profile_id: "provider-scoped-profile".into(),
|
||||
backend: TRANSPORT_BACKEND_REQWEST_RUSTLS.into(),
|
||||
http_mode: TRANSPORT_HTTP_MODE_AUTO.into(),
|
||||
pool_scope: "provider".into(),
|
||||
header_fingerprint: None,
|
||||
extra: None,
|
||||
};
|
||||
let cache_key = |key_id| {
|
||||
super::direct_reqwest_client_cache_key(
|
||||
"https://api.example.test/v1/messages",
|
||||
key_id,
|
||||
None,
|
||||
None,
|
||||
Some(&profile),
|
||||
ExecutionTransportControls::default(),
|
||||
)
|
||||
};
|
||||
|
||||
let first = cache_key("plain-key-identity-alpha");
|
||||
let second = cache_key("plain-key-identity-beta");
|
||||
|
||||
assert_eq!(first, second);
|
||||
assert!(first.pool_partition.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn direct_reqwest_client_cache_key_splits_origin_only_when_enabled() {
|
||||
let _guard = direct_reqwest_env_lock();
|
||||
@@ -4152,6 +4486,7 @@ mod tests {
|
||||
|
||||
let shared_left = super::direct_reqwest_client_cache_key(
|
||||
"http://127.0.0.1:18184/v1/chat/completions",
|
||||
"key-1",
|
||||
None,
|
||||
None,
|
||||
Some(&profile),
|
||||
@@ -4159,6 +4494,7 @@ mod tests {
|
||||
);
|
||||
let shared_right = super::direct_reqwest_client_cache_key(
|
||||
"http://127.0.0.1:18185/v1/chat/completions",
|
||||
"key-1",
|
||||
None,
|
||||
None,
|
||||
Some(&profile),
|
||||
@@ -4169,6 +4505,7 @@ mod tests {
|
||||
let _per_origin = set_test_env_var(super::DIRECT_REQWEST_CACHE_PER_ORIGIN_ENV, "true");
|
||||
let split_left = super::direct_reqwest_client_cache_key(
|
||||
"http://127.0.0.1:18184/v1/chat/completions",
|
||||
"key-1",
|
||||
None,
|
||||
None,
|
||||
Some(&profile),
|
||||
@@ -4176,6 +4513,7 @@ mod tests {
|
||||
);
|
||||
let split_right = super::direct_reqwest_client_cache_key(
|
||||
"http://127.0.0.1:18185/v1/chat/completions",
|
||||
"key-1",
|
||||
None,
|
||||
None,
|
||||
Some(&profile),
|
||||
@@ -4201,6 +4539,7 @@ mod tests {
|
||||
|
||||
let auto_key = super::direct_reqwest_client_cache_key(
|
||||
"http://127.0.0.1:18184/v1/chat/completions",
|
||||
"key-1",
|
||||
None,
|
||||
None,
|
||||
Some(&auto_profile),
|
||||
@@ -4208,6 +4547,7 @@ mod tests {
|
||||
);
|
||||
let h2c_key = super::direct_reqwest_client_cache_key(
|
||||
"http://127.0.0.1:18184/v1/chat/completions",
|
||||
"key-1",
|
||||
None,
|
||||
None,
|
||||
Some(&h2c_profile),
|
||||
@@ -4547,6 +4887,7 @@ mod tests {
|
||||
|
||||
let cache_key = super::direct_reqwest_client_cache_key(
|
||||
&plan.url,
|
||||
&plan.key_id,
|
||||
plan.timeouts.as_ref(),
|
||||
None,
|
||||
Some(&profile),
|
||||
@@ -4610,6 +4951,7 @@ mod tests {
|
||||
|
||||
let cache_key = super::direct_reqwest_client_cache_key(
|
||||
&plan.url,
|
||||
&plan.key_id,
|
||||
plan.timeouts.as_ref(),
|
||||
None,
|
||||
Some(&profile),
|
||||
@@ -4664,6 +5006,7 @@ mod tests {
|
||||
|
||||
let cache_key = super::direct_reqwest_client_cache_key(
|
||||
&plan.url,
|
||||
&plan.key_id,
|
||||
plan.timeouts.as_ref(),
|
||||
None,
|
||||
Some(&profile),
|
||||
@@ -4782,6 +5125,57 @@ mod tests {
|
||||
.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn response_body_mode_control_header_is_never_forwarded_upstream() {
|
||||
let headers = BTreeMap::from([
|
||||
("content-type".into(), "application/json".into()),
|
||||
(
|
||||
EXECUTION_RESPONSE_BODY_MODE_HEADER.into(),
|
||||
ExecutionResponseBodyMode::PreserveBytes
|
||||
.as_str()
|
||||
.to_string(),
|
||||
),
|
||||
]);
|
||||
|
||||
let forwarded = build_request_headers(&headers, None, true)
|
||||
.expect("headers should build after stripping internal controls");
|
||||
|
||||
assert!(forwarded.get("content-type").is_some());
|
||||
assert!(forwarded.get(EXECUTION_RESPONSE_BODY_MODE_HEADER).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn response_body_mode_requires_same_format_raw_sync_plan() {
|
||||
let mut plan = tunnel_timeout_plan(false);
|
||||
plan.headers.insert(
|
||||
EXECUTION_RESPONSE_BODY_MODE_HEADER.to_string(),
|
||||
ExecutionResponseBodyMode::PreserveBytes
|
||||
.as_str()
|
||||
.to_string(),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
execution_response_body_mode(&plan),
|
||||
ExecutionResponseBodyMode::StructuredJson
|
||||
);
|
||||
|
||||
plan.body = RequestBody {
|
||||
json_body: None,
|
||||
body_bytes_b64: Some("e30=".to_string()),
|
||||
body_ref: None,
|
||||
};
|
||||
assert_eq!(
|
||||
execution_response_body_mode(&plan),
|
||||
ExecutionResponseBodyMode::PreserveBytes
|
||||
);
|
||||
|
||||
plan.provider_api_format = "claude:messages".to_string();
|
||||
assert_eq!(
|
||||
execution_response_body_mode(&plan),
|
||||
ExecutionResponseBodyMode::StructuredJson
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tunnel_request_meta_uses_total_timeout_for_non_stream_requests() {
|
||||
let plan = tunnel_timeout_plan(false);
|
||||
@@ -6465,6 +6859,7 @@ mod tests {
|
||||
);
|
||||
let cache_key = super::direct_reqwest_client_cache_key(
|
||||
&plan.url,
|
||||
&plan.key_id,
|
||||
plan.timeouts.as_ref(),
|
||||
None,
|
||||
Some(&profile),
|
||||
@@ -6646,6 +7041,7 @@ mod tests {
|
||||
|
||||
let error = match build_client(
|
||||
"https://api.example.test/v1/chat/completions",
|
||||
"key-test",
|
||||
None,
|
||||
None,
|
||||
Some(&profile),
|
||||
@@ -6673,6 +7069,51 @@ mod tests {
|
||||
assert!(!response_body_is_json(&headers, &body));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn structured_json_response_does_not_duplicate_body_bytes() {
|
||||
let headers =
|
||||
BTreeMap::from([("content-type".to_string(), "application/json".to_string())]);
|
||||
let body_bytes = br#"{ "unknown": true, "ok": true }"#;
|
||||
|
||||
let body = build_execution_response_body(
|
||||
&headers,
|
||||
body_bytes,
|
||||
body_bytes,
|
||||
false,
|
||||
ExecutionResponseBodyMode::StructuredJson,
|
||||
)
|
||||
.expect("body should build")
|
||||
.expect("body should be present");
|
||||
|
||||
assert!(body.json_body.is_some());
|
||||
assert!(body.body_bytes_b64.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn preserve_bytes_json_response_keeps_parsed_and_wire_representations() {
|
||||
let headers =
|
||||
BTreeMap::from([("content-type".to_string(), "application/json".to_string())]);
|
||||
let body_bytes = br#"{ "unknown": true, "ok": true }"#;
|
||||
|
||||
let body = build_execution_response_body(
|
||||
&headers,
|
||||
body_bytes,
|
||||
body_bytes,
|
||||
false,
|
||||
ExecutionResponseBodyMode::PreserveBytes,
|
||||
)
|
||||
.expect("body should build")
|
||||
.expect("body should be present");
|
||||
|
||||
assert_eq!(body.json_body, Some(json!({"unknown": true, "ok": true})));
|
||||
assert_eq!(
|
||||
base64::engine::general_purpose::STANDARD
|
||||
.decode(body.body_bytes_b64.expect("wire bytes should be present"))
|
||||
.expect("wire body should decode"),
|
||||
body_bytes
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn connect_json_error_response_is_decoded_for_stream_sync_body() {
|
||||
let headers = BTreeMap::from([(
|
||||
@@ -6684,9 +7125,15 @@ mod tests {
|
||||
body_bytes.extend_from_slice(&(payload.len() as u32).to_be_bytes());
|
||||
body_bytes.extend_from_slice(payload);
|
||||
|
||||
let body = build_execution_response_body(&headers, &body_bytes, &body_bytes, true)
|
||||
.expect("body should build")
|
||||
.expect("body should be present");
|
||||
let body = build_execution_response_body(
|
||||
&headers,
|
||||
&body_bytes,
|
||||
&body_bytes,
|
||||
true,
|
||||
ExecutionResponseBodyMode::StructuredJson,
|
||||
)
|
||||
.expect("body should build")
|
||||
.expect("body should be present");
|
||||
|
||||
assert_eq!(
|
||||
body.json_body
|
||||
|
||||
Reference in New Issue
Block a user