From 9d2fe268ae68c1d5cffbaf96e01d59d5aa09484c Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Mon, 26 Oct 2026 08:04:40 +0000 Subject: [PATCH] fix(core): redact query credentials in logs and preserve JSON Schema property order Port QuantumNous/new-api 30748c6558bc97de49ff74871cfec717e50f3322 and d0b840e75ff36d731b04ff981ee361819acd593f into services/core. Access logs mask credential query keys, and tool schemas sent upstream keep the client's property order. Fixes #944 Co-authored-by: LIghtJUNction --- services/core/src/lib.rs | 1 + services/core/src/log_redact.rs | 471 ++++++++++++++ services/core/src/main.rs | 4 + services/core/src/relay/honeypot.rs | 4 +- services/core/src/relay/mod.rs | 5 +- services/core/src/relay/native_anthropic.rs | 3 + services/core/src/relay/native_gemini.rs | 3 + services/core/src/relay/request.rs | 63 +- services/core/src/relay/schema_order.rs | 657 ++++++++++++++++++++ 9 files changed, 1202 insertions(+), 9 deletions(-) create mode 100644 services/core/src/log_redact.rs create mode 100644 services/core/src/relay/schema_order.rs diff --git a/services/core/src/lib.rs b/services/core/src/lib.rs index fcc3d41ff..bd74b2c49 100644 --- a/services/core/src/lib.rs +++ b/services/core/src/lib.rs @@ -29,6 +29,7 @@ pub mod identity_http; pub mod internal_rpc; pub mod legacy_auth; pub mod lifecycle; +pub mod log_redact; pub mod market_oauth; pub mod model_http; pub mod native_key_read; diff --git a/services/core/src/log_redact.rs b/services/core/src/log_redact.rs new file mode 100644 index 000000000..11993b46d --- /dev/null +++ b/services/core/src/log_redact.rs @@ -0,0 +1,471 @@ +//! Query-credential redaction for request logs. +//! +//! Adapted from QuantumNous/new-api `30748c65` (`SanitizeURLForLog` and the +//! access-log formatter). The deny list is exact, case-insensitive, and +//! whitespace-trimmed. Any other query key is kept so operators can still see +//! `alt`, `api-version`, `model`, `stream`, `client_id`, and similar values. +//! +//! A key is also denied when its normalized form contains `token`, `secret`, +//! or `signature` (`session_token`, `nottoken`, `client_secret`). `code` and +//! `state` are not on that list. Paths under `/api/oauth/` and `/oauth/` drop +//! the whole query, matching the upstream logger. `/api/oauth2/` does not. +//! +//! Access logs re-encode the query before masking, which drops invalid +//! percent-escapes and semicolon pairs. `sanitize_url_for_log` itself follows +//! `url.URL.Query`: if nothing sensitive remains, the original string is kept. + +use axum::{extract::Request, middleware::Next, response::Response}; + +const MASK: &str = "***masked***"; + +const EXACT_KEYS: &[&str] = &[ + "key", + "api_key", + "api-key", + "apikey", + "x-api-key", + "access_token", + "refresh_token", + "id_token", + "token", + "authorization", + "auth", + "client_secret", + "secret", + "password", + "passwd", + "signature", + "sig", + "awsaccesskeyid", + "x-amz-credential", + "x-amz-security-token", + "x-amz-signature", +]; + +/// True when a decoded query key must not be written to a log. +pub fn is_sensitive_query_key(key: &str) -> bool { + let normalized = key.trim().to_lowercase(); + EXACT_KEYS.contains(&normalized.as_str()) + || normalized.contains("token") + || normalized.contains("secret") + || normalized.contains("signature") +} + +/// Mask sensitive query values in a URL. Non-sensitive queries, including +/// their original key order, are returned unchanged. +pub fn sanitize_url_for_log(raw: &str) -> String { + if raw.is_empty() || has_control(raw) { + return raw.to_owned(); + } + let (base, fragment) = match split_once_char(raw, '#') { + Some((base, fragment)) => (base, Some(fragment)), + None => (raw, None), + }; + let Some((prefix, query)) = split_once_char(base, '?') else { + return raw.to_owned(); + }; + let mut pairs = parse_query(query); + if pairs.is_empty() || !pairs.keys().any(|key| is_sensitive_query_key(key)) { + return raw.to_owned(); + } + for (key, values) in &mut pairs { + if is_sensitive_query_key(key) { + values.clear(); + values.push(MASK.to_owned()); + } + } + let mut out = format!("{prefix}?{}", encode_query(&pairs)); + if let Some(fragment) = fragment { + out.push('#'); + out.push_str(&escape(fragment, Escape::Fragment)); + } + out +} + +/// Request-target as it should appear in an access log. +/// +/// OAuth callback paths lose the query. Other targets keep a re-encoded query +/// with sensitive values replaced by `***masked***` (`*` is percent-encoded). +pub fn logged_request_target(target: &str) -> String { + let Some(parsed) = parse_request_target(target) else { + return strip_query(target); + }; + if parsed.decoded_path.starts_with("/api/oauth/") || parsed.decoded_path.starts_with("/oauth/") + { + return parsed.decoded_path; + } + if !parsed.had_query { + return parsed.decoded_path; + } + let encoded = encode_query(&parse_query(&parsed.raw_query)); + let suffix = sanitize_url_for_log(&format!("?{encoded}")); + format!("{}{suffix}", parsed.escaped_path) +} + +/// One stderr line for a request whose target carries a query. The original +/// target is not included. +pub fn access_log_line(method: &str, status: u16, target: &str) -> Option { + if !target.contains('?') { + return None; + } + let mut logged = logged_request_target(target); + if logged.chars().count() > 1024 { + logged = logged.chars().take(1024).collect(); + } + let logged: String = logged + .chars() + .map(|c| if c.is_control() { '_' } else { c }) + .collect(); + let method: String = method + .chars() + .filter(|c| !c.is_control()) + .take(16) + .collect(); + Some(format!( + "core_access method={method} status={status} target={logged}" + )) +} + +/// Outermost access note. Query-less requests stay quiet; credential-bearing +/// queries are masked before they reach stderr. +pub async fn log_redacted_request_target(request: Request, next: Next) -> Response { + let method = request.method().as_str().to_owned(); + let target = request + .uri() + .path_and_query() + .map(|value| value.as_str().to_owned()) + .unwrap_or_default(); + let response = next.run(request).await; + if let Some(line) = access_log_line(&method, response.status().as_u16(), &target) { + eprintln!("{line}"); + } + response +} + +fn strip_query(target: &str) -> String { + let path = split_once_char(target, '?') + .map(|(path, _)| path) + .unwrap_or(target); + path.chars() + .map(|c| if c.is_control() { '_' } else { c }) + .collect() +} + +struct RequestTarget { + decoded_path: String, + escaped_path: String, + raw_query: String, + had_query: bool, +} + +fn parse_request_target(target: &str) -> Option { + if target.is_empty() || has_control(target) { + return None; + } + // Request-targets are parsed like Go's ParseRequestURI: `#` stays in the + // query because the request line has no fragment. + let (path_raw, raw_query, had_query) = match split_once_char(target, '?') { + Some((path, query)) => (path, query.to_owned(), true), + None => (target, String::new(), false), + }; + if path_raw != "*" && !path_raw.starts_with('/') { + return None; + } + let decoded = unescape(path_raw, Escape::Path)?; + let escaped_path = escaped_path(path_raw, &decoded); + Some(RequestTarget { + decoded_path: decoded, + escaped_path, + raw_query, + had_query, + }) +} + +fn escaped_path(raw: &str, decoded: &str) -> String { + if raw != "*" + && valid_encoded(raw, Escape::Path) + && unescape(raw, Escape::Path).as_deref() == Some(decoded) + { + let escaped = escape(decoded, Escape::Path); + if escaped != raw { + return raw.to_owned(); + } + } + if decoded == "*" { + return "*".to_owned(); + } + escape(decoded, Escape::Path) +} + +fn parse_query(query: &str) -> std::collections::BTreeMap> { + let mut pairs = std::collections::BTreeMap::new(); + if query.is_empty() { + return pairs; + } + for segment in query.split('&') { + if segment.is_empty() || segment.contains(';') { + continue; + } + let (key, value) = segment.split_once('=').unwrap_or((segment, "")); + let (Some(key), Some(value)) = + (unescape(key, Escape::Query), unescape(value, Escape::Query)) + else { + continue; + }; + pairs.entry(key).or_default().push(value); + } + pairs +} + +fn encode_query(pairs: &std::collections::BTreeMap>) -> String { + let mut out = String::new(); + for (key, values) in pairs { + let key = escape(key, Escape::Query); + for value in values { + if !out.is_empty() { + out.push('&'); + } + out.push_str(&key); + out.push('='); + out.push_str(&escape(value, Escape::Query)); + } + } + out +} + +#[derive(Clone, Copy, PartialEq, Eq)] +enum Escape { + Path, + Query, + Fragment, +} + +fn has_control(value: &str) -> bool { + value.bytes().any(|byte| byte < 0x20 || byte == 0x7f) +} + +fn split_once_char(value: &str, separator: char) -> Option<(&str, &str)> { + value.split_once(separator) +} + +fn should_escape(byte: u8, mode: Escape) -> bool { + if byte.is_ascii_alphanumeric() { + return false; + } + if matches!(byte, b'-' | b'_' | b'.' | b'~') { + return false; + } + if matches!( + byte, + b'$' | b'&' | b'+' | b',' | b'/' | b':' | b';' | b'=' | b'?' | b'@' + ) { + return match mode { + Escape::Path => byte == b'?', + Escape::Query => true, + Escape::Fragment => false, + }; + } + if mode == Escape::Fragment && matches!(byte, b'!' | b'(' | b')' | b'*') { + return false; + } + true +} + +fn escape(value: &str, mode: Escape) -> String { + let bytes = value.as_bytes(); + if !bytes + .iter() + .any(|byte| should_escape(*byte, mode) || (*byte == b' ' && mode == Escape::Query)) + { + return value.to_owned(); + } + let mut out = String::with_capacity(bytes.len()); + for byte in bytes { + if *byte == b' ' && mode == Escape::Query { + out.push('+'); + } else if should_escape(*byte, mode) { + const HEX: &[u8; 16] = b"0123456789ABCDEF"; + out.push('%'); + out.push(HEX[(byte >> 4) as usize] as char); + out.push(HEX[(byte & 0x0f) as usize] as char); + } else { + out.push(*byte as char); + } + } + out +} + +fn unescape(value: &str, mode: Escape) -> Option { + let bytes = value.as_bytes(); + let mut out = Vec::with_capacity(bytes.len()); + let mut index = 0; + while index < bytes.len() { + match bytes[index] { + b'%' => { + if index + 2 >= bytes.len() { + return None; + } + let high = from_hex(bytes[index + 1])?; + let low = from_hex(bytes[index + 2])?; + out.push((high << 4) | low); + index += 3; + } + b'+' if mode == Escape::Query => { + out.push(b' '); + index += 1; + } + byte => { + out.push(byte); + index += 1; + } + } + } + String::from_utf8(out).ok() +} + +fn from_hex(byte: u8) -> Option { + match byte { + b'0'..=b'9' => Some(byte - b'0'), + b'a'..=b'f' => Some(byte - b'a' + 10), + b'A'..=b'F' => Some(byte - b'A' + 10), + _ => None, + } +} + +fn valid_encoded(value: &str, mode: Escape) -> bool { + value.bytes().all(|byte| { + matches!( + byte, + b'!' | b'$' + | b'&' + | b'\'' + | b'(' + | b')' + | b'*' + | b'+' + | b',' + | b';' + | b'=' + | b':' + | b'@' + | b'[' + | b']' + | b'%' + ) || !should_escape(byte, mode) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn access_log_matches_upstream_query_redaction_goldens() { + let cases = [ + ( + "/v1/models?%6bey=first&key=second&API-Key=third&session_token=fourth&alt=sse", + "/v1/models?API-Key=%2A%2A%2Amasked%2A%2A%2A&alt=sse&key=%2A%2A%2Amasked%2A%2A%2A&session_token=%2A%2A%2Amasked%2A%2A%2A", + ), + ( + "/v1/models%23test%25?key=#secret&alt=#sse", + "/v1/models%23test%25?alt=%23sse&key=%2A%2A%2Amasked%2A%2A%2A", + ), + ( + "//v1%2Fmodels?key=secret", + "//v1%2Fmodels?key=%2A%2A%2Amasked%2A%2A%2A", + ), + ( + "/v1/models?key=bad%ZZ&%ZZ=bad&token=bad;other=true&alt=sse", + "/v1/models?alt=sse", + ), + ("/v1/models?key=bad%ZZ", "/v1/models?"), + ( + "/v1/models?alt=sse&api-version=2024-02-01", + "/v1/models?alt=sse&api-version=2024-02-01", + ), + ("/v1/models", "/v1/models"), + ( + "/%61pi/oauth/test?code=secret&state=secret", + "/api/oauth/test", + ), + ("/oauth/callback?code=abc&state=xyz", "/oauth/callback"), + ( + "/v1?+key+=secret&model=gpt", + "/v1?+key+=%2A%2A%2Amasked%2A%2A%2A&model=gpt", + ), + ( + "/v1?author=ada&authentication=no&auth=yes&alt=sse", + "/v1?alt=sse&auth=%2A%2A%2Amasked%2A%2A%2A&authentication=no&author=ada", + ), + ( + "/api/oauth2/authorize?client_id=public&code=secret", + "/api/oauth2/authorize?client_id=public&code=secret", + ), + ( + "/v1/models?api_key=k&apikey=k2&x-api-key=k3&refresh_token=r&id_token=i&password=p&passwd=p2&sig=s&signature=s2&awsaccesskeyid=a&x-amz-security-token=t&model=keep", + "/v1/models?api_key=%2A%2A%2Amasked%2A%2A%2A&apikey=%2A%2A%2Amasked%2A%2A%2A&awsaccesskeyid=%2A%2A%2Amasked%2A%2A%2A&id_token=%2A%2A%2Amasked%2A%2A%2A&model=keep&passwd=%2A%2A%2Amasked%2A%2A%2A&password=%2A%2A%2Amasked%2A%2A%2A&refresh_token=%2A%2A%2Amasked%2A%2A%2A&sig=%2A%2A%2Amasked%2A%2A%2A&signature=%2A%2A%2Amasked%2A%2A%2A&x-amz-security-token=%2A%2A%2Amasked%2A%2A%2A&x-api-key=%2A%2A%2Amasked%2A%2A%2A", + ), + ]; + for (target, expected) in cases { + let logged = logged_request_target(target); + assert_eq!(logged, expected, "{target}"); + for secret in ["first", "second", "third", "fourth", "sk-secret", "#secret"] { + if target.contains(secret) && !expected.contains(secret) { + assert!(!logged.contains(secret), "{target} leaked {secret}"); + } + } + } + } + + #[test] + fn sanitize_url_for_log_masks_keys_and_keeps_operational_queries() { + let masked = sanitize_url_for_log( + "https://example.test/v1beta/models/gemini:streamGenerateContent?alt=sse&key=sk-secret&access_token=ya29-secret&api-version=2024-02-01", + ); + assert_eq!( + masked, + "https://example.test/v1beta/models/gemini:streamGenerateContent?access_token=%2A%2A%2Amasked%2A%2A%2A&alt=sse&api-version=2024-02-01&key=%2A%2A%2Amasked%2A%2A%2A" + ); + assert!(!masked.contains("sk-secret")); + assert!(!masked.contains("ya29-secret")); + let aws = sanitize_url_for_log( + "https://example.test/path?X-Amz-Credential=credential&X-Amz-Signature=signature&session_token=session&client_secret=secret&model=gpt-test", + ); + assert!(!aws.contains("=credential")); + assert!(!aws.contains("=signature")); + assert!(!aws.contains("=session")); + assert!(!aws.contains("=secret")); + assert!(aws.contains("X-Amz-Credential=%2A%2A%2Amasked%2A%2A%2A")); + assert!(aws.contains("model=gpt-test")); + let kept = "https://example.test/v1/chat/completions?api-version=2024-02-01&alt=sse"; + assert_eq!(sanitize_url_for_log(kept), kept); + assert_eq!( + sanitize_url_for_log("https://user:pass@example.test/v1?key=sek&alt=sse"), + "https://user:pass@example.test/v1?alt=sse&key=%2A%2A%2Amasked%2A%2A%2A" + ); + assert_eq!( + sanitize_url_for_log("https://example.test/v1?key=bad%ZZ"), + "https://example.test/v1?key=bad%ZZ" + ); + assert_eq!(sanitize_url_for_log("not a url %%"), "not a url %%"); + assert_eq!(sanitize_url_for_log(""), ""); + } + + #[test] + fn access_log_line_omits_plaintext_credentials_and_keeps_model() { + let line = access_log_line( + "GET", + 401, + "/v1/models?key=sk-secret&token=plain-token&api_key=abc&access_token=def&secret=ghi&model=gpt-test&alt=sse", + ) + .expect("query is logged"); + for secret in ["sk-secret", "plain-token", "abc", "def", "ghi"] { + assert!(!line.contains(secret), "{line}"); + } + assert!(line.contains("model=gpt-test")); + assert!(line.contains("alt=sse")); + assert!(line.contains("%2A%2A%2Amasked%2A%2A%2A")); + assert!(line.starts_with("core_access method=GET status=401 target=")); + assert!(access_log_line("GET", 200, "/health/live").is_none()); + } +} diff --git a/services/core/src/main.rs b/services/core/src/main.rs index 31ca9c18d..deee212b8 100644 --- a/services/core/src/main.rs +++ b/services/core/src/main.rs @@ -631,6 +631,10 @@ async fn run_core() -> Result<(), Box> { lmm_core::remote_control::preadmit, )); } + // Outermost: see the final status, and mask query credentials before stderr. + app = app.layer(axum::middleware::from_fn( + lmm_core::log_redact::log_redacted_request_target, + )); // All enabled components and listener binds have succeeded. The health body // lists capabilities; model readiness has its own explicit endpoint. assert!(lifecycle.mark_ready()); diff --git a/services/core/src/relay/honeypot.rs b/services/core/src/relay/honeypot.rs index b216a6d80..2be7bd38d 100644 --- a/services/core/src/relay/honeypot.rs +++ b/services/core/src/relay/honeypot.rs @@ -565,9 +565,7 @@ impl Worker { } else { None }; - let payload = self.request.upstream_to(route.protocol, &route.model)?; - let bytes = serde_json::to_vec(&payload) - .map_err(|_| RelayError::Invalid("request encoding failed"))?; + let bytes = self.request.upstream_body(route.protocol, &route.model)?; if bytes.len() > self.limits.request_bytes { return Err(RelayError::Limit("encoded request")); } diff --git a/services/core/src/relay/mod.rs b/services/core/src/relay/mod.rs index 4f08354be..964ea3c29 100644 --- a/services/core/src/relay/mod.rs +++ b/services/core/src/relay/mod.rs @@ -12,6 +12,7 @@ pub(crate) mod native_anthropic; pub(crate) mod native_gemini; pub(crate) mod provider_network; mod request; +mod schema_order; mod sse; #[cfg(test)] mod tests; @@ -680,9 +681,7 @@ impl Worker { .await?; self.report.route_id = Some(route.id.clone()); self.report.price_version = Some(route.price_version); - let payload = self.request.upstream(&route)?; - let body = serde_json::to_vec(&payload) - .map_err(|_| RelayError::Invalid("request encoding failed"))?; + let body = self.request.upstream_body(route.protocol, &route.model)?; if body.len() > self.inner.limits.request_bytes { return Err(RelayError::Limit("encoded request")); } diff --git a/services/core/src/relay/native_anthropic.rs b/services/core/src/relay/native_anthropic.rs index 4b14f1eb0..426d3ae21 100644 --- a/services/core/src/relay/native_anthropic.rs +++ b/services/core/src/relay/native_anthropic.rs @@ -258,6 +258,9 @@ pub(crate) fn parse(body: &[u8], limits: &Limits) -> Result &serde_json::to_vec(&canonical).map_err(|_| RelayError::Invalid("request encoding"))?, limits, )?; + if let Some(captured) = super::schema_order::anthropic_schemas(body) { + request.assign_ordered_schemas(&captured); + } let mut messages = Vec::new(); if let Some(s) = v.get("system") { messages.push(Message { diff --git a/services/core/src/relay/native_gemini.rs b/services/core/src/relay/native_gemini.rs index a73670263..dc7b87915 100644 --- a/services/core/src/relay/native_gemini.rs +++ b/services/core/src/relay/native_gemini.rs @@ -376,6 +376,9 @@ pub(crate) fn parse( .map_err(|_| RelayError::Invalid("Gemini conversion failed"))?; let mut request = Request::parse(ClientProtocol::Chat, &bytes, limits)?; request.continuation_token = continuation_token(v.get("continuationToken"), limits)?; + if let Some(captured) = super::schema_order::gemini_schemas(body) { + request.assign_ordered_schemas(&captured); + } Ok(request) } fn continuation_token(v: Option<&Value>, limits: &Limits) -> Result, RelayError> { diff --git a/services/core/src/relay/request.rs b/services/core/src/relay/request.rs index cf3682422..785c23ab8 100644 --- a/services/core/src/relay/request.rs +++ b/services/core/src/relay/request.rs @@ -1,5 +1,7 @@ //! Bounded, text/function-call request conversion. Unsupported inputs fail closed. -use super::{ClientProtocol, Limits, RelayError, Route, UpstreamProtocol}; +#[cfg(test)] +use super::Route; +use super::{ClientProtocol, Limits, RelayError, UpstreamProtocol}; use serde::{ Deserialize, de::{MapAccess, Visitor}, @@ -66,6 +68,9 @@ pub struct FunctionTool { pub description: Option, pub parameters: Value, pub strict: Option, + /// Compact JSON for `parameters`, using the client's object-key order when + /// that text still parses as the validated value. + pub(crate) ordered_parameters: String, } #[derive(Clone, Debug)] pub enum ToolChoice { @@ -372,11 +377,14 @@ impl Request { if !parameters.is_object() { return Err(RelayError::Invalid("function parameters must be an object")); } + let ordered_parameters = + serde_json::to_string(¶meters).unwrap_or_else(|_| "{}".to_owned()); tools.push(FunctionTool { name, parameters, description: f.get("description").map(string).transpose()?, strict: boolean(f.get("strict"))?, + ordered_parameters, }); } } @@ -638,7 +646,7 @@ impl Request { Ok(options.clone()) }) .transpose()?; - Ok(Self { + let mut request = Self { client, model, stream, @@ -654,9 +662,58 @@ impl Request { reasoning_effort, prompt_cache_options, continuation_token: None, - }) + }; + let captured = match client { + ClientProtocol::Chat => super::schema_order::chat_schemas(body), + ClientProtocol::Responses => super::schema_order::responses_schemas(body), + ClientProtocol::Anthropic | ClientProtocol::Gemini => None, + }; + if let Some(captured) = captured { + request.assign_ordered_schemas(&captured); + } + Ok(request) + } + + pub(crate) fn assign_ordered_schemas( + &mut self, + captured: &[Option], + ) { + if captured.len() != self.tools.len() { + return; + } + for (tool, item) in self.tools.iter_mut().zip(captured) { + let Some(item) = item else { continue }; + let ordered = match item { + super::schema_order::RawSchema::Json(raw) => { + super::schema_order::ordered_json_matching(raw, &tool.parameters) + } + super::schema_order::RawSchema::GeminiOpenApi(raw) => { + super::schema_order::ordered_gemini_schema(raw, &tool.parameters) + } + }; + if let Some(text) = ordered { + tool.ordered_parameters = text; + } + } + } + + /// Upstream JSON bytes. Tool schemas keep the client's property order when + /// that encoding is semantically identical to the validated value. + pub(crate) fn upstream_body( + &self, + protocol: UpstreamProtocol, + model: &str, + ) -> Result, RelayError> { + let payload = self.upstream_to(protocol, model)?; + let schemas: Vec = self + .tools + .iter() + .map(|tool| tool.ordered_parameters.clone()) + .collect(); + super::schema_order::splice_ordered_schemas(&payload, &schemas) } + #[cfg(test)] pub(crate) fn upstream(&self, route: &Route) -> Result { self.upstream_to(route.protocol, &route.model) } diff --git a/services/core/src/relay/schema_order.rs b/services/core/src/relay/schema_order.rs new file mode 100644 index 000000000..a7fc6999d --- /dev/null +++ b/services/core/src/relay/schema_order.rs @@ -0,0 +1,657 @@ +//! Keep JSON Schema object-key order when a tool schema is re-encoded. +//! +//! `serde_json::Map` is a `BTreeMap`, so a decode/encode cycle sorts keys. +//! Structured-output and tool callers that fill `properties` in listed order +//! then see `answer` before `reasoning`. This module reads the client's schema +//! bytes, writes them again with that order, and splices the result into the +//! upstream body. Global JSON encoding stays sorted: billing and policy +//! digests must not start following client key order. +//! +//! Adapted from QuantumNous/new-api `d0b840e7`, which keeps schema bytes +//! instead of round-tripping them through a map. + +use super::RelayError; +use serde::Deserialize; +use serde_json::{Value, value::RawValue}; + +#[derive(Clone, Debug)] +enum Node { + Null, + Bool(bool), + Number(String), + String(String), + Array(Vec), + Object(Vec<(String, Node)>), +} + +#[derive(Clone, Debug)] +pub(crate) enum RawSchema { + Json(String), + GeminiOpenApi(String), +} + +pub(crate) fn ordered_json_matching(raw: &str, expected: &Value) -> Option { + let node = parse_json(raw)?; + let text = emit(&node)?; + (serde_json::from_str::(&text).ok().as_ref() == Some(expected)).then_some(text) +} + +pub(crate) fn ordered_gemini_schema(raw: &str, expected: &Value) -> Option { + let node = convert_gemini(&parse_json(raw)?)?; + let text = emit(&node)?; + (serde_json::from_str::(&text).ok().as_ref() == Some(expected)).then_some(text) +} + +pub(crate) fn chat_schemas(body: &[u8]) -> Option>> { + #[derive(Deserialize)] + struct Body { + #[serde(default)] + tools: Vec, + } + #[derive(Deserialize)] + struct Tool { + #[serde(default)] + function: Option, + } + #[derive(Deserialize)] + struct Function { + #[serde(default)] + parameters: Option>, + } + let body: Body = serde_json::from_slice(body).ok()?; + Some( + body.tools + .into_iter() + .map(|tool| { + tool.function.and_then(|function| { + function + .parameters + .map(|raw| RawSchema::Json(raw.get().to_owned())) + }) + }) + .collect(), + ) +} + +pub(crate) fn responses_schemas(body: &[u8]) -> Option>> { + #[derive(Deserialize)] + struct Body { + #[serde(default)] + tools: Vec, + } + #[derive(Deserialize)] + struct Tool { + #[serde(default)] + parameters: Option>, + } + let body: Body = serde_json::from_slice(body).ok()?; + Some( + body.tools + .into_iter() + .map(|tool| { + tool.parameters + .map(|raw| RawSchema::Json(raw.get().to_owned())) + }) + .collect(), + ) +} + +pub(crate) fn anthropic_schemas(body: &[u8]) -> Option>> { + #[derive(Deserialize)] + struct Body { + #[serde(default)] + tools: Vec, + } + #[derive(Deserialize)] + struct Tool { + #[serde(default)] + input_schema: Option>, + } + let body: Body = serde_json::from_slice(body).ok()?; + Some( + body.tools + .into_iter() + .map(|tool| { + tool.input_schema + .map(|raw| RawSchema::Json(raw.get().to_owned())) + }) + .collect(), + ) +} + +pub(crate) fn gemini_schemas(body: &[u8]) -> Option>> { + #[derive(Deserialize)] + struct Body { + #[serde(default)] + tools: Vec, + } + #[derive(Deserialize)] + struct Tool { + #[serde(default, rename = "functionDeclarations")] + function_declarations: Vec, + } + #[derive(Deserialize)] + struct Function { + #[serde(default)] + parameters: Option>, + #[serde(default, rename = "parametersJsonSchema")] + parameters_json_schema: Option>, + } + let body: Body = serde_json::from_slice(body).ok()?; + Some( + body.tools + .into_iter() + .flat_map(|tool| tool.function_declarations) + .map( + |function| match (function.parameters, function.parameters_json_schema) { + (Some(raw), None) => Some(RawSchema::GeminiOpenApi(raw.get().to_owned())), + (None, Some(raw)) => Some(RawSchema::Json(raw.get().to_owned())), + (None, None) => None, + (Some(_), Some(_)) => None, + }, + ) + .collect(), + ) +} + +/// Replace each tool schema in `payload`'s compact encoding with `schemas`, +/// which use the client's object-key order. The parsed JSON value stays equal +/// to `payload`. A splice that would change that value is discarded. +pub(crate) fn splice_ordered_schemas( + payload: &Value, + schemas: &[String], +) -> Result, RelayError> { + let canonical = + serde_json::to_vec(payload).map_err(|_| RelayError::Invalid("request encoding failed"))?; + let mut body = canonical.clone(); + let mut from = 0usize; + for schema in schemas { + let Ok(parsed) = serde_json::from_str::(schema) else { + continue; + }; + let Ok(needle) = serde_json::to_vec(&parsed) else { + continue; + }; + let Some(relative) = find_subslice(&body[from..], &needle) else { + continue; + }; + let at = from + relative; + if needle.as_slice() == schema.as_bytes() { + from = at + needle.len(); + continue; + } + body.splice(at..at + needle.len(), schema.bytes()); + from = at + schema.len(); + } + match serde_json::from_slice::(&body) { + Ok(value) if value == *payload => Ok(body), + _ => Ok(canonical), + } +} + +fn find_subslice(haystack: &[u8], needle: &[u8]) -> Option { + haystack + .windows(needle.len()) + .position(|window| window == needle) +} + +const GEMINI_KEYS: &[&str] = &[ + "type", + "properties", + "required", + "description", + "items", + "enum", + "format", + "minimum", + "maximum", + "minItems", + "maxItems", +]; + +fn convert_gemini(node: &Node) -> Option { + let Node::Object(items) = node else { + return None; + }; + if items + .iter() + .any(|(key, _)| !GEMINI_KEYS.contains(&key.as_str())) + { + return None; + } + let mut out = Vec::with_capacity(items.len()); + for (key, value) in items { + let converted = match key.as_str() { + "type" => { + let Node::String(kind) = value else { + return None; + }; + let mapped = match kind.as_str() { + "OBJECT" => "object", + "ARRAY" => "array", + "STRING" => "string", + "NUMBER" => "number", + "INTEGER" => "integer", + "BOOLEAN" => "boolean", + _ => return None, + }; + Node::String(mapped.to_owned()) + } + "properties" => { + let Node::Object(properties) = value else { + return None; + }; + let mut converted = Vec::with_capacity(properties.len()); + for (name, schema) in properties { + converted.push((name.clone(), convert_gemini(schema)?)); + } + Node::Object(converted) + } + "items" => convert_gemini(value)?, + _ => value.clone(), + }; + out.push((key.clone(), converted)); + } + Some(Node::Object(out)) +} + +struct Parser<'a> { + bytes: &'a [u8], + index: usize, + depth: usize, +} + +fn parse_json(raw: &str) -> Option { + let mut parser = Parser { + bytes: raw.as_bytes(), + index: 0, + depth: 0, + }; + let node = parser.value()?; + parser.skip_ws(); + (parser.index == parser.bytes.len()).then_some(node) +} + +impl<'a> Parser<'a> { + fn value(&mut self) -> Option { + self.skip_ws(); + if self.depth > 64 { + return None; + } + match self.peek()? { + b'{' => self.object(), + b'[' => self.array(), + b'"' => self.string().map(Node::String), + b't' => self.literal(b"true").map(|()| Node::Bool(true)), + b'f' => self.literal(b"false").map(|()| Node::Bool(false)), + b'n' => self.literal(b"null").map(|()| Node::Null), + b'-' | b'0'..=b'9' => self.number(), + _ => None, + } + } + + fn object(&mut self) -> Option { + self.bump(b'{')?; + self.depth += 1; + let mut items = Vec::new(); + self.skip_ws(); + if self.peek() == Some(b'}') { + self.index += 1; + self.depth -= 1; + return Some(Node::Object(items)); + } + loop { + self.skip_ws(); + let key = self.string()?; + self.skip_ws(); + self.bump(b':')?; + let value = self.value()?; + if let Some(existing) = items.iter_mut().find(|(name, _)| name == &key) { + existing.1 = value; + } else { + items.push((key, value)); + } + self.skip_ws(); + match self.peek()? { + b',' => self.index += 1, + b'}' => { + self.index += 1; + self.depth -= 1; + return Some(Node::Object(items)); + } + _ => return None, + } + } + } + + fn array(&mut self) -> Option { + self.bump(b'[')?; + self.depth += 1; + let mut items = Vec::new(); + self.skip_ws(); + if self.peek() == Some(b']') { + self.index += 1; + self.depth -= 1; + return Some(Node::Array(items)); + } + loop { + items.push(self.value()?); + self.skip_ws(); + match self.peek()? { + b',' => self.index += 1, + b']' => { + self.index += 1; + self.depth -= 1; + return Some(Node::Array(items)); + } + _ => return None, + } + } + } + + fn string(&mut self) -> Option { + self.bump(b'"')?; + let mut out = String::new(); + while let Some(byte) = self.peek() { + match byte { + b'"' => { + self.index += 1; + return Some(out); + } + b'\\' => { + self.index += 1; + let escaped = self.peek()?; + self.index += 1; + match escaped { + b'"' | b'\\' | b'/' => out.push(escaped as char), + b'b' => out.push('\u{0008}'), + b'f' => out.push('\u{000c}'), + b'n' => out.push('\n'), + b'r' => out.push('\r'), + b't' => out.push('\t'), + b'u' => out.push(self.unicode_escape()?), + _ => return None, + } + } + 0x00..=0x1f => return None, + byte if byte < 0x80 => { + out.push(byte as char); + self.index += 1; + } + _ => { + let rest = self.bytes.get(self.index..)?; + let text = std::str::from_utf8(rest).ok()?; + let ch = text.chars().next()?; + out.push(ch); + self.index += ch.len_utf8(); + } + } + } + None + } + + fn unicode_escape(&mut self) -> Option { + let unit = self.hex4()?; + if (0xD800..=0xDBFF).contains(&unit) { + self.bump(b'\\')?; + self.bump(b'u')?; + let low = self.hex4()?; + if !(0xDC00..=0xDFFF).contains(&low) { + return None; + } + let code = 0x10000 + (((unit as u32 - 0xD800) << 10) | (low as u32 - 0xDC00)); + return char::from_u32(code); + } + if (0xDC00..=0xDFFF).contains(&unit) { + return None; + } + char::from_u32(unit as u32) + } + + fn hex4(&mut self) -> Option { + let mut value = 0u16; + for _ in 0..4 { + let byte = self.peek()?; + self.index += 1; + let digit = match byte { + b'0'..=b'9' => byte - b'0', + b'a'..=b'f' => byte - b'a' + 10, + b'A'..=b'F' => byte - b'A' + 10, + _ => return None, + }; + value = (value << 4) | u16::from(digit); + } + Some(value) + } + + fn number(&mut self) -> Option { + let start = self.index; + if self.peek() == Some(b'-') { + self.index += 1; + } + match self.peek()? { + b'0' => self.index += 1, + b'1'..=b'9' => { + self.index += 1; + while matches!(self.peek(), Some(b'0'..=b'9')) { + self.index += 1; + } + } + _ => return None, + } + if self.peek() == Some(b'.') { + self.index += 1; + let mut digits = false; + while matches!(self.peek(), Some(b'0'..=b'9')) { + digits = true; + self.index += 1; + } + if !digits { + return None; + } + } + if matches!(self.peek(), Some(b'e' | b'E')) { + self.index += 1; + if matches!(self.peek(), Some(b'+' | b'-')) { + self.index += 1; + } + let mut digits = false; + while matches!(self.peek(), Some(b'0'..=b'9')) { + digits = true; + self.index += 1; + } + if !digits { + return None; + } + } + let literal = std::str::from_utf8(&self.bytes[start..self.index]).ok()?; + Some(Node::Number(literal.to_owned())) + } + + fn literal(&mut self, expected: &[u8]) -> Option<()> { + if self.bytes.get(self.index..)?.starts_with(expected) { + self.index += expected.len(); + Some(()) + } else { + None + } + } + + fn bump(&mut self, expected: u8) -> Option<()> { + if self.peek() == Some(expected) { + self.index += 1; + Some(()) + } else { + None + } + } + + fn peek(&self) -> Option { + self.bytes.get(self.index).copied() + } + + fn skip_ws(&mut self) { + while matches!(self.peek(), Some(b' ' | b'\n' | b'\r' | b'\t')) { + self.index += 1; + } + } +} + +fn emit(node: &Node) -> Option { + let mut out = String::new(); + write_node(&mut out, node)?; + Some(out) +} + +fn write_node(out: &mut String, node: &Node) -> Option<()> { + match node { + Node::Null => out.push_str("null"), + Node::Bool(true) => out.push_str("true"), + Node::Bool(false) => out.push_str("false"), + Node::Number(literal) => out.push_str(literal), + Node::String(text) => out.push_str(&serde_json::to_string(text).ok()?), + Node::Array(items) => { + out.push('['); + for (index, item) in items.iter().enumerate() { + if index > 0 { + out.push(','); + } + write_node(out, item)?; + } + out.push(']'); + } + Node::Object(items) => { + out.push('{'); + for (index, (key, value)) in items.iter().enumerate() { + if index > 0 { + out.push(','); + } + out.push_str(&serde_json::to_string(key).ok()?); + out.push(':'); + write_node(out, value)?; + } + out.push('}'); + } + } + Some(()) +} + +#[cfg(test)] +mod tests { + use super::super::{ + ClientProtocol, Limits, UpstreamProtocol, native_anthropic, native_gemini, request, + }; + use super::*; + use serde_json::json; + + const SCHEMA: &str = r#"{"type":"object","properties":{"reasoning":{"type":"string","description":"why"},"answer":{"type":"string"}},"required":["reasoning","answer"]}"#; + + #[test] + fn schema_properties_keep_insertion_order_across_a_marshal_round_trip() { + let ordered = + ordered_json_matching(SCHEMA, &serde_json::from_str::(SCHEMA).unwrap()) + .expect("schema matches its parsed value"); + assert!( + ordered.contains( + r#""properties":{"reasoning":{"type":"string","description":"why"},"answer":{"type":"string"}}"# + ), + "{ordered}" + ); + let again = + ordered_json_matching(&ordered, &serde_json::from_str::(&ordered).unwrap()) + .unwrap(); + assert_eq!(again, ordered); + let sorted = + serde_json::to_string(&serde_json::from_str::(SCHEMA).unwrap()).unwrap(); + assert!(sorted.find("\"answer\"").unwrap() < sorted.find("\"reasoning\"").unwrap()); + assert_ne!(ordered, sorted); + } + + #[test] + fn chat_responses_and_anthropic_upstream_bodies_keep_property_order() { + let chat = br#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tools":[{"type":"function","function":{"name":"answer","description":"d","parameters":{"type":"object","properties":{"reasoning":{"type":"string"},"answer":{"type":"string"}},"required":["reasoning","answer"]}}}]}"#; + let request = + request::Request::parse(ClientProtocol::Chat, chat, &Limits::default()).unwrap(); + let body = String::from_utf8( + request + .upstream_body(UpstreamProtocol::Chat, "upstream-model") + .unwrap(), + ) + .unwrap(); + assert!( + body.contains( + r#""parameters":{"type":"object","properties":{"reasoning":{"type":"string"},"answer":{"type":"string"}},"required":["reasoning","answer"]}"# + ), + "{body}" + ); + let responses = br#"{"model":"m","input":"hi","tools":[{"type":"function","name":"answer","parameters":{"type":"object","properties":{"reasoning":{"type":"string"},"answer":{"type":"string"}}}}]}"#; + let request = + request::Request::parse(ClientProtocol::Responses, responses, &Limits::default()) + .unwrap(); + let body = String::from_utf8( + request + .upstream_body(UpstreamProtocol::Responses, "upstream-model") + .unwrap(), + ) + .unwrap(); + assert!( + body.contains( + r#""parameters":{"type":"object","properties":{"reasoning":{"type":"string"},"answer":{"type":"string"}}}"# + ), + "{body}" + ); + let anthropic = br#"{"model":"m","max_tokens":16,"messages":[{"role":"user","content":"hi"}],"tools":[{"name":"answer","input_schema":{"type":"object","properties":{"reasoning":{"type":"string"},"answer":{"type":"string"}}}}]}"#; + let request = native_anthropic::parse(anthropic, &Limits::default()).unwrap(); + let body = String::from_utf8( + request + .upstream_body(UpstreamProtocol::Anthropic, "upstream-model") + .unwrap(), + ) + .unwrap(); + assert!( + body.contains( + r#""input_schema":{"type":"object","properties":{"reasoning":{"type":"string"},"answer":{"type":"string"}}}"# + ), + "{body}" + ); + } + + #[test] + fn gemini_openapi_parameters_keep_property_order_while_types_are_lowercased() { + let body = br#"{"contents":[{"parts":[{"text":"hi"}]}],"tools":[{"functionDeclarations":[{"name":"answer","parameters":{"type":"OBJECT","properties":{"reasoning":{"type":"STRING"},"answer":{"type":"STRING"}}}}]}]}"#; + let request = native_gemini::parse("gemini-test", false, body, &Limits::default()).unwrap(); + let encoded = request + .upstream_body(UpstreamProtocol::Gemini, "gemini-test") + .unwrap(); + let text = String::from_utf8(encoded).unwrap(); + assert!( + text.contains( + r#""parametersJsonSchema":{"type":"object","properties":{"reasoning":{"type":"string"},"answer":{"type":"string"}}}"# + ), + "{text}" + ); + let value: Value = serde_json::from_slice(body).unwrap(); + assert_eq!( + value["tools"][0]["functionDeclarations"][0]["parameters"]["type"], + json!("OBJECT") + ); + } + + #[test] + fn gemini_json_schema_parameters_keep_client_key_order() { + let body = br#"{"contents":[{"parts":[{"text":"hi"}]}],"tools":[{"functionDeclarations":[{"name":"answer","parametersJsonSchema":{"type":"object","properties":{"reasoning":{"type":"string"},"answer":{"type":"string"}}}}]}]}"#; + let request = native_gemini::parse("gemini-test", false, body, &Limits::default()).unwrap(); + let text = String::from_utf8( + request + .upstream_body(UpstreamProtocol::Gemini, "gemini-test") + .unwrap(), + ) + .unwrap(); + assert!( + text.contains( + r#""parametersJsonSchema":{"type":"object","properties":{"reasoning":{"type":"string"},"answer":{"type":"string"}}}"# + ), + "{text}" + ); + } +}