Improve request trace upstream diagnostics

This commit is contained in:
fawney19
2026-05-09 10:46:04 +08:00
parent 4cf0de681a
commit 79d07ac79a
33 changed files with 2315 additions and 427 deletions

View File

@@ -112,6 +112,8 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
provider_request_method: Some(serde_json::Value::Null),
provider_request_headers: Some(&resolved.provider_request_headers),
original_headers: &parts.headers,
request_path: Some(parts.uri.path()),
request_query_string: parts.uri.query(),
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
original_request_body_json: Some(body_json),
original_request_body_base64: None,

View File

@@ -9,7 +9,11 @@ use aether_ai_serving::{
use aether_scheduler_core::{ClientSessionAffinity, SchedulerRankingOutcome};
use serde_json::{Map, Value};
use crate::ai_serving::{request_origin_from_headers, ExecutionRuntimeAuthContext, RequestOrigin};
use crate::ai_serving::{
request_origin_from_headers, request_path_implies_stream_request, sanitize_request_path,
sanitize_request_path_and_query, sanitize_request_query_string, ExecutionRuntimeAuthContext,
RequestOrigin,
};
use crate::client_session_affinity::{
client_session_affinity_report_context_value, CLIENT_SESSION_AFFINITY_REPORT_CONTEXT_FIELD,
};
@@ -40,6 +44,8 @@ pub(crate) struct LocalExecutionReportContextParts<'a> {
pub(crate) provider_request_method: Option<Value>,
pub(crate) provider_request_headers: Option<&'a BTreeMap<String, String>>,
pub(crate) original_headers: &'a http::HeaderMap,
pub(crate) request_path: Option<&'a str>,
pub(crate) request_query_string: Option<&'a str>,
pub(crate) request_origin: Option<RequestOrigin>,
pub(crate) original_request_body_json: Option<&'a Value>,
pub(crate) original_request_body_base64: Option<&'a str>,
@@ -80,6 +86,15 @@ pub(crate) fn build_local_execution_report_context(
{
merge_incoming_tls_fingerprint(&mut extra_fields, incoming_tls);
}
insert_request_path_fields(
&mut extra_fields,
parts.request_path,
parts.request_query_string,
);
let client_requested_stream = parts.client_requested_stream
|| parts
.request_path
.is_some_and(request_path_implies_stream_request);
build_ai_execution_report_context(AiExecutionReportContextParts {
auth_context: parts.auth_context,
@@ -113,7 +128,7 @@ pub(crate) fn build_local_execution_report_context(
client_ip,
user_agent,
},
client_requested_stream: parts.client_requested_stream,
client_requested_stream,
upstream_is_stream: parts.upstream_is_stream,
has_envelope: parts.has_envelope,
needs_conversion: parts.needs_conversion,
@@ -121,6 +136,30 @@ pub(crate) fn build_local_execution_report_context(
})
}
fn insert_request_path_fields(
extra_fields: &mut Map<String, Value>,
request_path: Option<&str>,
request_query_string: Option<&str>,
) {
let Some(path) = request_path.and_then(sanitize_request_path) else {
return;
};
let query = request_query_string.and_then(sanitize_request_query_string);
let path_and_query = sanitize_request_path_and_query(path.as_str(), query.as_deref())
.unwrap_or_else(|| path.clone());
extra_fields
.entry("request_path".to_string())
.or_insert_with(|| Value::String(path.clone()));
if let Some(query) = query.clone() {
extra_fields
.entry("request_query_string".to_string())
.or_insert_with(|| Value::String(query.to_string()));
}
extra_fields
.entry("request_path_and_query".to_string())
.or_insert_with(|| Value::String(path_and_query));
}
pub(crate) fn provider_stream_event_api_format_for_provider_type(
provider_type: &str,
) -> Option<&'static str> {
@@ -226,6 +265,8 @@ mod tests {
provider_request_method: None,
provider_request_headers: Some(&provider_request_headers),
original_headers: &original_headers,
request_path: Some("/v1/chat/completions"),
request_query_string: Some("debug=true&limit=10"),
request_origin: Some(RequestOrigin {
client_ip: Some("203.0.113.8".to_string()),
user_agent: Some("Claude-Code/1.0".to_string()),
@@ -255,6 +296,77 @@ mod tests {
"session_key": "account=account-1;session=session-1"
})
);
assert_eq!(report_context["request_path"], "/v1/chat/completions");
assert_eq!(report_context["request_query_string"], "limit=10");
assert_eq!(
report_context["request_path_and_query"],
"/v1/chat/completions?limit=10"
);
}
#[test]
fn local_execution_report_context_treats_stream_generate_content_path_as_client_stream() {
let auth_context = ExecutionRuntimeAuthContext {
user_id: "user-1".to_string(),
api_key_id: "api-key-1".to_string(),
username: None,
api_key_name: None,
balance_remaining: None,
access_allowed: true,
api_key_is_standalone: false,
};
let original_headers = http::HeaderMap::new();
let provider_request_headers = BTreeMap::new();
let report_context =
build_local_execution_report_context(LocalExecutionReportContextParts {
auth_context: &auth_context,
request_id: "trace-1",
candidate_id: "candidate-1",
attempt_identity: ExecutionAttemptIdentity::new(0, 0),
model: "gemini-3.1-flash-image-preview",
provider_name: "Gemini",
provider_id: "provider-1",
endpoint_id: "endpoint-1",
key_id: "key-1",
key_name: None,
model_id: None,
global_model_id: None,
global_model_name: None,
provider_api_format: "gemini:generate_content",
client_api_format: "gemini:generate_content",
mapped_model: None,
candidate_group_id: None,
ranking: None,
upstream_url: None,
header_rules: None,
body_rules: None,
provider_request_method: None,
provider_request_headers: Some(&provider_request_headers),
original_headers: &original_headers,
request_path: Some(
"/v1beta/models/gemini-3.1-flash-image-preview:streamGenerateContent",
),
request_query_string: Some("key=secret&alt=sse"),
request_origin: None,
original_request_body_json: Some(&json!({
"contents": [{"role": "user", "parts": [{"text": "hi"}]}]
})),
original_request_body_base64: None,
client_session_affinity: None,
client_requested_stream: false,
upstream_is_stream: true,
has_envelope: false,
needs_conversion: false,
extra_fields: Map::new(),
});
assert_eq!(report_context["client_requested_stream"], true);
assert_eq!(report_context["request_query_string"], "alt=sse");
assert_eq!(
report_context["request_path_and_query"],
"/v1beta/models/gemini-3.1-flash-image-preview:streamGenerateContent?alt=sse"
);
}
#[test]
@@ -300,6 +412,8 @@ mod tests {
provider_request_method: None,
provider_request_headers: Some(&provider_request_headers),
original_headers: &original_headers,
request_path: None,
request_query_string: None,
request_origin: None,
original_request_body_json: Some(&json!({"model": "gpt-5"})),
original_request_body_base64: None,

View File

@@ -94,6 +94,8 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
provider_request_method: None,
provider_request_headers: None,
original_headers: &parts.headers,
request_path: Some(parts.uri.path()),
request_query_string: parts.uri.query(),
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
original_request_body_json: Some(body_json),
original_request_body_base64: resolved.provider_request_body_base64.as_deref(),

View File

@@ -114,6 +114,8 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
provider_request_method: Some(serde_json::Value::String(parts.method.to_string())),
provider_request_headers: Some(&resolved.provider_request_headers),
original_headers: &parts.headers,
request_path: Some(parts.uri.path()),
request_query_string: parts.uri.query(),
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
original_request_body_json: Some(body_json),
original_request_body_base64: body_base64,

View File

@@ -75,6 +75,8 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat
provider_request_method: None,
provider_request_headers: None,
original_headers: &parts.headers,
request_path: Some(parts.uri.path()),
request_query_string: parts.uri.query(),
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
original_request_body_json: Some(body_json),
original_request_body_base64: None,

View File

@@ -120,6 +120,8 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
provider_request_method: Some(serde_json::Value::Null),
provider_request_headers: Some(&resolved.provider_request_headers),
original_headers: &parts.headers,
request_path: Some(parts.uri.path()),
request_query_string: parts.uri.query(),
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
original_request_body_json: Some(body_json),
original_request_body_base64: None,

View File

@@ -107,6 +107,8 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
provider_request_method: Some(serde_json::Value::Null),
provider_request_headers: Some(&resolved.provider_request_headers),
original_headers: &parts.headers,
request_path: Some(parts.uri.path()),
request_query_string: parts.uri.query(),
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
original_request_body_json: Some(body_json),
original_request_body_base64: None,

View File

@@ -105,6 +105,8 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
provider_request_method: Some(serde_json::Value::Null),
provider_request_headers: Some(&resolved.provider_request_headers),
original_headers: &parts.headers,
request_path: Some(parts.uri.path()),
request_query_string: parts.uri.query(),
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
original_request_body_json: Some(body_json),
original_request_body_base64: None,

View File

@@ -70,7 +70,8 @@ pub(crate) use aether_ai_formats::api::{
provider_adaptation_should_unwrap_stream_envelope,
provider_private_response_allows_sync_finalize, request_candidate_api_format_preference,
request_candidate_api_formats, request_conversion_kind,
request_conversion_requires_enable_flag, resolve_claude_stream_spec, resolve_claude_sync_spec,
request_conversion_requires_enable_flag, request_path_implies_stream_request,
resolve_claude_stream_spec, resolve_claude_sync_spec,
resolve_execution_runtime_stream_plan_kind, resolve_execution_runtime_sync_plan_kind,
resolve_finalize_stream_rewrite_mode, resolve_gemini_files_stream_spec,
resolve_gemini_files_sync_spec, resolve_gemini_stream_spec, resolve_gemini_sync_spec,
@@ -79,11 +80,12 @@ pub(crate) use aether_ai_formats::api::{
resolve_local_video_sync_spec, resolve_openai_chat_max_tokens,
resolve_openai_responses_stream_spec, resolve_openai_responses_sync_spec,
resolve_requested_gemini_image_model_for_request,
resolve_requested_openai_image_model_for_request, stream_body_contains_error_event,
supports_stream_execution_decision_kind, supports_sync_execution_decision_kind,
sync_chat_response_conversion_kind, sync_cli_response_conversion_kind,
transform_provider_private_stream_line, value_as_u64, AiControlPlanRequest,
AiSurfaceFinalizeError, AiSurfaceStreamRewriter, CanonicalStreamFrame,
resolve_requested_openai_image_model_for_request, sanitize_request_path,
sanitize_request_path_and_query, sanitize_request_query_string,
stream_body_contains_error_event, supports_stream_execution_decision_kind,
supports_sync_execution_decision_kind, sync_chat_response_conversion_kind,
sync_cli_response_conversion_kind, transform_provider_private_stream_line, value_as_u64,
AiControlPlanRequest, AiSurfaceFinalizeError, AiSurfaceStreamRewriter, CanonicalStreamFrame,
ChatGptWebImageRequestError, ClaudeClientEmitter, ClaudeProviderState,
ExecutionRuntimeAuthContext, FinalizeStreamRewriteMode, FormatContext, GeminiClientEmitter,
GeminiImageRequestForOpenAi, GeminiProviderState, KiroToClaudeCliStreamState,