mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +08:00
Improve request trace upstream diagnostics
This commit is contained in:
@@ -15,4 +15,5 @@ serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
sha1 = "0.10"
|
||||
sha2.workspace = true
|
||||
url.workspace = true
|
||||
uuid.workspace = true
|
||||
|
||||
@@ -96,7 +96,9 @@ pub use crate::formats::shared::response::{
|
||||
};
|
||||
pub use crate::formats::shared::routing::{
|
||||
is_matching_stream_http_request, is_matching_stream_request,
|
||||
resolve_execution_runtime_stream_plan_kind, resolve_execution_runtime_sync_plan_kind,
|
||||
request_path_implies_stream_request, resolve_execution_runtime_stream_plan_kind,
|
||||
resolve_execution_runtime_sync_plan_kind, sanitize_request_path,
|
||||
sanitize_request_path_and_query, sanitize_request_query_string,
|
||||
supports_stream_execution_decision_kind, supports_sync_execution_decision_kind,
|
||||
};
|
||||
pub use crate::formats::shared::sse::{encode_done_sse, encode_json_sse, map_claude_stop_reason};
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
use http::Method;
|
||||
use url::form_urlencoded;
|
||||
|
||||
use crate::contracts::{
|
||||
CLAUDE_CHAT_STREAM_PLAN_KIND, CLAUDE_CHAT_SYNC_PLAN_KIND, CLAUDE_CLI_STREAM_PLAN_KIND,
|
||||
@@ -304,6 +305,67 @@ fn resolve_gemini_generate_content_plan_kind(
|
||||
}
|
||||
}
|
||||
|
||||
pub fn request_path_implies_stream_request(path: &str) -> bool {
|
||||
let trimmed = path.trim();
|
||||
let path = trimmed
|
||||
.split_once('?')
|
||||
.map(|(path, _)| path)
|
||||
.unwrap_or(trimmed);
|
||||
path.ends_with(":streamGenerateContent")
|
||||
}
|
||||
|
||||
pub fn sanitize_request_path(path: &str) -> Option<String> {
|
||||
let path = path
|
||||
.trim()
|
||||
.split_once('?')
|
||||
.map(|(path, _)| path)
|
||||
.unwrap_or_else(|| path.trim())
|
||||
.trim();
|
||||
(!path.is_empty()).then(|| path.to_string())
|
||||
}
|
||||
|
||||
pub fn sanitize_request_query_string(query: &str) -> Option<String> {
|
||||
let query = query.trim().trim_start_matches('?').trim();
|
||||
if query.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut serializer = form_urlencoded::Serializer::new(String::new());
|
||||
for (key, value) in form_urlencoded::parse(query.as_bytes()) {
|
||||
if request_query_key_is_safe_to_trace(key.as_ref()) {
|
||||
serializer.append_pair(key.as_ref(), value.as_ref());
|
||||
}
|
||||
}
|
||||
let sanitized = serializer.finish();
|
||||
(!sanitized.is_empty()).then_some(sanitized)
|
||||
}
|
||||
|
||||
pub fn sanitize_request_path_and_query(path: &str, query: Option<&str>) -> Option<String> {
|
||||
let trimmed = path.trim();
|
||||
let (path, embedded_query) = trimmed
|
||||
.split_once('?')
|
||||
.map(|(path, query)| (path.trim(), Some(query)))
|
||||
.unwrap_or((trimmed, None));
|
||||
if path.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let sanitized_query = query
|
||||
.and_then(sanitize_request_query_string)
|
||||
.or_else(|| embedded_query.and_then(sanitize_request_query_string));
|
||||
Some(match sanitized_query {
|
||||
Some(query) => format!("{path}?{query}"),
|
||||
None => path.to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
fn request_query_key_is_safe_to_trace(key: &str) -> bool {
|
||||
matches!(
|
||||
key.to_ascii_lowercase().as_str(),
|
||||
"alt" | "view" | "pagesize" | "page_size" | "limit" | "offset"
|
||||
)
|
||||
}
|
||||
|
||||
pub fn is_matching_stream_request(
|
||||
plan_kind: &str,
|
||||
path: &str,
|
||||
@@ -320,7 +382,7 @@ pub fn is_matching_stream_request(
|
||||
.and_then(|value| value.as_bool())
|
||||
.unwrap_or(false),
|
||||
GEMINI_CHAT_STREAM_PLAN_KIND | GEMINI_CLI_STREAM_PLAN_KIND => {
|
||||
path.ends_with(":streamGenerateContent")
|
||||
request_path_implies_stream_request(path)
|
||||
}
|
||||
_ => true,
|
||||
}
|
||||
@@ -388,7 +450,9 @@ mod tests {
|
||||
|
||||
use super::{
|
||||
is_matching_stream_http_request, is_matching_stream_request,
|
||||
resolve_execution_runtime_stream_plan_kind, resolve_execution_runtime_sync_plan_kind,
|
||||
request_path_implies_stream_request, resolve_execution_runtime_stream_plan_kind,
|
||||
resolve_execution_runtime_sync_plan_kind, sanitize_request_path,
|
||||
sanitize_request_path_and_query, sanitize_request_query_string,
|
||||
supports_stream_execution_decision_kind, supports_sync_execution_decision_kind,
|
||||
};
|
||||
use crate::contracts::{
|
||||
@@ -609,6 +673,41 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stream_path_detection_handles_gemini_method_paths_with_query() {
|
||||
assert!(request_path_implies_stream_request(
|
||||
"/v1beta/models/gemini-2.5-pro:streamGenerateContent?alt=sse"
|
||||
));
|
||||
assert!(request_path_implies_stream_request(
|
||||
" /v1internal:streamGenerateContent?alt=sse "
|
||||
));
|
||||
assert!(!request_path_implies_stream_request(
|
||||
"/v1beta/models/gemini-2.5-pro:generateContent?alt=sse"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_path_metadata_sanitizer_drops_sensitive_query_parameters() {
|
||||
assert_eq!(
|
||||
sanitize_request_path("/v1beta/models/gemini-2.5-pro:generateContent?key=secret")
|
||||
.as_deref(),
|
||||
Some("/v1beta/models/gemini-2.5-pro:generateContent")
|
||||
);
|
||||
assert_eq!(
|
||||
sanitize_request_query_string("?key=secret&alt=sse&pageSize=10&token=hidden")
|
||||
.as_deref(),
|
||||
Some("alt=sse&pageSize=10")
|
||||
);
|
||||
assert_eq!(
|
||||
sanitize_request_path_and_query(
|
||||
"/v1beta/models/gemini-2.5-pro:streamGenerateContent?key=secret&alt=sse",
|
||||
None
|
||||
)
|
||||
.as_deref(),
|
||||
Some("/v1beta/models/gemini-2.5-pro:streamGenerateContent?alt=sse")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stream_matching_requires_openai_stream_flag() {
|
||||
assert!(!is_matching_stream_request(
|
||||
|
||||
Reference in New Issue
Block a user