From 0f48d95766f74bb8840bef08f8ecde027651087c Mon Sep 17 00:00:00 2001 From: rabi Date: Sat, 17 Jan 2026 21:38:12 +0530 Subject: [PATCH] fix(google): use parametersJsonSchema for full JSON Schema support Google's Schema spec for 'parameters' does not support $ref/$defs in tool definitions. Gemini 3 enforced stricter validation recently, causing silent failures. Switch to 'parametersJsonSchema' which supports full JSON Schema including $ref and $defs. This removes the need for process_map schema filtering. Removes technical debt (custom schema sanitizers), and fixes a functional regression with newer Gemini models. Signed-off-by: rabi --- .../goose/src/providers/formats/databricks.rs | 47 ++- crates/goose/src/providers/formats/google.rs | 363 +++--------------- 2 files changed, 81 insertions(+), 329 deletions(-) diff --git a/crates/goose/src/providers/formats/databricks.rs b/crates/goose/src/providers/formats/databricks.rs index e723cca2f0a..c43c77bf794 100644 --- a/crates/goose/src/providers/formats/databricks.rs +++ b/crates/goose/src/providers/formats/databricks.rs @@ -1,6 +1,5 @@ use crate::conversation::message::{Message, MessageContent}; use crate::model::ModelConfig; -use crate::providers::formats::google as gemini_schema; use crate::providers::utils::{ convert_image, detect_image_path, is_valid_function_name, load_image_file, safely_parse_json, sanitize_function_name, ImageFormat, @@ -232,19 +231,35 @@ pub fn format_tools(tools: &[Tool], model_name: &str) -> anyhow::Result Vec { let mut parameters = Map::new(); parameters.insert("name".to_string(), json!(tool.name)); parameters.insert("description".to_string(), json!(tool.description)); - let tool_input_schema = &tool.input_schema; - if tool_input_schema + // Use parametersJsonSchema which supports full JSON Schema including $ref/$defs + if tool + .input_schema .get("properties") .and_then(|v| v.as_object()) .is_some_and(|p| !p.is_empty()) { - parameters.insert( - "parameters".to_string(), - process_map(tool_input_schema, None), - ); + parameters.insert("parametersJsonSchema".to_string(), json!(tool.input_schema)); } json!(parameters) }) .collect() } -pub fn get_accepted_keys(parent_key: Option<&str>) -> Vec<&str> { - match parent_key { - Some("properties") => vec![ - "anyOf", - "allOf", - "type", - "description", - "nullable", - "enum", - "properties", - "required", - "items", - ], - Some("items") => vec!["type", "properties", "items", "required"], - _ => vec!["type", "properties", "required", "anyOf", "allOf"], - } -} - -pub fn process_value(value: &Value, parent_key: Option<&str>) -> Value { - match value { - Value::Object(map) => process_map(map, parent_key), - Value::Array(arr) if parent_key == Some("type") => arr - .iter() - .find(|v| v.as_str() != Some("null")) - .cloned() - .unwrap_or_else(|| json!("string")), - _ => value.clone(), - } -} - -/// Process a JSON map to filter out unsupported attributes, mirroring the logic -/// from the official Google Gemini CLI. -/// See: https://github.com/google-gemini/gemini-cli/blob/8a6509ffeba271a8e7ccb83066a9a31a5d72a647/packages/core/src/tools/tool-registry.ts#L356 -pub fn process_map(map: &Map, parent_key: Option<&str>) -> Value { - let accepted_keys = get_accepted_keys(parent_key); - - let filtered_map: Map = map - .iter() - .filter_map(|(key, value)| { - if !accepted_keys.contains(&key.as_str()) { - return None; - } - - let processed_value = match key.as_str() { - "properties" => { - if let Some(nested_map) = value.as_object() { - let processed_properties: Map = nested_map - .iter() - .map(|(prop_key, prop_value)| { - if let Some(prop_obj) = prop_value.as_object() { - (prop_key.clone(), process_map(prop_obj, Some("properties"))) - } else { - (prop_key.clone(), prop_value.clone()) - } - }) - .collect(); - Value::Object(processed_properties) - } else { - value.clone() - } - } - "items" => { - if let Some(items_map) = value.as_object() { - process_map(items_map, Some("items")) - } else { - value.clone() - } - } - "anyOf" | "allOf" => { - if let Some(arr) = value.as_array() { - let processed_arr: Vec = arr - .iter() - .map(|item| { - item.as_object().map_or_else( - || item.clone(), - |obj| process_map(obj, parent_key), - ) - }) - .collect(); - Value::Array(processed_arr) - } else { - value.clone() - } - } - _ => process_value(value, Some(key.as_str())), - }; - - Some((key.clone(), processed_value)) - }) - .collect(); - - Value::Object(filtered_map) -} - #[derive(Clone, Copy)] enum SignedTextHandling { SignedTextAsThinking, @@ -840,211 +744,23 @@ mod tests { #[test] fn test_tools_to_google_spec_with_valid_tools() { - let params1 = object!({ + let params = object!({ "properties": { "param1": { "type": "string", - "description": "A parameter", - "field_does_not_accept": ["value1", "value2"] - } - } - }); - let params2 = object!({ - "properties": { - "param2": { - "type": "string", - "description": "B parameter", + "description": "A parameter" } } }); - let params3 = object!({ - "properties": { - "body": { - "description": "Review comment text", - "type": "string" - }, - "comments": { - "description": "Line-specific comments array of objects to place comments on pull request changes. Requires path and body. For line comments use line or position. For multi-line comments use start_line and line with optional side parameters.", - "type": "array", - "items": { - "additionalProperties": false, - "properties": { - "body": { - "description": "comment body", - "type": "string" - }, - "line": { - "anyOf": [ - { "type": "number" }, - { "type": "null" } - ], - "description": "line number in the file to comment on. For multi-line comments, the end of the line range" - }, - "path": { - "description": "path to the file", - "type": "string" - }, - "position": { - "anyOf": [ - { "type": "number" }, - { "type": "null" } - ], - "description": "position of the comment in the diff" - }, - "side": { - "anyOf": [ - { "type": "string" }, - { "type": "null" } - ], - "description": "The side of the diff on which the line resides. For multi-line comments, this is the side for the end of the line range. (LEFT or RIGHT)" - }, - "start_line": { - "anyOf": [ - { "type": "number" }, - { "type": "null" } - ], - "description": "The first line of the range to which the comment refers. Required for multi-line comments." - }, - "start_side": { - "anyOf": [ - { "type": "string" }, - { "type": "null" } - ], - "description": "The side of the diff on which the start line resides for multi-line comments. (LEFT or RIGHT)" - } - }, - "required": ["path", "body", "position", "line", "side", "start_line", "start_side"], - "type": "object" - } - }, - "commitId": { - "description": "SHA of commit to review", - "type": "string" - }, - "event": { - "description": "Review action to perform", - "enum": ["APPROVE", "REQUEST_CHANGES", "COMMENT"], - "type": "string" - }, - "owner": { - "description": "Repository owner", - "type": "string" - }, - "pullNumber": { - "description": "Pull request number", - "type": "number" - } - } - }); - let tools = vec![ - Tool::new("tool1", "description1", params1), - Tool::new("tool2", "description2", params2), - Tool::new("tool3", "description3", params3), - ]; + let tools = vec![Tool::new("tool1", "description1", params.clone())]; let result = format_tools(&tools); - assert_eq!(result.len(), 3); + + assert_eq!(result.len(), 1); assert_eq!(result[0]["name"], "tool1"); assert_eq!(result[0]["description"], "description1"); - assert_eq!( - result[0]["parameters"]["properties"], - json!({"param1": json!({ - "type": "string", - "description": "A parameter" - })}) - ); - assert_eq!(result[1]["name"], "tool2"); - assert_eq!(result[1]["description"], "description2"); - assert_eq!( - result[1]["parameters"]["properties"], - json!({"param2": json!({ - "type": "string", - "description": "B parameter" - })}) - ); - - assert_eq!(result[2]["name"], "tool3"); - assert_eq!( - result[2]["parameters"]["properties"], - json!( - - { - "body": { - "description": "Review comment text", - "type": "string" - }, - "comments": { - "description": "Line-specific comments array of objects to place comments on pull request changes. Requires path and body. For line comments use line or position. For multi-line comments use start_line and line with optional side parameters.", - "type": "array", - "items": { - "properties": { - "body": { - "description": "comment body", - "type": "string" - }, - "line": { - "anyOf": [ - { "type": "number" }, - { "type": "null" } - ], - "description": "line number in the file to comment on. For multi-line comments, the end of the line range" - }, - "path": { - "description": "path to the file", - "type": "string" - }, - "position": { - "anyOf": [ - { "type": "number" }, - { "type": "null" } - ], - "description": "position of the comment in the diff" - }, - "side": { - "anyOf": [ - { "type": "string" }, - { "type": "null" } - ], - "description": "The side of the diff on which the line resides. For multi-line comments, this is the side for the end of the line range. (LEFT or RIGHT)" - }, - "start_line": { - "anyOf": [ - { "type": "number" }, - { "type": "null" } - ], - "description": "The first line of the range to which the comment refers. Required for multi-line comments." - }, - "start_side": { - "anyOf": [ - { "type": "string" }, - { "type": "null" } - ], - "description": "The side of the diff on which the start line resides for multi-line comments. (LEFT or RIGHT)" - } - }, - "required": ["path", "body", "position", "line", "side", "start_line", "start_side"], - "type": "object" - } - }, - "commitId": { - "description": "SHA of commit to review", - "type": "string" - }, - "event": { - "description": "Review action to perform", - "enum": ["APPROVE", "REQUEST_CHANGES", "COMMENT"], - "type": "string" - }, - "owner": { - "description": "Repository owner", - "type": "string" - }, - "pullNumber": { - "description": "Pull request number", - "type": "number" - } - } - ) - ); + assert!(result[0].get("parametersJsonSchema").is_some()); + assert!(result[0].get("parameters").is_none()); + assert_eq!(result[0]["parametersJsonSchema"], json!(params)); } #[test] @@ -1060,7 +776,7 @@ mod tests { assert_eq!(result.len(), 1); assert_eq!(result[0]["name"], "tool1"); assert_eq!(result[0]["description"], "description1"); - assert!(result[0]["parameters"].get("properties").is_none()); + assert!(result[0].get("parametersJsonSchema").is_none()); } #[test] @@ -1184,35 +900,22 @@ mod tests { } #[test] - fn test_tools_with_nullable_types_converted_to_single_type() { - // Test that type arrays like ["string", "null"] are converted to single types + fn test_tools_uses_parameters_json_schema() { let params = object!({ "properties": { - "nullable_field": { + "field": { "type": ["string", "null"], - "description": "A nullable string field" - }, - "regular_field": { - "type": "number", - "description": "A regular number field" + "description": "A field" } } }); - let tools = vec![Tool::new("test_tool", "test description", params)]; + let tools = vec![Tool::new("test_tool", "test description", params.clone())]; let result = format_tools(&tools); assert_eq!(result.len(), 1); assert_eq!(result[0]["name"], "test_tool"); - - // Verify that the type array was converted to a single string type - let nullable_field = &result[0]["parameters"]["properties"]["nullable_field"]; - assert_eq!(nullable_field["type"], "string"); - assert_eq!(nullable_field["description"], "A nullable string field"); - - // Verify that regular types are unchanged - let regular_field = &result[0]["parameters"]["properties"]["regular_field"]; - assert_eq!(regular_field["type"], "number"); - assert_eq!(regular_field["description"], "A regular number field"); + assert!(result[0].get("parametersJsonSchema").is_some()); + assert_eq!(result[0]["parametersJsonSchema"], json!(params)); } fn google_response(parts: Vec) -> Value { @@ -1509,4 +1212,32 @@ data: [DONE]"#; // Only "Complete" should be captured, stream should stop at [DONE] assert_eq!(text_parts, vec!["Complete"]); } + + #[test] + fn test_format_tools_uses_parameters_json_schema() { + let tool = Tool::new( + "test_tool", + "Test tool with $ref", + object!({ + "type": "object", + "$defs": { + "MyType": { "type": "string", "description": "A custom type" } + }, + "properties": { + "field": { "$ref": "#/$defs/MyType" } + } + }), + ); + + let result = format_tools(&[tool]); + + assert_eq!(result.len(), 1); + assert_eq!(result[0]["name"], "test_tool"); + assert!(result[0].get("parametersJsonSchema").is_some()); + assert!(result[0].get("parameters").is_none()); + + let schema = &result[0]["parametersJsonSchema"]; + assert_eq!(schema["properties"]["field"]["$ref"], "#/$defs/MyType"); + assert!(schema.get("$defs").is_some()); + } }