From a04673a90d0d0ae9d07ab1533d6c6c93e07cb5b7 Mon Sep 17 00:00:00 2001 From: elky Date: Thu, 30 Jul 2026 01:03:05 +0800 Subject: [PATCH] feat(gateway): harden failover and payload handling Retry pre-response transport failures across candidates with an explicit stop policy, and propagate end-to-end timing into usage records and UI diagnostics. Remove legacy body, import, cookie, PII, and tunnel replay caps while preserving optional operator-configured gateway limits. --- .env.example | 7 +- README.md | 6 +- .../src/ai_serving/planner/redaction.rs | 7 +- apps/aether-gateway/src/data/candidates.rs | 34 + .../src/data/state/integrations.rs | 43 +- .../execution_runtime/chatgpt_web_image.rs | 114 +++- .../src/execution_runtime/fallback.rs | 1 + .../src/execution_runtime/grok.rs | 126 ++-- .../src/execution_runtime/mod.rs | 5 + .../src/execution_runtime/server.rs | 3 + .../src/execution_runtime/stream/execution.rs | 588 +++++++++++++++++- .../stream/execution_failures.rs | 364 ++++++++++- .../src/execution_runtime/stream_pump.rs | 31 +- .../src/execution_runtime/sync/execution.rs | 297 ++++++++- .../src/execution_runtime/transport.rs | 2 + .../execution_runtime/transport_failure.rs | 139 +++++ .../src/execution_runtime/windsurf.rs | 114 +++- .../src/executor/candidate_loop.rs | 260 ++++++-- .../src/executor/orchestration.rs | 204 +++++- .../admin/provider/oauth/dispatch/cookie.rs | 24 +- .../provider/oauth/dispatch/cookie_task.rs | 31 +- .../handlers/admin/request/system/import.rs | 18 +- .../src/handlers/admin/request/system/mod.rs | 1 - .../handlers/public/support/user_me_usage.rs | 41 ++ apps/aether-gateway/src/headers.rs | 92 +-- .../src/orchestration/classifier.rs | 51 +- apps/aether-gateway/src/orchestration/mod.rs | 21 +- .../src/orchestration/policy.rs | 43 ++ .../src/orchestration/recovery.rs | 48 +- apps/aether-gateway/src/privacy/mod.rs | 210 ++----- .../aether-gateway/src/request_diagnostics.rs | 247 +++++++- .../ai_execute/sync/chat/pii_redaction.rs | 94 +-- .../src/tests/control/internal.rs | 60 +- apps/aether-gateway/src/tests/usage/local.rs | 179 +++++- apps/aether-gateway/src/tunnel/mod.rs | 116 +--- apps/aether-tunnel/.env.example | 3 - apps/aether-tunnel/README.md | 3 +- apps/aether-tunnel/src/app.rs | 6 +- apps/aether-tunnel/src/config.rs | 230 ++----- apps/aether-tunnel/src/setup/tui.rs | 33 +- apps/aether-tunnel/src/tunnel/mod.rs | 2 +- .../src/tunnel/stream_handler.rs | 151 ++--- .../aether-admin/src/observability/usage.rs | 54 ++ crates/aether-admin/src/system.rs | 4 +- crates/aether-admission-core/src/budget.rs | 2 +- crates/aether-admission-core/src/policy.rs | 13 + .../src/repository/candidates/types.rs | 14 +- crates/aether-gateway/frontdoor/src/body.rs | 70 ++- crates/aether-gateway/tunnel/src/admission.rs | 32 +- crates/aether-gateway/tunnel/src/lib.rs | 4 +- crates/aether-gateway/tunnel/src/relay.rs | 1 - .../src/provider/providers/claude_code.rs | 10 +- .../aether-usage/runtime/src/body_capture.rs | 8 +- .../runtime/src/request_metadata.rs | 26 +- crates/aether-usage/runtime/src/runtime.rs | 8 +- frontend/src/api/dashboard.ts | 2 + frontend/src/api/endpoints/types/provider.ts | 1 + frontend/src/api/me.ts | 4 + frontend/src/api/usage.ts | 6 + .../components/FailoverRulesDialog.vue | 34 + .../components/ProviderDetailDrawer.vue | 1 + .../components/ProviderFormDialog.vue | 6 +- ...iloverRulesDialog.transport-errors.spec.ts | 130 ++++ .../ProviderDetailDrawer.loading.spec.ts | 4 + ...ProviderFormDialog.transfer-limits.spec.ts | 11 +- .../components/HorizontalRequestTimeline.vue | 25 +- .../usage/components/RequestDetailDrawer.vue | 4 +- .../usage/components/UsageRecordsTable.vue | 26 +- .../HorizontalRequestTimeline.spec.ts | 40 ++ .../RequestDetailDrawer.pricing.spec.ts | 42 ++ .../__tests__/UsageRecordsTable.spec.ts | 61 +- frontend/src/features/usage/types.ts | 2 + frontend/src/i18n/__tests__/i18n.spec.ts | 1 + frontend/src/i18n/messages.ts | 11 +- frontend/src/views/admin/SystemSettings.vue | 6 - .../system-settings/RequestLogSection.vue | 46 +- .../__tests__/useConfigExportImport.spec.ts | 48 +- .../__tests__/useSystemConfig.spec.ts | 4 +- .../composables/useConfigExportImport.ts | 27 - .../composables/useSystemConfig.ts | 39 +- 80 files changed, 3640 insertions(+), 1236 deletions(-) create mode 100644 apps/aether-gateway/src/execution_runtime/transport_failure.rs create mode 100644 frontend/src/features/providers/components/__tests__/FailoverRulesDialog.transport-errors.spec.ts diff --git a/.env.example b/.env.example index b92e6b245..931d4109b 100644 --- a/.env.example +++ b/.env.example @@ -72,10 +72,11 @@ ADMIN_USERNAME=admin123456 # AETHER_GATEWAY_MAX_IN_FLIGHT_REQUESTS=2048 # AETHER_GATEWAY_REQUEST_BODY_BUFFER_BUDGET_MB=256 # AETHER_GATEWAY_REQUEST_BODY_READ_TIMEOUT_MS=120000 -# AETHER_MAX_REQUEST_BODY_MB=64 +# 可选的 Payload 上限（MiB）；默认及 0 均表示不限制。 +# AETHER_MAX_REQUEST_BODY_MB=0 # AETHER_GATEWAY_SECURITY_CACHE_TTL_MS=1000 -# AETHER_MAX_REDACTED_SYNC_RESPONSE_BODY_MB=64 -# AETHER_MAX_INTERNAL_BUFFERED_BODY_MB=128 +# AETHER_MAX_REDACTED_SYNC_RESPONSE_BODY_MB=0 +# AETHER_MAX_INTERNAL_BUFFERED_BODY_MB=0 # AETHER_TUNNEL_NODE_STATUS_QUEUE_CAPACITY=1024 # PostgreSQL 容器调优：docker-compose.yml 已内置通用默认值，通常不用配置。 diff --git a/README.md b/README.md index aceee2ef2..2f804363d 100644 --- a/README.md +++ b/README.md @@ -146,11 +146,11 @@ Aether Tunnel 是配套的正向代理节点，部署在海外 VPS 上，为墙 - `AETHER_GATEWAY_MAX_IN_FLIGHT_REQUESTS`：单实例请求并发上限；未配置时按 CPU 自动推导（基础范围 `512-65536`），低文件描述符预算时会进一步下调 - `AETHER_GATEWAY_REQUEST_BODY_BUFFER_BUDGET_MB`：单实例同时读取和解压请求体的加权内存预算，默认 `256MB` - `AETHER_GATEWAY_REQUEST_BODY_READ_TIMEOUT_MS`：请求体完整读取超时，默认 `120000ms` -- `AETHER_MAX_REQUEST_BODY_MB`：单请求解压后的最大请求体，默认 `64MB` -- `AETHER_MAX_INTERNAL_BUFFERED_BODY_MB`：heartbeat、管理探测等内部必须整包读取的响应体上限，默认 `128MB` +- `AETHER_MAX_REQUEST_BODY_MB`：可选的单请求解压后请求体上限；未配置或设为 `0` 时不限制 +- `AETHER_MAX_INTERNAL_BUFFERED_BODY_MB`：可选的 heartbeat、管理探测等内部整包响应体上限；未配置或设为 `0` 时不限制 - `AETHER_TUNNEL_NODE_STATUS_QUEUE_CAPACITY`：隧道节点状态上报队列容量，默认 `1024`；满载时拒绝新事件，避免控制面故障导致无界内存增长 - `AETHER_GATEWAY_SECURITY_CACHE_TTL_MS`：IP 黑白名单本地缓存时间，默认 `1000ms`，写操作会主动失效相关缓存 -- `AETHER_MAX_REDACTED_SYNC_RESPONSE_BODY_MB`：启用 PII 恢复时同步响应允许缓冲的最大大小，默认 `64MB` +- `AETHER_MAX_REDACTED_SYNC_RESPONSE_BODY_MB`：可选的 PII 恢复同步响应缓冲上限；未配置或设为 `0` 时不限制 - `REDIS_URL`：Redis 连接串；仅 Postgres + Redis 的 Docker Compose 部署需要配置 - `AETHER_RUNTIME_BACKEND=memory|redis`：运行时缓存/协调后端。SQLite 默认用 `memory`，不会连接 Redis - `AETHER_GATEWAY_AUTO_PREPARE_DATABASE`：常规启动前自动执行挂起的 schema migration 和 backfill；仓库自带的 `docker-compose.yml` 默认开启 diff --git a/apps/aether-gateway/src/ai_serving/planner/redaction.rs b/apps/aether-gateway/src/ai_serving/planner/redaction.rs index 99bd6d4e2..4b4814d21 100644 --- a/apps/aether-gateway/src/ai_serving/planner/redaction.rs +++ b/apps/aether-gateway/src/ai_serving/planner/redaction.rs @@ -216,12 +216,7 @@ async fn resolve_chat_pii_redaction_feature_settings( } fn redaction_mask_error_to_gateway_error(error: RedactionMaskError) -> GatewayError { - match error { - RedactionMaskError::Limit(limit) => GatewayError::Client { - status: limit.client_status(), - message: limit.safe_message().to_string(), - }, - } + match error {} } #[cfg(test)] diff --git a/apps/aether-gateway/src/data/candidates.rs b/apps/aether-gateway/src/data/candidates.rs index da2d2f3f0..7ed4854a2 100644 --- a/apps/aether-gateway/src/data/candidates.rs +++ b/apps/aether-gateway/src/data/candidates.rs @@ -127,4 +127,38 @@ mod tests { assert_eq!(trace.final_status, RequestCandidateFinalStatus::Failed); assert_eq!(trace.total_latency_ms, 33); } + + #[tokio::test] + async fn request_candidate_trace_recovers_missing_transport_latency_from_timestamps() { + let repository = Arc::new(InMemoryRequestCandidateRepository::seed(vec![ + sample_candidate( + "cand-timeout", + "req-failover", + 0, + RequestCandidateStatus::Failed, + Some(100), + None, + None, + ), + sample_candidate( + "cand-success", + "req-failover", + 1, + RequestCandidateStatus::Success, + Some(101), + Some(626), + Some(200), + ), + ])); + let state = GatewayDataState::with_request_candidate_reader_for_tests(repository); + + let trace = read_request_candidate_trace(&state, "req-failover", true) + .await + .expect("trace should succeed") + .expect("trace should exist"); + + assert_eq!(trace.total_candidates, 2); + assert_eq!(trace.final_status, RequestCandidateFinalStatus::Success); + assert_eq!(trace.total_latency_ms, 1_626); + } } diff --git a/apps/aether-gateway/src/data/state/integrations.rs b/apps/aether-gateway/src/data/state/integrations.rs index f0361bad9..5ad389012 100644 --- a/apps/aether-gateway/src/data/state/integrations.rs +++ b/apps/aether-gateway/src/data/state/integrations.rs @@ -20,8 +20,6 @@ use aether_runtime_state::RuntimeQueueStore; use aether_usage_runtime::{ UsageBillingEventEnricher, UsageBodyCapturePolicy, UsageEvent, UsageRecordWriter, UsageRequestRecordLevel, UsageRuntimeAccess, UsageSettlementWriter, - DEFAULT_USAGE_REQUEST_BODY_CAPTURE_LIMIT_BYTES, - DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES, }; use aether_video_tasks_core::StoredVideoTaskReadSide; use async_trait::async_trait; @@ -33,8 +31,6 @@ use crate::provider_transport::ProviderTransportSnapshotSource; const REQUEST_RECORD_LEVEL_KEY: &str = "request_record_level"; const LEGACY_REQUEST_LOG_LEVEL_KEY: &str = "request_log_level"; -const MAX_REQUEST_BODY_SIZE_KEY: &str = "max_request_body_size"; -const MAX_RESPONSE_BODY_SIZE_KEY: &str = "max_response_body_size"; fn usage_request_record_level_from_value(value: Option<&Value>) -> UsageRequestRecordLevel { let Some(value) = value.and_then(Value::as_str).map(str::trim) else { @@ -53,14 +49,6 @@ fn usage_request_record_level_from_value(value: Option<&Value>) -> UsageRequestR } } -fn usage_body_capture_limit_from_value(value: Option<&Value>, default: usize) -> Option { - match value.and_then(Value::as_u64) { - Some(0) => None, - Some(limit) => usize::try_from(limit).ok().filter(|limit| *limit > 0), - None => Some(default), - } -} - #[async_trait] impl RequestAuditReader for GatewayDataState { async fn find_request_usage_audit_by_request_id( @@ -282,20 +270,8 @@ impl UsageRuntimeAccess for GatewayDataState { .await? } }; - let max_request_body_size = - GatewayDataState::find_system_config_value(self, MAX_REQUEST_BODY_SIZE_KEY).await?; - let max_response_body_size = - GatewayDataState::find_system_config_value(self, MAX_RESPONSE_BODY_SIZE_KEY).await?; Ok(UsageBodyCapturePolicy { record_level: usage_request_record_level_from_value(value.as_ref()), - max_request_body_bytes: usage_body_capture_limit_from_value( - max_request_body_size.as_ref(), - DEFAULT_USAGE_REQUEST_BODY_CAPTURE_LIMIT_BYTES, - ), - max_response_body_bytes: usage_body_capture_limit_from_value( - max_response_body_size.as_ref(), - DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES, - ), }) } } @@ -502,7 +478,7 @@ mod tests { } #[tokio::test] - async fn usage_runtime_access_reads_body_capture_limits_from_system_config() { + async fn usage_runtime_access_ignores_legacy_body_capture_limits() { let state = GatewayDataState::disabled().with_system_config_values_for_tests([ ("max_request_body_size".to_string(), json!(1234)), ("max_response_body_size".to_string(), json!(5678)), @@ -513,22 +489,5 @@ mod tests { .expect("body capture policy should read"); assert_eq!(policy.record_level, UsageRequestRecordLevel::Full); - assert_eq!(policy.max_request_body_bytes, Some(1234)); - assert_eq!(policy.max_response_body_bytes, Some(5678)); - } - - #[tokio::test] - async fn usage_runtime_access_treats_zero_body_capture_limit_as_unbounded() { - let state = GatewayDataState::disabled().with_system_config_values_for_tests([ - ("max_request_body_size".to_string(), json!(0)), - ("max_response_body_size".to_string(), json!(0)), - ]); - - let policy = UsageRuntimeAccess::body_capture_policy(&state) - .await - .expect("body capture policy should read"); - - assert_eq!(policy.max_request_body_bytes, None); - assert_eq!(policy.max_response_body_bytes, None); } } diff --git a/apps/aether-gateway/src/execution_runtime/chatgpt_web_image.rs b/apps/aether-gateway/src/execution_runtime/chatgpt_web_image.rs index fee4f6eba..1f13401b6 100644 --- a/apps/aether-gateway/src/execution_runtime/chatgpt_web_image.rs +++ b/apps/aether-gateway/src/execution_runtime/chatgpt_web_image.rs @@ -115,7 +115,16 @@ pub(crate) async fn maybe_execute_chatgpt_web_image_sync( let result = match execute_chatgpt_web_image(state, plan, report_context, started_at).await { Ok(result) => result, - Err(err) => chatgpt_web_transport_error_execution_result(plan, started_at, &err), + Err(ExecutionRuntimeTransportError::UpstreamHttpStatus { + status_code, + message, + }) => chatgpt_web_http_error_execution_result( + plan, + started_at, + status_code, + message.as_str(), + ), + Err(error) => return Err(error), }; Ok(Some(result)) }) @@ -133,7 +142,13 @@ pub(crate) async fn maybe_execute_chatgpt_web_image_stream( let started_at = Instant::now(); let result = match execute_chatgpt_web_image(state, plan, report_context, started_at).await { Ok(result) => result, - Err(err) => chatgpt_web_transport_error_execution_result(plan, started_at, &err), + Err(ExecutionRuntimeTransportError::UpstreamHttpStatus { + status_code, + message, + }) => { + chatgpt_web_http_error_execution_result(plan, started_at, status_code, message.as_str()) + } + Err(error) => return Err(error), }; Ok(Some(ChatGptWebImageStream { frame_stream: execution_result_frame_stream(plan, &result, report_context), @@ -2567,19 +2582,20 @@ fn json_execution_result( } } -fn chatgpt_web_transport_error_execution_result( +fn chatgpt_web_http_error_execution_result( plan: &ExecutionPlan, started_at: Instant, - error: &ExecutionRuntimeTransportError, + status_code: u16, + message: &str, ) -> ExecutionResult { json_execution_result( plan, - 503, + status_code, json!({ "error": { "type": "upstream_error", "code": "chatgpt_web_image_execution_unavailable", - "message": error.to_string() + "message": message } }), started_at, @@ -2798,11 +2814,14 @@ fn ensure_success( return Ok(()); } let body = String::from_utf8_lossy(&execution_result_body_bytes_lossy(result)).to_string(); - Err(ExecutionRuntimeTransportError::UpstreamRequest(format!( - "{stage} returned {}: {}", - result.status_code, - body.chars().take(320).collect::() - ))) + Err(ExecutionRuntimeTransportError::UpstreamHttpStatus { + status_code: result.status_code, + message: format!( + "{stage} returned {}: {}", + result.status_code, + body.chars().take(320).collect::() + ), + }) } fn chatgpt_web_base_url_from_plan(plan: &ExecutionPlan) -> String { @@ -3992,10 +4011,14 @@ data: [DONE] Some(&json!({"chatgpt_web_image": true})), ) .await - .expect("executor should run") + .expect("executor should preserve the upstream HTTP response") .expect("plan should be intercepted"); - assert_ne!(result.status_code, 200); + assert_eq!(result.status_code, 500); + assert_eq!( + execution_result_json(&result).expect("error response should be json")["error"]["code"], + json!("chatgpt_web_image_execution_unavailable") + ); let metadata = reloaded_chatgpt_web_metadata(repository.as_ref()).await; assert_eq!(metadata["image_quota_remaining"], json!(25.0)); assert_eq!(metadata["image_quota_used"], json!(0.0)); @@ -4005,6 +4028,71 @@ data: [DONE] handle.abort(); } + #[tokio::test] + async fn chatgpt_web_image_sync_propagates_network_failure_without_synthetic_503() { + let listener = crate::test_support::bind_loopback_listener() + .await + .expect("listener should bind"); + let base_url = format!( + "http://{}", + listener.local_addr().expect("local addr should resolve") + ); + drop(listener); + let state = crate::AppState::new().expect("state should build"); + let plan = sample_plan( + base_url.as_str(), + json!({"prompt": "draw a small test image"}), + false, + ); + + let error = maybe_execute_chatgpt_web_image_sync( + &state, + &plan, + Some(&json!({"chatgpt_web_image": true})), + ) + .await + .expect_err("connection failure should propagate to the candidate loop"); + + assert!(matches!( + error, + ExecutionRuntimeTransportError::UpstreamRequest(_) + )); + } + + #[tokio::test] + async fn chatgpt_web_image_stream_propagates_network_failure_without_synthetic_503() { + let listener = crate::test_support::bind_loopback_listener() + .await + .expect("listener should bind"); + let base_url = format!( + "http://{}", + listener.local_addr().expect("local addr should resolve") + ); + drop(listener); + let state = crate::AppState::new().expect("state should build"); + let plan = sample_plan( + base_url.as_str(), + json!({"prompt": "draw a small test image"}), + true, + ); + + let error = match maybe_execute_chatgpt_web_image_stream( + &state, + &plan, + Some(&json!({"chatgpt_web_image": true})), + ) + .await + { + Err(error) => error, + Ok(_) => panic!("connection failure should propagate to the candidate loop"), + }; + + assert!(matches!( + error, + ExecutionRuntimeTransportError::UpstreamRequest(_) + )); + } + #[tokio::test] async fn chatgpt_web_image_stream_path_wraps_success_sse_as_ndjson_frames() { let (base_url, handle) = start_mock_chatgpt_web().await; diff --git a/apps/aether-gateway/src/execution_runtime/fallback.rs b/apps/aether-gateway/src/execution_runtime/fallback.rs index 6cba6676e..e475850e4 100644 --- a/apps/aether-gateway/src/execution_runtime/fallback.rs +++ b/apps/aether-gateway/src/execution_runtime/fallback.rs @@ -940,6 +940,7 @@ mod tests { max_transfer_timeout_seconds: 0, stop_status_codes: [503].into_iter().collect(), continue_status_codes: [409, 429].into_iter().collect(), + stop_on_transport_errors: false, success_failover_patterns: Vec::new(), error_stop_patterns: Vec::new(), stop_cyber_policy_errors: true, diff --git a/apps/aether-gateway/src/execution_runtime/grok.rs b/apps/aether-gateway/src/execution_runtime/grok.rs index 5ec901d43..48918c677 100644 --- a/apps/aether-gateway/src/execution_runtime/grok.rs +++ b/apps/aether-gateway/src/execution_runtime/grok.rs @@ -43,7 +43,6 @@ const GROK_MEDIA_POST_PATH: &str = "/rest/media/post/create"; const GROK_IMAGINE_WS_URL: &str = "wss://grok.com/ws/imagine/listen"; const GROK_STANDARD_PROVIDER_API_FORMAT: &str = "openai:responses"; const GROK_PROMPT_OVERHEAD_TOKENS: u64 = 4; -const GROK_MAX_ATTACHMENT_BYTES: usize = 25 * 1024 * 1024; const GROK_MAX_ATTACHMENT_REDIRECTS: usize = 5; const GROK_IMAGINE_STREAM_TIMEOUT_MS: u64 = 10_000; const GROK_IMAGINE_ROUND_TIMEOUT_MS: u64 = 120_000; @@ -642,7 +641,7 @@ fn grok_success_frame_stream( let chunk = match item { Ok(chunk) => chunk, Err(message) => { - match encode_grok_error_frame(status_code, message) { + match encode_grok_error_frame(message) { Ok(frame) => yield Ok(frame), Err(err) => { yield Err(err); @@ -872,17 +871,17 @@ fn encode_grok_telemetry_frame( }) } -fn encode_grok_error_frame(status_code: u16, message: String) -> Result { +fn encode_grok_error_frame(message: String) -> Result { encode_stream_frame_ndjson(&StreamFrame { frame_type: StreamFrameType::Error, payload: StreamFramePayload::Error { error: aether_contracts::ExecutionError { - kind: aether_contracts::ExecutionErrorKind::Internal, + kind: aether_contracts::ExecutionErrorKind::ProtocolError, phase: aether_contracts::ExecutionPhase::StreamRead, message, - upstream_status: Some(status_code), - retryable: false, - failover_recommended: false, + upstream_status: None, + retryable: true, + failover_recommended: true, }, }, }) @@ -896,7 +895,7 @@ fn encode_grok_first_byte_timeout_frame(timeout: Duration) -> Result(); - let decoded_len = base64::engine::general_purpose::STANDARD - .decode(&normalized_b64) - .map_err(|err| { - ExecutionRuntimeTransportError::UpstreamRequest(format!( - "Grok attachment data URI base64 is invalid: {err}" - )) - })? - .len(); - if decoded_len > GROK_MAX_ATTACHMENT_BYTES { - return Err(ExecutionRuntimeTransportError::UpstreamRequest(format!( - "Grok attachment exceeds {} byte limit", - GROK_MAX_ATTACHMENT_BYTES - ))); - } + drop( + base64::engine::general_purpose::STANDARD + .decode(&normalized_b64) + .map_err(|err| { + ExecutionRuntimeTransportError::UpstreamRequest(format!( + "Grok attachment data URI base64 is invalid: {err}" + )) + })?, + ); Ok(GrokAttachmentPayload { filename: input .filename @@ -1635,12 +1629,6 @@ async fn collect_grok_attachment_url_bytes( let chunk = chunk.map_err(|err| { ExecutionRuntimeTransportError::UpstreamRequest(format_upstream_request_error(&err)) })?; - if bytes.len().saturating_add(chunk.len()) > GROK_MAX_ATTACHMENT_BYTES { - return Err(ExecutionRuntimeTransportError::UpstreamRequest(format!( - "Grok attachment exceeds {} byte limit", - GROK_MAX_ATTACHMENT_BYTES - ))); - } bytes.extend_from_slice(&chunk); } Ok(bytes) @@ -3073,12 +3061,6 @@ async fn grok_download_image_asset( if bytes.is_empty() { return Ok(None); } - if bytes.len() > GROK_MAX_ATTACHMENT_BYTES { - return Err(ExecutionRuntimeTransportError::UpstreamRequest(format!( - "Grok image asset exceeds {} byte limit", - GROK_MAX_ATTACHMENT_BYTES - ))); - } Ok(Some(format!( "data:{content_type};base64,{}", base64::engine::general_purpose::STANDARD.encode(bytes) @@ -3236,7 +3218,10 @@ fn push_sse_event(body: &mut String, event: &str, data: &Value) { mod tests { use std::collections::BTreeMap; - use aether_contracts::{ExecutionPlan, RequestBody, StreamFrame, StreamFramePayload}; + use aether_contracts::{ + ExecutionErrorKind, ExecutionPhase, ExecutionPlan, RequestBody, StreamFrame, + StreamFramePayload, + }; use axum::body::{Body, Bytes}; use axum::extract::Request; use axum::routing::any; @@ -3246,16 +3231,18 @@ mod tests { use http::{Method, StatusCode}; use super::{ + encode_grok_error_frame, encode_grok_first_byte_timeout_frame, extract_grok_attachment_inputs, grok_aspect_ratio_from_provider_body, grok_asset_url, - grok_attachment_ip_is_public, grok_client_json_body, grok_client_stream_body, - grok_handle_imagine_ws_message, grok_image_count_from_provider_body, - grok_image_prompt_from_provider_body, grok_imagine_request_message, - grok_imagine_reset_message, grok_media_post_url, + grok_attachment_ip_is_public, grok_attachment_payload_from_data_uri, grok_client_json_body, + grok_client_stream_body, grok_handle_imagine_ws_message, + grok_image_count_from_provider_body, grok_image_prompt_from_provider_body, + grok_imagine_request_message, grok_imagine_reset_message, grok_media_post_url, grok_plan_uses_structured_image_generation, grok_should_collect_image_stream, grok_should_use_imagine_websocket, grok_success_frame_stream, grok_upload_url, grok_upstream_model_name, grok_usage_estimate, grok_user_id_from_cookie_header, materialize_grok_image_assets, openai_chat_body, openai_image_body, openai_responses_body, - set_grok_image_edit_config, GrokCollected, GrokImagineImage, GrokStreamAdapter, + set_grok_image_edit_config, GrokAttachmentInput, GrokCollected, GrokImagineImage, + GrokStreamAdapter, }; fn sample_plan(body: serde_json::Value, client_api_format: &str) -> ExecutionPlan { @@ -3326,6 +3313,45 @@ mod tests { out } + fn decode_encoded_frame(encoded: Bytes) -> StreamFrame { + let line = String::from_utf8(encoded.to_vec()).expect("frame should be utf8"); + serde_json::from_str(line.trim()).expect("frame should deserialize") + } + + #[test] + fn grok_stream_read_error_is_retryable_transport_without_upstream_status() { + let frame = decode_encoded_frame( + encode_grok_error_frame("connection reset while reading response body".to_string()) + .expect("error frame should encode"), + ); + let StreamFramePayload::Error { error } = frame.payload else { + panic!("encoded frame should contain an execution error"); + }; + + assert_eq!(error.kind, ExecutionErrorKind::ProtocolError); + assert_eq!(error.phase, ExecutionPhase::StreamRead); + assert_eq!(error.upstream_status, None); + assert!(error.retryable); + assert!(error.failover_recommended); + } + + #[test] + fn grok_first_byte_timeout_is_retryable_transport_without_upstream_status() { + let frame = decode_encoded_frame( + encode_grok_first_byte_timeout_frame(std::time::Duration::from_millis(250)) + .expect("timeout frame should encode"), + ); + let StreamFramePayload::Error { error } = frame.payload else { + panic!("encoded frame should contain an execution error"); + }; + + assert_eq!(error.kind, ExecutionErrorKind::FirstByteTimeout); + assert_eq!(error.phase, ExecutionPhase::FirstByte); + assert_eq!(error.upstream_status, None); + assert!(error.retryable); + assert!(error.failover_recommended); + } + #[tokio::test] async fn grok_success_stream_forwards_token_chunks_incrementally() { let plan = sample_plan( @@ -4062,6 +4088,26 @@ mod tests { assert_eq!(inputs[1].source.as_str(), "data:text/plain;base64,bm90ZXM="); } + #[test] + fn grok_data_uri_attachment_accepts_content_above_previous_size_cap() { + const PREVIOUS_CAP_BYTES: usize = 25 * 1024 * 1024; + let base64_blocks = PREVIOUS_CAP_BYTES / 3 + 1; + let mut source = String::from("data:application/octet-stream;base64,"); + source.extend(std::iter::repeat_n('A', base64_blocks * 4)); + let input = GrokAttachmentInput { + source, + filename: Some("large.bin".to_string()), + mime_type: None, + }; + + let payload = grok_attachment_payload_from_data_uri(&input, 0) + .expect("attachment above the previous size cap should be accepted"); + + assert_eq!(payload.filename, "large.bin"); + assert_eq!(payload.mime_type, "application/octet-stream"); + assert_eq!(payload.content_b64.len(), base64_blocks * 4); + } + #[test] fn extracts_responses_and_claude_attachment_inputs() { let responses = extract_grok_attachment_inputs( diff --git a/apps/aether-gateway/src/execution_runtime/mod.rs b/apps/aether-gateway/src/execution_runtime/mod.rs index 4f074b315..260660cd9 100644 --- a/apps/aether-gateway/src/execution_runtime/mod.rs +++ b/apps/aether-gateway/src/execution_runtime/mod.rs @@ -20,6 +20,7 @@ mod stream_pump; pub(crate) mod submission; pub(crate) mod sync; pub(crate) mod transport; +mod transport_failure; mod windsurf; pub(crate) use self::chatgpt_web_image::maybe_execute_chatgpt_web_image_sync; @@ -133,6 +134,10 @@ pub(crate) use transport::{ execute_sync_plan as execute_execution_runtime_sync_plan, DirectSyncExecutionRuntime, DirectUpstreamStreamExecution, ExecutionRuntimeTransportError, }; +pub(crate) use transport_failure::{ + build_transport_error_stop_response, mark_stream_candidate_watchdog_terminal_started, + StreamCandidateWatchdogProgress, +}; #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub(crate) struct ClientIntent { diff --git a/apps/aether-gateway/src/execution_runtime/server.rs b/apps/aether-gateway/src/execution_runtime/server.rs index ddc00cf04..8308e58ea 100644 --- a/apps/aether-gateway/src/execution_runtime/server.rs +++ b/apps/aether-gateway/src/execution_runtime/server.rs @@ -370,6 +370,9 @@ impl IntoResponse for ExecutionRuntimeAppError { | ExecutionRuntimeTransportError::UnsupportedTransportProfile(_) | ExecutionRuntimeTransportError::BodyEncode(_), ) => StatusCode::BAD_REQUEST, + ExecutionRuntimeServerError::Transport( + ExecutionRuntimeTransportError::UpstreamHttpStatus { status_code, .. }, + ) => StatusCode::from_u16(status_code).unwrap_or(StatusCode::BAD_GATEWAY), ExecutionRuntimeServerError::Transport( ExecutionRuntimeTransportError::ClientBuild(_) | ExecutionRuntimeTransportError::BrowserClientBuild(_) diff --git a/apps/aether-gateway/src/execution_runtime/stream/execution.rs b/apps/aether-gateway/src/execution_runtime/stream/execution.rs index 587b22f63..d7fd67e77 100644 --- a/apps/aether-gateway/src/execution_runtime/stream/execution.rs +++ b/apps/aether-gateway/src/execution_runtime/stream/execution.rs @@ -24,7 +24,7 @@ use aether_scheduler_core::{ use aether_usage_runtime::{ build_lifecycle_usage_seed, build_stream_terminal_usage_payload_seed, build_sync_terminal_usage_payload_seed, build_terminal_usage_context_seed, LifecycleUsageSeed, - SyncTerminalUsagePayloadSeed, TerminalUsageContextSeed, + SyncTerminalUsagePayloadSeed, TerminalUsageContextSeed, UsageRequestRecordLevel, DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES, }; use async_stream::stream; @@ -56,8 +56,9 @@ use super::error::{ mod execution_failures; use self::execution_failures::{ build_stream_failure_from_execution_error, build_stream_failure_from_provider_error_body, - build_stream_failure_report, handle_prefetch_provider_private_stream_error, - handle_prefetch_stream_failure, submit_midstream_stream_failure, StreamFailureReport, + build_stream_failure_report, build_stream_transport_failure_report, + handle_prefetch_provider_private_stream_error, handle_prefetch_stream_failure, + submit_midstream_stream_failure, StreamFailureReport, }; use crate::ai_serving::api::{ extract_provider_private_stream_error_body, maybe_bridge_standard_sync_json_to_stream, @@ -129,7 +130,8 @@ use crate::request_candidate_runtime::{ }; use crate::request_diagnostics::{ attach_current_request_diagnostics_to_report_context, - attach_request_diagnostics_to_report_context, current_request_diagnostics, RequestDiagnostics, + attach_request_diagnostics_and_candidate_start_timing_to_report_context, + current_request_diagnostics, RequestDiagnostics, }; use crate::stage_metrics::{ attach_stage_trace_to_report_context, observe_gateway_stage_ms, observe_gateway_stage_trace_ms, @@ -148,6 +150,7 @@ const SSE_CONTROL_FILTER_MAX_BUFFER_BYTES: usize = 1024 * 1024; const SSE_TERMINAL_DETECTOR_MAX_LINE_BYTES: usize = 1024 * 1024; const SSE_TERMINAL_DETECTOR_MAX_RECORD_BYTES: usize = SSE_TERMINAL_DETECTOR_MAX_LINE_BYTES; const PROVIDER_STREAM_ERROR_INSPECTION_MAX_BYTES: usize = SSE_TERMINAL_DETECTOR_MAX_LINE_BYTES; +const BASIC_STREAM_BODY_ANALYSIS_LIMIT_BYTES: usize = 5 * 1024 * 1024; const STREAM_IDLE_LOG_INTERVAL: Duration = Duration::from_secs(60); const STREAM_IDLE_LOG_INTERVAL_MS: u64 = 60_000; const REWRITTEN_STREAM_PREFETCH_TIMEOUT: Duration = Duration::from_millis(750); @@ -281,8 +284,15 @@ fn report_context_with_stage_trace( fn report_context_with_request_diagnostics( report_context: Option, diagnostics: Option<&Arc>, + candidate_started_at: Instant, + terminal_telemetry: Option<&ExecutionTelemetry>, ) -> Option { - attach_request_diagnostics_to_report_context(report_context, diagnostics) + attach_request_diagnostics_and_candidate_start_timing_to_report_context( + report_context, + diagnostics, + Some(candidate_started_at), + terminal_telemetry.and_then(|telemetry| telemetry.ttfb_ms), + ) } fn request_accepted_elapsed_ms(diagnostics: Option<&Arc>) -> Option { @@ -332,6 +342,37 @@ fn direct_passthrough_mode() -> DirectPassthroughMode { .unwrap_or(DirectPassthroughMode::Inline) } +fn stream_body_buffer_limit_for_record_level(record_level: UsageRequestRecordLevel) -> usize { + match record_level { + UsageRequestRecordLevel::Basic => BASIC_STREAM_BODY_ANALYSIS_LIMIT_BYTES, + UsageRequestRecordLevel::Full => DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES, + } +} + +async fn resolve_stream_body_buffer_limit(state: &AppState) -> usize { + if !state.usage_runtime.is_enabled() { + return BASIC_STREAM_BODY_ANALYSIS_LIMIT_BYTES; + } + + match state + .usage_runtime + .body_capture_policy_for(state.usage_lifecycle_data_state().as_ref()) + .await + { + Ok(policy) => stream_body_buffer_limit_for_record_level(policy.record_level), + Err(error) => { + warn!( + event_name = "stream_body_capture_policy_read_failed", + log_type = "ops", + error = %error, + fallback = "full", + "gateway could not resolve stream body capture policy" + ); + DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES + } + } +} + fn parse_direct_passthrough_mode(value: &str) -> DirectPassthroughMode { match value.trim().to_ascii_lowercase().as_str() { "legacy" | "pump" | "mpsc" => DirectPassthroughMode::Legacy, @@ -379,6 +420,7 @@ async fn record_sync_terminal_usage_with_handoff_after_spawn( ) where F: Future + Send + 'static, { + crate::execution_runtime::mark_stream_candidate_watchdog_terminal_started(); // Capture request task-local diagnostics before handing the work to a spawned task. Tokio // task-local values do not propagate across spawn boundaries. let (context_seed, payload_seed) = @@ -481,6 +523,7 @@ async fn record_stream_terminal_usage( payload: &GatewayStreamReportRequest, cancelled: bool, ) { + crate::execution_runtime::mark_stream_candidate_watchdog_terminal_started(); let context_seed = build_terminal_usage_context_seed(plan, report_context); let payload_seed = build_stream_terminal_usage_payload_seed(payload); state @@ -1719,6 +1762,7 @@ struct DirectPassthroughFinalizerCore { stream_usage_observer: Option, stream_usage_observer_buffered: Vec, provider_error_inspection: ProviderStreamErrorInspection, + max_stream_body_buffer_bytes: usize, provider_buffered_body: Vec, buffered_body: Vec, provider_body_truncated: bool, @@ -1869,7 +1913,7 @@ impl DirectPassthroughFinalizer { append_stream_capture_bytes( &mut core.provider_buffered_body, chunk.as_ref(), - DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES, + core.max_stream_body_buffer_bytes, &mut core.provider_body_truncated, ); if let (Some(observer), Some(report_context)) = ( @@ -1907,7 +1951,7 @@ impl DirectPassthroughFinalizer { append_stream_capture_bytes( &mut core.buffered_body, chunk.as_ref(), - DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES, + core.max_stream_body_buffer_bytes, &mut core.client_body_truncated, ); if !core.requires_anthropic_message_stop { @@ -2094,6 +2138,7 @@ impl DirectPassthroughFinalizerCore { stream_usage_observer: _, stream_usage_observer_buffered: _, provider_error_inspection: _, + max_stream_body_buffer_bytes: _, provider_buffered_body, buffered_body, provider_body_truncated, @@ -2150,6 +2195,8 @@ impl DirectPassthroughFinalizerCore { let report_context_for_payload = report_context_with_request_diagnostics( report_context_for_payload, request_diagnostics.as_ref(), + stream_started_at, + terminal_telemetry.as_ref(), ); submit_midstream_stream_failure( &state, @@ -2183,6 +2230,8 @@ impl DirectPassthroughFinalizerCore { let report_context_for_payload = report_context_with_request_diagnostics( report_context_for_payload, request_diagnostics.as_ref(), + stream_started_at, + terminal_telemetry.as_ref(), ); let usage_payload = build_stream_usage_payload( trace_id, @@ -2275,6 +2324,8 @@ impl DirectPassthroughFinalizerCore { let report_context_for_payload = report_context_with_request_diagnostics( report_context_for_payload, request_diagnostics.as_ref(), + stream_started_at, + terminal_telemetry.as_ref(), ); let usage_payload = build_stream_usage_payload( trace_id.clone(), @@ -2582,7 +2633,7 @@ impl DirectPassthroughInlineBodyState { Ok(item) => item, Err(timeout) => { if let Some(finalizer) = self.finalizer.as_mut() { - finalizer.set_terminal_failure(build_stream_failure_report( + finalizer.set_terminal_failure(build_stream_transport_failure_report( "first_byte_timeout", stream_first_byte_timeout_message(timeout), 504, @@ -2637,7 +2688,7 @@ impl DirectPassthroughInlineBodyState { error = %message, "gateway direct passthrough upstream body read failed" ); - finalizer.set_terminal_failure(build_stream_failure_report( + finalizer.set_terminal_failure(build_stream_transport_failure_report( "execution_runtime_stream_read_error", message, 502, @@ -2794,6 +2845,7 @@ async fn execute_stream_from_direct_passthrough( } let lifecycle_seed = build_lifecycle_usage_seed(&plan, report_context.as_ref()); + let max_stream_body_buffer_bytes = resolve_stream_body_buffer_limit(state).await; let request_candidate_status_snapshot = snapshot_local_request_candidate_status(&plan, report_context.as_ref()); let passthrough_mode = direct_passthrough_mode(); @@ -2889,6 +2941,7 @@ async fn execute_stream_from_direct_passthrough( stream_usage_observer, stream_usage_observer_buffered: Vec::new(), provider_error_inspection: ProviderStreamErrorInspection::default(), + max_stream_body_buffer_bytes, provider_buffered_body: Vec::new(), buffered_body: Vec::new(), provider_body_truncated: false, @@ -2958,7 +3011,6 @@ async fn execute_stream_from_direct_passthrough( StageElapsedGuard::from_started_at("stream_total", stream_started_at_for_report); let _provider_pool_in_flight_guard = provider_pool_in_flight_guard_for_report; let _upstream_target_permit = upstream_target_permit; - let max_stream_body_buffer_bytes = DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES; let stream_usage_report_context = normalized_stream_report_context_owned.clone().or_else(|| { Some(serde_json::json!({ @@ -3014,7 +3066,7 @@ async fn execute_stream_from_direct_passthrough( match result { Ok(item) => item, Err(timeout) => { - terminal_failure = Some(build_stream_failure_report( + terminal_failure = Some(build_stream_transport_failure_report( "first_byte_timeout", stream_first_byte_timeout_message(timeout), 504, @@ -3063,7 +3115,7 @@ async fn execute_stream_from_direct_passthrough( error = %message, "gateway direct passthrough upstream body read failed" ); - terminal_failure = Some(build_stream_failure_report( + terminal_failure = Some(build_stream_transport_failure_report( "execution_runtime_stream_read_error", message, 502, @@ -3308,6 +3360,8 @@ async fn execute_stream_from_direct_passthrough( let report_context_for_payload = report_context_with_request_diagnostics( report_context_for_payload, request_diagnostics_for_report.as_ref(), + stream_started_at_for_report, + terminal_telemetry.as_ref(), ); let usage_payload = build_stream_usage_payload( trace_id_owned, @@ -3368,6 +3422,8 @@ async fn execute_stream_from_direct_passthrough( let report_context_for_payload = report_context_with_request_diagnostics( report_context_for_payload, request_diagnostics_for_report.as_ref(), + stream_started_at_for_report, + terminal_telemetry.as_ref(), ); submit_midstream_stream_failure( &state_for_report, @@ -3433,6 +3489,8 @@ async fn execute_stream_from_direct_passthrough( let report_context_for_payload = report_context_with_request_diagnostics( report_context_for_payload, request_diagnostics_for_report.as_ref(), + stream_started_at_for_report, + terminal_telemetry.as_ref(), ); let usage_payload = build_stream_usage_payload( trace_id_owned.clone(), @@ -3627,6 +3685,41 @@ pub(crate) fn execute_execution_runtime_stream_with_retry_scope<'a>( }) } +async fn maybe_build_stream_transport_error_stop_response( + state: &AppState, + plan: &ExecutionPlan, + report_context: Option<&Value>, + trace_id: &str, + decision: &GatewayControlDecision, + error_type: &str, + error_message: &str, + elapsed_ms: u64, +) -> Result>, GatewayError> { + let analysis = crate::orchestration::resolve_local_transport_failover_analysis_for_attempt( + state, + plan, + report_context, + ) + .await; + if !matches!(analysis.decision, LocalFailoverDecision::StopLocalFailover) { + return Ok(None); + } + + crate::execution_runtime::build_transport_error_stop_response( + state, + plan, + report_context, + trace_id, + decision, + http::StatusCode::BAD_GATEWAY.as_u16(), + error_type, + error_message, + elapsed_ms, + ) + .await + .map(Some) +} + async fn execute_execution_runtime_stream_inner( state: &AppState, mut plan: ExecutionPlan, @@ -3731,6 +3824,7 @@ async fn execute_execution_runtime_stream_inner( } Ok(None) => {} Err(err) => { + let transport_error_message = err.to_string(); info!( event_name = "grok_execution_unavailable", log_type = "ops", @@ -3754,13 +3848,27 @@ async fn execute_execution_runtime_stream_inner( status: RequestCandidateStatus::Failed, status_code: None, error_type: Some("grok_execution_unavailable".to_string()), - error_message: Some(format!("{err:?}")), - latency_ms: None, + error_message: Some(transport_error_message.clone()), + latency_ms: Some(stream_elapsed_ms_since(stream_started_at)), started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; + if let Some(response) = maybe_build_stream_transport_error_stop_response( + state, + &plan, + report_context.as_ref(), + trace_id, + decision, + "grok_execution_unavailable", + transport_error_message.as_str(), + stream_elapsed_ms_since(stream_started_at), + ) + .await? + { + return Ok(Some(response)); + } return Ok(None); } } @@ -3788,6 +3896,7 @@ async fn execute_execution_runtime_stream_inner( } Ok(None) => {} Err(err) => { + let transport_error_message = err.to_string(); info!( event_name = "windsurf_native_execution_unavailable", log_type = "ops", @@ -3811,13 +3920,27 @@ async fn execute_execution_runtime_stream_inner( status: RequestCandidateStatus::Failed, status_code: None, error_type: Some("windsurf_native_execution_unavailable".to_string()), - error_message: Some(err.to_string()), - latency_ms: None, + error_message: Some(transport_error_message.clone()), + latency_ms: Some(stream_elapsed_ms_since(stream_started_at)), started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; + if let Some(response) = maybe_build_stream_transport_error_stop_response( + state, + &plan, + report_context.as_ref(), + trace_id, + decision, + "windsurf_native_execution_unavailable", + transport_error_message.as_str(), + stream_elapsed_ms_since(stream_started_at), + ) + .await? + { + return Ok(Some(response)); + } return Ok(None); } } @@ -3845,6 +3968,7 @@ async fn execute_execution_runtime_stream_inner( } Ok(None) => {} Err(err) => { + let transport_error_message = err.to_string(); info!( event_name = "kiro_web_search_mcp_unavailable", log_type = "ops", @@ -3868,13 +3992,27 @@ async fn execute_execution_runtime_stream_inner( status: RequestCandidateStatus::Failed, status_code: None, error_type: Some("kiro_web_search_mcp_unavailable".to_string()), - error_message: Some(format!("{err:?}")), - latency_ms: None, + error_message: Some(transport_error_message.clone()), + latency_ms: Some(stream_elapsed_ms_since(stream_started_at)), started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; + if let Some(response) = maybe_build_stream_transport_error_stop_response( + state, + &plan, + report_context.as_ref(), + trace_id, + decision, + "kiro_web_search_mcp_unavailable", + transport_error_message.as_str(), + stream_elapsed_ms_since(stream_started_at), + ) + .await? + { + return Ok(Some(response)); + } return Ok(None); } } @@ -3902,6 +4040,7 @@ async fn execute_execution_runtime_stream_inner( } Ok(None) => {} Err(err) => { + let transport_error_message = err.to_string(); info!( event_name = "chatgpt_web_image_execution_unavailable", log_type = "ops", @@ -3925,13 +4064,27 @@ async fn execute_execution_runtime_stream_inner( status: RequestCandidateStatus::Failed, status_code: None, error_type: Some("chatgpt_web_image_execution_unavailable".to_string()), - error_message: Some(format!("{err:?}")), - latency_ms: None, + error_message: Some(transport_error_message.clone()), + latency_ms: Some(stream_elapsed_ms_since(stream_started_at)), started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; + if let Some(response) = maybe_build_stream_transport_error_stop_response( + state, + &plan, + report_context.as_ref(), + trace_id, + decision, + "chatgpt_web_image_execution_unavailable", + transport_error_message.as_str(), + stream_elapsed_ms_since(stream_started_at), + ) + .await? + { + return Ok(Some(response)); + } return Ok(None); } } @@ -3961,6 +4114,7 @@ async fn execute_execution_runtime_stream_inner( return Err(err); } Err(InProcessStreamExecutionError::Transport(err)) => { + let transport_error_message = err.to_string(); info!( event_name = "stream_execution_runtime_unavailable", log_type = "ops", @@ -3984,13 +4138,27 @@ async fn execute_execution_runtime_stream_inner( status: RequestCandidateStatus::Failed, status_code: None, error_type: Some("execution_runtime_unavailable".to_string()), - error_message: Some(format!("{err:?}")), - latency_ms: None, + error_message: Some(transport_error_message.clone()), + latency_ms: Some(stream_elapsed_ms_since(stream_started_at)), started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; + if let Some(response) = maybe_build_stream_transport_error_stop_response( + state, + &plan, + report_context.as_ref(), + trace_id, + decision, + "execution_runtime_unavailable", + transport_error_message.as_str(), + stream_elapsed_ms_since(stream_started_at), + ) + .await? + { + return Ok(Some(response)); + } return Ok(None); } }; @@ -4076,6 +4244,7 @@ async fn execute_execution_runtime_stream_inner( return Err(err); } Err(InProcessStreamExecutionError::Transport(err)) => { + let transport_error_message = err.to_string(); info!( event_name = "stream_execution_runtime_unavailable", log_type = "ops", @@ -4099,13 +4268,27 @@ async fn execute_execution_runtime_stream_inner( status: RequestCandidateStatus::Failed, status_code: None, error_type: Some("execution_runtime_unavailable".to_string()), - error_message: Some(err.to_string()), - latency_ms: None, + error_message: Some(transport_error_message.clone()), + latency_ms: Some(stream_elapsed_ms_since(stream_started_at)), started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; + if let Some(response) = maybe_build_stream_transport_error_stop_response( + state, + &plan, + report_context.as_ref(), + trace_id, + decision, + "execution_runtime_unavailable", + transport_error_message.as_str(), + stream_elapsed_ms_since(stream_started_at), + ) + .await? + { + return Ok(Some(response)); + } return Ok(None); } }; @@ -4177,6 +4360,7 @@ async fn execute_execution_runtime_stream_inner( { Ok(response) => response, Err(err) => { + let transport_error_message = format!("{err:?}"); warn!( event_name = "stream_execution_runtime_remote_unavailable", log_type = "ops", @@ -4195,13 +4379,27 @@ async fn execute_execution_runtime_stream_inner( status: RequestCandidateStatus::Failed, status_code: None, error_type: Some("execution_runtime_unavailable".to_string()), - error_message: Some(format!("{err:?}")), - latency_ms: None, + error_message: Some(transport_error_message.clone()), + latency_ms: Some(stream_elapsed_ms_since(stream_started_at)), started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; + if let Some(response) = maybe_build_stream_transport_error_stop_response( + state, + &plan, + report_context.as_ref(), + trace_id, + decision, + "execution_runtime_unavailable", + transport_error_message.as_str(), + stream_elapsed_ms_since(stream_started_at), + ) + .await? + { + return Ok(Some(response)); + } return Ok(None); } }; @@ -5315,6 +5513,7 @@ async fn execute_stream_from_frame_stream_with_retry_scope( if !lifecycle_pending_recorded { record_stream_pending_lifecycle(state, &lifecycle_seed, &mut stage_trace).await; } + let max_stream_body_buffer_bytes = resolve_stream_body_buffer_limit(state).await; let request_candidate_status_snapshot = snapshot_local_request_candidate_status(&plan, report_context.as_ref()); let candidate_index = parse_request_candidate_report_context(report_context.as_ref()) @@ -5912,7 +6111,10 @@ async fn execute_stream_from_frame_stream_with_retry_scope( headers, prefetched_usage_telemetry.clone(), &provider_prefetched_body, + candidate_started_unix_secs, + stream_elapsed_ms_since(stream_started_at), failure, + None, ) .await; } @@ -5984,7 +6186,10 @@ async fn execute_stream_from_frame_stream_with_retry_scope( headers, prefetched_usage_telemetry.clone(), &prefetched_body, + candidate_started_unix_secs, + stream_elapsed_ms_since(stream_started_at), failure, + None, ) .await; } @@ -6110,10 +6315,17 @@ async fn execute_stream_from_frame_stream_with_retry_scope( provider_prefetched_body_bytes = provider_prefetched_body.len(), "gateway detected embedded error while prefetching execution runtime stream" ); + let request_diagnostics = current_request_diagnostics(); + let terminal_report_context = report_context_with_request_diagnostics( + report_context, + request_diagnostics.as_ref(), + stream_started_at, + prefetched_usage_telemetry.as_ref(), + ); let payload = build_stream_sync_payload( trace_id, report_kind.clone(), - report_context, + terminal_report_context, status_code, headers, Some(body_json), @@ -6187,7 +6399,10 @@ async fn execute_stream_from_frame_stream_with_retry_scope( headers, prefetched_usage_telemetry.clone(), &provider_prefetched_body, + candidate_started_unix_secs, + stream_elapsed_ms_since(stream_started_at), failure, + None, ) .await; } @@ -6220,7 +6435,10 @@ async fn execute_stream_from_frame_stream_with_retry_scope( headers, prefetched_usage_telemetry.clone(), &provider_prefetched_body, + candidate_started_unix_secs, + stream_elapsed_ms_since(stream_started_at), failure, + None, ) .await; } @@ -6251,7 +6469,10 @@ async fn execute_stream_from_frame_stream_with_retry_scope( headers, prefetched_usage_telemetry.clone(), &provider_prefetched_body, + candidate_started_unix_secs, + stream_elapsed_ms_since(stream_started_at), failure, + None, ) .await; } @@ -6335,7 +6556,10 @@ async fn execute_stream_from_frame_stream_with_retry_scope( headers, prefetched_usage_telemetry.clone(), &provider_prefetched_body, + candidate_started_unix_secs, + stream_elapsed_ms_since(stream_started_at), build_stream_failure_from_execution_error(&error), + retry_scope_out.as_deref_mut(), ) .await; } @@ -6429,7 +6653,6 @@ async fn execute_stream_from_frame_stream_with_retry_scope( let _stream_total_guard = StageElapsedGuard::from_started_at("stream_total", stream_started_at_for_report); let _provider_pool_in_flight_guard = provider_pool_in_flight_guard_for_report; - let max_stream_body_buffer_bytes = DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES; let mut provider_buffered_body = Vec::new(); let mut buffered_body = Vec::new(); let mut provider_body_truncated = false; @@ -7373,6 +7596,8 @@ async fn execute_stream_from_frame_stream_with_retry_scope( let report_context_for_payload = report_context_with_request_diagnostics( report_context_for_payload, request_diagnostics_for_report.as_ref(), + stream_started_at_for_report, + terminal_telemetry.as_ref(), ); let usage_payload = build_stream_usage_payload( trace_id_owned, @@ -7433,6 +7658,8 @@ async fn execute_stream_from_frame_stream_with_retry_scope( let report_context_for_payload = report_context_with_request_diagnostics( report_context_for_payload, request_diagnostics_for_report.as_ref(), + stream_started_at_for_report, + terminal_telemetry.as_ref(), ); submit_midstream_stream_failure( &state_for_report, @@ -7499,6 +7726,8 @@ async fn execute_stream_from_frame_stream_with_retry_scope( let report_context_for_payload = report_context_with_request_diagnostics( report_context_for_payload, request_diagnostics_for_report.as_ref(), + stream_started_at_for_report, + terminal_telemetry.as_ref(), ); let usage_payload = build_stream_usage_payload( trace_id_owned.clone(), @@ -7686,12 +7915,14 @@ mod tests { StoredUsageSettlement, UsageSettlementInput, }; use aether_data_contracts::repository::usage::UsageReadRepository; - use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UpsertUsageRecord}; + use aether_data_contracts::repository::usage::{ + StoredRequestUsageAudit, UpsertUsageRecord, UsageBodyCaptureState, + }; use aether_data_contracts::DataLayerError; use aether_usage_runtime::{ - UsageBillingEventEnricher, UsageBodyCapturePolicy, UsageEvent, UsageEventData, - UsageEventType, UsageRecordWriter, UsageRuntimeAccess, UsageRuntimeConfig, - UsageSettlementWriter, + apply_usage_body_capture_policy_to_event, UsageBillingEventEnricher, + UsageBodyCapturePolicy, UsageEvent, UsageEventData, UsageEventType, UsageRecordWriter, + UsageRequestRecordLevel, UsageRuntimeAccess, UsageRuntimeConfig, UsageSettlementWriter, }; use async_stream::stream; use async_trait::async_trait; @@ -8058,6 +8289,185 @@ mod tests { .expect("execution should succeed") } + async fn execute_prefetched_transport_failure( + stop_on_transport_errors: bool, + ) -> AiAttemptExecutionOutcome> { + let request_id = if stop_on_transport_errors { + "req-prefetch-transport-stop" + } else { + "req-prefetch-transport-retry" + }; + let plan = native_anthropic_stream_plan(request_id); + let provider_config = stop_on_transport_errors.then(|| { + json!({ + "failover_rules": { + "stop_on_transport_errors": true, + } + }) + }); + let provider_catalog = provider_catalog_for_plan(&plan, provider_config); + let data_state = crate::data::GatewayDataState::with_provider_transport_reader_for_tests( + Arc::new(provider_catalog), + "development-key", + ); + let state = AppState::new() + .expect("app state should build") + .with_data_state_for_tests(data_state); + let frame_stream = stream! { + yield Ok::(ndjson_frame(StreamFrame { + frame_type: StreamFrameType::Headers, + payload: StreamFramePayload::Headers { + status_code: 200, + headers: BTreeMap::from([( + "content-type".to_string(), + "text/event-stream".to_string(), + )]), + }, + })); + yield Ok::(ndjson_frame(StreamFrame { + frame_type: StreamFrameType::Error, + payload: StreamFramePayload::Error { + error: ExecutionError { + kind: ExecutionErrorKind::Internal, + phase: ExecutionPhase::StreamRead, + message: "connection reset before first body byte".to_string(), + upstream_status: None, + retryable: true, + failover_recommended: true, + }, + }, + })); + } + .boxed(); + let mut retry_scope = AiAttemptRetryScope::Provider; + let response = execute_stream_from_frame_stream_with_retry_scope( + &state, + plan, + &format!("trace-{request_id}"), + &test_decision(), + "claude_chat_stream", + Some("claude_chat_stream_success".to_string()), + Some(json!({ + "request_id": request_id, + "candidate_id": format!("candidate-{request_id}"), + "candidate_index": 0, + "retry_index": 0, + "provider_api_format": "claude:messages", + "client_api_format": "claude:messages" + })), + crate::clock::current_unix_ms(), + Instant::now(), + RequestStageTrace::from_env(), + true, + frame_stream, + false, + None, + Some(&mut retry_scope), + None, + ) + .await + .expect("prefetch transport execution should resolve"); + + match response { + Some(response) => AiAttemptExecutionOutcome::Responded(response), + None => AiAttemptExecutionOutcome::Retry { + scope: retry_scope, + fallback_response: None, + }, + } + } + + async fn execute_prefetched_http_status_failure( + continue_failover: bool, + ) -> AiAttemptExecutionOutcome> { + let request_id = if continue_failover { + "req-prefetch-http-continue" + } else { + "req-prefetch-http-stop" + }; + let plan = native_anthropic_stream_plan(request_id); + let failover_rules = if continue_failover { + json!({"continue_status_codes": [500]}) + } else { + json!({"stop_status_codes": [500]}) + }; + let provider_catalog = provider_catalog_for_plan( + &plan, + Some(json!({ + "failover_rules": failover_rules, + })), + ); + let data_state = crate::data::GatewayDataState::with_provider_transport_reader_for_tests( + Arc::new(provider_catalog), + "development-key", + ); + let state = AppState::new() + .expect("app state should build") + .with_data_state_for_tests(data_state); + let frame_stream = stream! { + yield Ok::(ndjson_frame(StreamFrame { + frame_type: StreamFrameType::Headers, + payload: StreamFramePayload::Headers { + status_code: 200, + headers: BTreeMap::from([( + "content-type".to_string(), + "text/event-stream".to_string(), + )]), + }, + })); + yield Ok::(ndjson_frame(StreamFrame { + frame_type: StreamFrameType::Error, + payload: StreamFramePayload::Error { + error: ExecutionError { + kind: ExecutionErrorKind::Internal, + phase: ExecutionPhase::StreamRead, + message: "upstream returned 500 before the first body byte".to_string(), + upstream_status: Some(500), + retryable: true, + failover_recommended: true, + }, + }, + })); + } + .boxed(); + let mut retry_scope = AiAttemptRetryScope::Provider; + let response = execute_stream_from_frame_stream_with_retry_scope( + &state, + plan, + &format!("trace-{request_id}"), + &test_decision(), + "claude_chat_stream", + Some("claude_chat_stream_success".to_string()), + Some(json!({ + "request_id": request_id, + "candidate_id": format!("candidate-{request_id}"), + "candidate_index": 0, + "retry_index": 0, + "provider_api_format": "claude:messages", + "client_api_format": "claude:messages" + })), + crate::clock::current_unix_ms(), + Instant::now(), + RequestStageTrace::from_env(), + true, + frame_stream, + false, + None, + Some(&mut retry_scope), + None, + ) + .await + .expect("prefetch HTTP status execution should resolve"); + + match response { + Some(response) => AiAttemptExecutionOutcome::Responded(response), + None => AiAttemptExecutionOutcome::Retry { + scope: retry_scope, + fallback_response: None, + }, + } + } + fn native_anthropic_stream_plan(request_id: &str) -> ExecutionPlan { ExecutionPlan { request_id: request_id.to_string(), @@ -8126,6 +8536,7 @@ mod tests { stream_usage_observer: None, stream_usage_observer_buffered: Vec::new(), provider_error_inspection: ProviderStreamErrorInspection::default(), + max_stream_body_buffer_bytes: super::DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES, provider_buffered_body: Vec::new(), buffered_body: Vec::new(), provider_body_truncated: false, @@ -9240,6 +9651,7 @@ mod tests { stream_usage_observer: None, stream_usage_observer_buffered: Vec::new(), provider_error_inspection: ProviderStreamErrorInspection::default(), + max_stream_body_buffer_bytes: super::DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES, provider_buffered_body: Vec::new(), buffered_body: Vec::new(), provider_body_truncated: false, @@ -9655,6 +10067,79 @@ mod tests { assert!(truncated); } + #[test] + fn stream_capture_policy_keeps_full_unbounded_and_caps_basic_analysis_buffer() { + let oversized_chunk = vec![b'x'; super::BASIC_STREAM_BODY_ANALYSIS_LIMIT_BYTES + 1]; + + let full_limit = + super::stream_body_buffer_limit_for_record_level(UsageRequestRecordLevel::Full); + assert_eq!(full_limit, usize::MAX); + let mut full_buffer = Vec::new(); + let mut full_truncated = false; + super::append_stream_capture_bytes( + &mut full_buffer, + &oversized_chunk, + full_limit, + &mut full_truncated, + ); + assert_eq!(full_buffer, oversized_chunk); + assert!(!full_truncated); + let (full_body, full_state) = + super::build_stream_body_capture(&full_buffer, full_truncated); + assert!(full_body.is_some()); + assert_eq!(full_state, Some(UsageBodyCaptureState::Inline)); + drop(full_body); + + let basic_limit = + super::stream_body_buffer_limit_for_record_level(UsageRequestRecordLevel::Basic); + assert_eq!(basic_limit, super::BASIC_STREAM_BODY_ANALYSIS_LIMIT_BYTES); + let mut basic_buffer = Vec::new(); + let mut basic_truncated = false; + super::append_stream_capture_bytes( + &mut basic_buffer, + &oversized_chunk, + basic_limit, + &mut basic_truncated, + ); + assert_eq!(basic_buffer.len(), basic_limit); + assert!(basic_truncated); + let (basic_body, basic_state) = + super::build_stream_body_capture(&basic_buffer, basic_truncated); + assert!(basic_body.is_some()); + assert_eq!(basic_state, Some(UsageBodyCaptureState::Truncated)); + + let mut event = UsageEvent::new( + UsageEventType::Completed, + "req-basic-stream-capture", + UsageEventData { + provider_name: "provider".to_string(), + model: "model".to_string(), + response_body: basic_body.map(Value::String), + response_body_state: basic_state, + client_response_body: Some(json!("captured client body")), + client_response_body_state: Some(UsageBodyCaptureState::Truncated), + ..UsageEventData::default() + }, + ); + apply_usage_body_capture_policy_to_event( + UsageBodyCapturePolicy { + record_level: UsageRequestRecordLevel::Basic, + }, + &mut event, + ); + + assert_eq!(event.data.response_body, None); + assert_eq!( + event.data.response_body_state, + Some(UsageBodyCaptureState::Disabled) + ); + assert_eq!(event.data.client_response_body, None); + assert_eq!( + event.data.client_response_body_state, + Some(UsageBodyCaptureState::Disabled) + ); + } + #[test] fn provider_error_inspection_detects_response_failed_at_every_chunk_boundary() { let body = concat!( @@ -9719,6 +10204,45 @@ mod tests { ); } + #[tokio::test] + async fn prefetched_transport_failure_retries_by_default() { + assert!(matches!( + execute_prefetched_transport_failure(false).await, + AiAttemptExecutionOutcome::Retry { + scope: AiAttemptRetryScope::Candidate, + fallback_response: None, + } + )); + } + + #[tokio::test] + async fn prefetched_transport_failure_can_stop_without_matching_http_status_rules() { + let AiAttemptExecutionOutcome::Responded(response) = + execute_prefetched_transport_failure(true).await + else { + panic!("transport stop policy should return a local response"); + }; + assert_eq!(response.status(), StatusCode::BAD_GATEWAY); + } + + #[tokio::test] + async fn prefetched_http_error_frame_honors_continue_status_codes() { + assert!(matches!( + execute_prefetched_http_status_failure(true).await, + AiAttemptExecutionOutcome::Retry { .. } + )); + } + + #[tokio::test] + async fn prefetched_http_error_frame_honors_stop_status_codes() { + let AiAttemptExecutionOutcome::Responded(response) = + execute_prefetched_http_status_failure(false).await + else { + panic!("HTTP stop policy should return the upstream error"); + }; + assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR); + } + fn tunnel_proxy_snapshot(base_url: String) -> aether_contracts::ProxySnapshot { aether_contracts::ProxySnapshot { enabled: Some(true), diff --git a/apps/aether-gateway/src/execution_runtime/stream/execution_failures.rs b/apps/aether-gateway/src/execution_runtime/stream/execution_failures.rs index 57e1df1b2..117bdb71d 100644 --- a/apps/aether-gateway/src/execution_runtime/stream/execution_failures.rs +++ b/apps/aether-gateway/src/execution_runtime/stream/execution_failures.rs @@ -1,5 +1,7 @@ use aether_ai_serving::AiAttemptRetryScope; -use aether_contracts::{ExecutionError, ExecutionPlan, ExecutionTelemetry}; +use aether_contracts::{ + ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionPlan, ExecutionTelemetry, +}; use aether_data_contracts::repository::candidates::RequestCandidateStatus; use aether_scheduler_core::SchedulerRequestCandidateStatusUpdate; use aether_usage_runtime::{ @@ -22,13 +24,14 @@ use crate::execution_runtime::submission::{ use crate::log_ids::short_request_id; use crate::orchestration::{ apply_local_execution_effect, classify_failure_disposition, - resolve_local_failover_analysis_for_attempt, with_upstream_response_report_context, + resolve_local_failover_analysis_for_attempt, + resolve_local_transport_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; +use crate::request_diagnostics::attach_current_request_diagnostics_and_candidate_timing_to_report_context; use crate::usage::submit_sync_report; use crate::{usage::GatewaySyncReportRequest, AppState, GatewayError}; @@ -37,6 +40,9 @@ pub(super) struct StreamFailureReport { pub(super) status_code: u16, pub(super) error_type: String, pub(super) error_message: String, + upstream_status_code: Option, + transport_error: bool, + honor_http_failover: bool, extra_error_fields: Map, provider_body_json: Option, } @@ -68,6 +74,9 @@ impl StreamFailureReport { status_code, error_type, error_message, + upstream_status_code: _, + transport_error: _, + honor_http_failover: _, mut extra_error_fields, provider_body_json, } = self; @@ -110,6 +119,26 @@ pub(super) fn build_stream_failure_report( status_code, error_type, error_message, + upstream_status_code: Some(status_code), + transport_error: false, + honor_http_failover: false, + extra_error_fields: Map::new(), + provider_body_json: None, + } +} + +pub(super) fn build_stream_transport_failure_report( + error_type: impl Into, + error_message: impl Into, + status_code: u16, +) -> StreamFailureReport { + StreamFailureReport { + status_code, + error_type: error_type.into(), + error_message: error_message.into(), + upstream_status_code: None, + transport_error: true, + honor_http_failover: false, extra_error_fields: Map::new(), provider_body_json: None, } @@ -118,7 +147,19 @@ pub(super) fn build_stream_failure_report( pub(super) fn build_stream_failure_from_execution_error( error: &ExecutionError, ) -> StreamFailureReport { - let status_code = error.upstream_status.unwrap_or(502); + let transport_error = execution_error_is_transport(error); + let status_code = error.upstream_status.unwrap_or_else(|| { + if matches!( + error.kind, + ExecutionErrorKind::ConnectTimeout + | ExecutionErrorKind::FirstByteTimeout + | ExecutionErrorKind::ReadTimeout + ) { + 504 + } else { + 502 + } + }); let error_type = serde_json::to_value(&error.kind) .ok() .and_then(|value| value.as_str().map(ToOwned::to_owned)) @@ -141,6 +182,9 @@ pub(super) fn build_stream_failure_from_execution_error( status_code, error_type, error_message, + upstream_status_code: error.upstream_status, + transport_error, + honor_http_failover: error.upstream_status.is_some(), extra_error_fields: error_object, provider_body_json: None, } @@ -168,11 +212,40 @@ pub(super) fn build_stream_failure_from_provider_error_body( status_code, error_type, error_message, + upstream_status_code: Some(status_code), + transport_error: false, + honor_http_failover: true, extra_error_fields: Map::new(), provider_body_json: Some(body_json.clone()), } } +fn execution_error_is_transport(error: &ExecutionError) -> bool { + if error.upstream_status.is_some() { + return false; + } + let explicit_transport_kind = matches!( + error.kind, + ExecutionErrorKind::ConnectTimeout + | ExecutionErrorKind::FirstByteTimeout + | ExecutionErrorKind::ReadTimeout + | ExecutionErrorKind::TlsError + | ExecutionErrorKind::ProxyError + | ExecutionErrorKind::ProtocolError + ); + let retryable_internal_transport_phase = matches!(error.kind, ExecutionErrorKind::Internal) + && (error.retryable || error.failover_recommended) + && matches!( + error.phase, + ExecutionPhase::Connect + | ExecutionPhase::Handshake + | ExecutionPhase::Write + | ExecutionPhase::FirstByte + | ExecutionPhase::StreamRead + ); + explicit_transport_kind || retryable_internal_transport_phase +} + fn first_non_empty_error_text( error_object: Option<&Map>, body_object: Option<&Map>, @@ -205,6 +278,8 @@ fn build_stream_failure_sync_payload( failure: StreamFailureReport, ) -> GatewaySyncReportRequest { let status_code = failure.status_code; + let upstream_status_code = failure.upstream_status_code; + let transport_error = failure.transport_error; let (body, client_body) = failure.into_body_jsons(); headers.retain(|name, _| { !name.eq_ignore_ascii_case("content-encoding") @@ -212,23 +287,31 @@ fn build_stream_failure_sync_payload( && !name.eq_ignore_ascii_case("content-type") }); headers.insert("content-type".to_string(), "application/json".to_string()); - let report_context = with_upstream_response_report_context( - report_context.as_ref(), - status_code, - Some(&headers), - Some(&body), - None, - None, - ) - .or(report_context); + let report_context = upstream_status_code + .and_then(|upstream_status_code| { + with_upstream_response_report_context( + report_context.as_ref(), + upstream_status_code, + Some(&headers), + Some(&body), + None, + None, + ) + }) + .or(report_context); let report_context = report_context.map(|mut context| { if let Some(object) = context.as_object_mut() { let response_headers = serde_json::to_value(&headers).unwrap_or(Value::Null); - object.insert( - "provider_response_headers".to_string(), - response_headers.clone(), - ); + if upstream_status_code.is_some() { + object.insert( + "provider_response_headers".to_string(), + response_headers.clone(), + ); + } object.insert("client_response_headers".to_string(), response_headers); + if transport_error { + object.insert("transport_error".to_string(), Value::Bool(true)); + } } context }); @@ -265,6 +348,7 @@ async fn record_stream_sync_failure( plan: &ExecutionPlan, report_context: Option<&Value>, payload: &GatewaySyncReportRequest, + candidate_status_code: Option, started_at_unix_ms: Option, handling: StreamFailureHandling, ) -> LocalFailoverAnalysis { @@ -361,8 +445,19 @@ async fn record_stream_sync_failure( LocalFailoverDecision::RetryNextCandidate ); if !matches!(handling, StreamFailureHandling::HonorLocalFailover) || !retrying_next_candidate { + crate::execution_runtime::mark_stream_candidate_watchdog_terminal_started(); let report_context_with_diagnostics = - attach_current_request_diagnostics_to_report_context(report_context); + attach_current_request_diagnostics_and_candidate_timing_to_report_context( + report_context, + payload + .telemetry + .as_ref() + .and_then(|telemetry| telemetry.elapsed_ms), + payload + .telemetry + .as_ref() + .and_then(|telemetry| telemetry.ttfb_ms), + ); let context_seed = build_terminal_usage_context_seed( plan, report_context_with_diagnostics.as_ref().or(report_context), @@ -383,7 +478,7 @@ async fn record_stream_sync_failure( report_context, SchedulerRequestCandidateStatusUpdate { status: RequestCandidateStatus::Failed, - status_code: Some(payload.status_code), + status_code: candidate_status_code, error_type: Some(error_type.to_string()), error_message: Some(error_message.to_string()), latency_ms: payload @@ -439,6 +534,7 @@ pub(super) async fn handle_prefetch_provider_private_stream_error( plan, payload.report_context.as_ref(), &payload, + Some(status_code), None, StreamFailureHandling::HonorLocalFailover, ) @@ -505,9 +601,15 @@ pub(super) async fn handle_prefetch_stream_failure( headers: std::collections::BTreeMap, telemetry: Option, buffered_body: &[u8], + candidate_started_unix_ms: u64, + candidate_elapsed_ms: u64, failure: StreamFailureReport, + retry_scope_out: Option<&mut AiAttemptRetryScope>, ) -> Result>, GatewayError> { - let payload = build_stream_failure_sync_payload( + let transport_error = failure.transport_error; + let candidate_status_code = failure.upstream_status_code; + let honor_http_failover = failure.honor_http_failover; + let mut payload = build_stream_failure_sync_payload( trace_id, report_kind.to_string(), report_context, @@ -516,15 +618,179 @@ pub(super) async fn handle_prefetch_stream_failure( buffered_body, failure, ); - record_stream_sync_failure( + if transport_error { + let telemetry = payload.telemetry.get_or_insert(ExecutionTelemetry { + ttfb_ms: None, + elapsed_ms: None, + upstream_bytes: None, + }); + telemetry.elapsed_ms.get_or_insert(candidate_elapsed_ms); + return handle_prefetch_transport_stream_failure( + state, + trace_id, + decision, + plan, + request_id, + candidate_id, + payload, + candidate_started_unix_ms, + candidate_elapsed_ms, + retry_scope_out, + ) + .await; + } + let failure_analysis = record_stream_sync_failure( state, plan, payload.report_context.as_ref(), &payload, + candidate_status_code, None, - StreamFailureHandling::Terminal, + if honor_http_failover { + StreamFailureHandling::HonorLocalFailover + } else { + StreamFailureHandling::Terminal + }, ) .await; + if honor_http_failover + && matches!( + failure_analysis.decision, + LocalFailoverDecision::RetryNextCandidate + ) + { + let failure_disposition = classify_failure_disposition( + &plan.provider_api_format, + failure_analysis.classification, + payload.status_code, + ); + if let Some(retry_scope) = retry_scope_out { + *retry_scope = ai_attempt_retry_scope_from_failure_disposition(failure_disposition); + } + warn!( + event_name = "local_stream_candidate_retry_scheduled", + log_type = "event", + trace_id = %trace_id, + request_id = %request_id, + candidate_id = ?candidate_id, + status_code = payload.status_code, + failover_classification = failure_analysis.classification.as_str(), + "gateway local stream decision retrying next candidate after prefetched execution error" + ); + return Ok(None); + } + + let response = + submit_local_core_error_or_sync_finalize(state, trace_id, decision, payload).await?; + Ok(Some(attach_control_metadata_headers( + response, + Some(request_id), + candidate_id, + )?)) +} + +#[allow(clippy::too_many_arguments)] +async fn handle_prefetch_transport_stream_failure( + state: &AppState, + trace_id: &str, + decision: &GatewayControlDecision, + plan: &ExecutionPlan, + request_id: &str, + candidate_id: Option<&str>, + payload: GatewaySyncReportRequest, + candidate_started_unix_ms: u64, + candidate_elapsed_ms: u64, + retry_scope_out: Option<&mut AiAttemptRetryScope>, +) -> Result>, GatewayError> { + let error_type = stream_failure_body_field(&payload, "type").unwrap_or("internal"); + let error_message = stream_failure_body_field(&payload, "message").unwrap_or_default(); + if matches!(error_type, "first_byte_timeout" | "read_timeout") { + apply_local_execution_effect( + state, + LocalExecutionEffectContext { + plan, + report_context: payload.report_context.as_ref(), + }, + LocalExecutionEffect::PoolStreamTimeout, + ) + .await; + } + + let analysis = resolve_local_transport_failover_analysis_for_attempt( + state, + plan, + payload.report_context.as_ref(), + ) + .await; + let retrying_next_candidate = + matches!(analysis.decision, LocalFailoverDecision::RetryNextCandidate); + if !retrying_next_candidate { + crate::execution_runtime::mark_stream_candidate_watchdog_terminal_started(); + let report_context_with_diagnostics = + attach_current_request_diagnostics_and_candidate_timing_to_report_context( + payload.report_context.as_ref(), + payload + .telemetry + .as_ref() + .and_then(|telemetry| telemetry.elapsed_ms) + .or(Some(candidate_elapsed_ms)), + payload + .telemetry + .as_ref() + .and_then(|telemetry| telemetry.ttfb_ms), + ); + let context_seed = build_terminal_usage_context_seed( + plan, + report_context_with_diagnostics + .as_ref() + .or(payload.report_context.as_ref()), + ); + let payload_seed = build_sync_terminal_usage_payload_seed(&payload); + state + .usage_runtime + .record_sync_terminal( + state.usage_lifecycle_data_state().as_ref(), + context_seed, + payload_seed, + ) + .await; + } + + let terminal_unix_ms = current_request_candidate_unix_ms(); + record_report_request_candidate_status( + state, + payload.report_context.as_ref(), + SchedulerRequestCandidateStatusUpdate { + status: RequestCandidateStatus::Failed, + status_code: None, + error_type: Some(error_type.to_string()), + error_message: Some(error_message.to_string()), + latency_ms: payload + .telemetry + .as_ref() + .and_then(|telemetry| telemetry.elapsed_ms) + .or(Some(candidate_elapsed_ms)), + started_at_unix_ms: Some(candidate_started_unix_ms), + finished_at_unix_ms: Some(terminal_unix_ms), + }, + ) + .await; + + if retrying_next_candidate { + if let Some(retry_scope) = retry_scope_out { + *retry_scope = AiAttemptRetryScope::Candidate; + } + warn!( + event_name = "local_stream_transport_retry_scheduled", + log_type = "event", + trace_id = %trace_id, + request_id = %request_id, + candidate_id = ?candidate_id, + transport_classification = analysis.classification.as_str(), + "gateway retrying next candidate after precommit transport failure" + ); + return Ok(None); + } let response = submit_local_core_error_or_sync_finalize(state, trace_id, decision, payload).await?; @@ -553,6 +819,7 @@ pub(super) async fn submit_midstream_stream_failure( return; }; + let candidate_status_code = failure.upstream_status_code; let payload = build_stream_failure_sync_payload( trace_id, report_kind, @@ -567,6 +834,7 @@ pub(super) async fn submit_midstream_stream_failure( plan, payload.report_context.as_ref(), &payload, + candidate_status_code, Some(started_at_unix_ms), StreamFailureHandling::Terminal, ) @@ -590,10 +858,60 @@ pub(super) async fn submit_midstream_stream_failure( mod tests { use std::collections::BTreeMap; + use aether_contracts::{ExecutionError, ExecutionErrorKind, ExecutionPhase}; use base64::Engine as _; use serde_json::json; - use super::{build_stream_failure_from_provider_error_body, build_stream_failure_sync_payload}; + use super::{ + build_stream_failure_from_execution_error, build_stream_failure_from_provider_error_body, + build_stream_failure_sync_payload, build_stream_transport_failure_report, + }; + + #[test] + fn committed_transport_failure_has_no_upstream_status() { + for status_code in [502, 504] { + let failure = build_stream_transport_failure_report( + "execution_runtime_stream_read_error", + "upstream disconnected", + status_code, + ); + + assert_eq!(failure.status_code, status_code); + assert_eq!(failure.upstream_status_code, None); + assert!(failure.transport_error); + assert!(!failure.honor_http_failover); + } + } + + #[test] + fn precommit_protocol_error_is_transport_without_upstream_status() { + let failure = build_stream_failure_from_execution_error(&ExecutionError { + kind: ExecutionErrorKind::ProtocolError, + phase: ExecutionPhase::StreamRead, + message: "connection reset".to_string(), + upstream_status: None, + retryable: true, + failover_recommended: true, + }); + + assert!(failure.transport_error); + assert_eq!(failure.upstream_status_code, None); + assert_eq!(failure.status_code, 502); + } + + #[test] + fn cancelled_stream_is_not_reclassified_as_transport_retry() { + let failure = build_stream_failure_from_execution_error(&ExecutionError { + kind: ExecutionErrorKind::Cancelled, + phase: ExecutionPhase::StreamRead, + message: "downstream cancelled".to_string(), + upstream_status: None, + retryable: true, + failover_recommended: true, + }); + + assert!(!failure.transport_error); + } #[test] fn midstream_failure_trace_uses_terminal_error_instead_of_buffered_sse() { diff --git a/apps/aether-gateway/src/execution_runtime/stream_pump.rs b/apps/aether-gateway/src/execution_runtime/stream_pump.rs index 3f24c10ed..ce7515ffd 100644 --- a/apps/aether-gateway/src/execution_runtime/stream_pump.rs +++ b/apps/aether-gateway/src/execution_runtime/stream_pump.rs @@ -163,7 +163,7 @@ pub(crate) fn build_direct_execution_frame_stream( let error_frame = if let Some(timeout) = first_byte_timeout { encode_first_byte_timeout_frame(timeout) } else { - encode_error_frame(status_code, message) + encode_error_frame(message) }; match error_frame { Ok(frame) => yield Ok(frame), @@ -245,7 +245,7 @@ pub(crate) fn build_direct_execution_frame_stream( error = %message, "upstream body stream read error" ); - match encode_error_frame(status_code, message) { + match encode_error_frame(message) { Ok(frame) => yield Ok(frame), Err(encode_err) => { yield Err(encode_err); @@ -329,7 +329,7 @@ pub(crate) fn build_direct_execution_frame_stream( error = %message, "upstream body stream read error" ); - match encode_error_frame(status_code, message) { + match encode_error_frame(message) { Ok(frame) => yield Ok(frame), Err(encode_err) => { yield Err(encode_err); @@ -411,7 +411,7 @@ pub(crate) fn build_direct_execution_frame_stream( error = %message, "upstream body stream read error" ); - match encode_error_frame(status_code, message) { + match encode_error_frame(message) { Ok(frame) => yield Ok(frame), Err(encode_err) => { yield Err(encode_err); @@ -493,7 +493,7 @@ pub(crate) fn build_direct_execution_frame_stream( error = %message, "upstream body stream read error" ); - match encode_error_frame(status_code, message) { + match encode_error_frame(message) { Ok(frame) => yield Ok(frame), Err(encode_err) => { yield Err(encode_err); @@ -570,7 +570,7 @@ pub(crate) fn build_direct_execution_frame_stream( error = %message, "upstream body stream read error" ); - match encode_error_frame(status_code, message) { + match encode_error_frame(message) { Ok(frame) => yield Ok(frame), Err(encode_err) => { yield Err(encode_err); @@ -648,17 +648,17 @@ fn encode_data_frame(chunk: &Bytes) -> Result { }) } -fn encode_error_frame(status_code: u16, message: String) -> Result { +fn encode_error_frame(message: String) -> Result { encode_stream_frame_ndjson(&StreamFrame { frame_type: StreamFrameType::Error, payload: StreamFramePayload::Error { error: ExecutionError { - kind: ExecutionErrorKind::Internal, + kind: ExecutionErrorKind::ProtocolError, phase: ExecutionPhase::StreamRead, message, - upstream_status: Some(status_code), - retryable: false, - failover_recommended: false, + upstream_status: None, + retryable: true, + failover_recommended: true, }, }, }) @@ -672,7 +672,7 @@ fn encode_first_byte_timeout_frame(timeout: Duration) -> Result kind: ExecutionErrorKind::FirstByteTimeout, phase: ExecutionPhase::FirstByte, message: stream_first_byte_timeout_message(timeout), - upstream_status: Some(504), + upstream_status: None, retryable: true, failover_recommended: true, }, @@ -1490,6 +1490,13 @@ mod tests { .is_some_and( |message| message.contains("provider stream first byte timeout after 50 ms") )); + let error = error_frame + .get("payload") + .and_then(|payload| payload.get("error")) + .expect("timeout error should exist"); + assert_eq!(error.get("upstream_status"), None); + assert_eq!(error.get("retryable"), Some(&Value::Bool(true))); + assert_eq!(error.get("failover_recommended"), Some(&Value::Bool(true))); } #[tokio::test] diff --git a/apps/aether-gateway/src/execution_runtime/sync/execution.rs b/apps/aether-gateway/src/execution_runtime/sync/execution.rs index 61d35b812..a5ada874c 100644 --- a/apps/aether-gateway/src/execution_runtime/sync/execution.rs +++ b/apps/aether-gateway/src/execution_runtime/sync/execution.rs @@ -82,7 +82,11 @@ use crate::request_candidate_runtime::{ record_local_request_candidate_status, record_local_request_candidate_status_snapshot, snapshot_local_request_candidate_status, }; -use crate::request_diagnostics::attach_current_request_diagnostics_to_report_context; +use crate::request_diagnostics::{ + attach_current_request_diagnostics_and_candidate_start_timing_to_report_context, + attach_request_diagnostics_to_report_context, calibrate_candidate_first_byte_elapsed_ms, + current_request_diagnostics, RequestDiagnostics, +}; use crate::usage::{spawn_sync_report, submit_sync_report}; use crate::video_tasks::VideoTaskSyncReportMode; use crate::{usage::GatewaySyncReportRequest, AppState, GatewayError}; @@ -108,6 +112,22 @@ const OPENAI_IMAGE_SYNC_JSON_HEARTBEAT_BYTES: &[u8] = b"\n"; const OPENAI_IMAGE_SYNC_PROGRESS_WRITE_INTERVAL: Duration = Duration::from_secs(5); const INVALID_GEMINI_PROVIDER_SUCCESS_MESSAGE: &str = "Provider returned HTTP 200 but the Gemini response did not contain visible model output; refusing to finalize it as a successful response."; +fn elapsed_ms_since(started_at: Instant) -> u64 { + started_at.elapsed().as_millis().min(u128::from(u64::MAX)) as u64 +} + +fn calibrated_sync_candidate_first_byte_elapsed_ms( + candidate_started_at: Instant, + result: &ExecutionResult, +) -> Option { + let telemetry = result.telemetry.as_ref()?; + calibrate_candidate_first_byte_elapsed_ms( + elapsed_ms_since(candidate_started_at), + telemetry.elapsed_ms, + telemetry.ttfb_ms, + ) +} + #[derive(Debug)] struct SyncExecutionFailure { error_type: &'static str, @@ -143,7 +163,9 @@ struct SyncAttemptTerminalGuard { state: AppState, plan: ExecutionPlan, report_context: Option, + request_diagnostics: Option>, candidate_started_unix_ms: u64, + candidate_started_at: Instant, armed: bool, } @@ -153,12 +175,15 @@ impl SyncAttemptTerminalGuard { plan: &ExecutionPlan, report_context: Option, candidate_started_unix_ms: u64, + candidate_started_at: Instant, ) -> Self { Self { state: state.clone(), plan: plan.clone(), report_context, + request_diagnostics: current_request_diagnostics(), candidate_started_unix_ms, + candidate_started_at, armed: true, } } @@ -176,7 +201,9 @@ impl SyncAttemptTerminalGuard { self.state.clone(), self.plan.clone(), self.report_context.clone(), + self.request_diagnostics.clone(), self.candidate_started_unix_ms, + self.candidate_started_at, UsageEventType::Failed, RequestCandidateStatus::Failed, StatusCode::INTERNAL_SERVER_ERROR.as_u16(), @@ -196,14 +223,18 @@ impl Drop for SyncAttemptTerminalGuard { let state = self.state.clone(); let plan = self.plan.clone(); let report_context = self.report_context.clone(); + let request_diagnostics = self.request_diagnostics.clone(); let candidate_started_unix_ms = self.candidate_started_unix_ms; + let candidate_started_at = self.candidate_started_at; if let Ok(handle) = tokio::runtime::Handle::try_current() { handle.spawn(async move { record_sync_attempt_forced_terminal_state( state, plan, report_context, + request_diagnostics, candidate_started_unix_ms, + candidate_started_at, UsageEventType::Cancelled, RequestCandidateStatus::Cancelled, 499, @@ -229,7 +260,9 @@ async fn record_sync_attempt_forced_terminal_state( state: AppState, plan: ExecutionPlan, report_context: Option, + request_diagnostics: Option>, candidate_started_unix_ms: u64, + candidate_started_at: Instant, usage_event_type: UsageEventType, candidate_status: RequestCandidateStatus, status_code: u16, @@ -237,8 +270,10 @@ async fn record_sync_attempt_forced_terminal_state( error_message: impl Into, ) { let error_message = error_message.into(); + let report_context = + attach_request_diagnostics_to_report_context(report_context, request_diagnostics.as_ref()); let terminal_unix_ms = current_request_candidate_unix_ms(); - let latency_ms = terminal_unix_ms.saturating_sub(candidate_started_unix_ms); + let latency_ms = elapsed_ms_since(candidate_started_at); record_local_request_candidate_status( &state, &plan, @@ -409,6 +444,41 @@ fn maybe_store_sync_execution_failure_fallback( Ok(()) } +async fn maybe_build_sync_transport_error_stop_response( + state: &AppState, + plan: &ExecutionPlan, + report_context: Option<&Value>, + trace_id: &str, + decision: &GatewayControlDecision, + error_type: &str, + error_message: &str, + elapsed_ms: u64, +) -> Result>, GatewayError> { + let analysis = crate::orchestration::resolve_local_transport_failover_analysis_for_attempt( + state, + plan, + report_context, + ) + .await; + if !matches!(analysis.decision, LocalFailoverDecision::StopLocalFailover) { + return Ok(None); + } + + crate::execution_runtime::build_transport_error_stop_response( + state, + plan, + report_context, + trace_id, + decision, + StatusCode::BAD_GATEWAY.as_u16(), + error_type, + error_message, + elapsed_ms, + ) + .await + .map(Some) +} + struct ImplicitSyncFinalizeOutcome { payload: GatewaySyncReportRequest, outcome: LocalCoreSyncFinalizeOutcome, @@ -496,9 +566,15 @@ async fn record_sync_terminal_usage( plan: &ExecutionPlan, report_context: Option<&serde_json::Value>, payload: &GatewaySyncReportRequest, + candidate_started_at: Instant, + candidate_first_byte_elapsed_ms: Option, ) { let report_context_with_diagnostics = - attach_current_request_diagnostics_to_report_context(report_context); + attach_current_request_diagnostics_and_candidate_start_timing_to_report_context( + report_context, + candidate_started_at, + candidate_first_byte_elapsed_ms, + ); let context_seed = build_terminal_usage_context_seed( plan, report_context_with_diagnostics.as_ref().or(report_context), @@ -519,9 +595,19 @@ async fn record_sync_terminal_usage_and_disarm_guard( plan: &ExecutionPlan, report_context: Option<&serde_json::Value>, payload: &GatewaySyncReportRequest, + candidate_started_at: Instant, + candidate_first_byte_elapsed_ms: Option, terminal_guard: &mut SyncAttemptTerminalGuard, ) { - record_sync_terminal_usage(state, plan, report_context, payload).await; + record_sync_terminal_usage( + state, + plan, + report_context, + payload, + candidate_started_at, + candidate_first_byte_elapsed_ms, + ) + .await; terminal_guard.disarm(); } @@ -1850,6 +1936,7 @@ async fn execute_execution_runtime_sync_impl( .and_then(|context| context.candidate_index) .map(|value| value.to_string()) .unwrap_or_else(|| "-".to_string()); + let candidate_started_at = Instant::now(); let candidate_started_unix_secs = current_request_candidate_unix_ms(); let lifecycle_seed = build_lifecycle_usage_seed(&plan, report_context.as_ref()); let usage_data = state.usage_lifecycle_data_state().as_ref().clone(); @@ -1877,6 +1964,7 @@ async fn execute_execution_runtime_sync_impl( &plan, report_context.clone(), candidate_started_unix_secs, + candidate_started_at, ); let result = (async { let _provider_pool_in_flight_guard = acquire_provider_pool_in_flight_guard( @@ -1922,6 +2010,11 @@ async fn execute_execution_runtime_sync_impl( { Ok(result) => result, Err(err) => { + let failure_error_type = err.error_type; + let failure_message = err.message.clone(); + let failure_latency_ms = err + .latency_ms + .unwrap_or_else(|| elapsed_ms_since(candidate_started_at)); maybe_store_sync_execution_failure_fallback( &err, &plan, @@ -1952,19 +2045,34 @@ async fn execute_execution_runtime_sync_impl( report_context.as_ref(), SchedulerRequestCandidateStatusUpdate { status: RequestCandidateStatus::Failed, - status_code: err.status_code, - error_type: Some(err.error_type.to_string()), + status_code: None, + error_type: Some(failure_error_type.to_string()), error_message: Some(err.message), - latency_ms: err.latency_ms, + latency_ms: Some(failure_latency_ms), started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; + if let Some(response) = maybe_build_sync_transport_error_stop_response( + state, + &plan, + report_context.as_ref(), + trace_id, + decision, + failure_error_type, + failure_message.as_str(), + failure_latency_ms, + ) + .await? + { + return Ok(Some(response)); + } return Ok(None); } }, Err(err) => { + let transport_error_message = err.to_string(); warn!( event_name = "chatgpt_web_image_execution_unavailable", log_type = "ops", @@ -1990,18 +2098,33 @@ async fn execute_execution_runtime_sync_impl( error_type: Some( "chatgpt_web_image_execution_unavailable".to_string(), ), - error_message: Some(err.to_string()), - latency_ms: None, + error_message: Some(transport_error_message.clone()), + latency_ms: Some(elapsed_ms_since(candidate_started_at)), started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; + if let Some(response) = maybe_build_sync_transport_error_stop_response( + state, + &plan, + report_context.as_ref(), + trace_id, + decision, + "chatgpt_web_image_execution_unavailable", + transport_error_message.as_str(), + elapsed_ms_since(candidate_started_at), + ) + .await? + { + return Ok(Some(response)); + } return Ok(None); } } } Err(err) => { + let transport_error_message = err.to_string(); warn!( event_name = "grok_execution_unavailable", log_type = "ops", @@ -2025,13 +2148,27 @@ async fn execute_execution_runtime_sync_impl( status: RequestCandidateStatus::Failed, status_code: None, error_type: Some("grok_execution_unavailable".to_string()), - error_message: Some(err.to_string()), - latency_ms: None, + error_message: Some(transport_error_message.clone()), + latency_ms: Some(elapsed_ms_since(candidate_started_at)), started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; + if let Some(response) = maybe_build_sync_transport_error_stop_response( + state, + &plan, + report_context.as_ref(), + trace_id, + decision, + "grok_execution_unavailable", + transport_error_message.as_str(), + elapsed_ms_since(candidate_started_at), + ) + .await? + { + return Ok(Some(response)); + } return Ok(None); } } @@ -2042,6 +2179,7 @@ async fn execute_execution_runtime_sync_impl( match (override_fn.0)(&plan) { Ok(result) => result, Err(err) => { + let transport_error_message = format!("{err:?}"); warn!( event_name = "sync_execution_runtime_test_override_failed", log_type = "ops", @@ -2065,13 +2203,27 @@ async fn execute_execution_runtime_sync_impl( status: RequestCandidateStatus::Failed, status_code: None, error_type: Some("execution_runtime_unavailable".to_string()), - error_message: Some(format!("{err:?}")), - latency_ms: None, + error_message: Some(transport_error_message.clone()), + latency_ms: Some(elapsed_ms_since(candidate_started_at)), started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; + if let Some(response) = maybe_build_sync_transport_error_stop_response( + state, + &plan, + report_context.as_ref(), + trace_id, + decision, + "execution_runtime_unavailable", + transport_error_message.as_str(), + elapsed_ms_since(candidate_started_at), + ) + .await? + { + return Ok(Some(response)); + } return Ok(None); } } @@ -2111,6 +2263,11 @@ async fn execute_execution_runtime_sync_impl( { Ok(result) => result, Err(err) => { + let failure_error_type = err.error_type; + let failure_message = err.message.clone(); + let failure_latency_ms = err + .latency_ms + .unwrap_or_else(|| elapsed_ms_since(candidate_started_at)); maybe_store_sync_execution_failure_fallback( &err, &plan, @@ -2141,19 +2298,34 @@ async fn execute_execution_runtime_sync_impl( report_context.as_ref(), SchedulerRequestCandidateStatusUpdate { status: RequestCandidateStatus::Failed, - status_code: err.status_code, - error_type: Some(err.error_type.to_string()), + status_code: None, + error_type: Some(failure_error_type.to_string()), error_message: Some(err.message), - latency_ms: err.latency_ms, + latency_ms: Some(failure_latency_ms), started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; + if let Some(response) = maybe_build_sync_transport_error_stop_response( + state, + &plan, + report_context.as_ref(), + trace_id, + decision, + failure_error_type, + failure_message.as_str(), + failure_latency_ms, + ) + .await? + { + return Ok(Some(response)); + } return Ok(None); } }, Err(err) => { + let transport_error_message = err.to_string(); warn!( event_name = "chatgpt_web_image_execution_unavailable", log_type = "ops", @@ -2179,17 +2351,32 @@ async fn execute_execution_runtime_sync_impl( error_type: Some( "chatgpt_web_image_execution_unavailable".to_string(), ), - error_message: Some(err.to_string()), - latency_ms: None, + error_message: Some(transport_error_message.clone()), + latency_ms: Some(elapsed_ms_since(candidate_started_at)), started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; + if let Some(response) = maybe_build_sync_transport_error_stop_response( + state, + &plan, + report_context.as_ref(), + trace_id, + decision, + "chatgpt_web_image_execution_unavailable", + transport_error_message.as_str(), + elapsed_ms_since(candidate_started_at), + ) + .await? + { + return Ok(Some(response)); + } return Ok(None); } }, Err(err) => { + let transport_error_message = err.to_string(); warn!( event_name = "grok_execution_unavailable", log_type = "ops", @@ -2213,13 +2400,27 @@ async fn execute_execution_runtime_sync_impl( status: RequestCandidateStatus::Failed, status_code: None, error_type: Some("grok_execution_unavailable".to_string()), - error_message: Some(err.to_string()), - latency_ms: None, + error_message: Some(transport_error_message.clone()), + latency_ms: Some(elapsed_ms_since(candidate_started_at)), started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, ) .await; + if let Some(response) = maybe_build_sync_transport_error_stop_response( + state, + &plan, + report_context.as_ref(), + trace_id, + decision, + "grok_execution_unavailable", + transport_error_message.as_str(), + elapsed_ms_since(candidate_started_at), + ) + .await? + { + return Ok(Some(response)); + } return Ok(None); } } @@ -2237,6 +2438,7 @@ async fn execute_execution_runtime_sync_impl( plan_candidate_id.as_deref(), report_context.as_ref(), candidate_started_unix_secs, + candidate_started_at, ) .await?; match remote_outcome { @@ -2246,6 +2448,8 @@ async fn execute_execution_runtime_sync_impl( } } }; + let mut candidate_first_byte_elapsed_ms = + calibrated_sync_candidate_first_byte_elapsed_ms(candidate_started_at, &result); let mut oauth_retry_attempted = false; let ( result_error_type, @@ -2331,6 +2535,11 @@ async fn execute_execution_runtime_sync_impl( .await { Ok(retry_result) => { + candidate_first_byte_elapsed_ms = + calibrated_sync_candidate_first_byte_elapsed_ms( + candidate_started_at, + &retry_result, + ); result = retry_result; continue; } @@ -2657,6 +2866,8 @@ async fn execute_execution_runtime_sync_impl( &plan, implicit_finalize.payload.report_context.as_ref(), usage_payload, + candidate_started_at, + candidate_first_byte_elapsed_ms, &mut terminal_guard, ) .await; @@ -2711,6 +2922,8 @@ async fn execute_execution_runtime_sync_impl( &plan, payload.report_context.as_ref(), usage_payload, + candidate_started_at, + candidate_first_byte_elapsed_ms, &mut terminal_guard, ) .await; @@ -2758,6 +2971,8 @@ async fn execute_execution_runtime_sync_impl( &plan, original_report_context.as_ref(), &report_payload, + candidate_started_at, + candidate_first_byte_elapsed_ms, &mut terminal_guard, ) .await; @@ -2793,6 +3008,8 @@ async fn execute_execution_runtime_sync_impl( &plan, payload.report_context.as_ref(), &payload, + candidate_started_at, + candidate_first_byte_elapsed_ms, &mut terminal_guard, ) .await; @@ -2840,6 +3057,8 @@ async fn execute_execution_runtime_sync_impl( &plan, payload.report_context.as_ref(), &payload, + candidate_started_at, + candidate_first_byte_elapsed_ms, &mut terminal_guard, ) .await; @@ -2867,6 +3086,8 @@ async fn execute_execution_runtime_sync_impl( &plan, payload.report_context.as_ref(), &payload, + candidate_started_at, + candidate_first_byte_elapsed_ms, &mut terminal_guard, ) .await; @@ -2903,6 +3124,8 @@ async fn execute_execution_runtime_sync_impl( &plan, usage_payload.report_context.as_ref(), &usage_payload, + candidate_started_at, + candidate_first_byte_elapsed_ms, &mut terminal_guard, ) .await; @@ -2995,6 +3218,7 @@ async fn execute_sync_via_remote_execution_runtime( plan_candidate_id: Option<&str>, report_context: Option<&serde_json::Value>, candidate_started_unix_secs: u64, + candidate_started_at: Instant, ) -> Result { let response = match post_sync_plan_to_remote_execution_runtime( state, @@ -3025,7 +3249,7 @@ async fn execute_sync_via_remote_execution_runtime( status_code: None, error_type: Some("execution_runtime_unavailable".to_string()), error_message: Some(format!("{err:?}")), - latency_ms: None, + latency_ms: Some(elapsed_ms_since(candidate_started_at)), started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, @@ -3049,7 +3273,7 @@ async fn execute_sync_via_remote_execution_runtime( "execution runtime returned HTTP {}", response.status() )), - latency_ms: None, + latency_ms: Some(elapsed_ms_since(candidate_started_at)), started_at_unix_ms: Some(candidate_started_unix_secs), finished_at_unix_ms: Some(terminal_unix_secs), }, @@ -3373,6 +3597,7 @@ mod tests { })); ensure_execution_request_candidate_slot(&state, &mut plan, &mut report_context).await; + let candidate_started_at = Instant::now(); let started_at = current_request_candidate_unix_ms(); state.usage_runtime.record_pending( state.usage_lifecycle_data_state().as_ref(), @@ -3394,10 +3619,19 @@ mod tests { ) .await; - { - let _guard = - SyncAttemptTerminalGuard::new(&state, &plan, report_context.clone(), started_at); - } + crate::request_diagnostics::scope_request_diagnostics(async { + crate::request_diagnostics::record_request_accepted_at( + Instant::now() - Duration::from_millis(25), + ); + let _guard = SyncAttemptTerminalGuard::new( + &state, + &plan, + report_context.clone(), + started_at, + candidate_started_at, + ); + }) + .await; let mut stored_usage = None; for _ in 0..50 { @@ -3418,6 +3652,17 @@ mod tests { assert_eq!(stored_usage.billing_status, "void"); assert_eq!(stored_usage.status_code, Some(499)); assert_eq!(stored_usage.error_category.as_deref(), Some("cancelled")); + let request_metadata = stored_usage + .request_metadata + .as_ref() + .expect("cancelled usage should retain request diagnostics"); + assert!(request_metadata + .get("end_to_end_time_ms") + .and_then(Value::as_u64) + .is_some()); + assert!(request_metadata + .get("end_to_end_first_byte_time_ms") + .is_none()); let stored_candidates = request_candidate_repository .list_by_request_id("sync-cancel-guard-request") diff --git a/apps/aether-gateway/src/execution_runtime/transport.rs b/apps/aether-gateway/src/execution_runtime/transport.rs index 3c2107515..2a3786592 100644 --- a/apps/aether-gateway/src/execution_runtime/transport.rs +++ b/apps/aether-gateway/src/execution_runtime/transport.rs @@ -575,6 +575,8 @@ pub(crate) enum ExecutionRuntimeTransportError { BrowserClientBuild(wreq::Error), #[error("browser impersonation response body failed: {0}")] BrowserBody(String), + #[error("{message}")] + UpstreamHttpStatus { status_code: u16, message: String }, #[error("failed to execute upstream request: {0}")] UpstreamRequest(String), #[error("upstream response {phase} body exceeds {limit_bytes} bytes")] diff --git a/apps/aether-gateway/src/execution_runtime/transport_failure.rs b/apps/aether-gateway/src/execution_runtime/transport_failure.rs new file mode 100644 index 000000000..ccd532adf --- /dev/null +++ b/apps/aether-gateway/src/execution_runtime/transport_failure.rs @@ -0,0 +1,139 @@ +use std::collections::BTreeMap; +use std::future::Future; +use std::sync::{ + atomic::{AtomicBool, Ordering}, + Arc, +}; + +use aether_usage_runtime::{build_usage_event_data_seed, UsageEvent, UsageEventType}; +use axum::body::Body; +use axum::http::Response; +use serde_json::{json, Value}; + +use crate::ai_serving::{build_core_error_body_for_client_format, LocalCoreSyncErrorKind}; +use crate::api::response::{attach_control_metadata_headers, build_client_response_from_parts}; +use crate::control::GatewayControlDecision; +use crate::request_diagnostics::attach_current_request_diagnostics_and_candidate_timing_to_report_context; +use crate::{AppState, GatewayError}; + +const TRANSPORT_ERROR_CLIENT_MESSAGE: &str = + "Upstream transport failed before an HTTP response was received"; + +#[derive(Debug, Default)] +pub(crate) struct StreamCandidateWatchdogProgress { + terminal_started: AtomicBool, +} + +tokio::task_local! { + static STREAM_CANDIDATE_WATCHDOG_PROGRESS: Arc; +} + +impl StreamCandidateWatchdogProgress { + pub(crate) fn shared() -> Arc { + Arc::new(Self::default()) + } + + pub(crate) fn terminal_started(&self) -> bool { + self.terminal_started.load(Ordering::Acquire) + } + + pub(crate) async fn scope(self: Arc, future: F) -> F::Output + where + F: Future, + { + STREAM_CANDIDATE_WATCHDOG_PROGRESS.scope(self, future).await + } +} + +pub(crate) fn mark_stream_candidate_watchdog_terminal_started() { + let _ = STREAM_CANDIDATE_WATCHDOG_PROGRESS.try_with(|progress| { + progress.terminal_started.store(true, Ordering::Release); + }); +} + +#[allow(clippy::too_many_arguments)] +pub(crate) async fn build_transport_error_stop_response( + state: &AppState, + plan: &aether_contracts::ExecutionPlan, + report_context: Option<&Value>, + trace_id: &str, + decision: &GatewayControlDecision, + client_status_code: u16, + error_type: &str, + error_message: &str, + elapsed_ms: u64, +) -> Result, GatewayError> { + mark_stream_candidate_watchdog_terminal_started(); + let client_body = build_core_error_body_for_client_format( + &plan.client_api_format, + TRANSPORT_ERROR_CLIENT_MESSAGE, + Some("upstream_transport_error"), + LocalCoreSyncErrorKind::ServerError, + ) + .unwrap_or_else(|| { + json!({ + "error": { + "type": "server_error", + "message": TRANSPORT_ERROR_CLIENT_MESSAGE, + "code": "upstream_transport_error", + } + }) + }); + let body_bytes = + serde_json::to_vec(&client_body).map_err(|err| GatewayError::Internal(err.to_string()))?; + let headers = BTreeMap::from([ + ("content-type".to_string(), "application/json".to_string()), + ("content-length".to_string(), body_bytes.len().to_string()), + ]); + + if state.usage_runtime.is_enabled() { + let report_context_with_diagnostics = + attach_current_request_diagnostics_and_candidate_timing_to_report_context( + report_context, + Some(elapsed_ms), + None, + ); + let mut usage_data = build_usage_event_data_seed( + plan, + report_context_with_diagnostics.as_ref().or(report_context), + ); + usage_data.status_code = Some(client_status_code); + usage_data.error_message = Some(error_message.to_string()); + usage_data.error_category = Some("server_error".to_string()); + usage_data.response_time_ms = Some(elapsed_ms); + usage_data.response_headers = None; + usage_data.response_body = None; + usage_data.client_response_headers = Some(json!({"content-type": "application/json"})); + usage_data.client_response_body = Some(client_body); + let mut request_metadata = match usage_data.request_metadata.take() { + Some(Value::Object(object)) => object, + Some(other) => serde_json::Map::from_iter([("seed".to_string(), other)]), + None => serde_json::Map::new(), + }; + request_metadata.insert("transport_error".to_string(), Value::Bool(true)); + request_metadata.insert( + "transport_error_type".to_string(), + Value::String(error_type.to_string()), + ); + usage_data.request_metadata = Some(Value::Object(request_metadata)); + state + .usage_runtime + .record_terminal_event_direct( + state.usage_lifecycle_data_state().as_ref(), + UsageEvent::new(UsageEventType::Failed, plan.request_id.clone(), usage_data), + ) + .await; + } + + attach_control_metadata_headers( + build_client_response_from_parts( + client_status_code, + &headers, + Body::from(body_bytes), + trace_id, + Some(decision), + )?, + Some(plan.request_id.as_str()), + plan.candidate_id.as_deref(), + ) +} diff --git a/apps/aether-gateway/src/execution_runtime/windsurf.rs b/apps/aether-gateway/src/execution_runtime/windsurf.rs index fa8d7b6ec..27aa3936a 100644 --- a/apps/aether-gateway/src/execution_runtime/windsurf.rs +++ b/apps/aether-gateway/src/execution_runtime/windsurf.rs @@ -700,18 +700,9 @@ fn windsurf_execution_error_from_transport_error( failover_recommended: false, }; } - if is_windsurf_cascade_transport_error(err) { - return ExecutionError { - kind: ExecutionErrorKind::Upstream5xx, - phase, - message: format!("{message}; Windsurf IDE language server is unavailable"), - upstream_status: Some(503), - retryable: true, - failover_recommended: true, - }; - } + let (kind, phase) = classify_windsurf_transport_execution_error(&lower, phase); ExecutionError { - kind: ExecutionErrorKind::ProtocolError, + kind, phase, message, upstream_status: None, @@ -720,6 +711,57 @@ fn windsurf_execution_error_from_transport_error( } } +fn classify_windsurf_transport_execution_error( + message: &str, + phase: ExecutionPhase, +) -> (ExecutionErrorKind, ExecutionPhase) { + if message.contains("proxy") { + return (ExecutionErrorKind::ProxyError, ExecutionPhase::Connect); + } + if ["tls", "ssl", "certificate", "handshake"] + .iter() + .any(|needle| message.contains(needle)) + { + return (ExecutionErrorKind::TlsError, ExecutionPhase::Handshake); + } + if message.contains("first byte") { + return ( + ExecutionErrorKind::FirstByteTimeout, + ExecutionPhase::FirstByte, + ); + } + if ["timed out", "timeout", "deadline exceeded"] + .iter() + .any(|needle| message.contains(needle)) + { + let kind = match &phase { + ExecutionPhase::Connect | ExecutionPhase::Handshake | ExecutionPhase::Write => { + ExecutionErrorKind::ConnectTimeout + } + ExecutionPhase::FirstByte => ExecutionErrorKind::FirstByteTimeout, + _ => ExecutionErrorKind::ReadTimeout, + }; + return (kind, phase); + } + if [ + "dns", + "failed to lookup address", + "name or service not known", + "no such host", + "connection refused", + "tcp connect", + "kind=connect", + "connect error", + "failed to connect", + ] + .iter() + .any(|needle| message.contains(needle)) + { + return (ExecutionErrorKind::ProtocolError, ExecutionPhase::Connect); + } + (ExecutionErrorKind::ProtocolError, phase) +} + async fn poll_windsurf_cascade_with_transport_recovery( prepared: &PreparedCascade, mut on_event: F, @@ -4554,6 +4596,56 @@ mod tests { assert!(execution_error.failover_recommended); } + #[test] + fn windsurf_transport_errors_do_not_synthesize_provider_status() { + let cases = [ + ( + "tcp connect error: connection refused", + ExecutionErrorKind::ProtocolError, + ExecutionPhase::Connect, + ), + ( + "dns lookup failed: no such host", + ExecutionErrorKind::ProtocolError, + ExecutionPhase::Connect, + ), + ( + "TLS certificate handshake failed", + ExecutionErrorKind::TlsError, + ExecutionPhase::Handshake, + ), + ( + "proxy connection failed", + ExecutionErrorKind::ProxyError, + ExecutionPhase::Connect, + ), + ( + "cascade polling timed out", + ExecutionErrorKind::ReadTimeout, + ExecutionPhase::StreamRead, + ), + ( + "connection reset by peer", + ExecutionErrorKind::ProtocolError, + ExecutionPhase::StreamRead, + ), + ]; + + for (message, expected_kind, expected_phase) in cases { + let err = ExecutionRuntimeTransportError::UpstreamRequest(message.to_string()); + let execution_error = super::windsurf_execution_error_from_transport_error( + &err, + ExecutionPhase::StreamRead, + ); + + assert_eq!(execution_error.kind, expected_kind, "{message}"); + assert_eq!(execution_error.phase, expected_phase, "{message}"); + assert_eq!(execution_error.upstream_status, None, "{message}"); + assert!(execution_error.retryable, "{message}"); + assert!(execution_error.failover_recommended, "{message}"); + } + } + #[test] fn windsurf_sanitizer_redacts_workspace_paths_in_text_and_tool_args() { let text = super::sanitize_windsurf_text( diff --git a/apps/aether-gateway/src/executor/candidate_loop.rs b/apps/aether-gateway/src/executor/candidate_loop.rs index 188e0c950..6149bc60b 100644 --- a/apps/aether-gateway/src/executor/candidate_loop.rs +++ b/apps/aether-gateway/src/executor/candidate_loop.rs @@ -20,8 +20,9 @@ use crate::ai_serving::LocalExecutionAttemptSource; use crate::clock::current_unix_ms; use crate::control::GatewayControlDecision; use crate::execution_runtime::{ - execute_execution_runtime_stream_with_retry_scope, + build_transport_error_stop_response, execute_execution_runtime_stream_with_retry_scope, execute_execution_runtime_sync_with_retry_scope, + mark_stream_candidate_watchdog_terminal_started, StreamCandidateWatchdogProgress, }; use crate::executor::{ build_local_execution_exhaustion, mark_deferred_upstream_response, LocalExecutionRequestOutcome, @@ -30,7 +31,9 @@ use crate::handlers::shared::provider_pool::release_admin_provider_pool_key_leas use crate::log_ids::short_request_id; use crate::orchestration::{ local_execution_candidate_metadata_from_report_context, - local_failover_policy_from_report_context, resolve_local_failover_policy, LocalFailoverPolicy, + local_failover_policy_from_report_context, resolve_local_failover_policy, + resolve_local_transport_failover_analysis_for_attempt, LocalFailoverDecision, + LocalFailoverPolicy, }; use crate::privacy::RedactionExecutionCandidateId; use crate::request_candidate_runtime::{ @@ -1013,12 +1016,24 @@ where let execution_decision = self.decision.clone(); let execution_report_kind = attempt.report_kind(); let execution_plan = plan.clone(); - let mut execution = execute_stream_candidate_with_watchdog( + let stop_on_transport_errors = matches!( + resolve_local_transport_failover_analysis_for_attempt( + self.state, + plan, + watchdog_report_context, + ) + .await + .decision, + LocalFailoverDecision::StopLocalFailover + ); + let watchdog_started_at = std::time::Instant::now(); + let execution = execute_stream_candidate_with_watchdog( self.state, self.trace_id, self.plan_kind, plan, watchdog_report_context, + stop_on_transport_errors, move || async move { execute_execution_runtime_stream_with_retry_scope( &execution_state, @@ -1033,6 +1048,25 @@ where }, ) .await?; + let mut execution = match execution { + StreamCandidateWatchdogOutcome::TransportTimeout => { + AiAttemptExecutionOutcome::Responded( + build_transport_error_stop_response( + self.state, + plan, + watchdog_report_context, + self.trace_id, + self.decision, + http::StatusCode::GATEWAY_TIMEOUT.as_u16(), + "local_stream_candidate_watchdog_timeout", + stream_candidate_watchdog_timeout_message(), + watchdog_started_at.elapsed().as_millis() as u64, + ) + .await?, + ) + } + StreamCandidateWatchdogOutcome::Executed(execution) => execution, + }; match &mut execution { AiAttemptExecutionOutcome::Responded(response) | AiAttemptExecutionOutcome::Retry { @@ -1297,20 +1331,28 @@ fn log_stream_candidate_admission_timeout( ); } +#[derive(Debug)] +enum StreamCandidateWatchdogOutcome { + Executed(AiAttemptExecutionOutcome>), + TransportTimeout, +} + async fn execute_stream_candidate_with_watchdog( state: &(impl RequestCandidateRuntimeWriter + UpstreamExecutionGateProvider + ?Sized), trace_id: &str, plan_kind: &str, plan: &aether_contracts::ExecutionPlan, report_context: Option<&serde_json::Value>, + stop_on_transport_errors: bool, execute: impl FnOnce() -> Fut, -) -> Result>, GatewayError> +) -> Result where Fut: std::future::Future< Output = Result>, GatewayError>, > + Send, { let timeout_duration = resolve_stream_candidate_watchdog_timeout(plan, report_context); + let candidate_started_at = std::time::Instant::now(); let candidate_started_unix_ms = current_unix_ms(); let permit = match acquire_upstream_execution_gate(state, trace_id).await { Ok(permit) => permit, @@ -1324,17 +1366,33 @@ where ) .await; log_stream_candidate_admission_timeout(trace_id, plan_kind, plan, report_context, &err); - return Ok(AiAttemptExecutionOutcome::retry( - AiAttemptRetryScope::Candidate, + return Ok(StreamCandidateWatchdogOutcome::Executed( + AiAttemptExecutionOutcome::retry(AiAttemptRetryScope::Candidate), )); } Err(err) => return Err(err), }; let permit_hold = permit.map(UpstreamExecutionPermitHold::new); let watchdog_started_at = std::time::Instant::now(); - let outcome = match timeout(timeout_duration, execute()).await { - Ok(result) => result, - Err(_) => { + let watchdog_progress = StreamCandidateWatchdogProgress::shared(); + let execution = watchdog_progress.clone().scope(execute()); + tokio::pin!(execution); + let deadline = tokio::time::sleep(timeout_duration); + tokio::pin!(deadline); + let execution_result = tokio::select! { + biased; + result = &mut execution => Some(result), + () = &mut deadline => { + if watchdog_progress.terminal_started() { + Some(execution.await) + } else { + None + } + } + }; + let outcome = match execution_result { + Some(result) => result.map(StreamCandidateWatchdogOutcome::Executed), + None => { let finished_at_unix_ms = current_unix_ms(); let request_id = short_request_id(plan.request_id.as_str()); let provider_name = plan.provider_name.as_deref().unwrap_or("-"); @@ -1350,10 +1408,10 @@ where report_context, SchedulerRequestCandidateStatusUpdate { status: RequestCandidateStatus::Failed, - status_code: Some(http::StatusCode::GATEWAY_TIMEOUT.as_u16()), + status_code: None, error_type: Some("local_stream_candidate_watchdog_timeout".to_string()), error_message: Some(stream_candidate_watchdog_timeout_message().to_string()), - latency_ms: None, + latency_ms: Some(candidate_started_at.elapsed().as_millis() as u64), started_at_unix_ms: Some(candidate_started_unix_ms), finished_at_unix_ms: Some(finished_at_unix_ms), }, @@ -1374,9 +1432,13 @@ where timeout_ms, "gateway local stream candidate watchdog timed out" ); - Ok(AiAttemptExecutionOutcome::retry( - AiAttemptRetryScope::Candidate, - )) + if stop_on_transport_errors { + Ok(StreamCandidateWatchdogOutcome::TransportTimeout) + } else { + Ok(StreamCandidateWatchdogOutcome::Executed( + AiAttemptExecutionOutcome::retry(AiAttemptRetryScope::Candidate), + )) + } } }; observe_gateway_stage_ms( @@ -1384,20 +1446,30 @@ where watchdog_started_at.elapsed().as_millis() as u64, ); match outcome { - Ok(AiAttemptExecutionOutcome::Responded(response)) => { + Ok(StreamCandidateWatchdogOutcome::Executed(AiAttemptExecutionOutcome::Responded( + response, + ))) => { let response = maybe_hold_upstream_execution_permit(Some(response), permit_hold) .expect("responded stream attempt must retain its response"); - Ok(AiAttemptExecutionOutcome::Responded(response)) + Ok(StreamCandidateWatchdogOutcome::Executed( + AiAttemptExecutionOutcome::Responded(response), + )) } - Ok(AiAttemptExecutionOutcome::Retry { + Ok(StreamCandidateWatchdogOutcome::Executed(AiAttemptExecutionOutcome::Retry { scope, fallback_response, - }) => { + })) => { drop(permit_hold); - Ok(AiAttemptExecutionOutcome::Retry { - scope, - fallback_response, - }) + Ok(StreamCandidateWatchdogOutcome::Executed( + AiAttemptExecutionOutcome::Retry { + scope, + fallback_response, + }, + )) + } + Ok(StreamCandidateWatchdogOutcome::TransportTimeout) => { + drop(permit_hold); + Ok(StreamCandidateWatchdogOutcome::TransportTimeout) } Err(err) if is_candidate_level_admission_timeout(&err) => { drop(permit_hold); @@ -1412,8 +1484,8 @@ where .await; } log_stream_candidate_admission_timeout(trace_id, plan_kind, plan, report_context, &err); - Ok(AiAttemptExecutionOutcome::retry( - AiAttemptRetryScope::Candidate, + Ok(StreamCandidateWatchdogOutcome::Executed( + AiAttemptExecutionOutcome::retry(AiAttemptRetryScope::Candidate), )) } Err(err) => { @@ -2329,6 +2401,7 @@ mod tests { "claude_cli_stream", &plan, Some(&report_context), + false, || { std::future::pending::< Result>, GatewayError>, @@ -2342,20 +2415,19 @@ mod tests { let result = task.await.expect("watchdog task should join"); assert!(matches!( result, - Ok(AiAttemptExecutionOutcome::Retry { - scope: AiAttemptRetryScope::Candidate, - fallback_response: None, - }) + Ok(StreamCandidateWatchdogOutcome::Executed( + AiAttemptExecutionOutcome::Retry { + scope: AiAttemptRetryScope::Candidate, + fallback_response: None, + } + )) )); let records = writer.records.lock().await; assert_eq!(records.len(), 1); let record = &records[0]; assert_eq!(record.status, RequestCandidateStatus::Failed); - assert_eq!( - record.status_code, - Some(http::StatusCode::GATEWAY_TIMEOUT.as_u16()) - ); + assert_eq!(record.status_code, None); assert_eq!( record.error_type.as_deref(), Some("local_stream_candidate_watchdog_timeout") @@ -2367,6 +2439,108 @@ mod tests { assert_eq!(record.candidate_index, 2); } + #[tokio::test] + async fn stream_candidate_watchdog_can_stop_on_transport_error() { + let writer = Arc::new(TestRequestCandidateWriter::default()); + let plan = test_plan(Some(ExecutionTimeouts { + first_byte_ms: Some(5), + ..ExecutionTimeouts::default() + })); + let report_context = test_report_context(); + + let result = execute_stream_candidate_with_watchdog( + writer.as_ref(), + "trace_watchdog_stop", + "claude_cli_stream", + &plan, + Some(&report_context), + true, + || { + std::future::pending::< + Result>, GatewayError>, + >() + }, + ) + .await; + + assert!(matches!( + result, + Ok(StreamCandidateWatchdogOutcome::TransportTimeout) + )); + let records = writer.records.lock().await; + assert_eq!(records.len(), 1); + assert_eq!(records[0].status_code, None); + assert_eq!( + records[0].error_type.as_deref(), + Some("local_stream_candidate_watchdog_timeout") + ); + } + + #[tokio::test] + async fn stream_candidate_watchdog_does_not_cancel_started_terminalization() { + let writer = Arc::new(TestRequestCandidateWriter::default()); + let plan = test_plan(Some(ExecutionTimeouts { + first_byte_ms: Some(5), + ..ExecutionTimeouts::default() + })); + let report_context = test_report_context(); + + let result = execute_stream_candidate_with_watchdog( + writer.as_ref(), + "trace_terminalization", + "claude_cli_stream", + &plan, + Some(&report_context), + true, + || async { + mark_stream_candidate_watchdog_terminal_started(); + tokio::time::sleep(Duration::from_millis(20)).await; + Ok(AiAttemptExecutionOutcome::Responded(Response::new( + Body::from("terminal response"), + ))) + }, + ) + .await; + + assert!(matches!( + result, + Ok(StreamCandidateWatchdogOutcome::Executed( + AiAttemptExecutionOutcome::Responded(_) + )) + )); + assert!(writer.records.lock().await.is_empty()); + } + + #[tokio::test] + async fn stream_candidate_watchdog_does_not_relabel_execution_error_as_timeout() { + let writer = Arc::new(TestRequestCandidateWriter::default()); + let plan = test_plan(None); + let report_context = test_report_context(); + + let result = execute_stream_candidate_with_watchdog( + writer.as_ref(), + "trace_execution_error", + "claude_cli_stream", + &plan, + Some(&report_context), + true, + || async { + Err(GatewayError::UpstreamUnavailable { + trace_id: "trace_execution_error".to_string(), + message: "upstream connect failed".to_string(), + }) + }, + ) + .await; + + assert!(matches!( + result, + Err(GatewayError::UpstreamUnavailable { message, .. }) + if message == "upstream connect failed" + )); + assert!(writer.records.lock().await.is_empty()); + } + #[tokio::test] async fn stream_candidate_upstream_execution_admission_timeout_marks_failed_and_continues() { let writer = Arc::new(TestRequestCandidateWriter::with_upstream_gate( @@ -2388,6 +2562,7 @@ mod tests { "claude_cli_stream", &plan, Some(&report_context), + false, || async { panic!("execute future should not run while upstream execution gate is saturated") }, @@ -2396,10 +2571,12 @@ mod tests { assert!(matches!( result, - Ok(AiAttemptExecutionOutcome::Retry { - scope: AiAttemptRetryScope::Candidate, - fallback_response: None, - }) + Ok(StreamCandidateWatchdogOutcome::Executed( + AiAttemptExecutionOutcome::Retry { + scope: AiAttemptRetryScope::Candidate, + fallback_response: None, + } + )) )); let records = writer.records.lock().await; assert_eq!(records.len(), 1); @@ -2432,6 +2609,7 @@ mod tests { "claude_cli_stream", &plan, Some(&report_context), + false, || async { Err(GatewayError::AdmissionTimeout { trace_id: "trace_target_admission".to_string(), @@ -2444,10 +2622,12 @@ mod tests { assert!(matches!( result, - Ok(AiAttemptExecutionOutcome::Retry { - scope: AiAttemptRetryScope::Candidate, - fallback_response: None, - }) + Ok(StreamCandidateWatchdogOutcome::Executed( + AiAttemptExecutionOutcome::Retry { + scope: AiAttemptRetryScope::Candidate, + fallback_response: None, + } + )) )); assert!(writer.records.lock().await.is_empty()); } diff --git a/apps/aether-gateway/src/executor/orchestration.rs b/apps/aether-gateway/src/executor/orchestration.rs index 379eab37f..1137a9943 100644 --- a/apps/aether-gateway/src/executor/orchestration.rs +++ b/apps/aether-gateway/src/executor/orchestration.rs @@ -55,6 +55,7 @@ use crate::executor::{ LocalExecutionRequestOutcome, }; use crate::handlers::shared::system_config_bool; +use crate::request_diagnostics::{current_request_diagnostics, scope_request_diagnostics_with}; use crate::stage_metrics::observe_gateway_stage_ms; use crate::{AiExecutionDecision, AppState, GatewayError}; @@ -844,15 +845,19 @@ where let decision_for_response = decision.clone(); let started_at = Instant::now(); let (tx, rx) = mpsc::channel::>(1); + let request_diagnostics = current_request_diagnostics(); tokio::spawn(async move { - let bytes = standard_text_sync_heartbeat_final_bytes( - client_api_format.as_str(), - redaction_slot.as_ref(), - execute(state, parts, trace_id, decision, plan_kind, started_at).await, - ) + scope_request_diagnostics_with(request_diagnostics, async move { + let bytes = standard_text_sync_heartbeat_final_bytes( + client_api_format.as_str(), + redaction_slot.as_ref(), + execute(state, parts, trace_id, decision, plan_kind, started_at).await, + ) + .await; + let _ = tx.send(Ok(Bytes::from(bytes))).await; + }) .await; - let _ = tx.send(Ok(Bytes::from(bytes))).await; }); let headers = BTreeMap::from([( @@ -1115,23 +1120,27 @@ fn build_openai_image_sync_heartbeat_shell_response( let decision_for_response = decision.clone(); let started_at = Instant::now(); let (tx, rx) = mpsc::channel::>(1); + let request_diagnostics = current_request_diagnostics(); tokio::spawn(async move { - let bytes = openai_image_sync_heartbeat_final_bytes( - execute_openai_image_sync_heartbeat_attempts( - state, - request_path, - trace_id, - decision, - plan_kind, - attempts, - transfer_tracker, - started_at, + scope_request_diagnostics_with(request_diagnostics, async move { + let bytes = openai_image_sync_heartbeat_final_bytes( + execute_openai_image_sync_heartbeat_attempts( + state, + request_path, + trace_id, + decision, + plan_kind, + attempts, + transfer_tracker, + started_at, + ) + .await, ) - .await, - ) + .await; + let _ = tx.send(Ok(Bytes::from(bytes))).await; + }) .await; - let _ = tx.send(Ok(Bytes::from(bytes))).await; }); let headers = BTreeMap::from([( @@ -1557,9 +1566,14 @@ pub(crate) fn decision_payload_is_direct_execution(payload: &AiExecutionDecision #[cfg(test)] mod tests { use super::*; + use aether_data::repository::candidates::InMemoryRequestCandidateRepository; + use aether_data::repository::usage::InMemoryUsageReadRepository; + use aether_data_contracts::repository::usage::UsageReadRepository; + use aether_usage_runtime::UsageRuntimeConfig; use futures_util::StreamExt; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; + use std::time::Duration; const TEST_OPENAI_IMAGE_SYNC_PLAN_KIND: &str = "openai_image_sync"; const TEST_STANDARD_TEXT_SYNC_PLAN_KIND: &str = "openai_responses_compact_sync"; @@ -1684,6 +1698,61 @@ mod tests { } } + fn heartbeat_usage_test_state( + response_body: Value, + ) -> (AppState, Arc) { + let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); + let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); + let state = AppState::new() + .expect("state should build") + .with_data_state_for_tests( + crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests( + request_candidate_repository, + Arc::clone(&usage_repository), + ), + ) + .with_usage_runtime_for_tests(UsageRuntimeConfig { + enabled: true, + ..UsageRuntimeConfig::default() + }) + .with_execution_runtime_sync_override_for_tests(move |plan| { + let mut result = test_openai_image_execution_result( + plan, + StatusCode::OK.as_u16(), + response_body.clone(), + ); + if let Some(telemetry) = result.telemetry.as_mut() { + telemetry.ttfb_ms = Some(5); + } + Ok(result) + }); + (state, usage_repository) + } + + async fn assert_usage_has_end_to_end_timings( + usage_repository: &InMemoryUsageReadRepository, + request_id: &str, + ) { + let usage = usage_repository + .find_by_request_id(request_id) + .await + .expect("usage should read") + .expect("terminal usage should be recorded"); + let request_metadata = usage + .request_metadata + .as_ref() + .expect("terminal usage should retain request diagnostics"); + let end_to_end_time_ms = request_metadata + .get("end_to_end_time_ms") + .and_then(Value::as_u64) + .expect("end-to-end time should be recorded"); + let end_to_end_first_byte_time_ms = request_metadata + .get("end_to_end_first_byte_time_ms") + .and_then(Value::as_u64) + .expect("end-to-end first-byte time should be recorded"); + assert!(end_to_end_first_byte_time_ms <= end_to_end_time_ms); + } + fn test_standard_text_heartbeat_decision() -> GatewayControlDecision { GatewayControlDecision::synthetic( "/v1/responses", @@ -1799,6 +1868,46 @@ mod tests { assert_eq!(body["error"]["upstream_status"], json!(503)); } + #[tokio::test] + async fn openai_image_sync_heartbeat_propagates_request_diagnostics_to_terminal_usage() { + let (state, usage_repository) = heartbeat_usage_test_state(json!({ + "data": [{"b64_json": "heartbeat-image"}] + })); + let response = crate::request_diagnostics::scope_request_diagnostics(async move { + crate::request_diagnostics::record_request_accepted_at( + Instant::now() - Duration::from_millis(25), + ); + build_openai_image_sync_heartbeat_shell_response( + state, + "/v1/images/generations".to_string(), + "trace-image-heartbeat-retry".to_string(), + test_openai_image_heartbeat_decision(), + TEST_OPENAI_IMAGE_SYNC_PLAN_KIND.to_string(), + vec![test_openai_image_heartbeat_attempt( + 0, + "endpoint-success", + "candidate-success", + )], + ProviderTransferTracker::default(), + ) + }) + .await + .expect("heartbeat shell should build"); + + let body = to_bytes( + response.into_body(), + crate::headers::max_internal_buffered_body_bytes(), + ) + .await + .expect("heartbeat response body should complete"); + assert!(!body.is_empty()); + assert_usage_has_end_to_end_timings( + usage_repository.as_ref(), + "trace-image-heartbeat-retry", + ) + .await; + } + #[tokio::test] async fn openai_image_sync_heartbeat_attempts_retry_first_candidate_then_return_second() { let call_count = Arc::new(AtomicUsize::new(0)); @@ -2065,6 +2174,63 @@ mod tests { let _ = release_tx.send(()); } + #[tokio::test] + async fn standard_text_sync_heartbeat_propagates_request_diagnostics_to_terminal_usage() { + let (state, usage_repository) = heartbeat_usage_test_state(json!({ + "id": "resp_heartbeat", + "output": [] + })); + let (parts, _) = http::Request::builder() + .method(http::Method::POST) + .uri("/v1/responses") + .body(()) + .expect("request should build") + .into_parts(); + let response = crate::request_diagnostics::scope_request_diagnostics(async move { + crate::request_diagnostics::record_request_accepted_at( + Instant::now() - Duration::from_millis(25), + ); + build_standard_text_sync_heartbeat_shell_response( + state, + parts, + "trace-standard-text-heartbeat-retry".to_string(), + test_standard_text_heartbeat_decision(), + TEST_STANDARD_TEXT_SYNC_PLAN_KIND.to_string(), + move |state, parts, trace_id, decision, plan_kind, _started_at| async move { + execute_sync_attempt_source::( + &state, + &parts, + trace_id.as_str(), + &decision, + plan_kind.as_str(), + TestSyncAttemptSource::new(vec![test_standard_text_heartbeat_attempt( + 0, + "endpoint-success", + "candidate-success", + "openai:responses:compact", + )]), + ) + .await + }, + ) + }) + .await + .expect("heartbeat shell should build"); + + let body = to_bytes( + response.into_body(), + crate::headers::max_internal_buffered_body_bytes(), + ) + .await + .expect("heartbeat response body should complete"); + assert!(!body.is_empty()); + assert_usage_has_end_to_end_timings( + usage_repository.as_ref(), + "trace-standard-text-heartbeat-retry", + ) + .await; + } + #[test] fn standard_text_sync_heartbeat_compact_non_json_error_body_is_wrapped_in_client_format() { let bytes = standard_text_sync_heartbeat_error_body_from_response( diff --git a/apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/cookie.rs b/apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/cookie.rs index 45d799808..7d90b20a1 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/cookie.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/cookie.rs @@ -9,9 +9,6 @@ use crate::handlers::admin::request::{AdminAppState, AdminRequestContext}; use crate::GatewayError; use axum::{body::Body, http, response::Response}; -pub(super) const MAX_CLAUDE_COOKIE_AUTHORIZE_BODY_BYTES: usize = 32 * 1024; -pub(super) const MAX_CLAUDE_SESSION_KEY_BYTES: usize = 16 * 1024; - struct ClaudeCookieAuthorizeRequest { session_key: String, name: Option, @@ -102,9 +99,6 @@ fn parse_claude_cookie_authorize_request( let Some(request_body) = request_body else { return Err(bad_cookie_request("请求体必须是合法的 JSON 对象")); }; - if request_body.len() > MAX_CLAUDE_COOKIE_AUTHORIZE_BODY_BYTES { - return Err(bad_cookie_request("Cookie 授权请求体过大")); - } let payload = serde_json::from_slice::(request_body) .ok() .and_then(|value| value.as_object().cloned()) @@ -126,7 +120,7 @@ fn parse_claude_cookie_authorize_request( pub(super) fn normalize_claude_session_key(raw: &str) -> Option { let raw = raw.trim(); - if raw.is_empty() || raw.len() > MAX_CLAUDE_SESSION_KEY_BYTES || raw.contains(['\r', '\n']) { + if raw.is_empty() || raw.contains(['\r', '\n']) { return None; } let cookie = raw @@ -155,7 +149,6 @@ pub(super) fn normalize_claude_session_key(raw: &str) -> Option { fn valid_session_key_value(value: &str) -> bool { !value.is_empty() - && value.len() <= MAX_CLAUDE_SESSION_KEY_BYTES && !value.contains(['\r', '\n', ';']) && http::HeaderValue::from_str(value).is_ok() } @@ -178,7 +171,9 @@ fn bad_cookie_request(detail: &'static str) -> Response { #[cfg(test)] mod tests { - use super::normalize_claude_session_key; + use super::{normalize_claude_session_key, parse_claude_cookie_authorize_request}; + use axum::body::Bytes; + use serde_json::json; #[test] fn normalizes_supported_claude_cookie_inputs() { @@ -211,4 +206,15 @@ mod tests { ); } } + + #[test] + fn accepts_authorize_body_and_session_key_above_previous_caps() { + let session_key = "x".repeat(40 * 1024); + let body = Bytes::from(json!({ "sessionKey": session_key }).to_string()); + assert!(body.len() > 32 * 1024); + + let parsed = parse_claude_cookie_authorize_request(Some(&body)) + .expect("large Cookie authorization payload should parse"); + assert_eq!(parsed.session_key.len(), 40 * 1024); + } } diff --git a/apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/cookie_task.rs b/apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/cookie_task.rs index e5a266c74..76e627aba 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/cookie_task.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/cookie_task.rs @@ -8,7 +8,7 @@ use super::super::state::{ build_admin_provider_oauth_backend_unavailable_response, }; use super::batch::build_admin_provider_oauth_batch_task_state; -use super::cookie::{normalize_claude_session_key, MAX_CLAUDE_SESSION_KEY_BYTES}; +use super::cookie::normalize_claude_session_key; use crate::handlers::admin::provider::shared::paths::admin_provider_oauth_cookie_task_provider_id; use crate::handlers::admin::request::{AdminAppState, AdminRequestContext}; use crate::task_runtime::{ @@ -35,7 +35,6 @@ use uuid::Uuid; const CLAUDE_COOKIE_TASK_IMPORT_KIND: &str = "cookie_authorize"; const CLAUDE_COOKIE_TASK_ID_PREFIX: &str = "claude-cookie-"; const MAX_CLAUDE_COOKIE_TASK_ENTRIES: usize = 20; -const MAX_CLAUDE_COOKIE_TASK_BODY_BYTES: usize = 768 * 1024; const CLAUDE_COOKIE_AUTHORIZATION_CONCURRENCY: usize = 3; const MAX_SAFE_ERROR_DETAIL_BYTES: usize = 512; @@ -476,9 +475,6 @@ fn parse_claude_cookie_task_request( let Some(request_body) = request_body else { return Err(bad_cookie_task_request("请求体必须是合法的 JSON 对象")); }; - if request_body.len() > MAX_CLAUDE_COOKIE_TASK_BODY_BYTES { - return Err(bad_cookie_task_request("Cookie 授权请求体过大")); - } let payload = serde_json::from_slice::(request_body) .ok() .and_then(|value| value.as_object().cloned()) @@ -571,8 +567,7 @@ fn current_unix_secs_or(fallback: u64) -> u64 { #[cfg(test)] mod tests { use super::{ - parse_claude_cookie_task_request, safe_error_detail, MAX_CLAUDE_COOKIE_TASK_BODY_BYTES, - MAX_CLAUDE_COOKIE_TASK_ENTRIES, MAX_CLAUDE_SESSION_KEY_BYTES, + parse_claude_cookie_task_request, safe_error_detail, MAX_CLAUDE_COOKIE_TASK_ENTRIES, }; use axum::body::{to_bytes, Bytes}; use serde_json::json; @@ -625,26 +620,23 @@ mod tests { } #[test] - fn accepts_twenty_maximum_length_session_keys_within_batch_body_limit() { + fn accepts_twenty_long_session_keys_above_previous_body_cap() { let cookies = (0..MAX_CLAUDE_COOKIE_TASK_ENTRIES) .map(|index| { let prefix = format!("{index:02}-"); - format!( - "{prefix}{}", - "x".repeat(MAX_CLAUDE_SESSION_KEY_BYTES - prefix.len()) - ) + format!("{prefix}{}", "x".repeat(40 * 1024)) }) .collect::>(); let body = Bytes::from(json!({"cookies": cookies}).to_string()); - assert!(body.len() < MAX_CLAUDE_COOKIE_TASK_BODY_BYTES); - let parsed = parse_claude_cookie_task_request(Some(&body)) - .expect("maximum valid batch should parse"); + assert!(body.len() > 768 * 1024); + let parsed = + parse_claude_cookie_task_request(Some(&body)).expect("large valid batch should parse"); assert_eq!(parsed.entries.len(), MAX_CLAUDE_COOKIE_TASK_ENTRIES); assert!(parsed.entries.iter().all(Result::is_ok)); } #[tokio::test] - async fn rejects_ambiguous_or_oversized_cookie_batches_without_echoing_secrets() { + async fn rejects_ambiguous_cookie_batches_without_echoing_secrets() { let too_many = vec!["sessionKey=value"; MAX_CLAUDE_COOKIE_TASK_ENTRIES + 1]; for payload in [ json!({"cookie": "sessionKey=secret", "cookies": ["sessionKey=other"]}), @@ -663,13 +655,6 @@ mod tests { assert!(!text.contains("secret")); assert!(!text.contains("other")); } - - let oversized_body = Bytes::from(vec![b'x'; MAX_CLAUDE_COOKIE_TASK_BODY_BYTES + 1]); - let response = match parse_claude_cookie_task_request(Some(&oversized_body)) { - Ok(_) => panic!("oversized request should fail"), - Err(response) => response, - }; - assert_eq!(response.status(), http::StatusCode::BAD_REQUEST); } #[test] diff --git a/apps/aether-gateway/src/handlers/admin/request/system/import.rs b/apps/aether-gateway/src/handlers/admin/request/system/import.rs index 71540a7c2..0c34cddd0 100644 --- a/apps/aether-gateway/src/handlers/admin/request/system/import.rs +++ b/apps/aether-gateway/src/handlers/admin/request/system/import.rs @@ -1,6 +1,4 @@ -use super::{ - AdminAppState, ADMIN_SYSTEM_DATA_EXPORT_VERSION, ADMIN_SYSTEM_DATA_IMPORT_MAX_SIZE_BYTES, -}; +use super::{AdminAppState, ADMIN_SYSTEM_DATA_EXPORT_VERSION}; use crate::ai_serving::build_provider_key_pool_score_upsert; use crate::api::ai::admin_endpoint_signature_parts; use crate::handlers::admin::admin_provider_pool_config; @@ -56,8 +54,6 @@ use std::collections::{BTreeMap, BTreeSet}; use std::time::{SystemTime, UNIX_EPOCH}; use uuid::Uuid; -const ADMIN_SYSTEM_IMPORT_MAX_SIZE_BYTES: usize = 500 * 1024 * 1024; - fn invalid_request(detail: impl Into) -> (http::StatusCode, Value) { ( http::StatusCode::BAD_REQUEST, @@ -1183,10 +1179,6 @@ impl<'a> AdminAppState<'a> { ))); } - if request_body.len() > ADMIN_SYSTEM_DATA_IMPORT_MAX_SIZE_BYTES { - return Ok(Err(invalid_request("请求体大小不能超过 500MB"))); - } - let root = match serde_json::from_slice::(request_body) { Ok(Value::Object(map)) => map, _ => return Ok(Err(invalid_request("请求数据验证失败"))), @@ -1280,10 +1272,6 @@ impl<'a> AdminAppState<'a> { json!({ "detail": "Admin system data unavailable" }), ))); } - if request_body.len() > ADMIN_SYSTEM_IMPORT_MAX_SIZE_BYTES { - return Ok(Err(invalid_request("请求体大小不能超过 500MB"))); - } - let parsed = routed!(parse_admin_system_config_import_request(request_body)); let root = parsed.root; let merge_mode = parsed.request.merge_mode; @@ -2307,10 +2295,6 @@ impl<'a> AdminAppState<'a> { json!({ "detail": "Admin system data unavailable" }), ))); } - if request_body.len() > ADMIN_SYSTEM_IMPORT_MAX_SIZE_BYTES { - return Ok(Err(invalid_request("请求体大小不能超过 500MB"))); - } - let root = match serde_json::from_slice::(request_body) { Ok(Value::Object(map)) => map, _ => return Ok(Err(invalid_request("请求数据验证失败"))), diff --git a/apps/aether-gateway/src/handlers/admin/request/system/mod.rs b/apps/aether-gateway/src/handlers/admin/request/system/mod.rs index 93a6bfb4a..7ab47130b 100644 --- a/apps/aether-gateway/src/handlers/admin/request/system/mod.rs +++ b/apps/aether-gateway/src/handlers/admin/request/system/mod.rs @@ -9,7 +9,6 @@ mod proxy_nodes; mod templates; const ADMIN_SYSTEM_DATA_EXPORT_VERSION: &str = "1.0"; -const ADMIN_SYSTEM_DATA_IMPORT_MAX_SIZE_BYTES: usize = 500 * 1024 * 1024; impl<'a> AdminAppState<'a> { pub(crate) async fn upsert_system_config_json_value( diff --git a/apps/aether-gateway/src/handlers/public/support/user_me_usage.rs b/apps/aether-gateway/src/handlers/public/support/user_me_usage.rs index ad3fe49ba..b113dd02c 100644 --- a/apps/aether-gateway/src/handlers/public/support/user_me_usage.rs +++ b/apps/aether-gateway/src/handlers/public/support/user_me_usage.rs @@ -338,6 +338,14 @@ fn users_me_usage_metadata_string<'a>( .filter(|value| !value.is_empty()) } +fn users_me_usage_metadata_u64(item: &StoredRequestUsageAudit, key: &str) -> Option { + item.request_metadata + .as_ref() + .and_then(serde_json::Value::as_object) + .and_then(|metadata| metadata.get(key)) + .and_then(serde_json::Value::as_u64) +} + fn infer_client_family_from_user_agent(user_agent: &str) -> Option<&'static str> { let normalized = user_agent.trim().to_ascii_lowercase(); if normalized.is_empty() { @@ -499,6 +507,11 @@ fn build_users_me_usage_record_payload( auth_api_key_reader_available, ), }); + payload["end_to_end_time_ms"] = json!(users_me_usage_metadata_u64(item, "end_to_end_time_ms")); + payload["end_to_end_first_byte_time_ms"] = json!(users_me_usage_metadata_u64( + item, + "end_to_end_first_byte_time_ms" + )); if item.target_model.is_some() { payload["target_model"] = json!(item.target_model.clone()); @@ -559,6 +572,11 @@ fn build_users_me_usage_active_payload(item: &StoredRequestUsageAudit) -> serde_ "target_model": item.target_model, "has_fallback": item.has_fallback(), }); + payload["end_to_end_time_ms"] = json!(users_me_usage_metadata_u64(item, "end_to_end_time_ms")); + payload["end_to_end_first_byte_time_ms"] = json!(users_me_usage_metadata_u64( + item, + "end_to_end_first_byte_time_ms" + )); if item.api_format.is_none() { payload .as_object_mut() @@ -1661,6 +1679,29 @@ mod tests { assert_eq!(payload["cache_creation_ephemeral_1h_input_tokens"], 6); } + #[test] + fn user_usage_payloads_project_end_to_end_timings_from_metadata() { + let item = StoredRequestUsageAudit { + response_time_ms: Some(626), + first_byte_time_ms: Some(120), + request_metadata: Some(json!({ + "end_to_end_time_ms": 10_626, + "end_to_end_first_byte_time_ms": 10_120, + })), + ..sample_usage("completed") + }; + + let record = build_users_me_usage_record_payload(&item, false, &BTreeMap::new(), false); + let active = build_users_me_usage_active_payload(&item); + + for payload in [&record, &active] { + assert_eq!(payload["response_time_ms"], 626); + assert_eq!(payload["first_byte_time_ms"], 120); + assert_eq!(payload["end_to_end_time_ms"], 10_626); + assert_eq!(payload["end_to_end_first_byte_time_ms"], 10_120); + } + } + #[test] fn user_usage_payloads_expose_requested_and_provider_reasoning_mapping() { let item = StoredRequestUsageAudit { diff --git a/apps/aether-gateway/src/headers.rs b/apps/aether-gateway/src/headers.rs index 7138e856d..9c4c8042c 100644 --- a/apps/aether-gateway/src/headers.rs +++ b/apps/aether-gateway/src/headers.rs @@ -13,43 +13,35 @@ use flate2::read::{DeflateDecoder, GzDecoder, ZlibDecoder}; use serde_json::{Map, Value}; use uuid::Uuid; -const DEFAULT_MAX_REQUEST_BODY_MB: u64 = 64; const MAX_REQUEST_BODY_MB_ENV: &str = "AETHER_MAX_REQUEST_BODY_MB"; -const DEFAULT_MAX_REDACTED_SYNC_RESPONSE_BODY_MB: u64 = 64; const MAX_REDACTED_SYNC_RESPONSE_BODY_MB_ENV: &str = "AETHER_MAX_REDACTED_SYNC_RESPONSE_BODY_MB"; -const DEFAULT_MAX_INTERNAL_BUFFERED_BODY_MB: u64 = 128; const MAX_INTERNAL_BUFFERED_BODY_MB_ENV: &str = "AETHER_MAX_INTERNAL_BUFFERED_BODY_MB"; const TRUSTED_PROXY_CIDRS_ENV: &str = "AETHER_TRUSTED_PROXY_CIDRS"; -/// Upper bound applied to a request body after Content-Encoding decoding, and to -/// uncompressed bodies as-is. Guards against decompression bombs and oversized -/// request allocations. Overridable via `AETHER_MAX_REQUEST_BODY_MB`. -static MAX_REQUEST_BODY_BYTES: LazyLock = LazyLock::new(|| { - std::env::var(MAX_REQUEST_BODY_MB_ENV) - .ok() - .and_then(|value| value.parse::().ok()) - .filter(|value| *value > 0) - .unwrap_or(DEFAULT_MAX_REQUEST_BODY_MB) - .saturating_mul(1024 * 1024) -}); +/// Optional operator cap applied after Content-Encoding decoding, and to +/// uncompressed bodies as-is. Unset, zero, or invalid values disable the cap. +static MAX_REQUEST_BODY_BYTES: LazyLock = + LazyLock::new(|| body_limit_bytes_from_env(MAX_REQUEST_BODY_MB_ENV)); -static MAX_REDACTED_SYNC_RESPONSE_BODY_BYTES: LazyLock = LazyLock::new(|| { - std::env::var(MAX_REDACTED_SYNC_RESPONSE_BODY_MB_ENV) - .ok() - .and_then(|value| value.parse::().ok()) - .filter(|value| *value > 0) - .unwrap_or(DEFAULT_MAX_REDACTED_SYNC_RESPONSE_BODY_MB) - .saturating_mul(1024 * 1024) -}); +static MAX_REDACTED_SYNC_RESPONSE_BODY_BYTES: LazyLock = + LazyLock::new(|| body_limit_bytes_from_env(MAX_REDACTED_SYNC_RESPONSE_BODY_MB_ENV)); -static MAX_INTERNAL_BUFFERED_BODY_BYTES: LazyLock = LazyLock::new(|| { - std::env::var(MAX_INTERNAL_BUFFERED_BODY_MB_ENV) - .ok() +static MAX_INTERNAL_BUFFERED_BODY_BYTES: LazyLock = + LazyLock::new(|| body_limit_bytes_from_env(MAX_INTERNAL_BUFFERED_BODY_MB_ENV)); + +fn body_limit_bytes_from_env(name: &str) -> u64 { + let value = std::env::var(name).ok(); + body_limit_bytes(value.as_deref()) +} + +fn body_limit_bytes(value: Option<&str>) -> u64 { + value + .map(str::trim) .and_then(|value| value.parse::().ok()) .filter(|value| *value > 0) - .unwrap_or(DEFAULT_MAX_INTERNAL_BUFFERED_BODY_MB) - .saturating_mul(1024 * 1024) -}); + .map(|value| value.saturating_mul(1024 * 1024)) + .unwrap_or(u64::MAX) +} static TRUSTED_PROXY_CIDRS: LazyLock> = LazyLock::new(|| { std::env::var(TRUSTED_PROXY_CIDRS_ENV) @@ -836,13 +828,14 @@ mod tests { } #[test] - fn decoded_request_body_bytes_rejects_oversized_uncompressed_body() { - let limit = *super::MAX_REQUEST_BODY_BYTES; + fn explicit_limit_rejects_oversized_uncompressed_body() { + let limit = 4; let oversized = vec![b'a'; limit as usize + 1]; let headers = HeaderMap::new(); - let err = decoded_request_body_bytes(&headers, oversized.as_slice()) - .expect_err("oversized uncompressed body should fail"); + let err = + super::decoded_request_body_bytes_with_limit(&headers, oversized.as_slice(), limit) + .expect_err("oversized uncompressed body should fail"); assert_eq!( err, @@ -851,15 +844,15 @@ mod tests { } #[test] - fn check_request_content_length_rejects_oversized_declared_length() { - let limit = *super::MAX_REQUEST_BODY_BYTES; + fn explicit_limit_rejects_oversized_declared_length() { + let limit = 4; let mut headers = HeaderMap::new(); headers.insert( http::header::CONTENT_LENGTH, HeaderValue::from_str(&(limit + 1).to_string()).expect("length header should build"), ); - let err = super::check_request_content_length(&headers) + let err = super::check_request_content_length_with_limit(&headers, limit) .expect_err("oversized declared length should fail"); assert_eq!( @@ -868,6 +861,35 @@ mod tests { ); } + #[test] + fn body_limits_default_to_unlimited() { + assert_eq!(super::body_limit_bytes(None), u64::MAX); + assert_eq!(super::body_limit_bytes(Some("0")), u64::MAX); + assert_eq!(super::body_limit_bytes(Some("invalid")), u64::MAX); + } + + #[test] + fn positive_body_limit_is_converted_from_mibibytes() { + assert_eq!(super::body_limit_bytes(Some(" 8 ")), 8 * 1024 * 1024); + } + + #[test] + fn unlimited_limit_accepts_declared_and_buffered_body() { + let mut headers = HeaderMap::new(); + headers.insert( + http::header::CONTENT_LENGTH, + HeaderValue::from_static("18446744073709551615"), + ); + super::check_request_content_length_with_limit(&headers, u64::MAX) + .expect("unlimited mode should accept every representable content length"); + + let body = b"body larger than the former default is admitted by the unlimited sentinel"; + let decoded = + super::decoded_request_body_bytes_with_limit(&HeaderMap::new(), body, u64::MAX) + .expect("unlimited mode should accept buffered bytes"); + assert_eq!(decoded.as_ref(), body); + } + #[test] fn request_body_normalization_error_maps_http_status() { assert_eq!( diff --git a/apps/aether-gateway/src/orchestration/classifier.rs b/apps/aether-gateway/src/orchestration/classifier.rs index 0245f511a..416e8c28a 100644 --- a/apps/aether-gateway/src/orchestration/classifier.rs +++ b/apps/aether-gateway/src/orchestration/classifier.rs @@ -54,6 +54,31 @@ impl LocalFailoverClassification { } } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum LocalTransportFailoverClassification { + StopTransportError, + RetryTransportError, +} + +impl LocalTransportFailoverClassification { + pub(crate) const fn as_str(self) -> &'static str { + match self { + Self::StopTransportError => "stop_transport_error", + Self::RetryTransportError => "retry_transport_error", + } + } +} + +pub(crate) const fn classify_local_transport_error( + policy: &LocalFailoverPolicy, +) -> LocalTransportFailoverClassification { + if policy.stop_on_transport_errors { + LocalTransportFailoverClassification::StopTransportError + } else { + LocalTransportFailoverClassification::RetryTransportError + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) enum FailureRetryAction { Stop, @@ -486,8 +511,9 @@ mod tests { use super::{ classify_anthropic_failure_disposition, classify_local_failover, - failure_disposition_from_local_classification, FailureDisposition, FailureRetryAction, - FailureScope, FailureTokenAction, LocalFailoverClassification, LocalFailoverInput, + classify_local_transport_error, failure_disposition_from_local_classification, + FailureDisposition, FailureRetryAction, FailureScope, FailureTokenAction, + LocalFailoverClassification, LocalFailoverInput, LocalTransportFailoverClassification, }; use crate::orchestration::{LocalFailoverPolicy, LocalFailoverRegexRule}; @@ -504,6 +530,27 @@ mod tests { ); } + #[test] + fn classifier_retries_transport_errors_by_default_and_honors_explicit_stop() { + assert_eq!( + classify_local_transport_error(&LocalFailoverPolicy::default()), + LocalTransportFailoverClassification::RetryTransportError + ); + + let stop_policy = LocalFailoverPolicy { + stop_on_transport_errors: true, + ..LocalFailoverPolicy::default() + }; + assert_eq!( + classify_local_transport_error(&stop_policy), + LocalTransportFailoverClassification::StopTransportError + ); + assert_eq!( + classify_local_transport_error(&stop_policy).as_str(), + "stop_transport_error" + ); + } + #[test] fn classifier_detects_success_failover_pattern() { let policy = LocalFailoverPolicy { diff --git a/apps/aether-gateway/src/orchestration/mod.rs b/apps/aether-gateway/src/orchestration/mod.rs index 369a2d9a5..f79b25e8e 100644 --- a/apps/aether-gateway/src/orchestration/mod.rs +++ b/apps/aether-gateway/src/orchestration/mod.rs @@ -27,9 +27,10 @@ pub(crate) use self::attempt::{ }; pub(crate) use self::classifier::{ classify_anthropic_failure_disposition, classify_failure_disposition, classify_local_failover, - failure_disposition_from_local_classification, local_failover_error_message, - FailureDisposition, FailureRetryAction, FailureScope, FailureTokenAction, - LocalFailoverClassification, LocalFailoverInput, + classify_local_transport_error, failure_disposition_from_local_classification, + local_failover_error_message, FailureDisposition, FailureRetryAction, FailureScope, + FailureTokenAction, LocalFailoverClassification, LocalFailoverInput, + LocalTransportFailoverClassification, }; pub(crate) use self::effects::{ apply_local_execution_effect, LocalAdaptiveRateLimitEffect, LocalAdaptiveSuccessEffect, @@ -51,8 +52,9 @@ pub(crate) use self::policy::{ LocalFailoverRegexRule, CYBER_CONTINUE_FAILOVER_CONFIG_KEY, }; pub(crate) use self::recovery::{ - analyze_local_failover, apply_provider_failure_disposition, recover_local_failover_decision, - LocalFailoverAnalysis, LocalFailoverDecision, + analyze_local_failover, analyze_local_transport_error, apply_provider_failure_disposition, + recover_local_failover_decision, LocalFailoverAnalysis, LocalFailoverDecision, + LocalTransportFailoverAnalysis, }; #[cfg(test)] pub(crate) use self::report_effects::clear_local_report_effect_caches_for_tests; @@ -95,6 +97,15 @@ pub(crate) async fn resolve_local_failover_decision_for_attempt( .decision } +pub(crate) async fn resolve_local_transport_failover_analysis_for_attempt( + state: &AppState, + plan: &ExecutionPlan, + report_context: Option<&serde_json::Value>, +) -> LocalTransportFailoverAnalysis { + let policy = resolve_local_failover_policy(state, plan, report_context).await; + analyze_local_transport_error(&policy) +} + pub(crate) fn build_local_error_flow_metadata( status_code: u16, response_text: Option<&str>, diff --git a/apps/aether-gateway/src/orchestration/policy.rs b/apps/aether-gateway/src/orchestration/policy.rs index 57900e653..8a5b9f166 100644 --- a/apps/aether-gateway/src/orchestration/policy.rs +++ b/apps/aether-gateway/src/orchestration/policy.rs @@ -16,6 +16,7 @@ pub(crate) struct LocalFailoverPolicy { pub(crate) max_transfer_timeout_seconds: u64, pub(crate) stop_status_codes: BTreeSet, pub(crate) continue_status_codes: BTreeSet, + pub(crate) stop_on_transport_errors: bool, pub(crate) success_failover_patterns: Vec, pub(crate) error_stop_patterns: Vec, pub(crate) stop_cyber_policy_errors: bool, @@ -30,6 +31,7 @@ impl Default for LocalFailoverPolicy { max_transfer_timeout_seconds: 0, stop_status_codes: BTreeSet::new(), continue_status_codes: BTreeSet::new(), + stop_on_transport_errors: false, success_failover_patterns: Vec::new(), error_stop_patterns: Vec::new(), stop_cyber_policy_errors: true, @@ -71,6 +73,7 @@ pub(crate) async fn resolve_local_failover_policy( max_transfer_timeout_seconds = policy.max_transfer_timeout_seconds, stop_status_code_count = policy.stop_status_codes.len(), continue_status_code_count = policy.continue_status_codes.len(), + stop_on_transport_errors = policy.stop_on_transport_errors, success_failover_pattern_count = policy.success_failover_patterns.len(), error_stop_pattern_count = policy.error_stop_patterns.len(), cyber_continue_failover, @@ -157,6 +160,10 @@ pub(crate) fn local_failover_policy_from_transport( ) }) .unwrap_or_default(), + stop_on_transport_errors: rules + .and_then(|value| value.get("stop_on_transport_errors")) + .and_then(Value::as_bool) + .unwrap_or(false), success_failover_patterns: rules .map(|value| parse_regex_rules(value, "success_failover_patterns")) .unwrap_or_default(), @@ -192,6 +199,10 @@ pub(crate) fn local_failover_policy_from_report_context( .get("continue_status_codes") .map(parse_status_code_list) .unwrap_or_default(), + stop_on_transport_errors: object + .get("stop_on_transport_errors") + .and_then(Value::as_bool) + .unwrap_or(false), success_failover_patterns: parse_regex_rules(object, "success_failover_patterns"), error_stop_patterns: parse_regex_rules(object, "error_stop_patterns"), stop_cyber_policy_errors: object @@ -235,6 +246,7 @@ fn local_failover_policy_to_value(policy: &LocalFailoverPolicy) -> Value { "max_transfer_timeout_seconds": policy.max_transfer_timeout_seconds, "stop_status_codes": policy.stop_status_codes.iter().copied().collect::>(), "continue_status_codes": policy.continue_status_codes.iter().copied().collect::>(), + "stop_on_transport_errors": policy.stop_on_transport_errors, "success_failover_patterns": policy.success_failover_patterns.iter().map(local_failover_regex_rule_to_value).collect::>(), "error_stop_patterns": policy.error_stop_patterns.iter().map(local_failover_regex_rule_to_value).collect::>(), "stop_cyber_policy_errors": policy.stop_cyber_policy_errors, @@ -416,6 +428,7 @@ mod tests { "max_retries": 2, "continue_status_codes": [429], "stop_status_codes": [400], + "stop_on_transport_errors": true, "success_failover_patterns": [{"pattern": "quota", "status_codes": [200]}], "error_stop_patterns": [{"pattern": "validation", "status_codes": [422]}] } @@ -431,6 +444,7 @@ mod tests { max_transfer_timeout_seconds: 60, stop_status_codes: [400].into_iter().collect(), continue_status_codes: [429].into_iter().collect(), + stop_on_transport_errors: true, success_failover_patterns: vec![LocalFailoverRegexRule { pattern: "quota".to_string(), status_codes: [200].into_iter().collect(), @@ -445,6 +459,35 @@ mod tests { ); } + #[test] + fn transport_error_failover_defaults_to_continue_and_accepts_explicit_stop() { + let default_policy = + local_failover_policy_from_transport(&sample_transport(None, None, None)); + assert!(!default_policy.stop_on_transport_errors); + + let stop_policy = local_failover_policy_from_transport(&sample_transport( + None, + None, + Some(json!({ + "failover_rules": { + "stop_on_transport_errors": true, + } + })), + )); + assert!(stop_policy.stop_on_transport_errors); + + let invalid_policy = local_failover_policy_from_transport(&sample_transport( + None, + None, + Some(json!({ + "failover_rules": { + "stop_on_transport_errors": "true", + } + })), + )); + assert!(!invalid_policy.stop_on_transport_errors); + } + #[test] fn transfer_limits_are_read_only_from_top_level_provider_config() { let top_level = local_failover_policy_from_transport(&sample_transport( diff --git a/apps/aether-gateway/src/orchestration/recovery.rs b/apps/aether-gateway/src/orchestration/recovery.rs index d6266eede..3cdcfb79e 100644 --- a/apps/aether-gateway/src/orchestration/recovery.rs +++ b/apps/aether-gateway/src/orchestration/recovery.rs @@ -1,6 +1,7 @@ use super::classifier::{ - classify_failure_disposition, classify_local_failover, FailureRetryAction, - LocalFailoverClassification, LocalFailoverInput, + classify_failure_disposition, classify_local_failover, classify_local_transport_error, + FailureRetryAction, LocalFailoverClassification, LocalFailoverInput, + LocalTransportFailoverClassification, }; use super::LocalFailoverPolicy; @@ -27,6 +28,12 @@ pub(crate) struct LocalFailoverAnalysis { pub(crate) decision: LocalFailoverDecision, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct LocalTransportFailoverAnalysis { + pub(crate) classification: LocalTransportFailoverClassification, + pub(crate) decision: LocalFailoverDecision, +} + impl LocalFailoverAnalysis { pub(crate) const fn use_default() -> Self { Self { @@ -47,6 +54,24 @@ pub(crate) fn analyze_local_failover( } } +pub(crate) fn analyze_local_transport_error( + policy: &LocalFailoverPolicy, +) -> LocalTransportFailoverAnalysis { + let classification = classify_local_transport_error(policy); + let decision = match classification { + LocalTransportFailoverClassification::StopTransportError => { + LocalFailoverDecision::StopLocalFailover + } + LocalTransportFailoverClassification::RetryTransportError => { + LocalFailoverDecision::RetryNextCandidate + } + }; + LocalTransportFailoverAnalysis { + classification, + decision, + } +} + pub(crate) fn apply_provider_failure_disposition( provider_api_format: &str, status_code: u16, @@ -105,7 +130,7 @@ const fn decision_from_classification( #[cfg(test)] mod tests { use super::{ - analyze_local_failover, apply_provider_failure_disposition, + analyze_local_failover, analyze_local_transport_error, apply_provider_failure_disposition, recover_local_failover_decision, LocalFailoverAnalysis, LocalFailoverDecision, }; use crate::orchestration::{ @@ -136,6 +161,23 @@ mod tests { ); } + #[test] + fn transport_error_recovery_defaults_to_retry_and_can_stop() { + assert_eq!( + analyze_local_transport_error(&LocalFailoverPolicy::default()).decision, + LocalFailoverDecision::RetryNextCandidate + ); + + let stop_policy = LocalFailoverPolicy { + stop_on_transport_errors: true, + ..LocalFailoverPolicy::default() + }; + assert_eq!( + analyze_local_transport_error(&stop_policy).decision, + LocalFailoverDecision::StopLocalFailover + ); + } + #[test] fn recovery_retries_default_client_error_without_custom_rule() { assert_eq!( diff --git a/apps/aether-gateway/src/privacy/mod.rs b/apps/aether-gateway/src/privacy/mod.rs index 14f4cc70f..3e5ddc858 100644 --- a/apps/aether-gateway/src/privacy/mod.rs +++ b/apps/aether-gateway/src/privacy/mod.rs @@ -1,4 +1,5 @@ use std::collections::{BTreeMap, HashMap, HashSet}; +use std::convert::Infallible; use std::fmt; use std::net::{Ipv4Addr, Ipv6Addr}; use std::sync::atomic::{AtomicU64, Ordering}; @@ -19,8 +20,6 @@ use crate::GatewayError; type HmacSha256 = hmac::Hmac; const DEFAULT_REDACTION_TTL_SECONDS: u64 = 300; -const DEFAULT_MAX_SCANNED_CHAT_TEXT_BYTES: usize = 2 * 1024 * 1024; -const DEFAULT_MAX_REDACTION_DETECTIONS: usize = 1024; const HMAC96_BYTES: usize = 12; const DEFAULT_SENTINEL_NAMESPACE: &str = "AETHER"; const MAX_SENTINEL_NAMESPACE_LEN: usize = 32; @@ -245,87 +244,17 @@ impl RedactionSessionConfig { } } -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -pub(crate) struct RedactionScanLimits { - pub(crate) max_scanned_text_bytes: usize, - pub(crate) max_detections: usize, -} +pub(crate) type RedactionMaskError = Infallible; -impl Default for RedactionScanLimits { - fn default() -> Self { - Self { - max_scanned_text_bytes: DEFAULT_MAX_SCANNED_CHAT_TEXT_BYTES, - max_detections: DEFAULT_MAX_REDACTION_DETECTIONS, - } - } -} - -#[derive(Clone, Debug, Eq, PartialEq)] -pub(crate) enum RedactionLimitError { - ScannedTextTooLarge { limit: usize }, - TooManyDetections { limit: usize }, -} - -impl RedactionLimitError { - pub(crate) const fn client_status(&self) -> http::StatusCode { - match self { - Self::ScannedTextTooLarge { .. } => http::StatusCode::PAYLOAD_TOO_LARGE, - Self::TooManyDetections { .. } => http::StatusCode::UNPROCESSABLE_ENTITY, - } - } - - pub(crate) const fn safe_message(&self) -> &'static str { - match self { - Self::ScannedTextTooLarge { .. } => "chat pii redaction scanned text limit exceeded", - Self::TooManyDetections { .. } => "chat pii redaction detection limit exceeded", - } - } -} - -#[derive(Debug)] -pub(crate) enum RedactionMaskError { - Limit(RedactionLimitError), -} - -impl From for RedactionMaskError { - fn from(error: RedactionLimitError) -> Self { - Self::Limit(error) - } -} - -#[derive(Clone, Copy)] -struct RedactionScanState { - limits: RedactionScanLimits, - scanned_text_bytes: usize, - detections: usize, -} +#[derive(Clone, Copy, Default)] +struct RedactionScanState; impl RedactionScanState { - fn new(limits: RedactionScanLimits) -> Self { - Self { - limits, - scanned_text_bytes: 0, - detections: 0, - } - } - - fn record_scan(&mut self, text: &str) -> Result<(), RedactionLimitError> { - self.scanned_text_bytes = self.scanned_text_bytes.saturating_add(text.len()); - if self.scanned_text_bytes > self.limits.max_scanned_text_bytes { - return Err(RedactionLimitError::ScannedTextTooLarge { - limit: self.limits.max_scanned_text_bytes, - }); - } + fn record_scan(&mut self, _text: &str) -> Result<(), Infallible> { Ok(()) } - fn record_detections(&mut self, count: usize) -> Result<(), RedactionLimitError> { - self.detections = self.detections.saturating_add(count); - if self.detections > self.limits.max_detections { - return Err(RedactionLimitError::TooManyDetections { - limit: self.limits.max_detections, - }); - } + fn record_detections(&mut self, _count: usize) -> Result<(), Infallible> { Ok(()) } } @@ -403,7 +332,7 @@ impl RedactionSession { &mut self, input: &str, scan_state: &mut RedactionScanState, - ) -> Result { + ) -> Result { scan_state.record_scan(input)?; self.redact_text_internal(input, Some(scan_state)) } @@ -460,7 +389,7 @@ impl RedactionSession { &mut self, input: &str, mut scan_state: Option<&mut RedactionScanState>, - ) -> Result { + ) -> Result { let candidates = select_non_overlapping(detect_candidates_for_session_config( input, &self.config, @@ -1043,21 +972,12 @@ impl ChatPiiRedactionRuntimeConfigCache { } } -pub(crate) struct MaskChatRequestOptions { - pub(crate) scan_limits: RedactionScanLimits, -} +#[derive(Clone, Copy, Default)] +pub(crate) struct MaskChatRequestOptions; impl MaskChatRequestOptions { pub(crate) fn runtime() -> Self { - Self { - scan_limits: RedactionScanLimits::default(), - } - } - - #[cfg(test)] - fn with_scan_limits(mut self, scan_limits: RedactionScanLimits) -> Self { - self.scan_limits = scan_limits; - self + Self } } @@ -1453,7 +1373,7 @@ pub(crate) fn mask_chat_request_json_with_options( options: MaskChatRequestOptions, ) -> MaskedChatRequest { try_mask_chat_request_json_with_options(body, config, options) - .expect("default chat redaction limits should not be exceeded") + .expect("chat redaction masking should be infallible") } pub(crate) fn mask_chat_request_json( @@ -1467,7 +1387,7 @@ pub(crate) fn try_mask_chat_request_json_with_options( body: &[u8], config: RedactionSessionConfig, options: MaskChatRequestOptions, -) -> Result { +) -> Result { try_mask_chat_pii_request_json_with_options( body, ChatPiiRedactionRequestFormat::OpenAiChat, @@ -1497,7 +1417,7 @@ pub(crate) fn try_mask_chat_pii_request_json_with_options( format: ChatPiiRedactionRequestFormat, config: RedactionSessionConfig, options: MaskChatRequestOptions, -) -> Result { +) -> Result { let mut session = RedactionSession::new(config); let Ok(mut value) = serde_json::from_slice::(body) else { return Ok(MaskedChatRequest { @@ -1508,7 +1428,7 @@ pub(crate) fn try_mask_chat_pii_request_json_with_options( }; session.set_collision_corpus(request_collision_corpus(format, &value)); - let mut scan_state = RedactionScanState::new(options.scan_limits); + let mut scan_state = RedactionScanState; let redacted = mask_request_value(format, &mut value, &mut session, &mut scan_state, options)?; if !redacted { @@ -1544,7 +1464,7 @@ pub(crate) async fn try_mask_chat_pii_request_json_with_cache_options( }; session.set_collision_corpus(request_collision_corpus(format, &value)); - let mut scan_state = RedactionScanState::new(options.scan_limits); + let mut scan_state = RedactionScanState; let redacted = mask_request_value_async( format, &mut value, @@ -1582,7 +1502,7 @@ pub(crate) async fn try_mask_chat_pii_request_value_with_cache_options( let mut value = body_json.clone(); session.set_collision_corpus(request_collision_corpus(format, &value)); - let mut scan_state = RedactionScanState::new(options.scan_limits); + let mut scan_state = RedactionScanState; let redacted = mask_request_value_async( format, &mut value, @@ -1619,7 +1539,7 @@ fn mask_request_value( session: &mut RedactionSession, scan_state: &mut RedactionScanState, options: MaskChatRequestOptions, -) -> Result { +) -> Result { match format { ChatPiiRedactionRequestFormat::OpenAiChat => { mask_openai_chat_request_value(value, session, scan_state, options) @@ -1665,7 +1585,7 @@ fn mask_openai_chat_request_value( session: &mut RedactionSession, scan_state: &mut RedactionScanState, options: MaskChatRequestOptions, -) -> Result { +) -> Result { let Some(messages) = value.get_mut("messages").and_then(Value::as_array_mut) else { return Ok(false); }; @@ -1912,7 +1832,7 @@ fn mask_chat_message_value( message: &mut Value, session: &mut RedactionSession, scan_state: &mut RedactionScanState, -) -> Result { +) -> Result { let Some(message) = message.as_object_mut() else { return Ok(false); }; @@ -1933,7 +1853,7 @@ fn mask_chat_content_value( content: &mut Value, session: &mut RedactionSession, scan_state: &mut RedactionScanState, -) -> Result { +) -> Result { match content { Value::String(text) => mask_json_string(text, session, scan_state), Value::Array(parts) => { @@ -1951,7 +1871,7 @@ fn mask_chat_content_part( part: &mut Value, session: &mut RedactionSession, scan_state: &mut RedactionScanState, -) -> Result { +) -> Result { let Some(part) = part.as_object_mut() else { return Ok(false); }; @@ -1968,7 +1888,7 @@ fn mask_tool_call_arguments( tool_call: &mut Value, session: &mut RedactionSession, scan_state: &mut RedactionScanState, -) -> Result { +) -> Result { let Some(function) = tool_call.get_mut("function").and_then(Value::as_object_mut) else { return Ok(false); }; @@ -1982,7 +1902,7 @@ fn mask_claude_messages_request_value( value: &mut Value, session: &mut RedactionSession, scan_state: &mut RedactionScanState, -) -> Result { +) -> Result { let mut redacted = false; if let Some(system) = value.get_mut("system") { redacted |= mask_claude_text_content_value(system, session, scan_state)?; @@ -2001,7 +1921,7 @@ fn mask_claude_text_content_value( value: &mut Value, session: &mut RedactionSession, scan_state: &mut RedactionScanState, -) -> Result { +) -> Result { match value { Value::String(text) => mask_json_string(text, session, scan_state), Value::Array(parts) => { @@ -2019,7 +1939,7 @@ fn mask_claude_content_part( part: &mut Value, session: &mut RedactionSession, scan_state: &mut RedactionScanState, -) -> Result { +) -> Result { let Some(part) = part.as_object_mut() else { return Ok(false); }; @@ -2050,7 +1970,7 @@ fn mask_openai_responses_request_value( value: &mut Value, session: &mut RedactionSession, scan_state: &mut RedactionScanState, -) -> Result { +) -> Result { let mut redacted = false; if let Some(Value::String(instructions)) = value.get_mut("instructions") { redacted |= mask_json_string(instructions, session, scan_state)?; @@ -2065,7 +1985,7 @@ fn mask_openai_search_request_value( value: &mut Value, session: &mut RedactionSession, scan_state: &mut RedactionScanState, -) -> Result { +) -> Result { let mut redacted = false; if let Some(input) = value.get_mut("input") { redacted |= mask_openai_responses_input_value(input, session, scan_state)?; @@ -2090,7 +2010,7 @@ fn mask_openai_responses_input_value( value: &mut Value, session: &mut RedactionSession, scan_state: &mut RedactionScanState, -) -> Result { +) -> Result { match value { Value::String(text) => mask_json_string(text, session, scan_state), Value::Array(items) => { @@ -2108,7 +2028,7 @@ fn mask_openai_responses_input_item( item: &mut Value, session: &mut RedactionSession, scan_state: &mut RedactionScanState, -) -> Result { +) -> Result { let Some(item) = item.as_object_mut() else { return Ok(false); }; @@ -2133,7 +2053,7 @@ fn mask_openai_responses_content_value( content: &mut Value, session: &mut RedactionSession, scan_state: &mut RedactionScanState, -) -> Result { +) -> Result { match content { Value::String(text) => mask_json_string(text, session, scan_state), Value::Array(parts) => { @@ -2151,7 +2071,7 @@ fn mask_openai_responses_content_part( part: &mut Value, session: &mut RedactionSession, scan_state: &mut RedactionScanState, -) -> Result { +) -> Result { let Some(part) = part.as_object_mut() else { return Ok(false); }; @@ -2168,7 +2088,7 @@ fn mask_json_string( text: &mut String, session: &mut RedactionSession, scan_state: &mut RedactionScanState, -) -> Result { +) -> Result { let redacted = session.redact_text_checked(text, scan_state)?; if redacted.matches.is_empty() { return Ok(false); @@ -4169,9 +4089,9 @@ mod tests { try_mask_chat_pii_request_value_with_cache_options, try_mask_chat_request_json_with_cache_options, try_mask_chat_request_json_with_options, ChatPiiRedactionRequestFormat, ChatPiiRedactionRuntimeConfig, DetectorProbe, MappingKey, - MaskChatRequestOptions, RedactionKind, RedactionLimitError, RedactionMapping, - RedactionScanLimits, RedactionSession, RedactionSessionConfig, RedactionSessionSlot, - RedisRedactionMappingCache, SentinelMatcher, StreamingResponseRestorer, + MaskChatRequestOptions, RedactionKind, RedactionMapping, RedactionSession, + RedactionSessionConfig, RedactionSessionSlot, RedisRedactionMappingCache, SentinelMatcher, + StreamingResponseRestorer, }; use std::collections::BTreeMap; use std::time::Duration; @@ -5624,59 +5544,39 @@ mod tests { } #[test] - fn pii_redaction_performance_limits_reject_scanned_text_and_detection_overflow() { - assert_eq!( - RedactionScanLimits::default(), - RedactionScanLimits { - max_scanned_text_bytes: 2 * 1024 * 1024, - max_detections: 1024, - } - ); - + fn pii_redaction_accepts_text_and_detection_counts_above_previous_caps() { let large_request = json!({ "model": "gpt-5", - "messages": [{"role": "user", "content": "x".repeat(2 * 1024 * 1024 + 1)}] + "messages": [{ + "role": "user", + "content": format!("alice@example.com {}", "x".repeat(2 * 1024 * 1024 + 1)) + }] }); - let large_err = try_mask_chat_request_json_with_options( + let large_masked = try_mask_chat_request_json_with_options( &serde_json::to_vec(&large_request).expect("request should serialize"), test_config(), MaskChatRequestOptions::runtime(), ) - .expect_err("oversized scan should fail closed"); - assert_eq!( - large_err, - RedactionLimitError::ScannedTextTooLarge { - limit: 2 * 1024 * 1024, - } - ); - assert_eq!( - large_err.client_status(), - http::StatusCode::PAYLOAD_TOO_LARGE - ); - assert!(!large_err.safe_message().contains("alice@example.com")); + .expect("text above the previous scan cap should be redacted"); + assert!(large_masked.redacted); + assert_eq!(large_masked.session.mapping_count(), 1); + let dense_content = (0..1025) + .map(|index| format!("user{index}@example.com")) + .collect::>() + .join(" "); let dense_request = json!({ "model": "gpt-5", - "messages": [{"role": "user", "content": "alice@example.com bob@example.net"}] + "messages": [{"role": "user", "content": dense_content}] }); - let dense_err = try_mask_chat_request_json_with_options( + let dense_masked = try_mask_chat_request_json_with_options( &serde_json::to_vec(&dense_request).expect("request should serialize"), test_config(), - MaskChatRequestOptions::runtime().with_scan_limits(RedactionScanLimits { - max_scanned_text_bytes: 1024, - max_detections: 1, - }), + MaskChatRequestOptions::runtime(), ) - .expect_err("too many detections should fail closed"); - assert_eq!( - dense_err, - RedactionLimitError::TooManyDetections { limit: 1 } - ); - assert_eq!( - dense_err.client_status(), - http::StatusCode::UNPROCESSABLE_ENTITY - ); - assert!(!dense_err.safe_message().contains("alice@example.com")); + .expect("detections above the previous cap should be redacted"); + assert!(dense_masked.redacted); + assert_eq!(dense_masked.session.mapping_count(), 1025); } #[tokio::test] diff --git a/apps/aether-gateway/src/request_diagnostics.rs b/apps/aether-gateway/src/request_diagnostics.rs index ad4de6458..ff349e632 100644 --- a/apps/aether-gateway/src/request_diagnostics.rs +++ b/apps/aether-gateway/src/request_diagnostics.rs @@ -53,6 +53,20 @@ impl RequestDiagnostics { inner.request_accepted_at } + pub(crate) fn request_accepted_elapsed_ms(&self) -> Option { + self.request_accepted_at() + .map(|accepted_at| accepted_at.elapsed().as_millis().min(u128::from(u64::MAX)) as u64) + } + + fn request_elapsed_ms_at(&self, observed_at: Instant) -> Option { + self.request_accepted_at().map(|accepted_at| { + observed_at + .saturating_duration_since(accepted_at) + .as_millis() + .min(u128::from(u64::MAX)) as u64 + }) + } + fn record_db_timing_ms(&self, operation: &'static str, elapsed_ms: u64) { let Ok(mut inner) = self.inner.lock() else { return; @@ -151,9 +165,20 @@ pub(crate) async fn scope_request_diagnostics(future: F) -> F::Output where F: Future, { - REQUEST_DIAGNOSTICS - .scope(Arc::new(RequestDiagnostics::default()), future) - .await + scope_request_diagnostics_with(Some(Arc::new(RequestDiagnostics::default())), future).await +} + +pub(crate) async fn scope_request_diagnostics_with( + diagnostics: Option>, + future: F, +) -> F::Output +where + F: Future, +{ + match diagnostics { + Some(diagnostics) => REQUEST_DIAGNOSTICS.scope(diagnostics, future).await, + None => future.await, + } } pub(crate) fn current_request_diagnostics() -> Option> { @@ -199,23 +224,231 @@ pub(crate) fn attach_request_diagnostics_to_report_context( report_context: Option, diagnostics: Option<&Arc>, ) -> Option { - let Some(db_timings_ms) = diagnostics.and_then(|diagnostics| diagnostics.db_timings_metadata()) - else { + let db_timings_ms = diagnostics.and_then(|diagnostics| diagnostics.db_timings_metadata()); + let end_to_end_time_ms = + diagnostics.and_then(|diagnostics| diagnostics.request_accepted_elapsed_ms()); + if db_timings_ms.is_none() && end_to_end_time_ms.is_none() { return report_context; - }; + } let mut object = match report_context { Some(Value::Object(object)) => object, Some(other) => Map::from_iter([("seed".to_string(), other)]), None => Map::new(), }; - object.insert("db_timings_ms".to_string(), db_timings_ms); + if let Some(db_timings_ms) = db_timings_ms { + object.insert("db_timings_ms".to_string(), db_timings_ms); + } + if let Some(end_to_end_time_ms) = end_to_end_time_ms { + object.insert( + "end_to_end_time_ms".to_string(), + Value::from(end_to_end_time_ms), + ); + } Some(Value::Object(object)) } +pub(crate) fn attach_request_diagnostics_and_candidate_timing_to_report_context( + report_context: Option, + diagnostics: Option<&Arc>, + candidate_elapsed_ms: Option, + candidate_ttfb_ms: Option, +) -> Option { + let end_to_end_time_ms = + diagnostics.and_then(|diagnostics| diagnostics.request_accepted_elapsed_ms()); + let mut report_context = + attach_request_diagnostics_to_report_context(report_context, diagnostics); + let Some(end_to_end_first_byte_time_ms) = + end_to_end_first_byte_time_ms(end_to_end_time_ms, candidate_elapsed_ms, candidate_ttfb_ms) + else { + return report_context; + }; + + let object = report_context + .get_or_insert_with(|| Value::Object(Map::new())) + .as_object_mut()?; + object.insert( + "end_to_end_first_byte_time_ms".to_string(), + Value::from(end_to_end_first_byte_time_ms), + ); + report_context +} + +pub(crate) fn attach_request_diagnostics_and_candidate_start_timing_to_report_context( + report_context: Option, + diagnostics: Option<&Arc>, + candidate_started_at: Option, + candidate_ttfb_ms: Option, +) -> Option { + let end_to_end_time_ms = + diagnostics.and_then(|diagnostics| diagnostics.request_accepted_elapsed_ms()); + let candidate_started_elapsed_ms = diagnostics + .zip(candidate_started_at) + .and_then(|(diagnostics, started_at)| diagnostics.request_elapsed_ms_at(started_at)); + let mut report_context = + attach_request_diagnostics_to_report_context(report_context, diagnostics); + let Some(end_to_end_first_byte_time_ms) = end_to_end_first_byte_time_ms_from_candidate_start( + end_to_end_time_ms, + candidate_started_elapsed_ms, + candidate_ttfb_ms, + ) else { + return report_context; + }; + + let object = report_context + .get_or_insert_with(|| Value::Object(Map::new())) + .as_object_mut()?; + object.insert( + "end_to_end_first_byte_time_ms".to_string(), + Value::from(end_to_end_first_byte_time_ms), + ); + report_context +} + +pub(crate) fn attach_current_request_diagnostics_and_candidate_timing_to_report_context( + report_context: Option<&Value>, + candidate_elapsed_ms: Option, + candidate_ttfb_ms: Option, +) -> Option { + let diagnostics = current_request_diagnostics(); + attach_request_diagnostics_and_candidate_timing_to_report_context( + report_context.cloned(), + diagnostics.as_ref(), + candidate_elapsed_ms, + candidate_ttfb_ms, + ) +} + +pub(crate) fn attach_current_request_diagnostics_and_candidate_start_timing_to_report_context( + report_context: Option<&Value>, + candidate_started_at: Instant, + candidate_ttfb_ms: Option, +) -> Option { + let diagnostics = current_request_diagnostics(); + attach_request_diagnostics_and_candidate_start_timing_to_report_context( + report_context.cloned(), + diagnostics.as_ref(), + Some(candidate_started_at), + candidate_ttfb_ms, + ) +} + +fn end_to_end_first_byte_time_ms( + end_to_end_time_ms: Option, + candidate_elapsed_ms: Option, + candidate_ttfb_ms: Option, +) -> Option { + let end_to_end_time_ms = end_to_end_time_ms?; + let candidate_elapsed_ms = candidate_elapsed_ms?; + let candidate_ttfb_ms = candidate_ttfb_ms?; + Some( + end_to_end_time_ms + .saturating_sub(candidate_elapsed_ms) + .saturating_add(candidate_ttfb_ms) + .min(end_to_end_time_ms), + ) +} + +fn end_to_end_first_byte_time_ms_from_candidate_start( + end_to_end_time_ms: Option, + candidate_started_elapsed_ms: Option, + candidate_ttfb_ms: Option, +) -> Option { + let end_to_end_time_ms = end_to_end_time_ms?; + let candidate_started_elapsed_ms = candidate_started_elapsed_ms?; + let candidate_ttfb_ms = candidate_ttfb_ms?; + Some( + candidate_started_elapsed_ms + .saturating_add(candidate_ttfb_ms) + .min(end_to_end_time_ms), + ) +} + +pub(crate) fn calibrate_candidate_first_byte_elapsed_ms( + candidate_elapsed_at_result_ms: u64, + execution_elapsed_ms: Option, + execution_ttfb_ms: Option, +) -> Option { + let execution_elapsed_ms = execution_elapsed_ms?; + let execution_ttfb_ms = execution_ttfb_ms?; + Some( + candidate_elapsed_at_result_ms + .saturating_sub(execution_elapsed_ms) + .saturating_add(execution_ttfb_ms) + .min(candidate_elapsed_at_result_ms), + ) +} + pub(crate) fn attach_current_request_diagnostics_to_report_context( report_context: Option<&Value>, ) -> Option { let diagnostics = current_request_diagnostics()?; attach_request_diagnostics_to_report_context(report_context.cloned(), Some(&diagnostics)) } + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::time::{Duration, Instant}; + + use serde_json::json; + + use super::{ + attach_request_diagnostics_and_candidate_start_timing_to_report_context, + calibrate_candidate_first_byte_elapsed_ms, + end_to_end_first_byte_time_ms_from_candidate_start, RequestDiagnostics, + }; + + #[test] + fn candidate_first_byte_calibration_includes_pre_transport_wait() { + assert_eq!( + calibrate_candidate_first_byte_elapsed_ms(826, Some(626), Some(120)), + Some(320) + ); + } + + #[test] + fn end_to_end_first_byte_matches_candidate_ttfb_without_prior_retry() { + assert_eq!( + end_to_end_first_byte_time_ms_from_candidate_start(Some(626), Some(0), Some(120)), + Some(120) + ); + } + + #[test] + fn end_to_end_first_byte_includes_time_spent_before_successful_retry() { + assert_eq!( + end_to_end_first_byte_time_ms_from_candidate_start( + Some(10_626), + Some(10_000), + Some(120), + ), + Some(10_120) + ); + } + + #[test] + fn explicit_diagnostics_add_end_to_end_and_first_byte_timing() { + let diagnostics = Arc::new(RequestDiagnostics::default()); + let accepted_at = Instant::now() - Duration::from_millis(1_000); + let candidate_started_at = accepted_at + Duration::from_millis(800); + diagnostics.record_request_accepted_at(accepted_at); + + let context = attach_request_diagnostics_and_candidate_start_timing_to_report_context( + Some(json!({"candidate_index": 1})), + Some(&diagnostics), + Some(candidate_started_at), + Some(50), + ) + .expect("diagnostics context should build"); + + let end_to_end = context["end_to_end_time_ms"] + .as_u64() + .expect("end-to-end timing should exist"); + let first_byte = context["end_to_end_first_byte_time_ms"] + .as_u64() + .expect("first-byte timing should exist"); + assert!(end_to_end >= 1_000); + assert_eq!(first_byte, 850); + } +} diff --git a/apps/aether-gateway/src/tests/ai_execute/sync/chat/pii_redaction.rs b/apps/aether-gateway/src/tests/ai_execute/sync/chat/pii_redaction.rs index ad5f6866b..586f6943d 100644 --- a/apps/aether-gateway/src/tests/ai_execute/sync/chat/pii_redaction.rs +++ b/apps/aether-gateway/src/tests/ai_execute/sync/chat/pii_redaction.rs @@ -782,82 +782,30 @@ async fn ai_execute_pii_redaction_restores_executed_candidate_session_after_late } large_stack_async_test!( - pii_redaction_performance_limits_do_not_forward_unredacted_body_upstream, - pii_redaction_performance_limits_do_not_forward_unredacted_body_upstream_impl + pii_redaction_forwards_text_above_previous_scan_cap_only_after_masking, + pii_redaction_forwards_text_above_previous_scan_cap_only_after_masking_impl ); -async fn pii_redaction_performance_limits_do_not_forward_unredacted_body_upstream_impl() { - let provider_hits = Arc::new(AtomicUsize::new(0)); - let provider_hits_clone = Arc::clone(&provider_hits); - let provider_app = Router::new().route( - "/v1/chat/completions", - any(move |_request: Request| { - let provider_hits_inner = Arc::clone(&provider_hits_clone); - async move { - provider_hits_inner.fetch_add(1, Ordering::SeqCst); - Json(json!({ - "id": "unexpected", - "choices": [{"message": {"role": "assistant", "content": "unexpected"}}] - })) - } - }), +async fn pii_redaction_forwards_text_above_previous_scan_cap_only_after_masking_impl() { + let mut request = rich_pii_request(); + request["messages"] + .as_array_mut() + .expect("messages should be an array") + .push(json!({ + "role": "user", + "content": "x".repeat(2 * 1024 * 1024 + 1), + })); + + let (response_json, seen) = + run_sync_redaction_case("pii-redaction-large-request", true, true, "known", request).await; + + let provider_body_text = serde_json::to_string(&seen.body).expect("body should serialize"); + assert!(!provider_body_text.contains("alice@example.com")); + assert!(provider_body_text.contains(")); + let captured_body_clone = Arc::clone(&captured_body); let owner = Router::new().route( "/api/internal/tunnel/relay/node-123", - post(move |_request: Request| { - let owner_hits_inner = Arc::clone(&owner_hits_clone); + post(move |body: Body| { + let captured_body_inner = Arc::clone(&captured_body_clone); async move { - *owner_hits_inner.lock().expect("mutex should lock") += 1; - (StatusCode::OK, Body::from("unexpected owner hit")) + let body = axum::body::to_bytes(body, usize::MAX) + .await + .expect("owner body should read"); + *captured_body_inner.lock().expect("mutex should lock") = Some(body); + StatusCode::OK } }), ); @@ -564,7 +572,10 @@ async fn gateway_rejects_owner_relay_body_above_configured_limit() { "observed_at_unix_secs": 4_102_444_800u64, }), ), - ("max_request_body_size".to_string(), json!(8)), + ( + "max_request_body_size".to_string(), + json!(RECORDING_LIMIT_BYTES), + ), ]); let gateway = build_router_with_state( AppState::new() @@ -574,15 +585,24 @@ async fn gateway_rejects_owner_relay_body_above_configured_limit() { ); let (gateway_url, gateway_handle) = start_server(gateway).await; + let request_payload = vec![b'x'; RECORDING_LIMIT_BYTES + 1]; + let envelope = relay_envelope( + &relay_request_meta(false, Some(60_000), None), + &request_payload, + ); + assert!(envelope.len() > RECORDING_LIMIT_BYTES); let response = reqwest::Client::new() .post(format!("{gateway_url}/api/internal/tunnel/relay/node-123")) - .body("relay-envelope") + .body(envelope.clone()) .send() .await .expect("request should succeed"); - assert_eq!(response.status(), StatusCode::PAYLOAD_TOO_LARGE); - assert_eq!(*owner_hits.lock().expect("mutex should lock"), 0); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!( + captured_body.lock().expect("mutex should lock").as_ref(), + Some(&Bytes::from(envelope)) + ); gateway_handle.abort(); owner_handle.abort(); diff --git a/apps/aether-gateway/src/tests/usage/local.rs b/apps/aether-gateway/src/tests/usage/local.rs index a64a3080d..b4e23bd58 100644 --- a/apps/aether-gateway/src/tests/usage/local.rs +++ b/apps/aether-gateway/src/tests/usage/local.rs @@ -171,6 +171,7 @@ async fn gateway_handles_local_openai_chat_sync_report_with_local_reporting_when } }, "telemetry": { + "ttfb_ms": 10, "elapsed_ms": 25 } })) @@ -241,6 +242,19 @@ async fn gateway_handles_local_openai_chat_sync_report_with_local_reporting_when assert_eq!(stored_usage.status, "completed"); assert_eq!(stored_usage.total_tokens, 5); assert_eq!(stored_usage.response_time_ms, Some(25)); + let end_to_end_time_ms = stored_usage + .request_metadata + .as_ref() + .and_then(|metadata| metadata.get("end_to_end_time_ms")) + .and_then(serde_json::Value::as_u64) + .expect("end-to-end latency should be persisted separately"); + let end_to_end_first_byte_time_ms = stored_usage + .request_metadata + .as_ref() + .and_then(|metadata| metadata.get("end_to_end_first_byte_time_ms")) + .and_then(serde_json::Value::as_u64) + .expect("end-to-end first-byte latency should be persisted separately"); + assert!(end_to_end_first_byte_time_ms <= end_to_end_time_ms); let stored_candidates = request_candidate_repository .list_by_request_id("trace-openai-chat-local-report-sync-123") @@ -422,14 +436,14 @@ async fn gateway_truncates_deep_request_echo_for_local_openai_chat_sync_usage_im } #[test] -fn gateway_applies_system_max_request_body_size_to_local_openai_chat_sync_usage() { +fn gateway_ignores_legacy_max_request_body_size_for_local_openai_chat_sync_usage() { run_async_test_on_large_stack( - "gateway_applies_system_max_request_body_size_to_local_openai_chat_sync_usage", - gateway_applies_system_max_request_body_size_to_local_openai_chat_sync_usage_impl(), + "gateway_ignores_legacy_max_request_body_size_for_local_openai_chat_sync_usage", + gateway_ignores_legacy_max_request_body_size_for_local_openai_chat_sync_usage_impl(), ); } -async fn gateway_applies_system_max_request_body_size_to_local_openai_chat_sync_usage_impl() { +async fn gateway_ignores_legacy_max_request_body_size_for_local_openai_chat_sync_usage_impl() { let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); @@ -546,13 +560,13 @@ async fn gateway_applies_system_max_request_body_size_to_local_openai_chat_sync_ assert_eq!(stored_usage.total_tokens, 5); assert_eq!( stored_usage.request_body_state, - Some(UsageBodyCaptureState::Truncated) + Some(UsageBodyCaptureState::Inline) ); assert_eq!( stored_usage.provider_request_body_state, - Some(UsageBodyCaptureState::Truncated) + Some(UsageBodyCaptureState::Inline) ); - assert_eq!( + assert_ne!( stored_usage .request_body .as_ref() @@ -905,6 +919,7 @@ async fn gateway_records_failed_usage_when_sync_runtime_transport_is_unavailable .expect("gateway should build") .with_execution_runtime_sync_override_for_tests(move |_plan| { *execution_hits_clone.lock().expect("mutex should lock") += 1; + std::thread::sleep(std::time::Duration::from_millis(5)); Err(crate::GatewayError::Internal( "simulated transport unavailable".to_string(), )) @@ -968,12 +983,148 @@ async fn gateway_records_failed_usage_when_sync_runtime_transport_is_unavailable .expect("request candidate trace should read"); assert_eq!(stored_candidates.len(), 1); assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Failed); + assert!(stored_candidates[0] + .latency_ms + .is_some_and(|value| value >= 5)); assert_eq!( stored_candidates[0].error_type.as_deref(), Some("execution_runtime_unavailable") ); } +#[test] +fn sync_transport_error_policy_stops_or_retries_candidates_end_to_end() { + run_async_test_on_large_stack( + "sync_transport_error_policy_stops_or_retries_candidates_end_to_end", + sync_transport_error_policy_stops_or_retries_candidates_end_to_end_impl(), + ); +} + +async fn sync_transport_error_policy_stops_or_retries_candidates_end_to_end_impl() { + async fn run_case(stop_on_transport_errors: bool) -> (StatusCode, usize, Vec>) { + let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); + let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); + let execution_hits = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let execution_hits_for_override = Arc::clone(&execution_hits); + + let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( + Some(hash_api_key("sk-client-transport-policy")), + sample_local_openai_auth_snapshot("api-key-transport-policy", "user-transport-policy"), + )])); + let mut second_candidate = sample_local_openai_candidate_row(); + second_candidate.key_id = "key-openai-usage-local-2".to_string(); + second_candidate.key_name = "secondary".to_string(); + second_candidate.key_internal_priority = second_candidate.key_internal_priority - 1; + let candidate_selection_repository = + Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![ + sample_local_openai_candidate_row(), + second_candidate, + ])); + + let mut provider = sample_local_openai_provider(); + provider.config = stop_on_transport_errors.then(|| { + json!({ + "failover_rules": { + "stop_on_transport_errors": true, + } + }) + }); + let mut second_key = sample_local_openai_key(); + second_key.id = "key-openai-usage-local-2".to_string(); + second_key.name = "secondary".to_string(); + let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( + vec![provider], + vec![sample_local_openai_endpoint()], + vec![sample_local_openai_key(), second_key], + )); + + let gateway_state = crate::AppState::new() + .expect("gateway should build") + .with_execution_runtime_sync_override_for_tests(move |plan| { + let hit = execution_hits_for_override + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + if hit == 0 { + return Err(crate::GatewayError::Internal( + "simulated transport unavailable".to_string(), + )); + } + Ok(aether_contracts::ExecutionResult { + request_id: plan.request_id.clone(), + candidate_id: plan.candidate_id.clone(), + status_code: 200, + headers: std::collections::BTreeMap::from([( + "content-type".to_string(), + "application/json".to_string(), + )]), + body: Some(aether_contracts::ResponseBody { + json_body: Some(json!({ + "id": "chatcmpl-transport-policy", + "choices": [{"message": {"role": "assistant", "content": "ok"}}] + })), + body_bytes_b64: None, + }), + telemetry: Some(aether_contracts::ExecutionTelemetry { + ttfb_ms: Some(1), + elapsed_ms: Some(1), + upstream_bytes: None, + }), + error: None, + }) + }) + .with_data_state_for_tests( + GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_and_usage_for_tests( + auth_repository, + candidate_selection_repository, + provider_catalog_repository, + Arc::clone(&request_candidate_repository), + usage_repository, + DEVELOPMENT_ENCRYPTION_KEY, + ), + ); + let gateway = build_router_with_state(gateway_state); + let trace_id = if stop_on_transport_errors { + "trace-transport-policy-stop" + } else { + "trace-transport-policy-retry" + }; + let request = Request::builder() + .method(http::Method::POST) + .uri("/v1/chat/completions") + .header(http::header::CONTENT_TYPE, "application/json") + .header( + http::header::AUTHORIZATION, + "Bearer sk-client-transport-policy", + ) + .header(TRACE_ID_HEADER, trace_id) + .body(Body::from("{\"model\":\"gpt-5\",\"messages\":[]}")) + .expect("request should build"); + let response = send_request(gateway, request).await; + let candidates = request_candidate_repository + .list_by_request_id(trace_id) + .await + .expect("request candidates should read"); + ( + response.status(), + execution_hits.load(std::sync::atomic::Ordering::SeqCst), + candidates + .into_iter() + .filter(|candidate| candidate.status == RequestCandidateStatus::Failed) + .map(|candidate| candidate.status_code) + .collect(), + ) + } + + let (retry_status, retry_hits, retry_failure_statuses) = run_case(false).await; + assert_eq!(retry_status, StatusCode::OK); + assert!(retry_hits >= 2); + assert_eq!(retry_failure_statuses.first(), Some(&None)); + + let (stop_status, stop_hits, stop_failure_statuses) = run_case(true).await; + assert_eq!(stop_status, StatusCode::BAD_GATEWAY); + assert_eq!(stop_hits, 1); + assert_eq!(stop_failure_statuses, vec![None]); +} + #[test] fn gateway_records_failed_usage_for_claude_runtime_miss_without_execution_exhaustion() { run_async_test_on_large_stack( @@ -1365,14 +1516,14 @@ async fn gateway_handles_local_openai_chat_stream_report_with_local_reporting_wh } #[test] -fn gateway_preserves_stream_usage_when_max_response_body_size_truncates_capture() { +fn gateway_ignores_legacy_max_response_body_size_for_stream_usage() { run_async_test_on_large_stack( - "gateway_preserves_stream_usage_when_max_response_body_size_truncates_capture", - gateway_preserves_stream_usage_when_max_response_body_size_truncates_capture_impl(), + "gateway_ignores_legacy_max_response_body_size_for_stream_usage", + gateway_ignores_legacy_max_response_body_size_for_stream_usage_impl(), ); } -async fn gateway_preserves_stream_usage_when_max_response_body_size_truncates_capture_impl() { +async fn gateway_ignores_legacy_max_response_body_size_for_stream_usage_impl() { let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); @@ -1524,13 +1675,13 @@ async fn gateway_preserves_stream_usage_when_max_response_body_size_truncates_ca assert_eq!(stored_usage.total_tokens, 6); assert_eq!( stored_usage.response_body_state, - Some(UsageBodyCaptureState::Truncated) + Some(UsageBodyCaptureState::Inline) ); assert_eq!( stored_usage.client_response_body_state, - Some(UsageBodyCaptureState::Truncated) + Some(UsageBodyCaptureState::Inline) ); - assert_eq!( + assert_ne!( stored_usage .response_body .as_ref() diff --git a/apps/aether-gateway/src/tunnel/mod.rs b/apps/aether-gateway/src/tunnel/mod.rs index fcddb1e70..b67b39418 100644 --- a/apps/aether-gateway/src/tunnel/mod.rs +++ b/apps/aether-gateway/src/tunnel/mod.rs @@ -3,13 +3,13 @@ mod embedded; use std::collections::HashMap; use std::fmt; use std::io; -use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use std::time::{Duration, SystemTime}; use aether_contracts::tunnel::{ resolve_tunnel_request_timeouts, try_decode_tunnel_relay_request_meta, RequestMeta, - TUNNEL_RELAY_FORWARDED_BY_HEADER, TUNNEL_RELAY_OWNER_INSTANCE_HEADER, + MAX_TUNNEL_RELAY_META_LEN, TUNNEL_RELAY_FORWARDED_BY_HEADER, + TUNNEL_RELAY_OWNER_INSTANCE_HEADER, }; use aether_data::repository::proxy_nodes::{ ProxyNodeHeartbeatMutation, ProxyNodeTunnelStatusMutation, StoredProxyNode, @@ -39,8 +39,8 @@ use super::AppState; pub(crate) use aether_gateway_tunnel::{ is_tunnel_heartbeat_path, is_tunnel_node_status_path, TunnelAttachmentRecord, - DEFAULT_OWNER_RELAY_BODY_LIMIT_BYTES, DEFAULT_TUNNEL_PROBE_BODY_LIMIT_BYTES, PROXY_TUNNEL_PATH, - TUNNEL_HEARTBEAT_PATH, TUNNEL_NODE_STATUS_PATH, TUNNEL_RELAY_PATH_PATTERN, TUNNEL_ROUTE_FAMILY, + DEFAULT_TUNNEL_PROBE_BODY_LIMIT_BYTES, PROXY_TUNNEL_PATH, TUNNEL_HEARTBEAT_PATH, + TUNNEL_NODE_STATUS_PATH, TUNNEL_RELAY_PATH_PATTERN, TUNNEL_ROUTE_FAMILY, }; pub(crate) use embedded::DirectRelayResponse; pub(crate) use embedded::ProxyConn as TunnelProxyConn; @@ -827,37 +827,17 @@ async fn forward_relay_request_to_owner( ) -> Result, GatewayError> { let owner_url = build_owner_relay_url(&owner.relay_base_url, node_id)?; let (parts, body) = request.into_parts(); - let body_limit = owner_relay_body_limit_bytes(state.data.as_ref()).await; - if request_content_length_exceeds_limit(&parts.headers, body_limit) { - return build_local_http_error_response( - trace_id, - None, - StatusCode::PAYLOAD_TOO_LARGE, - &format!("tunnel relay body exceeds {body_limit} bytes"), - ); - } - let limit_exceeded = Arc::new(AtomicBool::new(false)); - let prepared_body = - match prepare_owner_relay_request_body(body, body_limit, Arc::clone(&limit_exceeded)).await - { - Ok(prepared_body) => prepared_body, - Err(_) if limit_exceeded.load(Ordering::SeqCst) => { - return build_local_http_error_response( - trace_id, - None, - StatusCode::PAYLOAD_TOO_LARGE, - &format!("tunnel relay body exceeds {body_limit} bytes"), - ); - } - Err(error) => { - return build_local_http_error_response( - trace_id, - None, - StatusCode::BAD_REQUEST, - &error, - ); - } - }; + let prepared_body = match prepare_owner_relay_request_body(body).await { + Ok(prepared_body) => prepared_body, + Err(error) => { + return build_local_http_error_response( + trace_id, + None, + StatusCode::BAD_REQUEST, + &error, + ); + } + }; let mut upstream_request = state.owner_forward_client.post(owner_url); for (name, value) in &parts.headers { @@ -894,14 +874,6 @@ async fn forward_relay_request_to_owner( .await { Ok(response) => response, - Err(err) if limit_exceeded.load(Ordering::SeqCst) => { - return build_local_http_error_response( - trace_id, - None, - StatusCode::PAYLOAD_TOO_LARGE, - &format!("tunnel relay body exceeds {body_limit} bytes"), - ); - } Err(err) => { return Err(GatewayError::Internal(format!( "owner tunnel relay failed: {err}" @@ -929,26 +901,6 @@ fn build_owner_relay_url(relay_base_url: &str, node_id: &str) -> Result usize { - data.find_system_config_value("max_request_body_size") - .await - .ok() - .flatten() - .and_then(|value| value.as_u64()) - .and_then(|value| usize::try_from(value).ok()) - .filter(|value| *value > 0) - .unwrap_or(DEFAULT_OWNER_RELAY_BODY_LIMIT_BYTES) -} - -fn request_content_length_exceeds_limit(headers: &HeaderMap, body_limit: usize) -> bool { - headers - .get(http::header::CONTENT_LENGTH) - .and_then(|value| value.to_str().ok()) - .and_then(|value| value.parse::().ok()) - .and_then(|value| usize::try_from(value).ok()) - .is_some_and(|value| value > body_limit) -} - struct PreparedOwnerRelayRequestBody { body: reqwest::Body, meta: RequestMeta, @@ -956,13 +908,10 @@ struct PreparedOwnerRelayRequestBody { async fn prepare_owner_relay_request_body( body: Body, - body_limit: usize, - limit_exceeded: Arc, ) -> Result { let mut body_stream = body.into_data_stream(); let mut buffered_chunks = Vec::new(); let mut meta_buffer = BytesMut::new(); - let mut forwarded = 0usize; let mut meta = None; while meta.is_none() { @@ -971,17 +920,16 @@ async fn prepare_owner_relay_request_body( }; match next_chunk { Ok(chunk) => { - let next_forwarded = forwarded.saturating_add(chunk.len()); - if next_forwarded > body_limit { - limit_exceeded.store(true, Ordering::SeqCst); - return Err(format!("tunnel relay body exceeds {body_limit} bytes")); + let meta_buffer_limit = 4usize.saturating_add(MAX_TUNNEL_RELAY_META_LEN); + let remaining = meta_buffer_limit.saturating_sub(meta_buffer.len()); + if remaining > 0 { + meta_buffer.extend_from_slice(&chunk[..chunk.len().min(remaining)]); } - forwarded = next_forwarded; - meta_buffer.extend_from_slice(&chunk); buffered_chunks.push(chunk); match try_decode_tunnel_relay_request_meta(&meta_buffer) { Ok(Some((parsed, _))) => meta = Some(parsed), - Ok(None) => {} + Ok(None) if meta_buffer.len() < meta_buffer_limit => {} + Ok(None) => return Err("incomplete tunnel relay metadata".to_string()), Err(error) => return Err(error), } } @@ -999,15 +947,6 @@ async fn prepare_owner_relay_request_body( while let Some(next_chunk) = body_stream.next().await { match next_chunk { Ok(chunk) => { - forwarded = forwarded.saturating_add(chunk.len()); - if forwarded > body_limit { - limit_exceeded.store(true, Ordering::SeqCst); - yield Err::(io::Error::new( - io::ErrorKind::InvalidData, - format!("tunnel relay body exceeds {body_limit} bytes"), - )); - break; - } yield Ok::(chunk); } Err(error) => { @@ -1194,7 +1133,6 @@ mod tests { use axum::routing::post; use axum::Router; use serde_json::json; - use std::sync::atomic::AtomicBool; use std::sync::{Arc, Mutex}; fn sample_proxy_node(node_id: &str) -> StoredProxyNode { @@ -1333,14 +1271,10 @@ mod tests { envelope.extend_from_slice(&1u32.to_be_bytes()); envelope.push(b'{'); - let error = prepare_owner_relay_request_body( - Body::from(envelope), - 1024, - Arc::new(AtomicBool::new(false)), - ) - .await - .err() - .expect("invalid metadata should fail"); + let error = prepare_owner_relay_request_body(Body::from(envelope)) + .await + .err() + .expect("invalid metadata should fail"); assert!(error.contains("invalid relay metadata")); } diff --git a/apps/aether-tunnel/.env.example b/apps/aether-tunnel/.env.example index 8ab259e80..a3d5cff75 100644 --- a/apps/aether-tunnel/.env.example +++ b/apps/aether-tunnel/.env.example @@ -11,9 +11,6 @@ AETHER_TUNNEL_NODE_NAME=jp-proxy-01 AETHER_TUNNEL_SECURITY=off # AETHER_TUNNEL_ENCRYPTION_KEY=base64-32-bytes -# Maximum request body buffered for 307/308 replay (supports K/M/G, 0 disables body replay buffering) -AETHER_TUNNEL_REDIRECT_REPLAY_BUDGET_BYTES=5M - # Optional tunnel TCP address-family restriction (set at most one to true) AETHER_TUNNEL_IPV4_ONLY=false AETHER_TUNNEL_IPV6_ONLY=false diff --git a/apps/aether-tunnel/README.md b/apps/aether-tunnel/README.md index c49c0fd3d..6299bb611 100644 --- a/apps/aether-tunnel/README.md +++ b/apps/aether-tunnel/README.md @@ -159,7 +159,8 @@ sudo aether-tunnel uninstall | `--upstream-tcp-keepalive-secs` | `AETHER_TUNNEL_UPSTREAM_TCP_KEEPALIVE_SECS` | `60` | TCP keepalive（秒，0 关闭） | | `--upstream-tcp-nodelay` | `AETHER_TUNNEL_UPSTREAM_TCP_NODELAY` | `true` | 启用 TCP_NODELAY | | `--upstream-proxy-url` | `AETHER_TUNNEL_UPSTREAM_PROXY_URL` | 空 | 仅 provider 上游请求使用的出口代理 | -| `--redirect-replay-budget-bytes` | `AETHER_TUNNEL_REDIRECT_REPLAY_BUDGET_BYTES` | `5M` | 307/308 请求体重放的预读预算，支持 `K/M/G`，`0` 表示禁用 body replay buffering | + +启用 `follow_redirects` 后，307/308 请求体重放不设置累计大小上限。 出口代理支持 `http://`、`socks5://`、`socks5h://`。配合 WARP sidecar 时可填写： diff --git a/apps/aether-tunnel/src/app.rs b/apps/aether-tunnel/src/app.rs index 213252cfa..77e0a0517 100644 --- a/apps/aether-tunnel/src/app.rs +++ b/apps/aether-tunnel/src/app.rs @@ -1005,9 +1005,7 @@ mod tests { use axum::Router; use serde_json::json; - use crate::config::{ - TunnelLogDestinationArg, TunnelLogRotationArg, DEFAULT_REDIRECT_REPLAY_BUDGET_BYTES, - }; + use crate::config::{TunnelLogDestinationArg, TunnelLogRotationArg}; use crate::hardware::HardwareInfo; use crate::state::AppState as TunnelAppState; use crate::target_filter::DnsCache; @@ -1397,7 +1395,7 @@ mod tests { upstream_tcp_keepalive_secs: 60, upstream_tcp_nodelay: true, upstream_proxy_url: None, - redirect_replay_budget_bytes: DEFAULT_REDIRECT_REPLAY_BUDGET_BYTES, + legacy_redirect_replay_budget_bytes_ignored: None, emit_proxy_timing_header: true, log_level: "info".to_string(), log_destination: TunnelLogDestinationArg::Stdout, diff --git a/apps/aether-tunnel/src/config.rs b/apps/aether-tunnel/src/config.rs index a8df48590..f1acbde7e 100644 --- a/apps/aether-tunnel/src/config.rs +++ b/apps/aether-tunnel/src/config.rs @@ -27,6 +27,9 @@ const REMOVED_TUNNEL_SECONDS_KEYS: &[&str] = &[ "tunnel_stale_timeout_secs", ]; const REMOVED_SINGLE_SERVER_KEYS: &[&str] = &["aether_url", "management_token"]; +/// Configuration keys that no longer affect runtime behavior but are ignored +/// while loading so existing installations can upgrade without editing TOML. +const IGNORED_CONFIG_KEYS: &[&str] = &["redirect_replay_budget_bytes"]; /// Fields renamed from 0.1.x `delegate_*` to 0.2.0 `upstream_*`. const DELEGATE_TO_UPSTREAM: &[(&str, &str)] = &[ @@ -46,13 +49,7 @@ const DELEGATE_TO_UPSTREAM: &[(&str, &str)] = &[ ("delegate_tcp_nodelay", "upstream_tcp_nodelay"), ]; -/// Default bytes buffered before a tunnel request becomes non-replayable for -/// 307/308 redirects. Kept aligned with the current admin-side request size -/// default, but exposed as an independent proxy transport budget. pub const DEFAULT_HEARTBEAT_INTERVAL_SECS: u64 = 5; -#[allow(dead_code)] -pub const DEFAULT_REDIRECT_REPLAY_BUDGET_BYTES: usize = 5_242_880; -pub const DEFAULT_REDIRECT_REPLAY_BUDGET_HUMAN: &str = "5M"; pub const DEFAULT_LOG_RETENTION_DAYS: u64 = 7; pub const DEFAULT_LOG_MAX_FILES: usize = 30; pub const DEFAULT_LOG_DIR: &str = "logs"; @@ -87,139 +84,6 @@ pub struct TunnelPoolSizing { pub initial_connections: u32, pub max_connections: u32, } -#[derive(Debug, Clone, PartialEq, Eq)] -enum ByteSizeValue { - Text(String), - Integer(u64), -} - -impl<'de> Deserialize<'de> for ByteSizeValue { - fn deserialize(deserializer: D) -> Result - where - D: serde::Deserializer<'de>, - { - struct ByteSizeValueVisitor; - - impl serde::de::Visitor<'_> for ByteSizeValueVisitor { - type Value = ByteSizeValue; - - fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter.write_str("a byte-size string like 5M or an integer byte count") - } - - fn visit_u64(self, value: u64) -> Result - where - E: serde::de::Error, - { - Ok(ByteSizeValue::Integer(value)) - } - - fn visit_i64(self, value: i64) -> Result - where - E: serde::de::Error, - { - if value < 0 { - return Err(E::custom("byte size must be >= 0")); - } - Ok(ByteSizeValue::Integer(value as u64)) - } - - fn visit_str(self, value: &str) -> Result - where - E: serde::de::Error, - { - Ok(ByteSizeValue::Text(value.to_string())) - } - - fn visit_string(self, value: String) -> Result - where - E: serde::de::Error, - { - Ok(ByteSizeValue::Text(value)) - } - } - - deserializer.deserialize_any(ByteSizeValueVisitor) - } -} - -fn deserialize_optional_byte_size<'de, D>(deserializer: D) -> Result, D::Error> -where - D: serde::Deserializer<'de>, -{ - let value = Option::::deserialize(deserializer)?; - value - .map(|value| match value { - ByteSizeValue::Text(text) => { - normalize_byte_size_text(&text).map_err(serde::de::Error::custom) - } - ByteSizeValue::Integer(value) => usize::try_from(value) - .map(format_byte_size_human) - .map_err(|_| serde::de::Error::custom("byte size exceeds usize")), - }) - .transpose() -} - -pub fn parse_byte_size(input: &str) -> Result { - let trimmed = input.trim(); - if trimmed.is_empty() { - return Err("byte size must not be empty".to_string()); - } - - let digits_end = trimmed - .find(|ch: char| !ch.is_ascii_digit()) - .unwrap_or(trimmed.len()); - if digits_end == 0 { - return Err(format!("invalid byte size `{trimmed}`")); - } - - let number = trimmed[..digits_end] - .parse::() - .map_err(|_| format!("invalid byte size `{trimmed}`"))?; - let suffix = trimmed[digits_end..].trim().to_ascii_lowercase(); - let multiplier = match suffix.as_str() { - "" | "b" => 1u64, - "k" | "kb" | "kib" => 1024u64, - "m" | "mb" | "mib" => 1024u64.pow(2), - "g" | "gb" | "gib" => 1024u64.pow(3), - _ => { - return Err(format!( - "invalid byte size suffix `{}`; use B, K, M, or G", - &trimmed[digits_end..].trim() - )) - } - }; - - let total = number - .checked_mul(multiplier) - .ok_or_else(|| format!("byte size `{trimmed}` is too large"))?; - usize::try_from(total).map_err(|_| format!("byte size `{trimmed}` exceeds usize")) -} - -fn normalize_byte_size_text(input: &str) -> Result { - parse_byte_size(input).map(format_byte_size_human) -} - -pub fn format_byte_size_human(bytes: usize) -> String { - const KIB: usize = 1024; - const MIB: usize = 1024 * 1024; - const GIB: usize = 1024 * 1024 * 1024; - - if bytes == 0 { - return "0".to_string(); - } - if bytes.is_multiple_of(GIB) { - return format!("{}G", bytes / GIB); - } - if bytes.is_multiple_of(MIB) { - return format!("{}M", bytes / MIB); - } - if bytes.is_multiple_of(KIB) { - return format!("{}K", bytes / KIB); - } - bytes.to_string() -} - #[derive(clap::ValueEnum, Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] #[serde(rename_all = "lowercase")] pub enum TunnelLogDestinationArg { @@ -577,15 +441,14 @@ pub struct Config { #[arg(long, env = "AETHER_TUNNEL_UPSTREAM_PROXY_URL")] pub upstream_proxy_url: Option, - /// Maximum request body bytes buffered to support 307/308 redirect replay. - /// Accepts values like 5M / 512K / 1G. Set to 0 to disable request-body replay buffering. + /// Accepted only so older launch commands and environments keep working. + /// Redirect request bodies are always replayed without a cumulative size limit. #[arg( - long, + long = "redirect-replay-budget-bytes", env = "AETHER_TUNNEL_REDIRECT_REPLAY_BUDGET_BYTES", - value_parser = parse_byte_size, - default_value = DEFAULT_REDIRECT_REPLAY_BUDGET_HUMAN + hide = true )] - pub redirect_replay_budget_bytes: usize, + pub legacy_redirect_replay_budget_bytes_ignored: Option, /// Emit detailed x-proxy-timing headers on tunneled upstream responses. #[arg( @@ -1127,12 +990,6 @@ pub struct ConfigFile { pub upstream_tcp_nodelay: Option, #[serde(skip_serializing_if = "Option::is_none")] pub upstream_proxy_url: Option, - #[serde( - default, - skip_serializing_if = "Option::is_none", - deserialize_with = "deserialize_optional_byte_size" - )] - pub redirect_replay_budget_bytes: Option, #[serde(skip_serializing_if = "Option::is_none")] pub emit_proxy_timing_header: Option, #[serde(skip_serializing_if = "Option::is_none")] @@ -1324,10 +1181,6 @@ impl ConfigFile { self.upstream_tcp_nodelay ); set!("AETHER_TUNNEL_UPSTREAM_PROXY_URL", self.upstream_proxy_url); - set!( - "AETHER_TUNNEL_REDIRECT_REPLAY_BUDGET_BYTES", - self.redirect_replay_budget_bytes - ); set!( "AETHER_TUNNEL_EMIT_PROXY_TIMING_HEADER", self.emit_proxy_timing_header @@ -1418,10 +1271,20 @@ impl ConfigFile { fn parse_config_file_content(content: &str) -> anyhow::Result { reject_removed_config_keys(content)?; let mut value: toml::Value = toml::from_str(content)?; + discard_ignored_config_keys(&mut value); promote_server_scoped_upstream_proxy_url(&mut value)?; Ok(value.try_into()?) } +fn discard_ignored_config_keys(value: &mut toml::Value) { + let Some(root) = value.as_table_mut() else { + return; + }; + for key in IGNORED_CONFIG_KEYS { + root.remove(*key); + } +} + fn normalized_proxy_url(value: &Option) -> Option<&str> { value .as_deref() @@ -1534,31 +1397,46 @@ mod tests { use crate::hardware::HardwareInfo; #[test] - fn parse_byte_size_supports_human_units() { - assert_eq!( - parse_byte_size("5M").expect("5M should parse"), - 5 * 1024 * 1024 - ); - assert_eq!( - parse_byte_size("512K").expect("512K should parse"), - 512 * 1024 - ); - assert_eq!( - parse_byte_size("1G").expect("1G should parse"), - 1024 * 1024 * 1024 - ); - assert_eq!(parse_byte_size("0").expect("0 should parse"), 0); + fn config_file_load_ignores_removed_redirect_replay_budget() { + let config = parse_config_file_content("redirect_replay_budget_bytes = \"1K\"") + .expect("removed replay budget should not break existing config files"); + let serialized = toml::to_string(&config).expect("config should serialize"); + assert!(!serialized.contains("redirect_replay_budget_bytes")); } #[test] - fn config_file_deserializes_budget_from_integer_and_string() { - let numeric: ConfigFile = - toml::from_str("redirect_replay_budget_bytes = 5242880").expect("numeric toml"); - assert_eq!(numeric.redirect_replay_budget_bytes.as_deref(), Some("5M")); + fn cli_accepts_but_hides_legacy_redirect_replay_budget() { + let config = Config::parse_from([ + "aether-tunnel", + "--aether-url", + "https://example.com", + "--management-token", + "ae_test", + "--node-name", + "tunnel-test", + "--redirect-replay-budget-bytes", + "1K", + ]); - let stringy: ConfigFile = - toml::from_str("redirect_replay_budget_bytes = \"6m\"").expect("string toml"); - assert_eq!(stringy.redirect_replay_budget_bytes.as_deref(), Some("6M")); + assert_eq!( + config + .legacy_redirect_replay_budget_bytes_ignored + .as_deref(), + Some("1K") + ); + let mut command = Config::command(); + let legacy_arg = command + .get_arguments() + .find(|arg| arg.get_id() == "legacy_redirect_replay_budget_bytes_ignored") + .expect("legacy redirect replay argument"); + assert_eq!( + legacy_arg.get_env(), + Some(std::ffi::OsStr::new( + "AETHER_TUNNEL_REDIRECT_REPLAY_BUDGET_BYTES" + )) + ); + let help = command.render_long_help().to_string(); + assert!(!help.contains("redirect-replay-budget-bytes")); } #[test] diff --git a/apps/aether-tunnel/src/setup/tui.rs b/apps/aether-tunnel/src/setup/tui.rs index 5c0e008f9..f56066285 100644 --- a/apps/aether-tunnel/src/setup/tui.rs +++ b/apps/aether-tunnel/src/setup/tui.rs @@ -21,9 +21,8 @@ use ratatui::Frame; use ratatui::Terminal; use crate::config::{ - format_byte_size_human, parse_byte_size, ConfigFile, ServerEntry, TunnelLogDestinationArg, - TunnelLogRotationArg, DEFAULT_HEARTBEAT_INTERVAL_SECS, DEFAULT_LOG_MAX_FILES, - DEFAULT_LOG_RETENTION_DAYS, DEFAULT_REDIRECT_REPLAY_BUDGET_HUMAN, + ConfigFile, ServerEntry, TunnelLogDestinationArg, TunnelLogRotationArg, + DEFAULT_HEARTBEAT_INTERVAL_SECS, DEFAULT_LOG_MAX_FILES, DEFAULT_LOG_RETENTION_DAYS, }; use crate::egress_proxy::UpstreamProxyConfig; @@ -215,15 +214,6 @@ impl App { required: false, help: "Heartbeat interval in seconds; default is 5", }, - Field { - label: "Redirect Replay Budget", - key: "redirect_replay_budget_bytes", - value: DEFAULT_REDIRECT_REPLAY_BUDGET_HUMAN.to_string(), - kind: FieldKind::Text, - required: false, - help: - "Prebuffer budget for 307/308 replay, e.g. 5M; set 0 to disable buffering", - }, ], selected: 0, mode: Mode::Normal, @@ -297,7 +287,6 @@ impl App { }), "allow_private_targets" => cfg.allow_private_targets.map(|v| v.to_string()), "heartbeat_interval" => cfg.heartbeat_interval.map(|v| v.to_string()), - "redirect_replay_budget_bytes" => cfg.redirect_replay_budget_bytes.clone(), "upstream_proxy_url" => cfg.upstream_proxy_url.clone(), _ => None, }; @@ -369,15 +358,6 @@ impl App { Ok(Some(value)) } - fn parse_optional_redirect_replay_budget(&self) -> anyhow::Result> { - let Some(raw) = self.get_global("redirect_replay_budget_bytes") else { - return Ok(None); - }; - let bytes = parse_byte_size(raw.trim()) - .map_err(|err| anyhow::anyhow!("redirect replay budget invalid: {err}"))?; - Ok(Some(format_byte_size_human(bytes))) - } - fn parse_optional_upstream_proxy_url(&self) -> anyhow::Result> { let Some(raw) = self.get_global("upstream_proxy_url") else { return Ok(None); @@ -420,7 +400,6 @@ impl App { log_level: get_global("log_level"), allow_private_targets: Some(self.toggle_enabled("allow_private_targets")), heartbeat_interval: self.parse_optional_heartbeat_interval()?, - redirect_replay_budget_bytes: self.parse_optional_redirect_replay_budget()?, upstream_proxy_url: self.parse_optional_upstream_proxy_url()?, log_destination: Some(if save_logs_to_file { TunnelLogDestinationArg::Both @@ -745,12 +724,6 @@ impl App { Err(_) => false, } } - "redirect_replay_budget_bytes" => { - if trimmed.is_empty() { - return true; - } - parse_byte_size(trimmed).is_ok() - } "upstream_proxy_url" => { if trimmed.is_empty() { return true; @@ -1229,13 +1202,11 @@ mod tests { let mut app = sample_app(); set_global_field(&mut app, "allow_private_targets", "true"); set_global_field(&mut app, "heartbeat_interval", "45"); - set_global_field(&mut app, "redirect_replay_budget_bytes", "6m"); set_global_field(&mut app, "upstream_proxy_url", "socks5h://127.0.0.1:1080"); let cfg = app.to_config().expect("config should serialize"); assert_eq!(cfg.allow_private_targets, Some(true)); assert_eq!(cfg.heartbeat_interval, Some(45)); - assert_eq!(cfg.redirect_replay_budget_bytes.as_deref(), Some("6M")); assert_eq!( cfg.upstream_proxy_url.as_deref(), Some("socks5h://127.0.0.1:1080") diff --git a/apps/aether-tunnel/src/tunnel/mod.rs b/apps/aether-tunnel/src/tunnel/mod.rs index 70e995052..da2cf1ce2 100644 --- a/apps/aether-tunnel/src/tunnel/mod.rs +++ b/apps/aether-tunnel/src/tunnel/mod.rs @@ -536,7 +536,7 @@ mod tests { upstream_tcp_keepalive_secs: 60, upstream_tcp_nodelay: true, upstream_proxy_url: None, - redirect_replay_budget_bytes: crate::config::DEFAULT_REDIRECT_REPLAY_BUDGET_BYTES, + legacy_redirect_replay_budget_bytes_ignored: None, emit_proxy_timing_header: true, log_level: "info".to_string(), log_destination: crate::config::TunnelLogDestinationArg::Stdout, diff --git a/apps/aether-tunnel/src/tunnel/stream_handler.rs b/apps/aether-tunnel/src/tunnel/stream_handler.rs index 466cf6ab3..85f52c229 100644 --- a/apps/aether-tunnel/src/tunnel/stream_handler.rs +++ b/apps/aether-tunnel/src/tunnel/stream_handler.rs @@ -190,7 +190,6 @@ struct RequestTimeouts { #[derive(Debug)] struct RequestBodyReplayState { - budget_bytes: usize, state: Mutex, ready: Notify, } @@ -203,7 +202,6 @@ enum RequestBodyReplayStatus { }, Ready(Bytes), Empty, - NonReplayable, Error(String), } @@ -211,7 +209,6 @@ enum RequestBodyReplayStatus { enum ReplayBodyResolution { Empty, Replayable(Bytes), - NonReplayable, } #[derive(Debug)] @@ -400,7 +397,6 @@ async fn prepare_redirect_request_body( match state.wait_for_resolution(deadline).await? { ReplayBodyResolution::Empty => Ok(Some(empty_request_body())), ReplayBodyResolution::Replayable(body) => Ok(Some(buffered_request_body(body))), - ReplayBodyResolution::NonReplayable => Ok(None), } } ReplayableRequestBody::NonReplayable => Ok(None), @@ -409,9 +405,8 @@ async fn prepare_redirect_request_body( } impl RequestBodyReplayState { - fn new(budget_bytes: usize) -> Self { + fn new() -> Self { Self { - budget_bytes, state: Mutex::new(RequestBodyReplayStatus::Collecting { chunks: Vec::new(), buffered_len: 0, @@ -421,27 +416,14 @@ impl RequestBodyReplayState { } fn push_chunk(&self, payload: Bytes) { - let mut notify = false; + let mut state = self.state.lock().expect("request body replay state lock"); + if let RequestBodyReplayStatus::Collecting { + chunks, + buffered_len, + } = &mut *state { - let mut state = self.state.lock().expect("request body replay state lock"); - if let RequestBodyReplayStatus::Collecting { - chunks, - buffered_len, - } = &mut *state - { - let next_len = buffered_len.saturating_add(payload.len()); - if next_len > self.budget_bytes { - chunks.clear(); - *state = RequestBodyReplayStatus::NonReplayable; - notify = true; - } else { - *buffered_len = next_len; - chunks.push(payload); - } - } - } - if notify { - self.ready.notify_waiters(); + *buffered_len = buffered_len.saturating_add(payload.len()); + chunks.push(payload); } } @@ -492,9 +474,6 @@ impl RequestBodyReplayState { Some(Ok(ReplayBodyResolution::Replayable(body.clone()))) } RequestBodyReplayStatus::Empty => Some(Ok(ReplayBodyResolution::Empty)), - RequestBodyReplayStatus::NonReplayable => { - Some(Ok(ReplayBodyResolution::NonReplayable)) - } RequestBodyReplayStatus::Error(message) => Some(Err(message.clone())), } }; @@ -575,22 +554,18 @@ fn buffered_request_body(body: Bytes) -> upstream_client::UpstreamRequestBody { } // Drain tunnel body frames on a detached task so the shared dispatcher is no -// longer coupled to upstream body polling. Redirect replay still reuses a full -// in-memory copy when the request body completes within budget. +// longer coupled to upstream body polling. When requested, redirect replay +// retains a complete in-memory copy without a cumulative payload limit. fn prepare_request_body( stream_id: u32, body_rx: mpsc::Receiver, body_size: Arc, deadline: Instant, - replay_budget_bytes: usize, + capture_for_redirects: bool, frame_tx: FrameSender, ) -> PreparedRequestBody { let (spool_tx, spool_rx) = mpsc::channel(REQUEST_BODY_SPOOL_QUEUE_CAPACITY); - let replay_state = if replay_budget_bytes == 0 { - None - } else { - Some(Arc::new(RequestBodyReplayState::new(replay_budget_bytes))) - }; + let replay_state = capture_for_redirects.then(|| Arc::new(RequestBodyReplayState::new())); let replay_body = match replay_state.as_ref() { Some(state) => ReplayableRequestBody::Pending(Arc::clone(state)), None => ReplayableRequestBody::NonReplayable, @@ -632,7 +607,6 @@ async fn collect_request_body_for_replay( mut body_rx: mpsc::Receiver, body_size: Arc, deadline: Instant, - replay_budget_bytes: usize, frame_tx: &FrameSender, ) -> Result { let mut body = BytesMut::new(); @@ -650,13 +624,6 @@ async fn collect_request_body_for_replay( .map_err(|error| format!("gzip decompress failed: {error}"))?; if !payload.is_empty() { - if body.len().saturating_add(payload.len()) > replay_budget_bytes { - return Err(format!( - "request body exceeds redirect replay budget: {} > {}", - body.len().saturating_add(payload.len()), - replay_budget_bytes - )); - } body_size.fetch_add(payload.len(), Ordering::Relaxed); try_send_window_update(frame_tx, stream_id, payload.len()); body.extend_from_slice(&payload); @@ -675,10 +642,8 @@ async fn collect_request_body_for_replay( } } -fn replay_body_from_buffered(body: Bytes, replay_budget_bytes: usize) -> ReplayableRequestBody { - let state = Arc::new(RequestBodyReplayState::new( - replay_budget_bytes.max(body.len()).max(1), - )); +fn replay_body_from_buffered(body: Bytes) -> ReplayableRequestBody { + let state = Arc::new(RequestBodyReplayState::new()); if !body.is_empty() { state.push_chunk(body); } @@ -1484,8 +1449,7 @@ async fn handle_stream_inner( let first_byte_timeout = request_timeouts.first_byte_timeout; let request_body_size = Arc::new(AtomicUsize::new(0)); let request_has_body = request_likely_has_body(¤t_method, &meta.headers); - let replay_budget_bytes = state.config.redirect_replay_budget_bytes; - let can_buffer_redirect_body = request_has_body && follow_redirects && replay_budget_bytes > 0; + let can_buffer_redirect_body = request_has_body && follow_redirects; let request_body_mode = if can_buffer_redirect_body { "buffered_fixed" } else if request_has_body { @@ -1499,7 +1463,6 @@ async fn handle_stream_inner( body_rx, Arc::clone(&request_body_size), first_byte_deadline, - replay_budget_bytes, frame_tx, ) .await @@ -1524,7 +1487,7 @@ async fn handle_stream_inner( }; PreparedRequestBody { first_request_body: Some(buffered_request_body(buffered_body.clone())), - replay_body: replay_body_from_buffered(buffered_body, replay_budget_bytes), + replay_body: replay_body_from_buffered(buffered_body), } } else if request_has_body { prepare_request_body( @@ -1532,7 +1495,7 @@ async fn handle_stream_inner( body_rx, Arc::clone(&request_body_size), first_byte_deadline, - 0, + false, frame_tx.clone(), ) } else { @@ -1878,7 +1841,7 @@ mod tests { use crate::tunnel::client::build_tls_config; fn completed_replay_body(body: Bytes) -> ReplayableRequestBody { - let state = Arc::new(RequestBodyReplayState::new(body.len().max(1))); + let state = Arc::new(RequestBodyReplayState::new()); if !body.is_empty() { state.push_chunk(body); } @@ -1983,7 +1946,7 @@ mod tests { rx, Arc::clone(&body_size), Instant::now() + Duration::from_secs(1), - 1024, + true, frame_tx.clone(), ); let mut body = prepared @@ -2608,7 +2571,10 @@ mod tests { } #[tokio::test] - async fn preserves_redirect_response_when_replay_budget_is_zero() { + async fn follows_307_redirect_for_fragmented_body_larger_than_legacy_replay_budget() { + const BODY_LEN: usize = 5 * 1024 * 1024 + 1; + const REQUEST_FRAME_BYTES: usize = 32 * 1024; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") .await .expect("listener"); @@ -2617,7 +2583,8 @@ mod tests { .route( "/start", post(|body: Bytes| async move { - assert_eq!(body, Bytes::from_static(b"hello")); + assert_eq!(body.len(), BODY_LEN); + assert!(body.iter().all(|byte| *byte == b'x')); Response::builder() .status(StatusCode::TEMPORARY_REDIRECT) .header(header::LOCATION, "/final") @@ -2627,35 +2594,51 @@ mod tests { ) .route( "/final", - post(|| async { + post(|body: Bytes| async move { + assert_eq!(body.len(), BODY_LEN); + assert!(body.iter().all(|byte| *byte == b'x')); Response::builder() .status(StatusCode::OK) - .body(Body::from("unexpected")) + .body(Body::from("redirected")) .expect("final response") }), - ); + ) + .layer(axum::extract::DefaultBodyLimit::disable()); let server = tokio::spawn(async move { axum::serve(listener, app) .await .expect("test server should run"); }); - let host = "redirect-budget-zero.test"; - let state = sample_state_for_budget(addr.port(), 0); + let host = "redirect-unlimited.test"; + let mut config = sample_config(); + config.allowed_ports.push(addr.port()); + config.legacy_redirect_replay_budget_bytes_ignored = Some("1".to_string()); + let state = sample_state_with_config(config); cache_test_host(&state, host, addr).await; let server_ctx = sample_server(&state); let (frame_tx, sent, writer_handle) = spawn_test_writer(); let (body_tx, body_rx) = mpsc::channel(4); - body_tx - .send(TunnelFrame::new( - 1, - MsgType::RequestBody, - flags::END_STREAM, - Bytes::from_static(b"hello"), - )) - .await - .expect("send body"); - drop(body_tx); + let body_sender = tokio::spawn(async move { + let body = vec![b'x'; BODY_LEN]; + let chunk_count = body.len().div_ceil(REQUEST_FRAME_BYTES); + for (index, chunk) in body.chunks(REQUEST_FRAME_BYTES).enumerate() { + let frame_flags = if index + 1 == chunk_count { + flags::END_STREAM + } else { + 0 + }; + body_tx + .send(TunnelFrame::new( + 1, + MsgType::RequestBody, + frame_flags, + Bytes::copy_from_slice(chunk), + )) + .await + .expect("send request body chunk"); + } + }); let mut meta = sample_request_meta(); meta.method = "POST".to_string(); @@ -2672,6 +2655,7 @@ mod tests { test_response_window(), ) .await; + body_sender.await.expect("request body sender task"); let result = collect_stream_result(frame_tx, sent, writer_handle).await; server.abort(); @@ -2681,15 +2665,8 @@ mod tests { result.error ); let response = result.response.expect("response metadata"); - assert_eq!(response.status, 307); - assert_eq!( - response - .headers - .iter() - .find(|(name, _)| name.eq_ignore_ascii_case("location")) - .map(|(_, value)| value.as_str()), - Some("/final") - ); + assert_eq!(response.status, 200); + assert_eq!(result.body, Bytes::from_static(b"redirected")); } #[tokio::test] @@ -2825,14 +2802,6 @@ mod tests { sample_state_with_config(config) } - fn sample_state_for_budget(port: u16, redirect_replay_budget_bytes: usize) -> Arc { - ensure_rustls_provider(); - let mut config = sample_config(); - config.allowed_ports.push(port); - config.redirect_replay_budget_bytes = redirect_replay_budget_bytes; - sample_state_with_config(config) - } - fn sample_state_with_config(config: Config) -> Arc { let config = Arc::new(config); let dns_cache = Arc::new(DnsCache::new(Duration::from_secs(60), 128)); @@ -2912,7 +2881,7 @@ mod tests { upstream_tcp_keepalive_secs: 60, upstream_tcp_nodelay: true, upstream_proxy_url: None, - redirect_replay_budget_bytes: crate::config::DEFAULT_REDIRECT_REPLAY_BUDGET_BYTES, + legacy_redirect_replay_budget_bytes_ignored: None, emit_proxy_timing_header: true, log_level: "info".to_string(), log_destination: crate::config::TunnelLogDestinationArg::Stdout, diff --git a/crates/aether-admin/src/observability/usage.rs b/crates/aether-admin/src/observability/usage.rs index e703fdcc0..c4d6145f1 100644 --- a/crates/aether-admin/src/observability/usage.rs +++ b/crates/aether-admin/src/observability/usage.rs @@ -1076,6 +1076,14 @@ fn admin_usage_metadata_string<'a>( .filter(|value| !value.is_empty()) } +fn admin_usage_metadata_u64(item: &StoredRequestUsageAudit, key: &str) -> Option { + item.request_metadata + .as_ref() + .and_then(Value::as_object) + .and_then(|metadata| metadata.get(key)) + .and_then(Value::as_u64) +} + fn infer_client_family_from_user_agent(user_agent: &str) -> Option<&'static str> { let normalized = user_agent.trim().to_ascii_lowercase(); if normalized.is_empty() { @@ -1218,6 +1226,11 @@ fn admin_usage_active_request_json( "request_path_and_query": admin_usage_metadata_string(item, "request_path_and_query"), "has_fallback": admin_usage_has_fallback(item), }); + value["end_to_end_time_ms"] = json!(admin_usage_metadata_u64(item, "end_to_end_time_ms")); + value["end_to_end_first_byte_time_ms"] = json!(admin_usage_metadata_u64( + item, + "end_to_end_first_byte_time_ms" + )); if let Some(api_format) = item.api_format.as_ref() { value["api_format"] = json!(api_format); } @@ -1324,6 +1337,17 @@ pub fn admin_usage_record_json( let object = payload .as_object_mut() .expect("admin usage record payload should be an object"); + object.insert( + "end_to_end_time_ms".to_string(), + json!(admin_usage_metadata_u64(item, "end_to_end_time_ms")), + ); + object.insert( + "end_to_end_first_byte_time_ms".to_string(), + json!(admin_usage_metadata_u64( + item, + "end_to_end_first_byte_time_ms" + )), + ); object.insert("is_stream".to_string(), json!(item.is_stream)); object.insert( UPSTREAM_IS_STREAM_KEY.to_string(), @@ -2643,6 +2667,36 @@ mod tests { assert_eq!(record["client_is_stream"], false); } + #[test] + fn admin_usage_payloads_project_end_to_end_timings_from_metadata() { + let item = StoredRequestUsageAudit { + response_time_ms: Some(626), + first_byte_time_ms: Some(120), + request_metadata: Some(json!({ + "end_to_end_time_ms": 10_626, + "end_to_end_first_byte_time_ms": 10_120, + })), + ..sample_usage("completed", Some(200), None) + }; + + let record = admin_usage_record_json( + &item, + &BTreeMap::new(), + &BTreeMap::new(), + false, + false, + None, + ); + let active = admin_usage_active_request_json(&item, None, None, None); + + for payload in [&record, &active] { + assert_eq!(payload["response_time_ms"], 626); + assert_eq!(payload["first_byte_time_ms"], 120); + assert_eq!(payload["end_to_end_time_ms"], 10_626); + assert_eq!(payload["end_to_end_first_byte_time_ms"], 10_120); + } + } + #[test] fn admin_usage_record_infers_client_family_from_user_agent() { let item = StoredRequestUsageAudit { diff --git a/crates/aether-admin/src/system.rs b/crates/aether-admin/src/system.rs index 20837e96f..219822ed3 100644 --- a/crates/aether-admin/src/system.rs +++ b/crates/aether-admin/src/system.rs @@ -1691,8 +1691,8 @@ pub fn admin_system_config_default_value(key: &str) -> Option "default_user_initial_gift_usd" => Some(json!(10.0)), "password_policy_level" => Some(json!("weak")), REQUEST_RECORD_LEVEL_KEY => Some(json!("full")), - "max_request_body_size" => Some(json!(5_242_880)), - "max_response_body_size" => Some(json!(5_242_880)), + "max_request_body_size" => Some(json!(0)), + "max_response_body_size" => Some(json!(0)), "sensitive_headers" => Some(json!([ "authorization", "x-api-key", diff --git a/crates/aether-admission-core/src/budget.rs b/crates/aether-admission-core/src/budget.rs index 7b4fb38ee..08dd54a2c 100644 --- a/crates/aether-admission-core/src/budget.rs +++ b/crates/aether-admission-core/src/budget.rs @@ -62,7 +62,7 @@ impl ResourceBudget { ResourceClass::Interactive => Self::interactive(), ResourceClass::Streaming => Self::streaming(), ResourceClass::Upload => Self { - body_bytes: 64 * 1024 * 1024, + body_bytes: 0, ..Self::interactive() }, ResourceClass::Background => Self { diff --git a/crates/aether-admission-core/src/policy.rs b/crates/aether-admission-core/src/policy.rs index b7b04304c..8e0bfd9c8 100644 --- a/crates/aether-admission-core/src/policy.rs +++ b/crates/aether-admission-core/src/policy.rs @@ -61,4 +61,17 @@ mod tests { AdmissionDecision::Reject(AdmissionRejectReason::InvalidRequest) ); } + + #[test] + fn default_policy_admits_uploads_above_the_legacy_body_limit() { + let decision = DefaultAdmissionPolicy.decide(AdmissionRequest { + trace_id: "trace-upload", + class: ResourceClass::Upload, + body_bytes: 64 * 1024 * 1024 + 1, + }); + let AdmissionDecision::Admit(budget) = decision else { + panic!("upload body should be admitted without a size limit"); + }; + assert_eq!(budget.body_bytes, 0); + } } diff --git a/crates/aether-data/contracts/src/repository/candidates/types.rs b/crates/aether-data/contracts/src/repository/candidates/types.rs index aae46b7ad..f933ca18f 100644 --- a/crates/aether-data/contracts/src/repository/candidates/types.rs +++ b/crates/aether-data/contracts/src/repository/candidates/types.rs @@ -238,10 +238,18 @@ impl RequestCandidateTrace { RequestCandidateStatus::Success | RequestCandidateStatus::Failed | RequestCandidateStatus::Cancelled - ) && candidate.latency_ms.is_some() + ) }) - .map(|candidate| candidate.latency_ms.unwrap_or(0)) - .sum(); + .map(|candidate| { + candidate.latency_ms.unwrap_or_else(|| { + candidate + .finished_at_unix_ms + .zip(candidate.started_at_unix_ms) + .map(|(finished_at, started_at)| finished_at.saturating_sub(started_at)) + .unwrap_or(0) + }) + }) + .fold(0_u64, u64::saturating_add); let final_status_source = if attempted_only && candidates.is_empty() { &all_candidates } else { diff --git a/crates/aether-gateway/frontdoor/src/body.rs b/crates/aether-gateway/frontdoor/src/body.rs index 9a986c05f..1d7e8f465 100644 --- a/crates/aether-gateway/frontdoor/src/body.rs +++ b/crates/aether-gateway/frontdoor/src/body.rs @@ -78,7 +78,7 @@ impl BodyBufferPolicy { } pub fn reservation_bytes(&self, headers: &HeaderMap) -> usize { - reservation_bytes(headers, self.max_bytes) + reservation_bytes(headers, self.max_bytes, self.budget_bytes) } pub fn reservation_permits(&self, reservation_bytes: usize) -> u32 { @@ -264,19 +264,25 @@ fn declared_content_length(headers: &HeaderMap) -> Option { .and_then(|value| value.trim().parse::().ok()) } -fn reservation_bytes(headers: &HeaderMap, max_bytes: u64) -> usize { - let max_bytes = usize::try_from(max_bytes).unwrap_or(usize::MAX); +fn reservation_bytes(headers: &HeaderMap, max_bytes: u64, budget_bytes: usize) -> usize { + let reservation_ceiling = usize::try_from(max_bytes) + .unwrap_or(usize::MAX) + .min(budget_bytes); let encoded = headers .get(header::CONTENT_ENCODING) .and_then(|value| value.to_str().ok()) .map(str::trim) .is_some_and(|value| !value.is_empty() && !value.eq_ignore_ascii_case("identity")); if encoded { - return max_bytes; + return reservation_ceiling; } declared_content_length(headers) - .map(|value| usize::try_from(value).unwrap_or(usize::MAX).min(max_bytes)) - .unwrap_or(max_bytes) + .map(|value| { + usize::try_from(value) + .unwrap_or(usize::MAX) + .min(reservation_ceiling) + }) + .unwrap_or(reservation_ceiling) } fn reservation_permits(reservation_bytes: usize, permit_bytes: usize) -> u32 { @@ -334,6 +340,58 @@ mod tests { assert_eq!(error, BodyBufferError::TooLarge { limit_bytes: 5 }); } + #[tokio::test] + async fn unlimited_body_uses_budget_as_reservation_ceiling() { + let budget = Arc::new(Semaphore::new(5)); + let policy = BodyBufferPolicy::with_permit_bytes( + u64::MAX, + Duration::from_secs(1), + Duration::from_secs(1), + 5, + 1, + Arc::clone(&budget), + ); + let mut headers = HeaderMap::new(); + headers.insert(header::CONTENT_LENGTH, HeaderValue::from_static("10")); + + let reservation = policy + .reserve(&headers) + .await + .expect("unlimited body should reserve the available budget"); + assert_eq!(reservation.requested_bytes(), 5); + assert_eq!(budget.available_permits(), 0); + + let buffered = reservation + .collect(Body::from(Bytes::from_static(b"0123456789"))) + .await + .expect("unlimited body should collect beyond the reservation ceiling"); + assert_eq!(buffered.bytes().as_ref(), b"0123456789"); + drop(buffered); + assert_eq!(budget.available_permits(), 5); + } + + #[tokio::test] + async fn encoded_unlimited_body_reserves_the_full_budget() { + let budget = Arc::new(Semaphore::new(4)); + let policy = BodyBufferPolicy::with_permit_bytes( + u64::MAX, + Duration::from_secs(1), + Duration::from_secs(1), + 4, + 1, + Arc::clone(&budget), + ); + let mut headers = HeaderMap::new(); + headers.insert(header::CONTENT_ENCODING, HeaderValue::from_static("gzip")); + + let reservation = policy + .reserve(&headers) + .await + .expect("encoded unlimited body should reserve the available budget"); + assert_eq!(reservation.requested_bytes(), 4); + assert_eq!(budget.available_permits(), 0); + } + #[tokio::test] async fn holds_weighted_permit_through_normalization_callback() { let budget = Arc::new(Semaphore::new(1)); diff --git a/crates/aether-gateway/tunnel/src/admission.rs b/crates/aether-gateway/tunnel/src/admission.rs index 9214647cd..42cbc2491 100644 --- a/crates/aether-gateway/tunnel/src/admission.rs +++ b/crates/aether-gateway/tunnel/src/admission.rs @@ -3,7 +3,7 @@ use aether_admission_core::{ DefaultAdmissionPolicy, ResourceClass, }; -use crate::{DEFAULT_OWNER_RELAY_BODY_LIMIT_BYTES, DEFAULT_TUNNEL_PROBE_BODY_LIMIT_BYTES}; +use crate::DEFAULT_TUNNEL_PROBE_BODY_LIMIT_BYTES; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum TunnelAdmissionClass { @@ -27,7 +27,7 @@ pub struct TunnelAdmissionPolicy { impl TunnelAdmissionPolicy { pub fn decide(&self, request: TunnelAdmissionRequest<'_>) -> AdmissionDecision { let body_limit = body_limit(request.class); - if request.body_bytes > body_limit { + if body_limit.is_some_and(|limit| request.body_bytes > limit) { return AdmissionDecision::Reject(AdmissionRejectReason::BodyTooLarge); } @@ -37,7 +37,7 @@ impl TunnelAdmissionPolicy { body_bytes: request.body_bytes, }) { AdmissionDecision::Admit(mut budget) => { - budget.body_bytes = body_limit; + budget.body_bytes = body_limit.unwrap_or(0); AdmissionDecision::Admit(budget) } rejected => rejected, @@ -45,11 +45,11 @@ impl TunnelAdmissionPolicy { } } -const fn body_limit(class: TunnelAdmissionClass) -> usize { +const fn body_limit(class: TunnelAdmissionClass) -> Option { match class { - TunnelAdmissionClass::Connection => 0, - TunnelAdmissionClass::Relay { .. } => DEFAULT_OWNER_RELAY_BODY_LIMIT_BYTES, - TunnelAdmissionClass::Probe => DEFAULT_TUNNEL_PROBE_BODY_LIMIT_BYTES, + TunnelAdmissionClass::Connection => Some(0), + TunnelAdmissionClass::Relay { .. } => None, + TunnelAdmissionClass::Probe => Some(DEFAULT_TUNNEL_PROBE_BODY_LIMIT_BYTES), } } @@ -68,7 +68,7 @@ mod tests { use super::{TunnelAdmissionClass, TunnelAdmissionPolicy, TunnelAdmissionRequest}; use aether_admission_core::{AdmissionDecision, AdmissionRejectReason, DbClass}; - use crate::DEFAULT_OWNER_RELAY_BODY_LIMIT_BYTES; + use crate::DEFAULT_TUNNEL_PROBE_BODY_LIMIT_BYTES; #[test] fn stream_connection_reserves_stream_and_upstream_permits() { @@ -86,23 +86,27 @@ mod tests { } #[test] - fn relay_budget_exposes_and_enforces_body_limit() { + fn relay_budget_has_no_body_limit() { let policy = TunnelAdmissionPolicy::default(); let decision = policy.decide(TunnelAdmissionRequest { trace_id: "trace-2", class: TunnelAdmissionClass::Relay { streaming: false }, - body_bytes: DEFAULT_OWNER_RELAY_BODY_LIMIT_BYTES, + body_bytes: usize::MAX, }); let AdmissionDecision::Admit(budget) = decision else { - panic!("relay at the limit should be admitted"); + panic!("relay body should be admitted without a size limit"); }; - assert_eq!(budget.body_bytes, DEFAULT_OWNER_RELAY_BODY_LIMIT_BYTES); + assert_eq!(budget.body_bytes, 0); + } + #[test] + fn probe_budget_still_enforces_body_limit() { + let policy = TunnelAdmissionPolicy::default(); assert_eq!( policy.decide(TunnelAdmissionRequest { trace_id: "trace-3", - class: TunnelAdmissionClass::Relay { streaming: false }, - body_bytes: DEFAULT_OWNER_RELAY_BODY_LIMIT_BYTES + 1, + class: TunnelAdmissionClass::Probe, + body_bytes: DEFAULT_TUNNEL_PROBE_BODY_LIMIT_BYTES + 1, }), AdmissionDecision::Reject(AdmissionRejectReason::BodyTooLarge) ); diff --git a/crates/aether-gateway/tunnel/src/lib.rs b/crates/aether-gateway/tunnel/src/lib.rs index 8a13cba58..157b1f4ce 100644 --- a/crates/aether-gateway/tunnel/src/lib.rs +++ b/crates/aether-gateway/tunnel/src/lib.rs @@ -18,6 +18,6 @@ pub use hub::{ }; pub use relay::{ is_tunnel_heartbeat_path, is_tunnel_node_status_path, TunnelAttachmentRecord, - DEFAULT_OWNER_RELAY_BODY_LIMIT_BYTES, DEFAULT_TUNNEL_PROBE_BODY_LIMIT_BYTES, PROXY_TUNNEL_PATH, - TUNNEL_HEARTBEAT_PATH, TUNNEL_NODE_STATUS_PATH, TUNNEL_RELAY_PATH_PATTERN, TUNNEL_ROUTE_FAMILY, + DEFAULT_TUNNEL_PROBE_BODY_LIMIT_BYTES, PROXY_TUNNEL_PATH, TUNNEL_HEARTBEAT_PATH, + TUNNEL_NODE_STATUS_PATH, TUNNEL_RELAY_PATH_PATTERN, TUNNEL_ROUTE_FAMILY, }; diff --git a/crates/aether-gateway/tunnel/src/relay.rs b/crates/aether-gateway/tunnel/src/relay.rs index 1e6764171..ea649e8e3 100644 --- a/crates/aether-gateway/tunnel/src/relay.rs +++ b/crates/aether-gateway/tunnel/src/relay.rs @@ -6,7 +6,6 @@ pub const TUNNEL_NODE_STATUS_PATH: &str = "/api/internal/tunnel/node-status"; pub const TUNNEL_RELAY_PATH_PATTERN: &str = "/api/internal/tunnel/relay/{node_id}"; pub const TUNNEL_ROUTE_FAMILY: &str = "tunnel_manage"; -pub const DEFAULT_OWNER_RELAY_BODY_LIMIT_BYTES: usize = 5_242_880; pub const DEFAULT_TUNNEL_PROBE_BODY_LIMIT_BYTES: usize = 64 * 1024; #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] diff --git a/crates/aether-oauth/src/provider/providers/claude_code.rs b/crates/aether-oauth/src/provider/providers/claude_code.rs index 584b54a3a..bf8f6d9c9 100644 --- a/crates/aether-oauth/src/provider/providers/claude_code.rs +++ b/crates/aether-oauth/src/provider/providers/claude_code.rs @@ -35,7 +35,6 @@ pub const CLAUDE_CODE_OAUTH_SCOPES: &[&str] = &[ pub const CLAUDE_CODE_COOKIE_SCOPE: &str = "user:profile user:inference user:sessions:claude_code user:mcp_servers user:file_upload"; -const MAX_CLAUDE_SESSION_KEY_BYTES: usize = 16 * 1024; const CLAUDE_CODE_BROWSER_USER_AGENT: &str = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/136.0.0.0 Safari/537.36"; #[derive(Debug, Clone)] @@ -85,7 +84,6 @@ impl ClaudeCodeProviderOAuthAdapter { fn session_cookie(session_key: &str) -> Result { let session_key = session_key.trim(); if session_key.is_empty() - || session_key.len() > MAX_CLAUDE_SESSION_KEY_BYTES || session_key.contains(['\r', '\n', ';']) || http::HeaderValue::from_str(session_key).is_err() { @@ -669,6 +667,14 @@ mod tests { assert!(token_body.get("code_verifier").is_some()); } + #[test] + fn session_cookie_accepts_values_above_previous_length_cap() { + let session_key = "x".repeat(20 * 1024); + let cookie = ClaudeCodeProviderOAuthAdapter::session_cookie(&session_key) + .expect("long sessionKey should remain valid"); + assert_eq!(cookie.len(), "sessionKey=".len() + session_key.len()); + } + #[tokio::test] async fn rejects_wrong_state_and_hostile_authorize_redirects() { for redirect_mode in [RedirectMode::WrongState, RedirectMode::HostileHost] { diff --git a/crates/aether-usage/runtime/src/body_capture.rs b/crates/aether-usage/runtime/src/body_capture.rs index d8e88383a..aca62f851 100644 --- a/crates/aether-usage/runtime/src/body_capture.rs +++ b/crates/aether-usage/runtime/src/body_capture.rs @@ -165,7 +165,7 @@ impl UsageBodyCaptureEngine { apply_usage_body_capture_limit( UsageBodyField::RequestBody, "request", - self.policy.max_request_body_bytes, + None, payload.request_body, payload.request_body_ref, payload.request_body_state, @@ -174,7 +174,7 @@ impl UsageBodyCaptureEngine { apply_usage_body_capture_limit( UsageBodyField::ProviderRequestBody, "provider_request", - self.policy.max_request_body_bytes, + None, payload.provider_request_body, payload.provider_request_body_ref, payload.provider_request_body_state, @@ -183,7 +183,7 @@ impl UsageBodyCaptureEngine { apply_usage_body_capture_limit( UsageBodyField::ResponseBody, "response", - self.policy.max_response_body_bytes, + None, payload.response_body, payload.response_body_ref, payload.response_body_state, @@ -192,7 +192,7 @@ impl UsageBodyCaptureEngine { apply_usage_body_capture_limit( UsageBodyField::ClientResponseBody, "client_response", - self.policy.max_response_body_bytes, + None, payload.client_response_body, payload.client_response_body_ref, payload.client_response_body_state, diff --git a/crates/aether-usage/runtime/src/request_metadata.rs b/crates/aether-usage/runtime/src/request_metadata.rs index fa2797e28..df62aa09b 100644 --- a/crates/aether-usage/runtime/src/request_metadata.rs +++ b/crates/aether-usage/runtime/src/request_metadata.rs @@ -360,6 +360,10 @@ fn copy_allowed_metadata_fields(source: &Map, target: &mut Map, target: &mut Map remove_number(&mut source, target, "client_response_body_base64_bytes"); remove_non_null_value(&mut source, target, "body_size"); remove_number(&mut source, target, "client_response_status_code"); + remove_number(&mut source, target, "end_to_end_time_ms"); + remove_number(&mut source, target, "end_to_end_first_byte_time_ms"); + remove_bool(&mut source, target, "transport_error"); + remove_non_empty_string(&mut source, target, "transport_error_type"); remove_non_null_value(&mut source, target, "billing_snapshot"); remove_non_empty_string(&mut source, target, "billing_snapshot_schema_version"); remove_non_empty_string(&mut source, target, "billing_snapshot_status"); @@ -776,6 +784,10 @@ mod tests { "provider_request_body_base64_bytes": 512, "provider_response_body_base64_bytes": 1024, "client_response_body_base64_bytes": 2048, + "end_to_end_time_ms": 10626, + "end_to_end_first_byte_time_ms": 10120, + "transport_error": true, + "transport_error_type": "connect_timeout", "body_size": { "client_request_body": "1 KB", "provider_request_body": "4 KB", @@ -816,6 +828,10 @@ mod tests { "provider_request_body_base64_bytes": 512, "provider_response_body_base64_bytes": 1024, "client_response_body_base64_bytes": 2048, + "end_to_end_time_ms": 10626, + "end_to_end_first_byte_time_ms": 10120, + "transport_error": true, + "transport_error_type": "connect_timeout", "body_size": { "client_request_body": "1 KB", "provider_request_body": "4 KB", @@ -965,6 +981,10 @@ mod tests { "client_requested_stream": false, "upstream_is_stream": true, "api_key_is_standalone": true, + "end_to_end_time_ms": 10626, + "end_to_end_first_byte_time_ms": 10120, + "transport_error": true, + "transport_error_type": "connect_timeout", "provider_id": "provider-1", "model_id": "model-1", "global_model_id": "global-model-1", @@ -998,6 +1018,10 @@ mod tests { "client_requested_stream": false, "upstream_is_stream": true, "api_key_is_standalone": true, + "end_to_end_time_ms": 10626, + "end_to_end_first_byte_time_ms": 10120, + "transport_error": true, + "transport_error_type": "connect_timeout", "model_id": "model-1", "global_model_id": "global-model-1", "global_model_name": "gpt-5", @@ -1142,8 +1166,6 @@ mod tests { apply_usage_body_capture_policy_to_event( UsageBodyCapturePolicy { record_level: UsageRequestRecordLevel::Basic, - max_request_body_bytes: Some(1024), - max_response_body_bytes: Some(1024), }, &mut event, ); diff --git a/crates/aether-usage/runtime/src/runtime.rs b/crates/aether-usage/runtime/src/runtime.rs index f7c15f3b1..60d59217a 100644 --- a/crates/aether-usage/runtime/src/runtime.rs +++ b/crates/aether-usage/runtime/src/runtime.rs @@ -44,22 +44,18 @@ pub enum UsageRequestRecordLevel { Full, } -pub const DEFAULT_USAGE_REQUEST_BODY_CAPTURE_LIMIT_BYTES: usize = 5 * 1024 * 1024; -pub const DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES: usize = 5 * 1024 * 1024; +pub const DEFAULT_USAGE_REQUEST_BODY_CAPTURE_LIMIT_BYTES: usize = usize::MAX; +pub const DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES: usize = usize::MAX; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct UsageBodyCapturePolicy { pub record_level: UsageRequestRecordLevel, - pub max_request_body_bytes: Option, - pub max_response_body_bytes: Option, } impl Default for UsageBodyCapturePolicy { fn default() -> Self { Self { record_level: UsageRequestRecordLevel::Full, - max_request_body_bytes: Some(DEFAULT_USAGE_REQUEST_BODY_CAPTURE_LIMIT_BYTES), - max_response_body_bytes: Some(DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES), } } } diff --git a/frontend/src/api/dashboard.ts b/frontend/src/api/dashboard.ts index 2f0e28372..e83b5857b 100644 --- a/frontend/src/api/dashboard.ts +++ b/frontend/src/api/dashboard.ts @@ -266,6 +266,8 @@ export interface RequestDetail { scheduling_failure?: RequestSchedulingFailure | null response_time_ms: number first_byte_time_ms?: number | null + end_to_end_time_ms?: number | null + end_to_end_first_byte_time_ms?: number | null created_at: string updated_at?: string | null request_headers?: Record diff --git a/frontend/src/api/endpoints/types/provider.ts b/frontend/src/api/endpoints/types/provider.ts index 2182dca9b..fe1c49ce1 100644 --- a/frontend/src/api/endpoints/types/provider.ts +++ b/frontend/src/api/endpoints/types/provider.ts @@ -846,6 +846,7 @@ export interface FailoverRuleItem { export interface FailoverRulesConfig { max_retries?: number + stop_on_transport_errors?: boolean stop_status_codes?: number[] stop_on_status_codes?: number[] early_stop_status_codes?: number[] diff --git a/frontend/src/api/me.ts b/frontend/src/api/me.ts index a5c18cfdf..e4a312cb5 100644 --- a/frontend/src/api/me.ts +++ b/frontend/src/api/me.ts @@ -68,6 +68,8 @@ export interface UsageRecordDetail { rate_multiplier?: number // 成本倍率（仅管理员可见） response_time_ms?: number | null first_byte_time_ms?: number | null + end_to_end_time_ms?: number | null + end_to_end_first_byte_time_ms?: number | null updated_at?: string | null response_time_updated_at?: string | null is_stream: boolean @@ -357,6 +359,8 @@ export const meApi = { rate_multiplier?: number | null response_time_ms: number | null first_byte_time_ms: number | null + end_to_end_time_ms?: number | null + end_to_end_first_byte_time_ms?: number | null updated_at?: string | null response_time_updated_at?: string | null status_code?: number | null diff --git a/frontend/src/api/usage.ts b/frontend/src/api/usage.ts index 29cd87cba..425a78c26 100644 --- a/frontend/src/api/usage.ts +++ b/frontend/src/api/usage.ts @@ -29,6 +29,10 @@ export interface UsageRecord { total_tokens: number cost?: number response_time?: number + response_time_ms?: number | null + first_byte_time_ms?: number | null + end_to_end_time_ms?: number | null + end_to_end_first_byte_time_ms?: number | null created_at: string updated_at?: string | null response_time_updated_at?: string | null @@ -554,6 +558,8 @@ export const usageApi = { rate_multiplier?: number | null response_time_ms: number | null first_byte_time_ms: number | null + end_to_end_time_ms?: number | null + end_to_end_first_byte_time_ms?: number | null updated_at?: string | null response_time_updated_at?: string | null status_code?: number | null diff --git a/frontend/src/features/providers/components/FailoverRulesDialog.vue b/frontend/src/features/providers/components/FailoverRulesDialog.vue index 21a3adb6e..36b5fea16 100644 --- a/frontend/src/features/providers/components/FailoverRulesDialog.vue +++ b/frontend/src/features/providers/components/FailoverRulesDialog.vue @@ -8,6 +8,33 @@ @update:model-value="handleClose" >
+ +
+
+

