Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

agent: Improve compatibility when using MCP servers with Gemini models (#28700)

Browse source 
WIP

Release Notes:

- agent: Improve compatibility when using MCPs with Gemini models
This commit is contained in:
Bennet Bo Fenner 2025-04-14 13:55:25 -06:00 • committed by GitHub
parent 6c93d107c2
commit 2603f36737
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
27 changed files with 307 additions and 103 deletions
Download patch file Download diff file Expand all files Collapse all files

6
crates/assistant_tool/src/assistant_tool.rs
View file

@ -1,5 +1,6 @@
mod action_log;
mod tool_registry;
mod tool_schema;
mod tool_working_set;
use std::fmt;
@ -16,6 +17,7 @@ use project::Project;
pub use crate::action_log::*;
pub use crate::tool_registry::*;
pub use crate::tool_schema::*;
pub use crate::tool_working_set::*;
pub fn init(cx: &mut App) {
@ -51,8 +53,8 @@ pub trait Tool: 'static + Send + Sync {
fn needs_confirmation(&self, input: &serde_json::Value, cx: &App) -> bool;
/// Returns the JSON schema that describes the tool's input.
fn input_schema(&self, _: LanguageModelToolSchemaFormat) -> serde_json::Value {
serde_json::Value::Object(serde_json::Map::default())
fn input_schema(&self, _: LanguageModelToolSchemaFormat) -> Result<serde_json::Value> {
Ok(serde_json::Value::Object(serde_json::Map::default()))
}
/// Returns markdown to be displayed in the UI for this tool.

236
crates/assistant_tool/src/tool_schema.rs Normal file
View file

@ -0,0 +1,236 @@
use anyhow::Result;
use serde_json::Value;
use crate::LanguageModelToolSchemaFormat;
/// Tries to adapt a JSON schema representation to be compatible with the specified format.
///
/// If the json cannot be made compatible with the specified format, an error is returned.
pub fn adapt_schema_to_format(
json: &mut Value,
format: LanguageModelToolSchemaFormat,
) -> Result<()> {
match format {
LanguageModelToolSchemaFormat::JsonSchema => Ok(()),
LanguageModelToolSchemaFormat::JsonSchemaSubset => adapt_to_json_schema_subset(json),
}
}
/// Tries to adapt the json schema so that it is compatible with https://ai.google.dev/api/caching#Schema
fn adapt_to_json_schema_subset(json: &mut Value) -> Result<()> {
if let Value::Object(obj) = json {
const UNSUPPORTED_KEYS: [&str; 4] = ["if", "then", "else", "$ref"];
for key in UNSUPPORTED_KEYS {
if obj.contains_key(key) {
return Err(anyhow::anyhow!(
"Schema cannot be made compatible because it contains \"{}\" ",
key
));
}
}
const KEYS_TO_REMOVE: [&str; 2] = ["format", "$schema"];
for key in KEYS_TO_REMOVE {
obj.remove(key);
}
if let Some(default) = obj.get("default") {
let is_null = default.is_null();
// Default is not supported, so we need to remove it
obj.remove("default");
if is_null {
obj.insert("nullable".to_string(), Value::Bool(true));
}
}
// If a type is not specified for an input parameter, add a default type
if obj.contains_key("description")
&& !obj.contains_key("type")
&& !(obj.contains_key("anyOf")
|| obj.contains_key("oneOf")
|| obj.contains_key("allOf"))
{
obj.insert("type".to_string(), Value::String("string".to_string()));
}
// Handle oneOf -> anyOf conversion
if let Some(subschemas) = obj.get_mut("oneOf") {
if subschemas.is_array() {
let subschemas_clone = subschemas.clone();
obj.remove("oneOf");
obj.insert("anyOf".to_string(), subschemas_clone);
}
}
// Recursively process all nested objects and arrays
for (_, value) in obj.iter_mut() {
if let Value::Object(_) | Value::Array(_) = value {
adapt_to_json_schema_subset(value)?;
}
}
} else if let Value::Array(arr) = json {
for item in arr.iter_mut() {
adapt_to_json_schema_subset(item)?;
}
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn test_transform_default_null_to_nullable() {
let mut json = json!({
"description": "A test field",
"type": "string",
"default": null
});
adapt_to_json_schema_subset(&mut json).unwrap();
assert_eq!(
json,
json!({
"description": "A test field",
"type": "string",
"nullable": true
})
);
}
#[test]
fn test_transform_adds_type_when_missing() {
let mut json = json!({
"description": "A test field without type"
});
adapt_to_json_schema_subset(&mut json).unwrap();
assert_eq!(
json,
json!({
"description": "A test field without type",
"type": "string"
})
);
}
#[test]
fn test_transform_removes_format() {
let mut json = json!({
"description": "A test field",
"type": "integer",
"format": "uint32"
});
adapt_to_json_schema_subset(&mut json).unwrap();
assert_eq!(
json,
json!({
"description": "A test field",
"type": "integer"
})
);
}
#[test]
fn test_transform_one_of_to_any_of() {
let mut json = json!({
"description": "A test field",
"oneOf": [
{ "type": "string" },
{ "type": "integer" }
]
});
adapt_to_json_schema_subset(&mut json).unwrap();
assert_eq!(
json,
json!({
"description": "A test field",
"anyOf": [
{ "type": "string" },
{ "type": "integer" }
]
})
);
}
#[test]
fn test_transform_nested_objects() {
let mut json = json!({
"type": "object",
"properties": {
"nested": {
"oneOf": [
{ "type": "string" },
{ "type": "null" }
],
"format": "email"
}
}
});
adapt_to_json_schema_subset(&mut json).unwrap();
assert_eq!(
json,
json!({
"type": "object",
"properties": {
"nested": {
"anyOf": [
{ "type": "string" },
{ "type": "null" }
]
}
}
})
);
}
#[test]
fn test_transform_fails_if_unsupported_keys_exist() {
let mut json = json!({
"type": "object",
"properties": {
"$ref": "#/definitions/User",
}
});
assert!(adapt_to_json_schema_subset(&mut json).is_err());
let mut json = json!({
"type": "object",
"properties": {
"if": "...",
}
});
assert!(adapt_to_json_schema_subset(&mut json).is_err());
let mut json = json!({
"type": "object",
"properties": {
"then": "...",
}
});
assert!(adapt_to_json_schema_subset(&mut json).is_err());
let mut json = json!({
"type": "object",
"properties": {
"else": "...",
}
});
assert!(adapt_to_json_schema_subset(&mut json).is_err());
}
}