From 5a55116b62e7157447ab9afee35e927429a16e18 Mon Sep 17 00:00:00 2001 From: wangpengxiang Date: Thu, 17 Sep 2026 10:46:09 +0800 Subject: [PATCH 1/2] fix(antigravity): harden tool schemas and Claude thought replay Preserve tool schemas across provider-scoped format conversion and OpenAI gateway planners until the Antigravity boundary. Bound reference expansion, safely merge schema constraints, and lower unsupported Claude unions. Drop unsigned historical Claude thinking without changing Gemini behavior. Add cross-format regression fixtures and retain the project's original error output policy. Validation: 945 format tests and 527 transport tests passed; gateway cargo check passed. --- .../planner/standard/normalize/chat.rs | 60 +- .../planner/standard/normalize/responses.rs | 63 +- .../aether-ai/formats/src/formats/context.rs | 3 + .../gemini/generate_content/request.rs | 149 ++-- .../src/formats/shared/standard_matrix.rs | 29 +- .../src/formats/shared/standard_normalize.rs | 2 +- .../src/antigravity/fabric_exec_schema.json | 24 + .../transport/src/antigravity/mod.rs | 1 + .../transport/src/antigravity/request.rs | 690 ++++++++++++++- .../transport/src/antigravity/schema.rs | 789 ++++++++++++++++++ 10 files changed, 1730 insertions(+), 80 deletions(-) create mode 100644 crates/aether-provider/transport/src/antigravity/fabric_exec_schema.json create mode 100644 crates/aether-provider/transport/src/antigravity/schema.rs diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/normalize/chat.rs b/apps/aether-gateway/src/ai_serving/planner/standard/normalize/chat.rs index c113e88f6..e0d44d283 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/normalize/chat.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/normalize/chat.rs @@ -70,13 +70,23 @@ pub(crate) fn build_cross_format_openai_chat_request_body( request_headers: &http::HeaderMap, enable_model_directives: bool, ) -> Option { - let provider_request_body = surface_build_cross_format_openai_chat_request_body( - body_json, - mapped_model, - provider_api_format, - upstream_is_stream, - enable_model_directives, - )?; + let provider_request_body = if provider_type.trim().eq_ignore_ascii_case("antigravity") + && aether_ai_formats::normalize_api_format_alias(provider_api_format) + == "gemini:generate_content" + { + aether_ai_formats::formats::shared::standard_matrix::build_standard_request_body_with_model_directives( + body_json, "openai:chat", mapped_model, provider_type, provider_api_format, + "", upstream_is_stream, None, user_api_key_id, enable_model_directives, + )? + } else { + surface_build_cross_format_openai_chat_request_body( + body_json, + mapped_model, + provider_api_format, + upstream_is_stream, + enable_model_directives, + )? + }; let mut provider_request_body = apply_standard_provider_request_body_rules_with_request_headers( provider_request_body, @@ -110,6 +120,42 @@ pub(crate) fn build_cross_format_openai_chat_request_body( Some(provider_request_body) } +#[cfg(test)] +mod antigravity_schema_tests { + use super::*; + use serde_json::json; + + #[test] + fn antigravity_chat_route_preserves_tool_schema_and_alternate_responses_shape() { + let schema = json!({"type": "object", "properties": {"mode": {"const": "fast"}}}); + let body = json!({"model": "client", "messages": [{"role": "user", "content": "hi"}], + "tools": [{"type": "function", "function": {"name": "probe", "parameters": schema}}]}); + let responses_body = json!({"model": "client", "input": "hi", + "tools": [{"type": "function", "name": "probe", "parameters": schema}]}); + for input in [body, responses_body] { + for provider in ["antigravity", "gemini"] { + let output = build_cross_format_openai_chat_request_body( + &input, + "claude-test", + provider, + "gemini:generate_content", + true, + false, + None, + None, + &http::HeaderMap::new(), + false, + ) + .unwrap(); + let parameters = &output["tools"][0]["functionDeclarations"][0]["parameters"]; + assert_eq!(parameters == &schema, provider == "antigravity"); + assert!(output.get("stream").is_none()); + assert_eq!(output["contents"][0]["parts"][0]["text"], "hi"); + } + } + } +} + pub(crate) fn build_cross_format_openai_chat_upstream_url( parts: &http::request::Parts, transport: &GatewayProviderTransportSnapshot, diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/normalize/responses.rs b/apps/aether-gateway/src/ai_serving/planner/standard/normalize/responses.rs index 47c91f1b9..724ec4a94 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/normalize/responses.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/normalize/responses.rs @@ -214,15 +214,25 @@ pub(crate) fn build_cross_format_openai_responses_request_body_with_codex_model_ model_capabilities: Option<&crate::ai_serving::CodexResponsesModelCapabilities>, enable_model_directives: bool, ) -> Option { - let provider_request_body = surface_build_cross_format_openai_responses_request_body( - body_json, - mapped_model, - client_api_format, - provider_api_format, - upstream_is_stream, - enable_model_directives, - history_scope, - )?; + let provider_request_body = if provider_type.trim().eq_ignore_ascii_case("antigravity") + && aether_ai_formats::normalize_api_format_alias(provider_api_format) + == "gemini:generate_content" + { + aether_ai_formats::formats::shared::standard_matrix::build_standard_request_body_with_model_directives( + body_json, client_api_format, mapped_model, provider_type, provider_api_format, + "", upstream_is_stream, None, history_scope, enable_model_directives, + )? + } else { + surface_build_cross_format_openai_responses_request_body( + body_json, + mapped_model, + client_api_format, + provider_api_format, + upstream_is_stream, + enable_model_directives, + history_scope, + )? + }; let mut provider_request_body = apply_standard_provider_request_body_rules_with_request_headers( provider_request_body, @@ -274,6 +284,41 @@ pub(crate) fn build_local_openai_responses_upstream_url( ) } +#[cfg(test)] +mod antigravity_schema_tests { + use super::*; + use serde_json::json; + + #[test] + fn antigravity_responses_route_preserves_tool_schema_without_changing_public_gemini() { + let schema = json!({"type": "object", "properties": {"mode": {"const": "fast"}}}); + let input = json!({"model": "client", "input": "hi", + "tools": [{"type": "function", "name": "probe", "parameters": schema}]}); + for provider in ["antigravity", "gemini"] { + let output = + build_cross_format_openai_responses_request_body_with_codex_model_capabilities( + &input, + "claude-test", + "openai:responses", + "gemini:generate_content", + true, + false, + provider, + None, + &http::HeaderMap::new(), + Some("antigravity-schema-test"), + None, + false, + ) + .unwrap(); + let parameters = &output["tools"][0]["functionDeclarations"][0]["parameters"]; + assert_eq!(parameters == &schema, provider == "antigravity"); + assert!(output.get("stream").is_none()); + assert_eq!(output["contents"][0]["parts"][0]["text"], "hi"); + } + } +} + pub(crate) fn build_cross_format_openai_responses_upstream_url( parts: &http::request::Parts, transport: &GatewayProviderTransportSnapshot, diff --git a/crates/aether-ai/formats/src/formats/context.rs b/crates/aether-ai/formats/src/formats/context.rs index 02b189156..be7dee222 100644 --- a/crates/aether-ai/formats/src/formats/context.rs +++ b/crates/aether-ai/formats/src/formats/context.rs @@ -10,6 +10,8 @@ pub struct FormatContext { pub upstream_is_stream: bool, pub report_context: Option, pub history_scope: Option, + /// Defer tool schema lowering to the private provider transport boundary. + pub preserve_gemini_tool_schemas: bool, } impl FormatContext { @@ -45,6 +47,7 @@ impl FormatContext { upstream_is_stream: false, report_context: self.report_context.clone(), history_scope: self.history_scope.clone(), + preserve_gemini_tool_schemas: false, } } diff --git a/crates/aether-ai/formats/src/formats/gemini/generate_content/request.rs b/crates/aether-ai/formats/src/formats/gemini/generate_content/request.rs index 7b95dcc1f..c6d51c4e6 100644 --- a/crates/aether-ai/formats/src/formats/gemini/generate_content/request.rs +++ b/crates/aether-ai/formats/src/formats/gemini/generate_content/request.rs @@ -31,10 +31,11 @@ pub fn from(body: &Value, ctx: &FormatContext) -> Option { } pub fn to(request: &CanonicalRequest, ctx: &FormatContext) -> Option { - to_raw( + to_raw_with_schema_policy( request, ctx.mapped_model_or(request.model.as_str()), ctx.upstream_is_stream, + ctx.preserve_gemini_tool_schemas, ) } @@ -187,7 +188,21 @@ pub fn to_raw( mapped_model: &str, upstream_is_stream: bool, ) -> Option { - let mut output = canonical_to_gemini_request_body(canonical, mapped_model, upstream_is_stream)?; + to_raw_with_schema_policy(canonical, mapped_model, upstream_is_stream, false) +} + +fn to_raw_with_schema_policy( + canonical: &CanonicalRequest, + mapped_model: &str, + upstream_is_stream: bool, + preserve_tool_schemas: bool, +) -> Option { + let mut output = canonical_to_gemini_request_body( + canonical, + mapped_model, + upstream_is_stream, + preserve_tool_schemas, + )?; apply_gemini_request_extensions(&mut output, &canonical.extensions)?; if !canonical_has_raw_gemini_tools(canonical) { enable_server_side_tool_invocations_for_mixed_tools(&mut output, mapped_model)?; @@ -244,7 +259,8 @@ pub fn ensure_server_side_tool_invocations_for_mixed_tools(output: &mut Value) - } pub(crate) fn canonical_has_mixed_gemini_tools(canonical: &CanonicalRequest) -> bool { - canonical_tools_to_gemini(canonical) + // Only tool kinds matter here; do not lower/expand schemas just to count them. + canonical_tools_to_gemini(canonical, true) .and_then(|tools| tools.as_array().cloned()) .is_some_and(|tools| gemini_tools_are_mixed(&tools)) } @@ -275,6 +291,7 @@ fn canonical_to_gemini_request_body( canonical: &CanonicalRequest, mapped_model: &str, _upstream_is_stream: bool, + preserve_tool_schemas: bool, ) -> Option { let mut output = Map::new(); if !mapped_model.trim().is_empty() { @@ -297,7 +314,7 @@ fn canonical_to_gemini_request_body( { output.insert("generationConfig".to_string(), generation_config); } - if let Some(tools) = canonical_tools_to_gemini(canonical) { + if let Some(tools) = canonical_tools_to_gemini(canonical, preserve_tool_schemas) { output.insert("tools".to_string(), tools); } if let Some(tool_config) = canonical_tool_choice_to_gemini(canonical.tool_choice.as_ref()) { @@ -701,7 +718,10 @@ fn apply_response_format_to_gemini_generation_config( } } -fn canonical_tools_to_gemini(canonical: &CanonicalRequest) -> Option { +fn canonical_tools_to_gemini( + canonical: &CanonicalRequest, + preserve_tool_schemas: bool, +) -> Option { let mut declarations = Vec::new(); let mut tools = Vec::new(); let mut google_search = canonical @@ -747,7 +767,10 @@ fn canonical_tools_to_gemini(canonical: &CanonicalRequest) -> Option { google_search = true; continue; } - declarations.push(canonical_tool_to_gemini_declaration(tool)); + declarations.push(canonical_tool_to_gemini_declaration( + tool, + preserve_tool_schemas, + )); } let mut emitted_google_search = false; let mut emitted_code_execution = false; @@ -877,7 +900,10 @@ fn gemini_unhandled_builtin_tool_portion(tool_object: &Map) -> Op (!builtin.is_empty()).then_some(Value::Object(builtin)) } -fn canonical_tool_to_gemini_declaration(tool: &CanonicalToolDefinition) -> Value { +fn canonical_tool_to_gemini_declaration( + tool: &CanonicalToolDefinition, + preserve_tool_schema: bool, +) -> Value { let mut declaration = Map::new(); declaration.insert("name".to_string(), Value::String(tool.name.clone())); if let Some(description) = &tool.description { @@ -898,7 +924,7 @@ fn canonical_tool_to_gemini_declaration(tool: &CanonicalToolDefinition) -> Value .clone() .or_else(|| tool.parameters.clone()) .map(|mut schema| { - if raw_parameters.is_none() { + if raw_parameters.is_none() && !preserve_tool_schema { clean_gemini_schema(&mut schema); } schema @@ -1217,43 +1243,46 @@ mod tests { #[test] fn canonical_tool_declaration_sanitizes_json_schema_for_gemini() { - let declaration = canonical_tool_to_gemini_declaration(&CanonicalToolDefinition { - name: "inspect".to_string(), - description: None, - parameters: Some(json!({ - "$defs": { - "Target": { - "type": "object", - "properties": { - "secret": { - "type": "string", - "encrypted": true - } + let declaration = canonical_tool_to_gemini_declaration( + &CanonicalToolDefinition { + name: "inspect".to_string(), + description: None, + parameters: Some(json!({ + "$defs": { + "Target": { + "type": "object", + "properties": { + "secret": { + "type": "string", + "encrypted": true + } + }, + "required": ["secret"], + "additionalProperties": false + } + }, + "type": "object", + "properties": { + "target": { + "oneOf": [ + {"$ref": "#/$defs/Target"}, + {"type": "null"} + ] }, - "required": ["secret"], - "additionalProperties": false + "mode": { + "type": ["string", "null"], + "enum": [1, "fast"] + }, + "value": { + "type": ["string", "integer"] + } } - }, - "type": "object", - "properties": { - "target": { - "oneOf": [ - {"$ref": "#/$defs/Target"}, - {"type": "null"} - ] - }, - "mode": { - "type": ["string", "null"], - "enum": [1, "fast"] - }, - "value": { - "type": ["string", "integer"] - } - } - })), - strict: None, - extensions: BTreeMap::new(), - }); + })), + strict: None, + extensions: BTreeMap::new(), + }, + false, + ); assert_eq!( declaration["parameters"], @@ -1361,17 +1390,20 @@ mod tests { ..CanonicalRequest::default() }; - let tools = canonical_tools_to_gemini(&canonical).expect("tools should be emitted"); - let tools = tools.as_array().expect("tools should be an array"); + for preserve_tool_schemas in [false, true] { + let tools = canonical_tools_to_gemini(&canonical, preserve_tool_schemas) + .expect("tools should be emitted"); + let tools = tools.as_array().expect("tools should be an array"); - assert!( - tools.iter().all(|tool| tool.get("googleSearch").is_none()), - "a client tool named WebSearch must not become server-side grounding: {tools:?}" - ); - assert_eq!( - tools[0]["functionDeclarations"][0]["name"], "WebSearch", - "the client declaration must survive: {tools:?}" - ); + assert!( + tools.iter().all(|tool| tool.get("googleSearch").is_none()), + "a client tool named WebSearch must not become server-side grounding: {tools:?}" + ); + assert_eq!( + tools[0]["functionDeclarations"][0]["name"], "WebSearch", + "the client declaration must survive: {tools:?}" + ); + } } #[test] @@ -1388,10 +1420,13 @@ mod tests { ..CanonicalRequest::default() }; - let tools = canonical_tools_to_gemini(&canonical).expect("tools should be emitted"); - let tools = tools.as_array().expect("tools should be an array"); + for preserve_tool_schemas in [false, true] { + let tools = canonical_tools_to_gemini(&canonical, preserve_tool_schemas) + .expect("tools should be emitted"); + let tools = tools.as_array().expect("tools should be an array"); - assert_eq!(tools.len(), 1, "{tools:?}"); - assert_eq!(tools[0]["googleSearch"], json!({})); + assert_eq!(tools.len(), 1, "{tools:?}"); + assert_eq!(tools[0]["googleSearch"], json!({})); + } } } diff --git a/crates/aether-ai/formats/src/formats/shared/standard_matrix.rs b/crates/aether-ai/formats/src/formats/shared/standard_matrix.rs index 3a862f407..29d5fbc93 100644 --- a/crates/aether-ai/formats/src/formats/shared/standard_matrix.rs +++ b/crates/aether-ai/formats/src/formats/shared/standard_matrix.rs @@ -131,6 +131,10 @@ pub fn build_standard_request_body_with_model_directives_and_request_headers_and .with_mapped_model(mapped_model) .with_request_path(request_path) .with_upstream_stream(upstream_is_stream); + format_context.preserve_gemini_tool_schemas = + provider_type.trim().eq_ignore_ascii_case("antigravity") + && aether_ai_formats::normalize_api_format_alias(provider_api_format) + == "gemini:generate_content"; if let Some(history_scope) = user_api_key_id { format_context = format_context.with_history_scope(history_scope); } @@ -139,6 +143,23 @@ pub fn build_standard_request_body_with_model_directives_and_request_headers_and client_api_format, provider_api_format, ); + // Keep the specialized OpenAI builders' compatibility/history preprocessing + // when routing them through the provider-aware schema-preserving path. + let antigravity_chat_body = if format_context.preserve_gemini_tool_schemas + && matches!( + aether_ai_formats::normalize_api_format_alias(source_api_format.as_ref()).as_str(), + "openai:chat" | "openai:responses" | "openai:responses:compact" + ) { + Some( + crate::formats::shared::standard_normalize::chat_compatible_body_for_standard_source( + body_json, + source_api_format.as_ref(), + user_api_key_id, + )?, + ) + } else { + None + }; // DeepSeek and xAI replay opaque provider state. Preserve their native // Responses input items: canonical conversion can lose reasoning IDs and // encrypted-only items even when source and destination formats are equal. @@ -152,9 +173,13 @@ pub fn build_standard_request_body_with_model_directives_and_request_headers_and Value::Object(object) } else { convert_request( - source_api_format.as_ref(), + if antigravity_chat_body.is_some() { + "openai:chat" + } else { + source_api_format.as_ref() + }, provider_api_format, - body_json, + antigravity_chat_body.as_deref().unwrap_or(body_json), &format_context, ) .ok()? diff --git a/crates/aether-ai/formats/src/formats/shared/standard_normalize.rs b/crates/aether-ai/formats/src/formats/shared/standard_normalize.rs index de6d1f070..c0e0062a6 100644 --- a/crates/aether-ai/formats/src/formats/shared/standard_normalize.rs +++ b/crates/aether-ai/formats/src/formats/shared/standard_normalize.rs @@ -65,7 +65,7 @@ fn chat_compatible_body_for_openai_chat_endpoint(body_json: &Value) -> Option( +pub(crate) fn chat_compatible_body_for_standard_source<'a>( body_json: &'a Value, client_api_format: &str, history_scope: Option<&str>, diff --git a/crates/aether-provider/transport/src/antigravity/fabric_exec_schema.json b/crates/aether-provider/transport/src/antigravity/fabric_exec_schema.json new file mode 100644 index 000000000..1b98d673e --- /dev/null +++ b/crates/aether-provider/transport/src/antigravity/fabric_exec_schema.json @@ -0,0 +1,24 @@ +{ + "type": "object", + "required": ["code"], + "properties": { + "code": {"type": "string", "description": "TypeScript function body."}, + "payloads": {"type": "object", "patternProperties": {"^.*$": {"type": "string"}}}, + "resultFormat": {"anyOf": [ + {"type": "string", "const": "auto"}, + {"type": "string", "const": "yaml"}, + {"type": "string", "const": "json"}, + {"type": "string", "const": "text"} + ]}, + "tokenBudget": {"type": "number", "minimum": 1}, + "agentBudget": {"type": "number", "minimum": 1}, + "timeoutMs": {"type": "number", "minimum": 1}, + "display": {"anyOf": [ + {"type": "object", "properties": { + "name": {"type": "string"}, + "description": {"type": "string"} + }}, + {"type": "string"} + ]} + } +} diff --git a/crates/aether-provider/transport/src/antigravity/mod.rs b/crates/aether-provider/transport/src/antigravity/mod.rs index 4d09239e2..ac1c91ffb 100644 --- a/crates/aether-provider/transport/src/antigravity/mod.rs +++ b/crates/aether-provider/transport/src/antigravity/mod.rs @@ -1,6 +1,7 @@ mod auth; mod policy; mod request; +mod schema; mod url; pub use auth::{ diff --git a/crates/aether-provider/transport/src/antigravity/request.rs b/crates/aether-provider/transport/src/antigravity/request.rs index c6002f295..2f009b7f3 100644 --- a/crates/aether-provider/transport/src/antigravity/request.rs +++ b/crates/aether-provider/transport/src/antigravity/request.rs @@ -1,6 +1,7 @@ use serde_json::{Map, Value}; use super::auth::{AntigravityRequestAuth, ANTIGRAVITY_REQUEST_USER_AGENT}; +use super::schema::{normalize_claude_unions, normalize_tool_parameters, SchemaBudget}; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum AntigravityEnvelopeRequestType { @@ -31,6 +32,7 @@ pub enum AntigravityRequestEnvelopeUnsupportedReason { MissingContents, MissingRequestId, MissingModel, + ToolSchemaBudgetExceeded, } pub fn classify_antigravity_safe_request_body( @@ -78,8 +80,14 @@ pub fn build_antigravity_safe_v1internal_request( inner_request.remove("model"); inner_request.remove("safetySettings"); inner_request.remove("safety_settings"); + normalize_antigravity_claude_thought_history(&mut inner_request, model); normalize_antigravity_builtin_tool_names(&mut inner_request); - normalize_antigravity_function_declaration_parameters(&mut inner_request); + if normalize_antigravity_function_declaration_parameters(&mut inner_request, model).is_err() + { + return AntigravityRequestEnvelopeSupport::Unsupported( + AntigravityRequestEnvelopeUnsupportedReason::ToolSchemaBudgetExceeded, + ); + } let request_id = non_empty_string_field(source, "requestId").unwrap_or(request_id); let user_agent = non_empty_string_field(source, "userAgent").unwrap_or(ANTIGRAVITY_REQUEST_USER_AGENT); @@ -104,8 +112,13 @@ pub fn build_antigravity_safe_v1internal_request( inner_request.remove("model"); inner_request.remove("safetySettings"); inner_request.remove("safety_settings"); + normalize_antigravity_claude_thought_history(&mut inner_request, model); normalize_antigravity_builtin_tool_names(&mut inner_request); - normalize_antigravity_function_declaration_parameters(&mut inner_request); + if normalize_antigravity_function_declaration_parameters(&mut inner_request, model).is_err() { + return AntigravityRequestEnvelopeSupport::Unsupported( + AntigravityRequestEnvelopeUnsupportedReason::ToolSchemaBudgetExceeded, + ); + } let mut envelope = serde_json::json!({ "project": auth.project_id, @@ -150,10 +163,47 @@ fn normalize_antigravity_builtin_tool_names(request: &mut Map) { } } -fn normalize_antigravity_function_declaration_parameters(request: &mut Map) { - let Some(tools) = request.get_mut("tools").and_then(Value::as_array_mut) else { +/// Claude requires a replayable signature on historical thinking blocks. In +/// particular, Responses reasoning summaries are not signed thinking. Omit that +/// non-replayable metadata instead of inventing a signature or promoting private +/// reasoning into ordinary assistant text. Keep Gemini's native policy unchanged. +fn normalize_antigravity_claude_thought_history(request: &mut Map, model: &str) { + if !model.trim().to_ascii_lowercase().starts_with("claude-") { + return; + } + let Some(contents) = request.get_mut("contents").and_then(Value::as_array_mut) else { return; }; + contents.retain_mut(|message| { + if message.get("role").and_then(Value::as_str) != Some("model") { + return true; + } + let Some(parts) = message.get_mut("parts").and_then(Value::as_array_mut) else { + return true; + }; + let previous_len = parts.len(); + parts.retain(|part| { + part.get("thought").and_then(Value::as_bool) != Some(true) + || ["thoughtSignature", "thought_signature"].iter().any(|key| { + part.get(*key) + .and_then(Value::as_str) + .is_some_and(|signature| !signature.trim().is_empty()) + }) + }); + // Do not introduce empty messages when a turn contained only a summary. + !parts.is_empty() || parts.len() == previous_len + }); +} + +fn normalize_antigravity_function_declaration_parameters( + request: &mut Map, + model: &str, +) -> Result<(), ()> { + let Some(tools) = request.get_mut("tools").and_then(Value::as_array_mut) else { + return Ok(()); + }; + let mut budget = SchemaBudget::default(); + let claude = model.trim().to_ascii_lowercase().starts_with("claude-"); for tool in tools { let Some(tool_object) = tool.as_object_mut() else { @@ -177,9 +227,16 @@ fn normalize_antigravity_function_declaration_parameters(request: &mut Map) -> Option<&Map> { @@ -216,6 +273,282 @@ mod tests { }; use crate::antigravity::ANTIGRAVITY_REQUEST_USER_AGENT; + #[test] + fn antigravity_claude_fabric_union_regression_across_client_formats() { + use aether_ai_formats::formats::shared::standard_matrix::build_standard_request_body; + let schema: serde_json::Value = + serde_json::from_str(include_str!("fabric_exec_schema.json")).unwrap(); + let clients = [ + ( + "gemini:generate_content", + json!({"contents": [{"role":"user","parts":[{"text":"hi"}]}], + "tools":[{"functionDeclarations":[{"name":"fabric_exec","parametersJsonSchema":schema}]}]}), + ), + ( + "claude:messages", + json!({"messages":[{"role":"user","content":"hi"}],"max_tokens":128, + "tools":[{"name":"fabric_exec","input_schema":schema}]}), + ), + ( + "openai:chat", + json!({"messages":[{"role":"user","content":"hi"}], + "tools":[{"type":"function","function":{"name":"fabric_exec","parameters":schema}}]}), + ), + ( + "openai:responses", + json!({"input":"hi", + "tools":[{"type":"function","name":"fabric_exec","parameters":schema}]}), + ), + ]; + for (format, body) in clients { + for model in ["claude-opus-4-6-thinking", "gemini-test"] { + let converted = build_standard_request_body( + &body, + format, + model, + "antigravity", + "gemini:generate_content", + "", + true, + None, + None, + ) + .unwrap(); + for request in [converted.clone(), json!({"request":converted})] { + let AntigravityRequestEnvelopeSupport::Supported(output) = + build_antigravity_safe_v1internal_request( + &sample_auth(), + "union-regression", + model, + &request, + AntigravityEnvelopeRequestType::Agent, + ) + else { + panic!("failed {format} {model}"); + }; + let s = &output["request"]["tools"][0]["functionDeclarations"][0]["parameters"]; + assert_eq!(s["required"], json!(["code"])); + assert_eq!( + s["properties"]["payloads"]["additionalProperties"], + json!({"type":"string"}) + ); + assert_eq!(s["properties"]["tokenBudget"]["minimum"], 1); + if model.starts_with("claude-") { + assert_eq!( + s["properties"]["resultFormat"], + json!({"type":"string","enum":["auto","json","text","yaml"]}) + ); + let display = &s["properties"]["display"]; + assert!( + display.get("type").is_none(), + "do not select one union branch" + ); + assert!(display.get("anyOf").is_none()); + assert!(display["description"].as_str().unwrap().contains("object")); + assert!(display["description"].as_str().unwrap().contains("string")); + } else { + assert!(s["properties"]["resultFormat"]["anyOf"].is_array()); + assert!(s["properties"]["display"]["anyOf"].is_array()); + } + let AntigravityRequestEnvelopeSupport::Supported(twice) = + build_antigravity_safe_v1internal_request( + &sample_auth(), + "union-regression", + model, + &output, + AntigravityEnvelopeRequestType::Agent, + ) + else { + panic!("idempotence"); + }; + assert_eq!(twice, output); + } + } + } + } + + #[test] + fn antigravity_claude_omits_only_unsigned_thought_parts() { + let signed = json!({"text":"signed plan","thought":true,"thoughtSignature":"signed-value"}); + let signed_alias = + json!({"text":"signed alias","thought":true,"thought_signature":"alias-value"}); + let call = json!({"functionCall":{"id":"call_1","name":"lookup","args":{}},"thoughtSignature":"skip_thought_signature_validator"}); + let result = json!({"role":"user","parts":[{"functionResponse":{"id":"call_1","name":"lookup","response":{"result":"ok"}}}]}); + let body = json!({ + "contents":[ + {"role":"user","parts":[{"text":"hello"}]}, + {"role":"model","parts":[{"text":"unsigned-only summary","thought":true}]}, + {"role":"model","parts":[ + {"text":"unsigned summary","thought":true}, + {"text":"empty signature","thought":true,"thoughtSignature":""}, + {"text":"blank signature","thought":true,"thoughtSignature":" "}, + {"text":"non-string signature","thought":true,"thoughtSignature":12}, + signed,signed_alias,{"text":"visible answer"},call + ]}, + result + ], + "generationConfig":{"maxOutputTokens":64000,"thinkingConfig":{"includeThoughts":true,"thinkingBudget":4096}} + }); + for model in [ + "claude-sonnet-4-6", + "claude-opus-4-6-thinking", + "gemini-3.8-flash-high", + ] { + for wrapped in [false, true] { + let input = if wrapped { + json!({"request":body}) + } else { + body.clone() + }; + let original = input.clone(); + let AntigravityRequestEnvelopeSupport::Supported(output) = + build_antigravity_safe_v1internal_request( + &sample_auth(), + "thought-test", + model, + &input, + AntigravityEnvelopeRequestType::Agent, + ) + else { + panic!("expected supported envelope"); + }; + let expected = if model.starts_with("claude-") { + json!([ + {"role":"user","parts":[{"text":"hello"}]}, + {"role":"model","parts":[signed,signed_alias,{"text":"visible answer"},call]}, + result + ]) + } else { + body["contents"].clone() + }; + assert_eq!( + output["request"]["contents"], expected, + "{model} wrapped={wrapped}" + ); + assert_eq!( + output["request"]["generationConfig"], + body["generationConfig"] + ); + assert_eq!(input, original, "do not mutate caller-owned input"); + let rebuilt = build_antigravity_safe_v1internal_request( + &sample_auth(), + "thought-test", + model, + &output, + AntigravityEnvelopeRequestType::Agent, + ); + assert_eq!( + rebuilt, + AntigravityRequestEnvelopeSupport::Supported(output) + ); + } + } + } + + #[test] + fn antigravity_claude_cross_format_unsigned_reasoning_history() { + use aether_ai_formats::formats::shared::standard_matrix::build_standard_request_body; + let clients = [ + ( + "openai:responses", + json!({ + "max_output_tokens":64000, + "input":[ + {"role":"user","content":"hello"}, + {"type":"reasoning","id":"rs_history","status":"completed","summary":[{"type":"summary_text","text":"historical summary"}],"content":[]}, + {"role":"assistant","content":[{"type":"output_text","text":"visible answer"}]}, + {"type":"function_call","call_id":"call_1","name":"lookup","arguments":"{}"}, + {"type":"function_call_output","call_id":"call_1","output":"ok"}, + {"role":"user","content":"continue"} + ] + }), + ), + ( + "claude:messages", + json!({ + "max_tokens":64000, + "messages":[ + {"role":"user","content":"hello"}, + {"role":"assistant","content":[ + {"type":"thinking","thinking":"historical summary"}, + {"type":"text","text":"visible answer"}, + {"type":"tool_use","id":"call_1","name":"lookup","input":{}} + ]}, + {"role":"user","content":[{"type":"tool_result","tool_use_id":"call_1","content":"ok"},{"type":"text","text":"continue"}]} + ] + }), + ), + ]; + for (format, body) in clients { + for model in [ + "claude-sonnet-4-6", + "claude-opus-4-6-thinking", + "gemini-3.8-flash-high", + ] { + let converted = build_standard_request_body( + &body, + format, + model, + "antigravity", + "gemini:generate_content", + "", + true, + None, + None, + ) + .unwrap(); + assert!( + converted["contents"] + .as_array() + .unwrap() + .iter() + .flat_map(|m| m["parts"].as_array().unwrap()) + .any(|p| p["thought"] == true), + "fixture must exercise unsigned thoughts: {format}" + ); + let AntigravityRequestEnvelopeSupport::Supported(output) = + build_antigravity_safe_v1internal_request( + &sample_auth(), + "reasoning-history", + model, + &converted, + AntigravityEnvelopeRequestType::Agent, + ) + else { + panic!("expected supported envelope"); + }; + let messages = output["request"]["contents"].as_array().unwrap(); + assert!(messages + .iter() + .all(|m| !m["parts"].as_array().unwrap().is_empty())); + let parts: Vec<_> = messages + .iter() + .flat_map(|m| m["parts"].as_array().unwrap()) + .collect(); + assert_eq!( + parts.iter().any(|p| p["thought"] == true), + !model.starts_with("claude-") + ); + assert!(parts.iter().any(|p| p["text"] == "visible answer")); + assert!(parts.iter().any(|p| p["text"] == "continue")); + assert!(parts.iter().any(|p| p["functionCall"]["id"] == "call_1" + && p["functionCall"]["name"] == "lookup")); + assert!(parts.iter().any(|p| p["functionResponse"]["id"] == "call_1" + && p["functionResponse"]["name"] == "lookup")); + if model.starts_with("claude-") { + assert!( + !parts.iter().any(|p| p["text"] == "historical summary"), + "do not promote private reasoning to visible text" + ); + } + assert_eq!( + output["request"]["generationConfig"]["maxOutputTokens"], + 64000 + ); + } + } + } + fn sample_auth() -> AntigravityRequestAuth { AntigravityRequestAuth { project_id: "project-ant-123".to_string(), @@ -224,6 +557,199 @@ mod tests { } } + #[test] + fn antigravity_combined_client_conversion_preserves_schemas_until_transport() { + use aether_ai_formats::formats::shared::standard_matrix::build_standard_request_body; + let schema = json!({"type": "object", "properties": { + "mode": {"const": "fast"}, + "payloads": {"type": "object", "patternProperties": {"^.*$": {"type": "string"}}}, + "name": {"type": "string", "minLength": 1} + }, "required": ["mode"]}); + let clients = [ + ( + "claude:messages", + json!({"model": "client-model", "max_tokens": 128, + "messages": [{"role": "user", "content": "hi"}], + "tools": [{"name": "probe", "input_schema": schema}]}), + ), + ( + "openai:chat", + json!({"model": "client-model", + "messages": [{"role": "user", "content": "hi"}], + "tools": [{"type": "function", "function": {"name": "probe", "parameters": schema}}]}), + ), + ( + "openai:responses", + json!({"model": "client-model", "input": "hi", + "tools": [{"type": "function", "name": "probe", "parameters": schema}]}), + ), + ( + "gemini:generate_content", + json!({"model": "client-model", + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "tools": [{"functionDeclarations": [{"name": "probe", "parameters": schema}]}]}), + ), + ]; + for (source, original) in clients { + for model in ["claude-sonnet-test", "gemini-test"] { + let converted = build_standard_request_body( + &original, + source, + model, + " AnTiGrAvItY ", + "gemini:generate_content", + "", + true, + None, + None, + ) + .expect("Antigravity conversion"); + assert_eq!( + converted["tools"][0]["functionDeclarations"][0]["parameters"], schema, + "{source}" + ); + let AntigravityRequestEnvelopeSupport::Supported(envelope) = + build_antigravity_safe_v1internal_request( + &sample_auth(), + "review-test", + model, + &converted, + AntigravityEnvelopeRequestType::Agent, + ) + else { + panic!("schema should fit budget"); + }; + let parameters = + &envelope["request"]["tools"][0]["functionDeclarations"][0]["parameters"]; + assert_eq!( + parameters["properties"]["mode"], + json!({"type": "string", "enum": ["fast"]}) + ); + assert_eq!( + parameters["properties"]["payloads"]["additionalProperties"], + json!({"type": "string"}) + ); + assert_eq!(parameters["properties"]["name"]["minLength"], 1); + assert_eq!(envelope["model"], model); + } + // The default public Gemini policy must remain unchanged. + let public = build_standard_request_body( + &original, + source, + "gemini-test", + "gemini", + "gemini:generate_content", + "", + true, + None, + None, + ) + .unwrap(); + let parameters = &public["tools"][0]["functionDeclarations"][0]["parameters"]; + if source == "gemini:generate_content" { + assert_eq!(parameters, &schema); + } else { + assert!(parameters["properties"]["mode"].get("const").is_none()); + assert_eq!(parameters["properties"]["name"]["minLength"], "1"); + } + } + } + + #[test] + fn antigravity_responses_conversion_keeps_scoped_previous_response_history() { + use aether_ai_formats::{ + api::record_converted_response_history, + formats::shared::standard_matrix::build_standard_request_body, + }; + let schema = json!({"type": "object", "properties": {"mode": {"const": "fast"}}}); + let response_id = "resp_antigravity_schema_history_test"; + let scope = "antigravity-schema-history"; + record_converted_response_history(&json!({ + "needs_conversion": true, "client_api_format": "openai:responses", + "provider_api_format": "openai:chat", "api_key_id": scope, + "original_request_body": {"model": "client", "input": "first"} + }), &json!({"id": response_id, "status": "completed", "output": [{ + "type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "remembered"}] + }]})).expect("seed scoped history"); + let input = json!({"model": "client", "previous_response_id": response_id, + "input": "second", "tools": [{"type": "function", "name": "probe", "parameters": schema}]}); + let output = build_standard_request_body( + &input, + "openai:responses", + "claude-test", + "antigravity", + "gemini:generate_content", + "", + true, + None, + Some(scope), + ) + .expect("expand history"); + assert_eq!(output["contents"][0]["parts"][0]["text"], "first"); + assert_eq!(output["contents"][1]["parts"][0]["text"], "remembered"); + assert_eq!(output["contents"][2]["parts"][0]["text"], "second"); + assert_eq!( + output["tools"][0]["functionDeclarations"][0]["parameters"], + schema + ); + assert!(build_standard_request_body( + &input, + "openai:responses", + "claude-test", + "antigravity", + "gemini:generate_content", + "", + true, + None, + Some("different-key") + ) + .is_none()); + } + + #[test] + fn antigravity_rejects_shared_schema_budget_exhaustion_on_both_envelope_paths() { + use super::AntigravityRequestEnvelopeUnsupportedReason; + let declaration = json!({"name": "probe", "parameters": { + "type": "object", "description": "x".repeat(600_000) + }}); + let mut body = json!({"contents": [], "tools": [{"functionDeclarations": [declaration]}]}); + assert!(matches!( + build_antigravity_safe_v1internal_request( + &sample_auth(), + "test", + "claude-test", + &body, + AntigravityEnvelopeRequestType::Agent, + ), + AntigravityRequestEnvelopeSupport::Supported(_) + )); + body["tools"][0]["functionDeclarations"] + .as_array_mut() + .unwrap() + .push(declaration); + for wrapped in [false, true] { + let input = if wrapped { + json!({"request": body}) + } else { + body.clone() + }; + let snapshot = input.clone(); + assert_eq!( + build_antigravity_safe_v1internal_request( + &sample_auth(), + "test", + "claude-test", + &input, + AntigravityEnvelopeRequestType::Agent, + ), + AntigravityRequestEnvelopeSupport::Unsupported( + AntigravityRequestEnvelopeUnsupportedReason::ToolSchemaBudgetExceeded + ) + ); + assert_eq!(input, snapshot); + } + } + #[test] fn search_only_request_keeps_the_modern_google_search_spelling() { // Reproduces the live failure: a grounding-only request (no function @@ -516,6 +1042,162 @@ mod tests { ); } + #[test] + fn antigravity_envelope_downgrades_fabric_schema_on_all_input_paths() { + let schema = json!({ + "$schema": "http://json-schema.org/draft-07/schema#", + "type": "object", + "properties": { + "code": { "type": "string", "description": "TypeScript function body" }, + "payloads": { + "type": "object", + "patternProperties": { "^.*$": { "type": "string" } } + }, + "resultFormat": { + "anyOf": [ + { "type": "string", "const": "auto" }, + { "type": "string", "const": "yaml" }, + { "type": "string", "const": "json" }, + { "type": "string", "const": "text" } + ] + }, + "display": { + "anyOf": [ + { "type": "object", "properties": { "name": { "type": "string" } } }, + { "type": "string" } + ] + }, + "tokenBudget": { "type": "number", "minimum": 1 } + }, + "required": ["code"], + "minProperties": "1" + }); + let mut expected = schema.clone(); + expected.as_object_mut().unwrap().remove("$schema"); + expected["minProperties"] = json!(1); + expected["properties"]["payloads"] = json!({ + "type": "object", "additionalProperties": { "type": "string" } + }); + for (index, format) in ["auto", "yaml", "json", "text"].iter().enumerate() { + expected["properties"]["resultFormat"]["anyOf"][index] = + json!({ "type": "string", "enum": [format] }); + } + + for declarations_key in ["functionDeclarations", "function_declarations"] { + for parameters_key in [ + "parametersJsonSchema", + "parameters_json_schema", + "parameters", + ] { + for wrapped in [false, true] { + let mut body = json!({ + "contents": [{ "role": "user", "parts": [{ "text": "hello" }] }], + "tools": [{ "googleSearch": {} }, {}], + "generationConfig": { "maxOutputTokens": 4096 }, + "labels": { "const": "not a schema" } + }); + body["tools"][1][declarations_key] = json!([ + { "name": "fabric_exec", "description": "Execute TypeScript" }, + { "name": "other", "parameters": { "type": "object" } } + ]); + body["tools"][1][declarations_key][0][parameters_key] = schema.clone(); + if wrapped { + body = json!({ "request": body, "requestId": "client-id" }); + } + let original = body.clone(); + let AntigravityRequestEnvelopeSupport::Supported(envelope) = + build_antigravity_safe_v1internal_request( + &sample_auth(), + "trace-id", + "gemini-test", + &body, + AntigravityEnvelopeRequestType::Agent, + ) + else { + panic!("Fabric request should be supported"); + }; + let declaration = &envelope["request"]["tools"][1][declarations_key][0]; + assert_eq!( + declaration["parameters"], expected, + "{declarations_key}/{parameters_key}/wrapped={wrapped}" + ); + assert_eq!(declaration["name"], "fabric_exec"); + assert!(declaration.get("parametersJsonSchema").is_none()); + assert!(declaration.get("parameters_json_schema").is_none()); + assert_eq!( + envelope["request"]["generationConfig"]["maxOutputTokens"], + 4096 + ); + assert_eq!(envelope["request"]["labels"]["const"], "not a schema"); + assert_eq!( + envelope["request"]["tools"][0], + json!({ "googleSearch": {} }) + ); + assert_eq!( + envelope["requestId"], + if wrapped { "client-id" } else { "trace-id" } + ); + assert_eq!(envelope["project"], "project-ant-123"); + assert_eq!(body, original, "caller-owned input must remain unchanged"); + + let AntigravityRequestEnvelopeSupport::Supported(rebuilt) = + build_antigravity_safe_v1internal_request( + &sample_auth(), + "trace-id", + "gemini-test", + &envelope, + AntigravityEnvelopeRequestType::Agent, + ) + else { + panic!("converted envelope should be supported"); + }; + assert_eq!(rebuilt, envelope, "normalization must be idempotent"); + } + } + } + } + + #[test] + fn antigravity_schema_aliases_do_not_override_existing_parameters() { + let body = json!({ + "contents": [], + "tools": [{ "functionDeclarations": [{ + "name": "select", + "parameters": { "type": "string", "const": "existing" }, + "parametersJsonSchema": { "type": "string", "const": "camel" }, + "parameters_json_schema": { "type": "string", "const": "snake" } + }, { + "name": "select_alias", + "parametersJsonSchema": { "type": "string", "const": "camel" }, + "parameters_json_schema": { "type": "string", "const": "snake" } + }] }] + }); + let AntigravityRequestEnvelopeSupport::Supported(envelope) = + build_antigravity_safe_v1internal_request( + &sample_auth(), + "trace-id", + "gemini-test", + &body, + AntigravityEnvelopeRequestType::Agent, + ) + else { + panic!("request should be supported"); + }; + let declarations = &envelope["request"]["tools"][0]["functionDeclarations"]; + assert_eq!( + declarations[0]["parameters"], + json!({ "type": "string", "enum": ["existing"] }) + ); + assert_eq!( + declarations[1]["parameters"], + json!({ "type": "string", "enum": ["camel"] }) + ); + for declaration in declarations.as_array().unwrap() { + assert!(declaration.get("parametersJsonSchema").is_none()); + assert!(declaration.get("parameters_json_schema").is_none()); + } + } + #[test] fn antigravity_envelope_normalizes_json_schema_parameter_spellings() { let request_body = json!({ diff --git a/crates/aether-provider/transport/src/antigravity/schema.rs b/crates/aether-provider/transport/src/antigravity/schema.rs new file mode 100644 index 000000000..fee8790f3 --- /dev/null +++ b/crates/aether-provider/transport/src/antigravity/schema.rs @@ -0,0 +1,789 @@ +use std::collections::BTreeSet; + +use serde_json::{json, Map, Value}; + +/// Cloud Code's `parameters` uses a protobuf-shaped schema, while its Claude +/// backend validates the translated `tools.*.custom.input_schema` as JSON Schema +/// draft 2020-12. Keep this lowering at the Antigravity boundary: it preserves the +/// common subset and converts protobuf JSON int64 strings back to JSON numbers. +/// Callers must still validate tool arguments against their original schema. +pub(super) fn normalize_tool_parameters( + parameters: &mut Value, + budget: &mut SchemaBudget, +) -> Result<(), ()> { + let lowered = lower_schema(parameters, parameters, &mut BTreeSet::new(), 0, budget); + if budget.exhausted { + return Err(()); + } + *parameters = lowered; + Ok(()) +} + +/// Cloud Code accepts these unions but its Claude bridge rejects typed anyOf +/// branches (verified with the real fabric_exec schema). Do not apply this to +/// Gemini models. Fold string literal alternatives exactly; otherwise relax the +/// union and retain its constraints as guidance, never choose an arbitrary branch. +/// The caller must validate generated arguments against the original schema. +pub(super) fn normalize_claude_unions( + schema: &mut Value, + budget: &mut SchemaBudget, +) -> Result<(), ()> { + lower_claude_unions(schema, budget) +} + +fn lower_claude_unions(value: &mut Value, budget: &mut SchemaBudget) -> Result<(), ()> { + let Some(schema) = value.as_object_mut() else { + return Ok(()); + }; + if let Some(Value::Array(branches)) = schema.remove("anyOf") { + // Only intersect an existing sibling enum when both sides are strings; + // an empty intersection cannot be represented by protobuf enum (omitted). + let literals = string_union_literals(&branches); + let folded = literals.and_then(|mut literals| { + if let Some(Value::Array(existing)) = schema.get("enum") { + literals.retain(|literal| existing.contains(&Value::String(literal.clone()))); + } + (!literals.is_empty()).then_some(literals) + }); + if let Some(literals) = folded { + schema.entry("type").or_insert(json!("string")); + schema.insert("enum".into(), json!(literals)); + } else { + let guidance = serde_json::to_string(&branches).expect("JSON value serializes"); + let description = schema.entry("description").or_insert(json!("")); + let original = description.as_str().unwrap_or_default(); + *description = json!(format!("{original}\nAccepted alternatives (validate against the original tool schema): {guidance}").trim()); + if !budget.charge(description) { + return Err(()); + } + } + } + // Traverse schema positions only: descriptions/defaults/examples and property + // names such as `anyOf` are data, not keywords to rewrite. + if let Some(Value::Object(properties)) = schema.get_mut("properties") { + for child in properties.values_mut() { + lower_claude_unions(child, budget)?; + } + } + for key in ["items", "additionalProperties"] { + if let Some(child) = schema.get_mut(key) { + lower_claude_unions(child, budget)?; + } + } + Ok(()) +} + +fn string_union_literals(branches: &[Value]) -> Option> { + if branches.is_empty() { + return None; + } + let mut literals = BTreeSet::new(); + for branch in branches { + let branch = branch.as_object()?; + if branch + .keys() + .any(|key| !matches!(key.as_str(), "type" | "enum" | "description" | "title")) + || branch.get("type").is_some_and(|ty| ty != "string") + { + return None; + } + let values = branch.get("enum")?.as_array()?; + if values.is_empty() { + return None; + } + for literal in values { + literals.insert(literal.as_str()?.to_owned()); + } + } + Some(literals) +} + +/// Shared across all tool schemas in one request. Counting serialized input at +/// every expansion conservatively bounds cloning work, including literal data +/// and repeated acyclic references, without allocating serialized copies. +pub(super) struct SchemaBudget { + nodes_left: usize, + bytes_left: usize, + exhausted: bool, +} + +impl Default for SchemaBudget { + fn default() -> Self { + Self { + nodes_left: 4096, + bytes_left: 1024 * 1024, + exhausted: false, + } + } +} + +impl SchemaBudget { + fn charge(&mut self, value: &Value) -> bool { + if self.exhausted || self.nodes_left == 0 { + self.exhausted = true; + return false; + } + self.nodes_left -= 1; + if serde_json::to_writer(&mut *self, value).is_err() { + self.exhausted = true; + return false; + } + true + } +} + +impl std::io::Write for SchemaBudget { + fn write(&mut self, bytes: &[u8]) -> std::io::Result { + if bytes.len() > self.bytes_left { + return Err(std::io::Error::other( + "tool schema expansion budget exceeded", + )); + } + self.bytes_left -= bytes.len(); + Ok(bytes.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } +} + +/// JSON Schema $ref siblings are conjunctive, not an object-spread override. +/// Merge properties and required sets; retain the referenced constraint when +/// two validation keywords cannot be intersected in the wire subset. This +/// relaxes validation rather than manufacturing a contradictory schema. +fn merge_ref_siblings(mut target: Value, siblings: Value) -> Value { + let target_object = target.as_object_mut().expect("lowered schema is an object"); + let Value::Object(siblings) = siblings else { + unreachable!("lowered schema is an object") + }; + for (key, value) in siblings { + match (key.as_str(), target_object.get_mut(&key), value) { + ("properties", Some(Value::Object(properties)), Value::Object(children)) => { + for (name, child) in children { + if let Some(existing) = properties.get_mut(&name) { + *existing = merge_ref_siblings(existing.take(), child); + } else { + properties.insert(name, child); + } + } + } + ("required", Some(Value::Array(required)), Value::Array(names)) => { + let mut seen: BTreeSet = required + .iter() + .filter_map(Value::as_str) + .map(str::to_owned) + .collect(); + for name in names { + if seen.insert(name.as_str().expect("required contains strings").to_owned()) { + required.push(name); + } + } + } + ("title" | "description" | "default" | "example", _, value) | (_, None, value) => { + target_object.insert(key, value); + } + _ => {} + } + } + target +} + +fn lower_schema( + value: &Value, + root: &Value, + resolving: &mut BTreeSet, + depth: usize, + budget: &mut SchemaBudget, +) -> Value { + // Bound both recursive references and expansion of deeply nested schemas. + if !budget.charge(value) || depth >= 64 { + return json!({}); + } + let Some(source) = value.as_object() else { + // Boolean schemas have no protobuf equivalent (false is relaxed). + return json!({}); + }; + let mut source = source.clone(); + if let Some(Value::String(reference)) = source.remove("$ref") { + if let Some(target) = reference + .strip_prefix('#') + .and_then(|pointer| root.pointer(pointer)) + .filter(|target| target.is_object()) + { + if resolving.insert(reference.clone()) { + let target = lower_schema(target, root, resolving, depth + 1, budget); + resolving.remove(&reference); + let siblings = + lower_schema(&Value::Object(source), root, resolving, depth + 1, budget); + return merge_ref_siblings(target, siblings); + } + } + // Unresolved, external or cyclic refs retain only their sibling fields. + } + + let mut schema = Map::new(); + // Keep only a typed JSON-Schema subset. Gemini's protobuf JSON mapping renders + // int64 constraints as strings, but Claude rejects those at custom.input_schema. + match source.get("type") { + Some(Value::String(schema_type)) => { + if let Some(schema_type) = json_schema_type(schema_type) { + schema.insert("type".to_string(), Value::String(schema_type.to_string())); + } + } + Some(Value::Array(types)) => { + schema.insert("type".to_string(), Value::Array(types.clone())); + } + _ => {} + } + for key in ["format", "title", "description", "pattern"] { + if let Some(value) = source.get(key).filter(|value| value.is_string()) { + schema.insert(key.to_string(), value.clone()); + } + } + if let Some(value) = source.get("nullable").filter(|value| value.is_boolean()) { + schema.insert("nullable".to_string(), value.clone()); + } + if let Some(values) = json_schema_string_array(source.get("enum"), true) { + schema.insert("enum".to_string(), values); + } + for key in ["minimum", "maximum"] { + if let Some(value) = source.get(key).filter(|value| value.is_number()) { + schema.insert(key.to_string(), value.clone()); + } + } + for key in [ + "minItems", + "maxItems", + "minLength", + "maxLength", + "minProperties", + "maxProperties", + ] { + if let Some(value) = source.get(key).and_then(json_schema_nonnegative_integer) { + schema.insert(key.to_string(), value); + } + } + if let Some(values) = json_schema_string_array(source.get("required"), false) { + schema.insert("required".to_string(), values); + } + if let Some(values) = json_schema_string_array(source.get("propertyOrdering"), false) { + schema.insert("propertyOrdering".to_string(), values); + } + for key in ["default", "example"] { + if let Some(value) = source.get(key) { + schema.insert(key.to_string(), value.clone()); + } + } + + if let Some(constant) = source.get("const") { + if constant.is_string() { + schema.insert("enum".to_string(), json!([constant])); + schema.entry("type").or_insert(json!("string")); + } else { + // The protobuf enum field is repeated string. Never stringify numeric + // or boolean literals into enums: that changes the argument's type. + let description = schema.entry("description").or_insert(json!("")); + let prefix = description.as_str().unwrap_or_default(); + *description = json!(format!("{prefix}\nMust equal: {constant}").trim()); + } + } + if let Some(values) = schema.get_mut("enum").and_then(Value::as_array_mut) { + // Mixed/non-string enums cannot be represented without changing types. + if !values.iter().all(Value::is_string) { + schema.remove("enum"); + } + } + + if let Some(properties) = source.get("properties").and_then(Value::as_object) { + schema.insert( + "properties".to_string(), + Value::Object( + properties + .iter() + .map(|(name, child)| { + ( + name.clone(), + lower_schema(child, root, resolving, depth + 1, budget), + ) + }) + .collect(), + ), + ); + } + if let Some(items) = source.get("items") { + schema.insert( + "items".to_string(), + lower_schema(items, root, resolving, depth + 1, budget), + ); + } + // oneOf's exclusivity is not supported; anyOf retains the alternatives. + if let Some(branches) = source + .get("anyOf") + .or_else(|| source.get("any_of")) + .or_else(|| source.get("oneOf")) + .and_then(Value::as_array) + { + if !branches.is_empty() { + schema.insert( + "anyOf".to_string(), + Value::Array( + branches + .iter() + .map(|child| lower_schema(child, root, resolving, depth + 1, budget)) + .collect(), + ), + ); + } + } + + let patterns = source.get("patternProperties").and_then(Value::as_object); + let wildcard = patterns + .filter(|patterns| patterns.len() == 1) + .and_then(|patterns| { + patterns + .iter() + .next() + .filter(|(pattern, _)| matches!(pattern.as_str(), ".*" | "^.*$")) + }) + .map(|(_, child)| child); + // For a catch-all, JSON Schema's additionalProperties applies to no unmatched + // keys, so even an explicit `false` must not suppress the dictionary values. + let additional = wildcard.or_else(|| { + // Without regex matching, an additional-properties constraint could + // incorrectly reject keys formerly accepted by a pattern. Relax it too. + patterns + .is_none_or(Map::is_empty) + .then(|| source.get("additionalProperties")) + .flatten() + }); + if let Some(additional) = additional { + if additional.is_boolean() { + // Dropping non-wildcard patterns must not turn their allowed keys into + // forbidden additional properties. General regex maps are relaxed. + if additional != &Value::Bool(false) || patterns.is_none_or(Map::is_empty) { + schema.insert("additionalProperties".to_string(), additional.clone()); + } + } else { + schema.insert( + "additionalProperties".to_string(), + lower_schema(additional, root, resolving, depth + 1, budget), + ); + } + } + + if let Some(Value::Array(types)) = schema.get("type").cloned() { + schema.remove("type"); + let nullable = types.iter().any(|value| value.as_str() == Some("null")); + let types: BTreeSet<_> = types + .iter() + .filter_map(Value::as_str) + .filter_map(json_schema_type) + .filter(|value| *value != "null") + .collect(); + if types.len() == 1 { + schema.insert("type".to_string(), json!(types.first().unwrap())); + } else if !types.is_empty() { + schema.entry("anyOf").or_insert_with(|| { + Value::Array(types.iter().map(|ty| json!({ "type": ty })).collect()) + }); + } else if nullable { + schema.insert("type".to_string(), json!("null")); + } + if nullable && !types.is_empty() { + schema.insert("nullable".to_string(), Value::Bool(true)); + } + } + Value::Object(schema) +} + +fn json_schema_type(value: &str) -> Option<&'static str> { + match value.trim().to_ascii_lowercase().as_str() { + "array" => Some("array"), + "boolean" => Some("boolean"), + "integer" => Some("integer"), + "null" => Some("null"), + "number" => Some("number"), + "object" => Some("object"), + "string" => Some("string"), + _ => None, + } +} + +fn json_schema_nonnegative_integer(value: &Value) -> Option { + let value = match value { + Value::Number(value) => value.as_u64(), + // protobuf JSON encodes int64 fields as decimal strings. + Value::String(value) => value.parse::().ok(), + _ => None, + }?; + Some(Value::from(value)) +} + +fn json_schema_string_array(value: Option<&Value>, require_non_empty: bool) -> Option { + let values = value?.as_array()?; + let mut seen = BTreeSet::new(); + let values = values + .iter() + .map(Value::as_str) + .collect::>>()? + .into_iter() + .filter(|value| seen.insert(*value)) + .map(|value| Value::String(value.to_string())) + .collect::>(); + (!require_non_empty || !values.is_empty()).then_some(Value::Array(values)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn claude_unions_keep_siblings_literals_and_nested_schema_positions() { + let union = json!({"anyOf":[{"type":"string"},{"type":"object"}]}); + let mut schema = json!({"type":"object", "default":union, "properties":{ + "anyOf":{"type":"array","items":union}, + "map":{"type":"object","additionalProperties":union}, + "literal":{"enum":["b"],"anyOf":[{"type":"string","enum":["a"]},{"type":"string","enum":["b"]}]}, + "constrained":{"description":"Original","type":"string","minLength":2,"anyOf":[{"type":"string","pattern":"a+"},{"type":"number"}]} + }}); + normalize_claude_unions(&mut schema, &mut SchemaBudget::default()).unwrap(); + assert_eq!( + schema["default"], union, + "literal data must not be rewritten" + ); + assert!(schema["properties"]["anyOf"]["items"]["description"].is_string()); + assert!(schema["properties"]["map"]["additionalProperties"]["description"].is_string()); + assert_eq!( + schema["properties"]["literal"], + json!({"type":"string","enum":["b"]}) + ); + let constrained = &schema["properties"]["constrained"]; + assert_eq!(constrained["minLength"], 2); + assert_eq!(constrained["type"], "string"); + assert!(constrained["description"] + .as_str() + .unwrap() + .starts_with("Original")); + assert!(constrained["description"] + .as_str() + .unwrap() + .contains("pattern")); + let once = schema.clone(); + normalize_claude_unions(&mut schema, &mut SchemaBudget::default()).unwrap(); + assert_eq!(schema, once); + } + + #[test] + fn claude_union_guidance_is_charged_to_shared_budget() { + let mut schema = json!({"anyOf":[{"type":"string"},{"type":"number"}]}); + let mut budget = SchemaBudget { + nodes_left: 4096, + bytes_left: 8, + exhausted: false, + }; + assert!(normalize_claude_unions(&mut schema, &mut budget).is_err()); + assert!(budget.exhausted); + } + + #[test] + fn claude_union_folding_does_not_drop_branch_constraints() { + let mut schema = json!({"anyOf":[ + {"type":"string","enum":["a"],"minLength":2}, + {"type":"string","enum":["b"]} + ]}); + normalize_claude_unions(&mut schema, &mut SchemaBudget::default()).unwrap(); + assert!(schema.get("enum").is_none()); + assert!(schema["description"] + .as_str() + .unwrap() + .contains("minLength")); + } + + fn lowered(mut schema: Value) -> Value { + normalize_tool_parameters(&mut schema, &mut SchemaBudget::default()).unwrap(); + let once = schema.clone(); + normalize_tool_parameters(&mut schema, &mut SchemaBudget::default()).unwrap(); + assert_eq!(schema, once, "lowering must be idempotent"); + schema + } + + #[test] + fn reference_siblings_preserve_properties_and_required_without_contradictions() { + let result = lowered(json!({ + "$defs": {"Base": { + "type": "object", "properties": {"a": {"type": "string", "minLength": 1}}, + "required": ["a"], "additionalProperties": false + }}, + "$ref": "#/$defs/Base", + "properties": {"a": {"maxLength": 8}, "b": {"type": "string"}}, + "required": ["a"] + })); + assert_eq!( + result, + json!({ + "type": "object", "properties": { + "a": {"type": "string", "minLength": 1, "maxLength": 8}, + "b": {"type": "string"} + }, "required": ["a"], "additionalProperties": false + }) + ); + // A property schema named "additionalProperties" is data, not a keyword. + let merged = merge_ref_siblings( + json!({"properties": {"additionalProperties": {"type": "string"}}, "required": ["a"]}), + json!({"properties": {"b": {"type": "number"}}, "required": ["b", "a"]}), + ); + assert_eq!(merged["required"], json!(["a", "b"])); + assert_eq!( + merged["properties"]["additionalProperties"]["type"], + "string" + ); + } + + #[test] + fn sibling_reference_to_completed_target_is_not_a_cycle() { + let result = lowered(json!({ + "$defs": {"Base": {"type": "object", "properties": {"mode": {"const": "fast"}}}}, + "$ref": "#/$defs/Base", "properties": {"nested": {"$ref": "#/$defs/Base"}} + })); + assert_eq!( + result["properties"]["nested"]["properties"]["mode"], + json!({"type": "string", "enum": ["fast"]}) + ); + } + + #[test] + fn acyclic_branching_references_exhaust_budget_without_mutating_input() { + let mut schema = json!({"$defs": {"D0": {"type": "string"}}, "$ref": "#/$defs/D24"}); + for index in 1..=24 { + let reference = format!("#/$defs/D{}", index - 1); + schema["$defs"][format!("D{index}")] = json!({"type": "object", "properties": { + "left": {"$ref": reference}, "right": {"$ref": reference} + }}); + } + let original = schema.clone(); + let mut budget = SchemaBudget::default(); + assert!(normalize_tool_parameters(&mut schema, &mut budget).is_err()); + assert!(budget.exhausted); + assert_eq!(schema, original); + } + + #[test] + fn limits_nodes_and_literal_bytes_not_only_reference_depth() { + let mut wide = json!({"type": "object", "properties": {}}); + for index in 0..5000 { + wide["properties"][format!("p{index}")] = json!({}); + } + let mut budget = SchemaBudget::default(); + assert!(normalize_tool_parameters(&mut wide, &mut budget).is_err()); + assert_eq!(budget.nodes_left, 0); + let mut literal = json!({"type": "object", "default": "x".repeat(1024 * 1024)}); + assert!(normalize_tool_parameters(&mut literal, &mut SchemaBudget::default()).is_err()); + } + + #[test] + fn recursively_lowers_schema_nodes_without_touching_property_names_or_literal_data() { + let literal = json!({ "const": "data", "patternProperties": { "^.*$": 1 } }); + let result = lowered(json!({ + "$schema": "draft", "$id": "id", "x-custom": true, + "type": "object", "additionalProperties": false, + "properties": { + "const": { "const": "value", "readOnly": true }, + "patternProperties": { + "type": "array", "uniqueItems": true, + "items": { "oneOf": [{ "const": "a" }, { "const": "b" }] } + }, + "dictionary": { + "type": "object", + "additionalProperties": { + "any_of": [{ "const": "nested" }], "deprecated": true + } + } + }, + "default": literal, "example": literal, + "required": ["const"], "propertyOrdering": ["const", "patternProperties"] + })); + assert_eq!( + result, + json!({ + "type": "object", "additionalProperties": false, + "properties": { + "const": { "type": "string", "enum": ["value"] }, + "patternProperties": { + "type": "array", + "items": { "anyOf": [ + { "type": "string", "enum": ["a"] }, + { "type": "string", "enum": ["b"] } + ] } + }, + "dictionary": { + "type": "object", "additionalProperties": { + "anyOf": [{ "type": "string", "enum": ["nested"] }] + } + } + }, + "default": literal, "example": literal, + "required": ["const"], "propertyOrdering": ["const", "patternProperties"] + }) + ); + } + + #[test] + fn catch_all_patterns_preserve_dictionary_values_even_with_additional_properties_false() { + for pattern in [".*", "^.*$"] { + for additional in [json!(false), json!(true), json!({ "type": "number" })] { + let mut schema = json!({ "type": "object", "additionalProperties": additional }); + schema["patternProperties"][pattern] = json!({ "const": "value" }); + assert_eq!( + lowered(schema), + json!({ + "type": "object", "additionalProperties": { "type": "string", "enum": ["value"] } + }) + ); + } + } + } + + #[test] + fn general_patterns_are_not_mistaken_for_catch_all_dictionaries() { + for patterns in [ + json!({ "^x-": { "type": "string" } }), + json!({ "^.*$": { "type": "string" }, "^x-": { "maxLength": 5 } }), + ] { + for additional in [json!(false), json!({ "type": "number" })] { + assert_eq!( + lowered(json!({ + "type": "object", "patternProperties": patterns, + "additionalProperties": additional + })), + json!({ "type": "object" }) + ); + } + } + assert_eq!( + lowered(json!({ "type": "object", "additionalProperties": true })), + json!({ "type": "object", "additionalProperties": true }) + ); + } + + #[test] + fn resolves_local_references_with_siblings_and_terminates_cycles() { + let result = lowered(json!({ + "$defs": { + "Mode": { "const": "fast", "description": "original" }, + "Node": { "type": "object", "properties": { "next": { "$ref": "#/$defs/Node" } } } + }, + "type": "object", + "properties": { + "mode": { "$ref": "#/$defs/Mode", "description": "override" }, + "again": { "$ref": "#/$defs/Mode" }, + "node": { "$ref": "#/$defs/Node" }, + "missing": { "$ref": "#/$defs/Missing", "type": "string" }, + "external": { "$ref": "https://example.test/schema", "description": "external" } + } + })); + assert_eq!( + result["properties"]["mode"], + json!({ "type": "string", "enum": ["fast"], "description": "override" }) + ); + assert_eq!(result["properties"]["again"]["enum"], json!(["fast"])); + assert_eq!( + result["properties"]["node"]["properties"]["next"], + json!({}) + ); + assert_eq!(result["properties"]["missing"], json!({ "type": "string" })); + assert_eq!( + result["properties"]["external"], + json!({ "description": "external" }) + ); + assert!(result.get("$defs").is_none()); + } + + #[test] + fn lowers_nullable_type_unions_and_non_string_constants_without_invalid_enums() { + assert_eq!( + lowered(json!({ "type": ["string", "null", "string"] })), + json!({ "type": "string", "nullable": true }) + ); + assert_eq!( + lowered(json!({ "type": ["string", "number"] })), + json!({ "anyOf": [{ "type": "number" }, { "type": "string" }] }) + ); + for (ty, value) in [ + ("integer", json!(42)), + ("boolean", json!(true)), + ("null", Value::Null), + ] { + let result = lowered(json!({ "type": ty, "const": value })); + assert_eq!(result["type"], ty); + assert_eq!(result["description"], format!("Must equal: {value}")); + assert!(result.get("enum").is_none()); + assert!(result.get("const").is_none()); + } + assert_eq!( + lowered(json!({ "type": "integer", "enum": [1, 2] })), + json!({ "type": "integer" }) + ); + } + + #[test] + fn converts_protobuf_integer_strings_to_draft_2020_numbers() { + let result = lowered(json!({ + "type": ["OBJECT", "not-a-json-schema-type"], + "minProperties": "1", + "maxProperties": "invalid", + "required": ["args", "args"], + "enum": [], + "anyOf": [], + "properties": { + "args": { + "type": "ARRAY", + "minItems": "2", + "maxItems": "3", + "minLength": "4", + "maxLength": 8, + "minimum": "invalid", + "maximum": 1, + "items": { "type": "STRING", "minLength": "0" } + } + } + })); + assert_eq!( + result, + json!({ + "type": "object", + "minProperties": 1, + "required": ["args"], + "properties": { + "args": { + "type": "array", + "minItems": 2, + "maxItems": 3, + "minLength": 4, + "maxLength": 8, + "maximum": 1, + "items": { "type": "string", "minLength": 0 } + } + } + }) + ); + } + + #[test] + fn handles_boolean_and_deep_schemas_without_panicking() { + for input in [json!(true), json!(false), Value::Null] { + assert_eq!(lowered(input), json!({})); + } + let mut nested = json!({ "const": "deep" }); + for _ in 0..70 { + nested = json!({ "type": "array", "items": nested }); + } + let result = lowered(nested); + let mut child = &result; + for _ in 0..64 { + assert_eq!(child["type"], "array"); + child = &child["items"]; + } + assert_eq!(child, &json!({})); + } +} From 4124749a7d70e407062248ca186331dec1989cd9 Mon Sep 17 00:00:00 2001 From: wangpengxiang Date: Thu, 17 Sep 2026 11:28:42 +0800 Subject: [PATCH 2/2] fix(ai-serving): route provider-aware normalization through root seams Move Antigravity schema-preservation policy into the format crate and expose provider-aware Chat and Responses builders through the existing gateway root seam. Preserve legacy conversion and scoped Responses history behavior without weakening architecture tests. Validated: 61 standalone architecture tests, 947 format tests, 527 transport tests, 2 actual-source planner tests, and gateway cargo check. --- .../planner/standard/normalize/chat.rs | 28 +-- .../planner/standard/normalize/responses.rs | 31 +-- .../aether-gateway/src/ai_serving/pure/mod.rs | 2 + crates/aether-ai/formats/src/api.rs | 2 + .../src/formats/shared/standard_matrix.rs | 15 +- .../src/formats/shared/standard_normalize.rs | 211 ++++++++++++++++++ 6 files changed, 248 insertions(+), 41 deletions(-) diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/normalize/chat.rs b/apps/aether-gateway/src/ai_serving/planner/standard/normalize/chat.rs index e0d44d283..5c2fbfe31 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/normalize/chat.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/normalize/chat.rs @@ -4,7 +4,7 @@ use crate::ai_serving::transport::apply_standard_provider_request_body_rules_wit use crate::ai_serving::{ apply_codex_openai_responses_chat_body_edits, apply_openai_responses_compact_special_body_edits, - build_cross_format_openai_chat_request_body_with_model_directives as surface_build_cross_format_openai_chat_request_body, + build_cross_format_openai_chat_request_body_with_provider_context as surface_build_cross_format_openai_chat_request_body, build_local_openai_chat_request_body_with_model_directives as surface_build_local_openai_chat_request_body, GatewayProviderTransportSnapshot, }; @@ -70,23 +70,15 @@ pub(crate) fn build_cross_format_openai_chat_request_body( request_headers: &http::HeaderMap, enable_model_directives: bool, ) -> Option { - let provider_request_body = if provider_type.trim().eq_ignore_ascii_case("antigravity") - && aether_ai_formats::normalize_api_format_alias(provider_api_format) - == "gemini:generate_content" - { - aether_ai_formats::formats::shared::standard_matrix::build_standard_request_body_with_model_directives( - body_json, "openai:chat", mapped_model, provider_type, provider_api_format, - "", upstream_is_stream, None, user_api_key_id, enable_model_directives, - )? - } else { - surface_build_cross_format_openai_chat_request_body( - body_json, - mapped_model, - provider_api_format, - upstream_is_stream, - enable_model_directives, - )? - }; + let provider_request_body = surface_build_cross_format_openai_chat_request_body( + body_json, + mapped_model, + provider_type, + provider_api_format, + upstream_is_stream, + enable_model_directives, + user_api_key_id, + )?; let mut provider_request_body = apply_standard_provider_request_body_rules_with_request_headers( provider_request_body, diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/normalize/responses.rs b/apps/aether-gateway/src/ai_serving/planner/standard/normalize/responses.rs index 724ec4a94..5a138ac83 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/normalize/responses.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/normalize/responses.rs @@ -3,7 +3,7 @@ use serde_json::Value; use crate::ai_serving::transport::apply_standard_provider_request_body_rules_with_request_headers; use crate::ai_serving::{ apply_openai_responses_compact_special_body_edits, - build_cross_format_openai_responses_request_body_with_model_directives_and_history_scope as surface_build_cross_format_openai_responses_request_body, + build_cross_format_openai_responses_request_body_with_provider_context as surface_build_cross_format_openai_responses_request_body, build_local_openai_responses_request_body_with_model_directives as surface_build_local_openai_responses_request_body, GatewayProviderTransportSnapshot, }; @@ -214,25 +214,16 @@ pub(crate) fn build_cross_format_openai_responses_request_body_with_codex_model_ model_capabilities: Option<&crate::ai_serving::CodexResponsesModelCapabilities>, enable_model_directives: bool, ) -> Option { - let provider_request_body = if provider_type.trim().eq_ignore_ascii_case("antigravity") - && aether_ai_formats::normalize_api_format_alias(provider_api_format) - == "gemini:generate_content" - { - aether_ai_formats::formats::shared::standard_matrix::build_standard_request_body_with_model_directives( - body_json, client_api_format, mapped_model, provider_type, provider_api_format, - "", upstream_is_stream, None, history_scope, enable_model_directives, - )? - } else { - surface_build_cross_format_openai_responses_request_body( - body_json, - mapped_model, - client_api_format, - provider_api_format, - upstream_is_stream, - enable_model_directives, - history_scope, - )? - }; + let provider_request_body = surface_build_cross_format_openai_responses_request_body( + body_json, + mapped_model, + client_api_format, + provider_type, + provider_api_format, + upstream_is_stream, + enable_model_directives, + history_scope, + )?; let mut provider_request_body = apply_standard_provider_request_body_rules_with_request_headers( provider_request_body, diff --git a/apps/aether-gateway/src/ai_serving/pure/mod.rs b/apps/aether-gateway/src/ai_serving/pure/mod.rs index 42c45f186..41542a1f2 100644 --- a/apps/aether-gateway/src/ai_serving/pure/mod.rs +++ b/apps/aether-gateway/src/ai_serving/pure/mod.rs @@ -17,9 +17,11 @@ pub(crate) use aether_ai_formats::api::{ build_codex_model_catalog_metadata, build_codex_openai_image_api_provider_request_body, build_core_error_body_for_client_format, build_cross_format_openai_chat_request_body, build_cross_format_openai_chat_request_body_with_model_directives, + build_cross_format_openai_chat_request_body_with_provider_context, build_cross_format_openai_responses_request_body, build_cross_format_openai_responses_request_body_with_model_directives, build_cross_format_openai_responses_request_body_with_model_directives_and_history_scope, + build_cross_format_openai_responses_request_body_with_provider_context, build_gemini_image_request_body_from_openai_image_request, build_gemini_image_response_from_openai_image_response, build_gemini_image_response_from_openai_responses_image_response, build_generated_tool_call_id, diff --git a/crates/aether-ai/formats/src/api.rs b/crates/aether-ai/formats/src/api.rs index 1f0a310bd..fbbd51933 100644 --- a/crates/aether-ai/formats/src/api.rs +++ b/crates/aether-ai/formats/src/api.rs @@ -224,9 +224,11 @@ pub use crate::formats::{ standard_normalize::{ build_cross_format_openai_chat_request_body, build_cross_format_openai_chat_request_body_with_model_directives, + build_cross_format_openai_chat_request_body_with_provider_context, build_cross_format_openai_responses_request_body, build_cross_format_openai_responses_request_body_with_model_directives, build_cross_format_openai_responses_request_body_with_model_directives_and_history_scope, + build_cross_format_openai_responses_request_body_with_provider_context, build_local_openai_chat_request_body, build_local_openai_chat_request_body_with_model_directives, build_local_openai_responses_request_body, diff --git a/crates/aether-ai/formats/src/formats/shared/standard_matrix.rs b/crates/aether-ai/formats/src/formats/shared/standard_matrix.rs index 29d5fbc93..9bb764655 100644 --- a/crates/aether-ai/formats/src/formats/shared/standard_matrix.rs +++ b/crates/aether-ai/formats/src/formats/shared/standard_matrix.rs @@ -23,6 +23,17 @@ use crate::formats::shared::standard_normalize::{ is_claude_messages_shaped_body_on_openai_chat_endpoint, }; +/// Tool schema preservation is a format-conversion policy, shared by the +/// standard matrix and provider-aware Chat/Responses entry points. +pub(super) fn preserves_gemini_tool_schemas( + provider_type: &str, + provider_api_format: &str, +) -> bool { + provider_type.trim().eq_ignore_ascii_case("antigravity") + && aether_ai_formats::normalize_api_format_alias(provider_api_format) + == "gemini:generate_content" +} + #[allow(clippy::too_many_arguments)] pub fn build_standard_request_body( body_json: &Value, @@ -132,9 +143,7 @@ pub fn build_standard_request_body_with_model_directives_and_request_headers_and .with_request_path(request_path) .with_upstream_stream(upstream_is_stream); format_context.preserve_gemini_tool_schemas = - provider_type.trim().eq_ignore_ascii_case("antigravity") - && aether_ai_formats::normalize_api_format_alias(provider_api_format) - == "gemini:generate_content"; + preserves_gemini_tool_schemas(provider_type, provider_api_format); if let Some(history_scope) = user_api_key_id { format_context = format_context.with_history_scope(history_scope); } diff --git a/crates/aether-ai/formats/src/formats/shared/standard_normalize.rs b/crates/aether-ai/formats/src/formats/shared/standard_normalize.rs index c0e0062a6..45c344161 100644 --- a/crates/aether-ai/formats/src/formats/shared/standard_normalize.rs +++ b/crates/aether-ai/formats/src/formats/shared/standard_normalize.rs @@ -167,6 +167,40 @@ pub fn build_cross_format_openai_chat_request_body( ) } +/// Provider-aware entry point for gateway Chat planners. Keep private schema +/// conversion policy in the format crate while retaining legacy behavior elsewhere. +pub fn build_cross_format_openai_chat_request_body_with_provider_context( + body_json: &Value, + mapped_model: &str, + provider_type: &str, + provider_api_format: &str, + upstream_is_stream: bool, + enable_model_directives: bool, + history_scope: Option<&str>, +) -> Option { + if super::standard_matrix::preserves_gemini_tool_schemas(provider_type, provider_api_format) { + return super::standard_matrix::build_standard_request_body_with_model_directives( + body_json, + "openai:chat", + mapped_model, + provider_type, + provider_api_format, + "", + upstream_is_stream, + None, + history_scope, + enable_model_directives, + ); + } + build_cross_format_openai_chat_request_body_with_model_directives( + body_json, + mapped_model, + provider_api_format, + upstream_is_stream, + enable_model_directives, + ) +} + pub fn build_cross_format_openai_chat_request_body_with_model_directives( body_json: &Value, mapped_model: &str, @@ -342,6 +376,44 @@ pub fn build_cross_format_openai_responses_request_body_with_model_directives( ) } +/// Provider-aware Responses entry point; preserve history scoping and defer +/// private tool schema lowering without exposing provider policy to the gateway. +#[allow(clippy::too_many_arguments)] +pub fn build_cross_format_openai_responses_request_body_with_provider_context( + body_json: &Value, + mapped_model: &str, + client_api_format: &str, + provider_type: &str, + provider_api_format: &str, + upstream_is_stream: bool, + enable_model_directives: bool, + history_scope: Option<&str>, +) -> Option { + if super::standard_matrix::preserves_gemini_tool_schemas(provider_type, provider_api_format) { + return super::standard_matrix::build_standard_request_body_with_model_directives( + body_json, + client_api_format, + mapped_model, + provider_type, + provider_api_format, + "", + upstream_is_stream, + None, + history_scope, + enable_model_directives, + ); + } + build_cross_format_openai_responses_request_body_with_model_directives_and_history_scope( + body_json, + mapped_model, + client_api_format, + provider_api_format, + upstream_is_stream, + enable_model_directives, + history_scope, + ) +} + pub fn build_cross_format_openai_responses_request_body_with_model_directives_and_history_scope( body_json: &Value, mapped_model: &str, @@ -433,6 +505,145 @@ mod tests { }; use serde_json::{json, Value}; + #[test] + fn provider_context_builders_preserve_private_schemas_and_legacy_routes() { + use crate::api::{ + build_cross_format_openai_chat_request_body_with_provider_context as chat, + build_cross_format_openai_responses_request_body_with_provider_context as responses, + }; + let schema = json!({"type":"object", "properties":{"mode":{"const":"fast"}}}); + let chat_input = json!({"model":"client", "messages":[{"role":"user","content":"hi"}], + "tools":[{"type":"function","function":{"name":"probe","parameters":schema}}]}); + let responses_input = json!({"model":"client", "input":"hi", + "tools":[{"type":"function","name":"probe","parameters":schema}]}); + for provider in ["antigravity", " AnTiGrAvItY ", "gemini", "openai"] { + for target in [ + "gemini:generate_content", + "claude:messages", + "openai:responses", + ] { + for stream in [false, true] { + for directives in [false, true] { + for input in [&chat_input, &responses_input] { + let actual = chat( + input, + "claude-test", + provider, + target, + stream, + directives, + Some("seam-test"), + ); + let expected = + if super::super::standard_matrix::preserves_gemini_tool_schemas( + provider, target, + ) { + super::super::standard_matrix::build_standard_request_body_with_model_directives( + input, "openai:chat", "claude-test", provider, target, "", stream, None, Some("seam-test"), directives) + } else { + super::build_cross_format_openai_chat_request_body_with_model_directives( + input, "claude-test", target, stream, directives) + }; + assert!(actual.is_some(), "chat {provider} {target}"); + assert_eq!(actual, expected); + if target == "gemini:generate_content" { + assert_eq!( + actual.unwrap()["tools"][0]["functionDeclarations"][0] + ["parameters"] + == schema, + provider.trim().eq_ignore_ascii_case("antigravity") + ); + } + } + let actual = responses( + &responses_input, + "claude-test", + "openai:responses", + provider, + target, + stream, + directives, + Some("seam-test"), + ); + let expected = + if super::super::standard_matrix::preserves_gemini_tool_schemas( + provider, target, + ) { + super::super::standard_matrix::build_standard_request_body_with_model_directives( + &responses_input, "openai:responses", "claude-test", provider, target, "", stream, None, Some("seam-test"), directives) + } else { + super::build_cross_format_openai_responses_request_body_with_model_directives_and_history_scope( + &responses_input, "claude-test", "openai:responses", target, stream, directives, Some("seam-test")) + }; + // Same-format Responses uses the local builder, not this cross-format API. + assert_eq!( + actual.is_some(), + target != "openai:responses", + "responses {provider} {target}" + ); + assert_eq!(actual, expected); + } + } + } + } + } + + #[test] + fn provider_context_builders_keep_scoped_responses_history() { + use crate::api::{ + build_cross_format_openai_chat_request_body_with_provider_context as chat, + build_cross_format_openai_responses_request_body_with_provider_context as responses, + record_converted_response_history, + }; + let response_id = "resp_provider_context_seam_history"; + let scope = "provider-context-seam-history"; + record_converted_response_history(&json!({ + "needs_conversion":true, "client_api_format":"openai:responses", + "provider_api_format":"openai:chat", "api_key_id":scope, + "original_request_body":{"model":"client", "input":"first"} + }), &json!({"id":response_id, "status":"completed", "output":[{ + "type":"message", "role":"assistant", "content":[{"type":"output_text", "text":"remembered"}] + }]})).expect("seed scoped history"); + let input = json!({"model":"client", "previous_response_id":response_id, "input":"second"}); + for use_chat in [false, true] { + let build = |history_scope| { + if use_chat { + chat( + &input, + "claude-test", + "antigravity", + "gemini:generate_content", + true, + false, + history_scope, + ) + } else { + responses( + &input, + "claude-test", + "openai:responses", + "antigravity", + "gemini:generate_content", + true, + false, + history_scope, + ) + } + }; + if use_chat { + // The legacy Chat alternate-shape path does not hydrate scoped + // Responses history. Preserve that behavior during this refactor. + assert!(build(Some(scope)).is_none()); + continue; + } + let output = build(Some(scope)).expect("expand scoped history"); + assert_eq!(output["contents"][0]["parts"][0]["text"], "first"); + assert_eq!(output["contents"][1]["parts"][0]["text"], "remembered"); + assert_eq!(output["contents"][2]["parts"][0]["text"], "second"); + assert!(build(Some("different-seam-key")).is_none()); + } + } + fn object_keys(value: &Value) -> Vec<&str> { value .as_object()