+ 传输错误 +

+

+ 这类错误没有可用于故障转移判断的上游 HTTP 状态码，因此不会命中下方的状态码规则 +

+
+ +
+
+ 继续尝试下一候选 +

+ 适用于 DNS 解析、TCP 连接、TLS 握手、代理/WARP 连接，以及响应提交前的连接重置或超时。默认开启；关闭后会立即返回网关错误。响应开始发送后无法切换候选。 +

+
+ +
+
+
@@ -243,6 +270,7 @@ import { Dialog, Button, Input, + Switch, Textarea, } from '@/components/ui' import { AlignLeft, Code2, GitBranch, Plus, Trash2 } from 'lucide-vue-next' @@ -263,6 +291,7 @@ const emit = defineEmits<{ const { success, error: showError } = useToast() const saving = ref(false) +const continueOnTransportErrors = ref(true) const successPatterns = ref([]) const errorPatterns = ref([]) @@ -284,6 +313,7 @@ const TOP_LEVEL_STOP_STATUS_CODE_KEYS = [ ] as const const MANAGED_FAILOVER_RULE_KEYS = [ + 'stop_on_transport_errors', 'success_failover_patterns', 'error_stop_patterns', ...TOP_LEVEL_STOP_STATUS_CODE_KEYS, @@ -334,6 +364,9 @@ function buildNextFailoverRules( if (filteredError.length > 0) { nextRules.error_stop_patterns = filteredError } + if (!continueOnTransportErrors.value) { + nextRules.stop_on_transport_errors = true + } return Object.values(nextRules).some(hasPersistableFailoverValue) ? nextRules as FailoverRulesConfig @@ -343,6 +376,7 @@ function buildNextFailoverRules( watch(() => [props.open, props.provider], () => { if (props.open && props.provider) { const rules = props.provider.failover_rules + continueOnTransportErrors.value = rules?.stop_on_transport_errors !== true successPatterns.value = (rules?.success_failover_patterns || []).map(r => ({ ...r, pattern: r.pattern || '', diff --git a/frontend/src/features/providers/components/ProviderDetailDrawer.vue b/frontend/src/features/providers/components/ProviderDetailDrawer.vue index af54646e4..c77ccb2bf 100644 --- a/frontend/src/features/providers/components/ProviderDetailDrawer.vue +++ b/frontend/src/features/providers/components/ProviderDetailDrawer.vue @@ -1188,6 +1188,7 @@ const hasFailoverRules = computed(() => { if (!rules) return false return FAILOVER_RULE_ARRAY_KEYS.some(key => (rules[key]?.length || 0) > 0) || typeof rules.max_retries === 'number' + || rules.stop_on_transport_errors === true }) // Provider 级别代理配置状态 diff --git a/frontend/src/features/providers/components/ProviderFormDialog.vue b/frontend/src/features/providers/components/ProviderFormDialog.vue index 6380b3627..2060ea37c 100644 --- a/frontend/src/features/providers/components/ProviderFormDialog.vue +++ b/frontend/src/features/providers/components/ProviderFormDialog.vue @@ -216,10 +216,11 @@
@@ -233,10 +234,11 @@
diff --git a/frontend/src/features/providers/components/__tests__/FailoverRulesDialog.transport-errors.spec.ts b/frontend/src/features/providers/components/__tests__/FailoverRulesDialog.transport-errors.spec.ts new file mode 100644 index 000000000..f5ffc9140 --- /dev/null +++ b/frontend/src/features/providers/components/__tests__/FailoverRulesDialog.transport-errors.spec.ts @@ -0,0 +1,130 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { createApp, nextTick, type App } from 'vue' + +import type { FailoverRulesConfig, ProviderWithEndpointsSummary } from '@/api/endpoints/types' +import FailoverRulesDialog from '../FailoverRulesDialog.vue' + +const endpointMocks = vi.hoisted(() => ({ + updateProvider: vi.fn(), +})) + +vi.mock('@/api/endpoints', () => ({ + updateProvider: endpointMocks.updateProvider, +})) + +vi.mock('@/composables/useToast', () => ({ + useToast: () => ({ + success: vi.fn(), + error: vi.fn(), + }), +})) + +const mountedApps: Array<{ app: App, root: HTMLElement }> = [] + +function makeProvider(failoverRules: FailoverRulesConfig | null): ProviderWithEndpointsSummary { + return { + id: 'provider-1', + failover_rules: failoverRules, + } as ProviderWithEndpointsSummary +} + +function mountDialog(failoverRules: FailoverRulesConfig | null = null) { + const root = document.createElement('div') + document.body.appendChild(root) + const app = createApp(FailoverRulesDialog, { + open: true, + provider: makeProvider(failoverRules), + 'onUpdate:open': vi.fn(), + }) + app.mount(root) + mountedApps.push({ app, root }) +} + +async function settle() { + for (let index = 0; index < 4; index += 1) { + await Promise.resolve() + await nextTick() + } +} + +function transportErrorSwitch(): HTMLButtonElement { + const control = document.body.querySelector( + '[role="switch"][aria-label="传输错误时继续尝试下一候选"]', + ) + if (!control) throw new Error('Missing transport-error failover switch') + return control +} + +function clickSave() { + const button = [...document.body.querySelectorAll('button')] + .find(candidate => candidate.textContent?.trim() === '保存') + if (!button) throw new Error('Missing save button') + button.click() +} + +beforeEach(() => { + endpointMocks.updateProvider.mockReset() + endpointMocks.updateProvider.mockResolvedValue(makeProvider(null)) +}) + +afterEach(() => { + for (const { app, root } of mountedApps.splice(0)) { + app.unmount() + root.remove() + } + document.body.innerHTML = '' +}) + +describe('FailoverRulesDialog transport errors', () => { + it('continues failover by default and explains why HTTP status rules do not apply', async () => { + mountDialog() + await settle() + + expect(transportErrorSwitch().getAttribute('aria-checked')).toBe('true') + expect(document.body.textContent).toContain('没有可用于故障转移判断的上游 HTTP 状态码') + expect(document.body.textContent).toContain('DNS 解析') + expect(document.body.textContent).toContain('TCP 连接') + expect(document.body.textContent).toContain('TLS 握手') + expect(document.body.textContent).toContain('响应提交前的连接重置或超时') + expect(document.body.textContent).toContain('响应开始发送后无法切换候选') + }) + + it('persists stop_on_transport_errors when continuing is disabled', async () => { + mountDialog() + await settle() + + transportErrorSwitch().click() + await nextTick() + expect(transportErrorSwitch().getAttribute('aria-checked')).toBe('false') + + clickSave() + await settle() + + expect(endpointMocks.updateProvider).toHaveBeenCalledWith('provider-1', { + failover_rules: { + stop_on_transport_errors: true, + }, + }) + }) + + it('loads a stop rule and removes only that rule when continuing is enabled', async () => { + mountDialog({ + max_retries: 2, + stop_on_transport_errors: true, + }) + await settle() + + expect(transportErrorSwitch().getAttribute('aria-checked')).toBe('false') + transportErrorSwitch().click() + await nextTick() + + clickSave() + await settle() + + expect(endpointMocks.updateProvider).toHaveBeenCalledWith('provider-1', { + failover_rules: { + max_retries: 2, + }, + }) + }) +}) diff --git a/frontend/src/features/providers/components/__tests__/ProviderDetailDrawer.loading.spec.ts b/frontend/src/features/providers/components/__tests__/ProviderDetailDrawer.loading.spec.ts index 7fbe87d56..7d90d6b63 100644 --- a/frontend/src/features/providers/components/__tests__/ProviderDetailDrawer.loading.spec.ts +++ b/frontend/src/features/providers/components/__tests__/ProviderDetailDrawer.loading.spec.ts @@ -54,4 +54,8 @@ describe('ProviderDetailDrawer loading priorities', () => { expect(source).toContain('v-if="open && batchAssignDialogOpen && provider"') expect(source).toContain('v-if="open && failoverRulesDialogOpen"') }) + + it('marks a transport-error stop policy as a configured failover rule', () => { + expect(source).toContain('rules.stop_on_transport_errors === true') + }) }) diff --git a/frontend/src/features/providers/components/__tests__/ProviderFormDialog.transfer-limits.spec.ts b/frontend/src/features/providers/components/__tests__/ProviderFormDialog.transfer-limits.spec.ts index 5f0fa1ea1..0a56a3e7c 100644 --- a/frontend/src/features/providers/components/__tests__/ProviderFormDialog.transfer-limits.spec.ts +++ b/frontend/src/features/providers/components/__tests__/ProviderFormDialog.transfer-limits.spec.ts @@ -177,15 +177,20 @@ describe('ProviderFormDialog transfer limits', () => { ) }) - it('defaults missing legacy values to explicit zero', async () => { + it('shows zero limits as unlimited placeholders while submitting explicit zero', async () => { mountDialog(makeProvider({ max_transfer_count: undefined, max_transfer_timeout_seconds: undefined, })) await settle() - expect(document.body.querySelector('#max-transfer-count')?.value).toBe('0') - expect(document.body.querySelector('#max-transfer-timeout-seconds')?.value).toBe('0') + const countInput = document.body.querySelector('#max-transfer-count') + const timeoutInput = document.body.querySelector('#max-transfer-timeout-seconds') + + expect(countInput?.value).toBe('') + expect(countInput?.placeholder).toBe('0 (不限制)') + expect(timeoutInput?.value).toBe('') + expect(timeoutInput?.placeholder).toBe('0 (不限制)') clickButton('保存') await settle() diff --git a/frontend/src/features/usage/components/HorizontalRequestTimeline.vue b/frontend/src/features/usage/components/HorizontalRequestTimeline.vue index 6e41929fb..2ea131f0f 100644 --- a/frontend/src/features/usage/components/HorizontalRequestTimeline.vue +++ b/frontend/src/features/usage/components/HorizontalRequestTimeline.vue @@ -1023,25 +1023,26 @@ const conversionBoundaryIndex = computed(() => { return idx }) -// 计算链路总耗时（使用成功候选的 latency_ms 字段） -// 优先使用 latency_ms，因为它与 Usage.response_time_ms 使用相同的时间基准 -// 避免 finished_at - started_at 带来的额外延迟（数据库操作时间） +// The trace aggregate includes every attempted candidate, including failed +// failover attempts. Per-candidate latency remains provider-scoped. const totalTraceLatency = computed(() => { if (!rawTimeline.value || rawTimeline.value.length === 0) return 0 - // 查找成功的候选，使用其 latency_ms - const successCandidate = rawTimeline.value.find(c => c.status === 'success') - if (successCandidate?.latency_ms != null) { - return successCandidate.latency_ms + const aggregateLatency = trace.value?.total_latency_ms + if (typeof aggregateLatency === 'number' && Number.isFinite(aggregateLatency) && aggregateLatency > 0) { + return aggregateLatency } - // 如果没有成功的候选，查找失败但有 latency_ms 的候选 - const failedWithLatency = rawTimeline.value.find(c => c.status === 'failed' && c.latency_ms != null) - if (failedWithLatency?.latency_ms != null) { - return failedWithLatency.latency_ms + const attemptedLatency = rawTimeline.value.reduce((sum, candidate) => { + const latency = normalizeLatencyMs(candidate.latency_ms) + return sum + (latency ?? 0) + }, 0) + if (attemptedLatency > 0) { + return attemptedLatency } - // 回退：使用 finished_at - started_at 计算 + // Historical transport failures may not have latency_ms. Recover the wall + // clock span from candidate timestamps for those records. let earliestStart: number | null = null let latestEnd: number | null = null diff --git a/frontend/src/features/usage/components/RequestDetailDrawer.vue b/frontend/src/features/usage/components/RequestDetailDrawer.vue index 258c33f26..15da9a818 100644 --- a/frontend/src/features/usage/components/RequestDetailDrawer.vue +++ b/frontend/src/features/usage/components/RequestDetailDrawer.vue @@ -175,7 +175,7 @@ 耗时 - {{ formatDurationMs(detail.first_byte_time_ms) }} / {{ formatDurationMs(detail.response_time_ms) }} + {{ formatDurationMs(detail.end_to_end_first_byte_time_ms ?? detail.first_byte_time_ms) }} / {{ formatDurationMs(detail.end_to_end_time_ms ?? detail.response_time_ms) }} | @@ -197,7 +197,7 @@ 耗时 - {{ formatDurationMs(detail.first_byte_time_ms) }} / {{ formatDurationMs(detail.response_time_ms) }} + {{ formatDurationMs(detail.end_to_end_first_byte_time_ms ?? detail.first_byte_time_ms) }} / {{ formatDurationMs(detail.end_to_end_time_ms ?? detail.response_time_ms) }} | diff --git a/frontend/src/features/usage/components/UsageRecordsTable.vue b/frontend/src/features/usage/components/UsageRecordsTable.vue index bfc02f1fa..ea849110a 100644 --- a/frontend/src/features/usage/components/UsageRecordsTable.vue +++ b/frontend/src/features/usage/components/UsageRecordsTable.vue @@ -358,7 +358,7 @@ {{ formatOutputRate(getRecordDisplayOutputRate(record)) }} {{ formatRecordLatencyPair(record) }} / {{ formatOutputRate(getRecordDisplayOutputRate(record)) }} -
+
+ 端到端 首字/总耗时 输出速度
@@ -925,7 +926,7 @@
@@ -1445,8 +1446,10 @@ function getRecordCacheTokensTitle(record: UsageRecord): string { } function formatRecordLatencyPair(record: UsageRecord): string { - const firstByte = formatRecordDurationSeconds(record.first_byte_time_ms) - const total = formatRecordDurationSeconds(record.response_time_ms) + const firstByte = formatRecordDurationSeconds( + record.end_to_end_first_byte_time_ms ?? record.first_byte_time_ms, + ) + const total = formatRecordDurationSeconds(record.end_to_end_time_ms ?? record.response_time_ms) return `${firstByte} / ${total}` } @@ -1455,6 +1458,13 @@ function formatRecordDurationSeconds(ms: number | null | undefined): string { return `${(ms / 1000).toFixed(2)}s` } +function hasRecordDisplayLatency(record: UsageRecord): boolean { + return record.end_to_end_time_ms != null + || record.end_to_end_first_byte_time_ms != null + || record.response_time_ms != null + || record.first_byte_time_ms != null +} + function getRecordDisplayOutputRate(record: UsageRecord): number | null { return getDisplayOutputRate({ output_tokens: record.output_tokens, @@ -1468,8 +1478,10 @@ function getRecordDisplayOutputRate(record: UsageRecord): number | null { function getRecordPerformanceTitle(record: UsageRecord): string { const outputRate = getRecordDisplayOutputRate(record) return [ - `首字: ${formatRecordDurationSeconds(record.first_byte_time_ms)}`, - `总耗时: ${formatRecordDurationSeconds(record.response_time_ms)}`, + `端到端首字: ${formatRecordDurationSeconds(record.end_to_end_first_byte_time_ms ?? record.first_byte_time_ms)}`, + `端到端总耗时: ${formatRecordDurationSeconds(record.end_to_end_time_ms ?? record.response_time_ms)}`, + `成功候选首字: ${formatRecordDurationSeconds(record.first_byte_time_ms)}`, + `成功候选耗时: ${formatRecordDurationSeconds(record.response_time_ms)}`, `生成耗时: ${formatRecordDurationSeconds(getGenerationTimeMs(record))}`, `输出速度: ${formatOutputRateTokensPerSecond(outputRate)}`, ].join('\n') diff --git a/frontend/src/features/usage/components/__tests__/HorizontalRequestTimeline.spec.ts b/frontend/src/features/usage/components/__tests__/HorizontalRequestTimeline.spec.ts index 5a00a5de1..9cdd5f45b 100644 --- a/frontend/src/features/usage/components/__tests__/HorizontalRequestTimeline.spec.ts +++ b/frontend/src/features/usage/components/__tests__/HorizontalRequestTimeline.spec.ts @@ -181,6 +181,46 @@ afterEach(() => { }) describe('HorizontalRequestTimeline', () => { + it('uses the trace aggregate latency instead of the successful candidate latency', async () => { + const trace = buildTrace([ + buildCandidate({ + id: 'cand-transport-timeout', + provider_id: 'provider-timeout', + provider_name: 'Provider Timeout', + key_id: 'key-timeout', + key_name: 'Timeout Key', + candidate_index: 0, + status: 'failed', + latency_ms: 10_000, + started_at: '2026-05-06T12:00:00.000Z', + finished_at: '2026-05-06T12:00:10.000Z', + }), + buildCandidate({ + id: 'cand-success-after-failover', + provider_id: 'provider-success', + provider_name: 'Provider Success', + key_id: 'key-success', + key_name: 'Success Key', + candidate_index: 1, + status: 'success', + latency_ms: 626, + started_at: '2026-05-06T12:00:10.000Z', + finished_at: '2026-05-06T12:00:10.626Z', + }), + ]) + trace.total_latency_ms = 10_626 + + const root = mountTimeline(trace) + await nextTick() + + const heading = [...root.querySelectorAll('h4')] + .find(element => element.textContent?.trim() === '请求链路追踪') + const overview = heading?.parentElement?.parentElement + const displayedLatency = overview?.lastElementChild?.textContent?.trim() + expect(displayedLatency).toBe('10.63s') + expect(displayedLatency).not.toBe('626ms') + }) + it('keeps attempted keys visible for ordinary provider groups that are not selected', async () => { const trace = buildTrace([ buildCandidate({ diff --git a/frontend/src/features/usage/components/__tests__/RequestDetailDrawer.pricing.spec.ts b/frontend/src/features/usage/components/__tests__/RequestDetailDrawer.pricing.spec.ts index f04c07df7..c0c9a45e4 100644 --- a/frontend/src/features/usage/components/__tests__/RequestDetailDrawer.pricing.spec.ts +++ b/frontend/src/features/usage/components/__tests__/RequestDetailDrawer.pricing.spec.ts @@ -113,6 +113,48 @@ function buildFastTierDetail(): RequestDetail { } describe('RequestDetailDrawer settlement pricing', () => { + it('shows end-to-end latency while keeping output TPS scoped to candidate timing', async () => { + apiMocks.getRequestDetail.mockResolvedValue({ + ...buildEmbeddingDetail(), + tokens: { input: 100, output: 50, total: 150 }, + input_tokens: 100, + output_tokens: 50, + total_tokens: 150, + is_stream: true, + upstream_is_stream: true, + response_time_ms: 626, + first_byte_time_ms: 100, + end_to_end_time_ms: 10_626, + end_to_end_first_byte_time_ms: 10_120, + } satisfies RequestDetail) + + let isOpen!: Ref + const Host = defineComponent({ + setup() { + isOpen = ref(false) + return () => h(RequestDetailDrawer, { + isOpen: isOpen.value, + requestId: 'usage-embedding-1', + }) + }, + }) + + const root = document.createElement('div') + document.body.appendChild(root) + const app = createApp(Host) + app.mount(root) + mountedApps.push({ app, root }) + + isOpen.value = true + await nextTick() + + await vi.waitFor(() => { + expect(document.body.textContent).toContain('10.12s / 10.63s') + expect(document.body.textContent).toContain('95.1tps') + expect(document.body.textContent).not.toContain('98.8tps') + }) + }) + it('renders an input-only embedding tier without treating the missing output price as zero', async () => { apiMocks.getRequestDetail.mockResolvedValue(buildEmbeddingDetail()) diff --git a/frontend/src/features/usage/components/__tests__/UsageRecordsTable.spec.ts b/frontend/src/features/usage/components/__tests__/UsageRecordsTable.spec.ts index 7cd4f13f1..9dc5834cc 100644 --- a/frontend/src/features/usage/components/__tests__/UsageRecordsTable.spec.ts +++ b/frontend/src/features/usage/components/__tests__/UsageRecordsTable.spec.ts @@ -195,8 +195,10 @@ describe('UsageRecordsTable', () => { const titles = [...root.querySelectorAll('[title]')] .map((element) => element.getAttribute('title')) expect(titles).toContain([ - '首字: 0.50s', - '总耗时: 1.00s', + '端到端首字: 0.50s', + '端到端总耗时: 1.00s', + '成功候选首字: 0.50s', + '成功候选耗时: 1.00s', '生成耗时: 0.50s', '输出速度: 100 tokens/s', ].join('\n')) @@ -204,6 +206,55 @@ describe('UsageRecordsTable', () => { expect(titles.join('\n')).not.toContain('首字后生成耗时') }) + it('shows end-to-end latency while keeping output TPS scoped to the successful candidate', () => { + const root = mountUsageRecordsTable([buildRecord({ + output_tokens: 50, + response_time_ms: 626, + first_byte_time_ms: 100, + end_to_end_time_ms: 10_626, + end_to_end_first_byte_time_ms: 10_120, + })]) + + const performanceCell = root.querySelector('table tbody tr td:last-child') as HTMLElement + expect(performanceCell.textContent).toContain('10.12s / 10.63s') + expect(performanceCell.textContent).toContain('95.1 tps') + expect(performanceCell.textContent).not.toContain('98.8 tps') + + const titles = [...root.querySelectorAll('[title]')] + .map((element) => element.getAttribute('title')) + expect(titles).toContain([ + '端到端首字: 10.12s', + '端到端总耗时: 10.63s', + '成功候选首字: 0.10s', + '成功候选耗时: 0.63s', + '生成耗时: 0.53s', + '输出速度: 95.1 tokens/s', + ].join('\n')) + }) + + it('shows end-to-end latency when candidate timing fields are unavailable', () => { + const root = mountUsageRecordsTable([buildRecord({ + response_time_ms: null, + first_byte_time_ms: null, + end_to_end_time_ms: 10_626, + end_to_end_first_byte_time_ms: 10_120, + })]) + + const performanceCell = root.querySelector('table tbody tr td:last-child') as HTMLElement + expect(performanceCell.textContent).toContain('10.12s / 10.63s') + + const titles = [...root.querySelectorAll('[title]')] + .map((element) => element.getAttribute('title')) + expect(titles).toContain([ + '端到端首字: 10.12s', + '端到端总耗时: 10.63s', + '成功候选首字: -', + '成功候选耗时: -', + '生成耗时: -', + '输出速度: -', + ].join('\n')) + }) + it('shows an output speed placeholder when the rate is unavailable', () => { const root = mountUsageRecordsTable([buildRecord({ output_tokens: 0, @@ -218,8 +269,10 @@ describe('UsageRecordsTable', () => { const titles = [...root.querySelectorAll('[title]')].map((element) => element.title) expect(titles).toContain([ - '首字: 0.50s', - '总耗时: 1.00s', + '端到端首字: 0.50s', + '端到端总耗时: 1.00s', + '成功候选首字: 0.50s', + '成功候选耗时: 1.00s', '生成耗时: 0.50s', '输出速度: -', ].join('\n')) diff --git a/frontend/src/features/usage/types.ts b/frontend/src/features/usage/types.ts index 890f81558..cb7433876 100644 --- a/frontend/src/features/usage/types.ts +++ b/frontend/src/features/usage/types.ts @@ -117,6 +117,8 @@ export interface UsageRecord { actual_cost?: number response_time_ms?: number | null first_byte_time_ms?: number | null // 首字时间 (TTFB) + end_to_end_time_ms?: number | null // 客户端从请求进入网关到完成的总耗时 + end_to_end_first_byte_time_ms?: number | null // 客户端从请求进入网关到首字节的耗时 is_stream: boolean upstream_is_stream?: boolean client_requested_stream?: boolean diff --git a/frontend/src/i18n/__tests__/i18n.spec.ts b/frontend/src/i18n/__tests__/i18n.spec.ts index 8ddf2a0f1..c8398aa49 100644 --- a/frontend/src/i18n/__tests__/i18n.spec.ts +++ b/frontend/src/i18n/__tests__/i18n.spec.ts @@ -114,6 +114,7 @@ describe('i18n infrastructure', () => { expect(translateLegacyText('请求记录清理策略', 'en-US')).toBe('Request log cleanup policy') expect(translateLegacyText('最大转移次数', 'en-US')).toBe('Max transfers') expect(translateLegacyText('最大转移超时', 'en-US')).toBe('Max transfer timeout') + expect(translateLegacyText('0 (不限制)', 'en-US')).toBe('0 (unlimited)') expect(translateLegacyText(' 发布于 2026-01-01 ', 'en-US')).toBe(' Published at 2026-01-01 ') expect(translateLegacyText('git clone https://github.com/fawney19/Aether.git', 'en-US')).toBe('git clone https://github.com/fawney19/Aether.git') }) diff --git a/frontend/src/i18n/messages.ts b/frontend/src/i18n/messages.ts index 3444c3d35..cab25966b 100644 --- a/frontend/src/i18n/messages.ts +++ b/frontend/src/i18n/messages.ts @@ -1412,6 +1412,7 @@ const legacyExactEnglishMessages: Record = { '默认 2': 'Default 2', '最大转移次数': 'Max transfers', '最大转移超时': 'Max transfer timeout', + '0 (不限制)': '0 (unlimited)', '流式首字节超时': 'Streaming first-byte timeout', '非流式请求超时': 'Non-streaming request timeout', '(秒)': '(seconds)', @@ -2409,12 +2410,8 @@ const legacyPhraseEnglishMessages: Array<[string, string]> = [ ['记录详细程度', 'Record detail level'], ['BASIC - 基本信息 (~1KB/条)', 'BASIC - basic information (~1KB/item)'], ['HEADERS - 含请求头 (~2-3KB/条)', 'HEADERS - includes request headers (~2-3KB/item)'], - ['FULL - 完整请求响应 (~50KB/条)', 'FULL - full request and response (~50KB/item)'], + ['FULL - 完整请求响应', 'FULL - full request and response'], ['敏感信息会自动脱敏', 'Sensitive information is automatically redacted'], - ['最大请求体大小 (KB)', 'Maximum request body size (KB)'], - ['超过此大小的请求体将被截断记录', 'Request bodies larger than this are truncated in records'], - ['最大响应体大小 (KB)', 'Maximum response body size (KB)'], - ['超过此大小的响应体将被截断记录', 'Response bodies larger than this are truncated in records'], ['敏感请求头', 'Sensitive request headers'], ['逗号分隔，这些请求头会被脱敏处理', 'Comma-separated request headers that will be redacted'], ['0 表示默认不限制；未单独配置的用户和独立 Key 会跟随这里', '0 means no default limit. Users and standalone keys without separate configuration follow this setting.'], @@ -2750,10 +2747,6 @@ const legacyFallbackTokens: Array<[string, string]> = [ ['完整请求响应', 'full request and response'], ['自动脱敏', 'automatically redacted'], ['脱敏', 'redacted'], - ['最大请求体', 'maximum request body'], - ['最大响应体', 'maximum response body'], - ['超过此大小', 'larger than this'], - ['截断记录', 'truncated in records'], ['逗号分隔', 'comma-separated'], ['这些请求头', 'these request headers'], ['脱敏处理', 'redacted'], diff --git a/frontend/src/views/admin/SystemSettings.vue b/frontend/src/views/admin/SystemSettings.vue index 5cfab5227..c87e96aa5 100644 --- a/frontend/src/views/admin/SystemSettings.vue +++ b/frontend/src/views/admin/SystemSettings.vue @@ -116,15 +116,11 @@ @@ -359,8 +355,6 @@ const { hasBasicConfigChanges, hasLogConfigChanges, hasCleanupConfigChanges, - maxRequestBodySizeKB, - maxResponseBodySizeKB, sensitiveHeadersStr, turnstileAllowedHostnamesStr, loadSystemConfig, diff --git a/frontend/src/views/admin/system-settings/RequestLogSection.vue b/frontend/src/views/admin/system-settings/RequestLogSection.vue index 3520427f7..3fe0743bd 100644 --- a/frontend/src/views/admin/system-settings/RequestLogSection.vue +++ b/frontend/src/views/admin/system-settings/RequestLogSection.vue @@ -38,7 +38,7 @@ HEADERS - 含请求头 (~2-3KB/条) - FULL - 完整请求响应 (~50KB/条) + FULL - 完整请求响应 @@ -47,46 +47,6 @@

-
- - -

- 超过此大小的请求体将被截断记录 -

-
- -
- - -

- 超过此大小的响应体将被截断记录 -

-
-