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!({})); + } +}