mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
feat(usage): record request origin metadata (#362)
Record client IP and User-Agent in usage request_metadata for local execution requests without adding schema, API, or table fields. References #357 and supersedes #358. Co-authored-by: RWDai <27391645+RWDai@users.noreply.github.com>
This commit is contained in:
@@ -109,6 +109,7 @@ 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_origin: Some(crate::headers::request_origin_from_parts(parts)),
|
||||
original_request_body_json: Some(body_json),
|
||||
original_request_body_base64: None,
|
||||
client_requested_stream: body_json
|
||||
|
||||
@@ -5,6 +5,7 @@ use serde_json::{Map, Value};
|
||||
|
||||
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
|
||||
use crate::ai_pipeline::planner::candidate_metadata::append_ranking_metadata_to_object;
|
||||
use crate::headers::RequestOrigin;
|
||||
use crate::orchestration::ExecutionAttemptIdentity;
|
||||
|
||||
pub(crate) struct LocalExecutionReportContextParts<'a> {
|
||||
@@ -32,6 +33,7 @@ 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_origin: Option<RequestOrigin>,
|
||||
pub(crate) original_request_body_json: Option<&'a Value>,
|
||||
pub(crate) original_request_body_base64: Option<&'a str>,
|
||||
pub(crate) client_requested_stream: bool,
|
||||
@@ -131,6 +133,18 @@ pub(crate) fn build_local_execution_report_context(
|
||||
)
|
||||
.unwrap_or(Value::Null),
|
||||
);
|
||||
let RequestOrigin {
|
||||
client_ip,
|
||||
user_agent,
|
||||
} = parts
|
||||
.request_origin
|
||||
.unwrap_or_else(|| crate::headers::request_origin_from_headers(parts.original_headers));
|
||||
if let Some(client_ip) = client_ip {
|
||||
object.insert("client_ip".to_string(), Value::String(client_ip));
|
||||
}
|
||||
if let Some(user_agent) = user_agent {
|
||||
object.insert("user_agent".to_string(), Value::String(user_agent));
|
||||
}
|
||||
object.insert(
|
||||
"client_requested_stream".to_string(),
|
||||
Value::Bool(parts.client_requested_stream),
|
||||
@@ -237,7 +251,17 @@ pub(crate) fn insert_provider_stream_event_api_format(
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::provider_stream_event_api_format_for_provider_type;
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use serde_json::{json, Map, Value};
|
||||
|
||||
use super::{
|
||||
build_local_execution_report_context, provider_stream_event_api_format_for_provider_type,
|
||||
LocalExecutionReportContextParts,
|
||||
};
|
||||
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
|
||||
use crate::headers::RequestOrigin;
|
||||
use crate::orchestration::ExecutionAttemptIdentity;
|
||||
|
||||
#[test]
|
||||
fn codex_provider_uses_openai_responses_stream_event_format() {
|
||||
@@ -262,4 +286,67 @@ mod tests {
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn local_execution_report_context_records_request_origin() {
|
||||
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: "gpt-5",
|
||||
provider_name: "OpenAI",
|
||||
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: "openai:chat",
|
||||
client_api_format: "openai:chat",
|
||||
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_origin: Some(RequestOrigin {
|
||||
client_ip: Some("203.0.113.8".to_string()),
|
||||
user_agent: Some("Claude-Code/1.0".to_string()),
|
||||
}),
|
||||
original_request_body_json: Some(&json!({"model": "gpt-5"})),
|
||||
original_request_body_base64: None,
|
||||
client_requested_stream: false,
|
||||
upstream_is_stream: false,
|
||||
has_envelope: false,
|
||||
needs_conversion: false,
|
||||
extra_fields: Map::new(),
|
||||
});
|
||||
|
||||
assert_eq!(
|
||||
report_context["client_ip"],
|
||||
Value::String("203.0.113.8".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
report_context["user_agent"],
|
||||
Value::String("Claude-Code/1.0".to_string())
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -90,6 +90,7 @@ 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_origin: Some(crate::headers::request_origin_from_parts(parts)),
|
||||
original_request_body_json: Some(body_json),
|
||||
original_request_body_base64: resolved.provider_request_body_base64.as_deref(),
|
||||
client_requested_stream: spec_metadata.require_streaming,
|
||||
|
||||
@@ -87,6 +87,7 @@ 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_origin: Some(crate::headers::request_origin_from_parts(parts)),
|
||||
original_request_body_json: Some(body_json),
|
||||
original_request_body_base64: body_base64,
|
||||
client_requested_stream: spec_metadata.require_streaming,
|
||||
|
||||
@@ -73,6 +73,7 @@ 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_origin: Some(crate::headers::request_origin_from_parts(parts)),
|
||||
original_request_body_json: Some(body_json),
|
||||
original_request_body_base64: None,
|
||||
client_requested_stream: false,
|
||||
|
||||
@@ -89,6 +89,7 @@ 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_origin: Some(crate::headers::request_origin_from_parts(parts)),
|
||||
original_request_body_json: Some(body_json),
|
||||
original_request_body_base64: None,
|
||||
client_requested_stream: body_json
|
||||
|
||||
@@ -107,6 +107,7 @@ 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_origin: Some(crate::headers::request_origin_from_parts(parts)),
|
||||
original_request_body_json: Some(body_json),
|
||||
original_request_body_base64: None,
|
||||
client_requested_stream: body_json
|
||||
|
||||
@@ -105,6 +105,7 @@ 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_origin: Some(crate::headers::request_origin_from_parts(parts)),
|
||||
original_request_body_json: Some(body_json),
|
||||
original_request_body_base64: None,
|
||||
client_requested_stream: body_json
|
||||
|
||||
@@ -49,7 +49,10 @@ use crate::handlers::shared::{
|
||||
local_proxy_route_requires_buffered_body, request_enables_control_execute,
|
||||
should_strip_forwarded_provider_credential_header, should_strip_forwarded_trusted_admin_header,
|
||||
};
|
||||
use crate::headers::{extract_or_generate_trace_id, should_skip_request_header};
|
||||
use crate::headers::{
|
||||
extract_or_generate_trace_id, request_origin_from_headers_and_remote_addr,
|
||||
should_skip_request_header,
|
||||
};
|
||||
use crate::router::RequestAdmissionError;
|
||||
use crate::{
|
||||
AppState, FrontdoorUserRpmOutcome, GatewayError, GatewayFallbackMetricKind,
|
||||
@@ -707,7 +710,13 @@ pub(crate) async fn proxy_request(
|
||||
)) => return Err(GatewayError::Internal(message)),
|
||||
};
|
||||
let request_admission_ms = started_at.elapsed().as_millis() as u64;
|
||||
let (parts, body) = request.into_parts();
|
||||
let (mut parts, body) = request.into_parts();
|
||||
parts
|
||||
.extensions
|
||||
.insert(request_origin_from_headers_and_remote_addr(
|
||||
&parts.headers,
|
||||
&remote_addr,
|
||||
));
|
||||
let trace_id = extract_or_generate_trace_id(&parts.headers);
|
||||
state.clear_local_execution_runtime_miss_diagnostic(&trace_id);
|
||||
if request_hits_execution_loop_guard(&parts) {
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::{collections::BTreeMap, net::SocketAddr};
|
||||
|
||||
use crate::constants::*;
|
||||
use uuid::Uuid;
|
||||
@@ -20,6 +20,61 @@ pub(crate) fn header_value_u64(headers: &http::HeaderMap, key: &str) -> Option<u
|
||||
header_value_str(headers, key).and_then(|value| value.parse::<u64>().ok())
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub(crate) struct RequestOrigin {
|
||||
pub(crate) client_ip: Option<String>,
|
||||
pub(crate) user_agent: Option<String>,
|
||||
}
|
||||
|
||||
pub(crate) fn request_origin_from_headers(headers: &http::HeaderMap) -> RequestOrigin {
|
||||
RequestOrigin {
|
||||
client_ip: client_ip_from_headers(headers),
|
||||
user_agent: header_value_str(headers, http::header::USER_AGENT.as_str())
|
||||
.map(|value| truncate_chars(value.as_str(), 1_000)),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn request_origin_from_headers_and_remote_addr(
|
||||
headers: &http::HeaderMap,
|
||||
remote_addr: &SocketAddr,
|
||||
) -> RequestOrigin {
|
||||
let mut origin = request_origin_from_headers(headers);
|
||||
if origin.client_ip.is_none() {
|
||||
origin.client_ip = Some(remote_addr.ip().to_string());
|
||||
}
|
||||
origin
|
||||
}
|
||||
|
||||
pub(crate) fn request_origin_from_parts(parts: &http::request::Parts) -> RequestOrigin {
|
||||
parts
|
||||
.extensions
|
||||
.get::<RequestOrigin>()
|
||||
.cloned()
|
||||
.unwrap_or_else(|| request_origin_from_headers(&parts.headers))
|
||||
}
|
||||
|
||||
fn client_ip_from_headers(headers: &http::HeaderMap) -> Option<String> {
|
||||
header_value_str(headers, "x-forwarded-for")
|
||||
.and_then(|value| {
|
||||
value
|
||||
.split(',')
|
||||
.map(str::trim)
|
||||
.find(|segment| !segment.is_empty() && !segment.eq_ignore_ascii_case("unknown"))
|
||||
.map(|segment| truncate_chars(segment, 45))
|
||||
})
|
||||
.or_else(|| {
|
||||
header_value_str(headers, "x-real-ip").and_then(|value| {
|
||||
let value = value.trim();
|
||||
(!value.is_empty() && !value.eq_ignore_ascii_case("unknown"))
|
||||
.then(|| truncate_chars(value, 45))
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
fn truncate_chars(value: &str, max_chars: usize) -> String {
|
||||
value.chars().take(max_chars).collect()
|
||||
}
|
||||
|
||||
pub(crate) fn should_skip_request_header(name: &str) -> bool {
|
||||
crate::provider_transport::should_skip_request_header(name)
|
||||
}
|
||||
@@ -74,3 +129,59 @@ pub(crate) fn header_equals(
|
||||
.map(|value| value.eq_ignore_ascii_case(expected))
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
request_origin_from_headers, request_origin_from_headers_and_remote_addr, RequestOrigin,
|
||||
};
|
||||
use http::{HeaderMap, HeaderValue};
|
||||
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
|
||||
|
||||
#[test]
|
||||
fn request_origin_prefers_first_forwarded_for_ip() {
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert(
|
||||
"x-forwarded-for",
|
||||
HeaderValue::from_static(" 203.0.113.8, 10.0.0.1 "),
|
||||
);
|
||||
headers.insert("x-real-ip", HeaderValue::from_static("198.51.100.4"));
|
||||
headers.insert(
|
||||
http::header::USER_AGENT,
|
||||
HeaderValue::from_static("Claude-Code/1.0"),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
request_origin_from_headers(&headers),
|
||||
RequestOrigin {
|
||||
client_ip: Some("203.0.113.8".to_string()),
|
||||
user_agent: Some("Claude-Code/1.0".to_string()),
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_origin_uses_real_ip_after_empty_forwarded_for_segments() {
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert("x-forwarded-for", HeaderValue::from_static(" , unknown "));
|
||||
headers.insert("x-real-ip", HeaderValue::from_static("198.51.100.4"));
|
||||
|
||||
assert_eq!(
|
||||
request_origin_from_headers(&headers).client_ip.as_deref(),
|
||||
Some("198.51.100.4")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_origin_falls_back_to_remote_addr() {
|
||||
let headers = HeaderMap::new();
|
||||
let remote_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(192, 0, 2, 10)), 443);
|
||||
|
||||
assert_eq!(
|
||||
request_origin_from_headers_and_remote_addr(&headers, &remote_addr)
|
||||
.client_ip
|
||||
.as_deref(),
|
||||
Some("192.0.2.10")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user