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

WIP

Release Notes:

- agent: Improve compatibility when using MCPs with Gemini models

crates/assistant_tools/src/assistant_tools.rs

@@ -84,7 +84,7 @@ mod tests {
     use super::*;
 
     #[gpui::test]
-    fn test_tool_schema_compatibility(cx: &mut App) {
+    fn test_builtin_tool_schema_compatibility(cx: &mut App) {
         crate::init(
             Arc::new(http_client::HttpClientWithUrl::new(
                 FakeHttpClient::with_200_response(),
@@ -95,18 +95,23 @@ mod tests {
         );
 
         for tool in ToolRegistry::global(cx).tools() {
-            let schema =
-                tool.input_schema(language_model::LanguageModelToolSchemaFormat::JsonSchemaSubset);
-            assert!(schema.is_object());
-            if schema.as_object().unwrap().contains_key("$schema") {
-                let error_message = format!(
-                    "Tool schema for `{}` is not compatible with `language_model::LanguageModelToolSchemaFormat::JsonSchemaSubset` (Gemini Models).\n\
-                    Are you using `schema::json_schema_for<T>(format)` to generate the schema?",
-                    tool.name()
-                );
+            let actual_schema = tool
+                .input_schema(language_model::LanguageModelToolSchemaFormat::JsonSchemaSubset)
+                .unwrap();
+            let mut expected_schema = actual_schema.clone();
+            assistant_tool::adapt_schema_to_format(
+                &mut expected_schema,
+                language_model::LanguageModelToolSchemaFormat::JsonSchemaSubset,
+            )
+            .unwrap();
 
-                panic!("{}", error_message)
-            }
+            let error_message = format!(
+                "Tool schema for `{}` is not compatible with `language_model::LanguageModelToolSchemaFormat::JsonSchemaSubset` (Gemini Models).\n\
+                Are you using `schema::json_schema_for<T>(format)` to generate the schema?",
+                tool.name(),
+            );
+
+            assert_eq!(actual_schema, expected_schema, "{}", error_message)
         }
     }
 }

crates/assistant_tools/src/batch_tool.rs

@@ -172,7 +172,7 @@ impl Tool for BatchTool {
         IconName::Cog
     }
 
-    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
+    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> Result<serde_json::Value> {
         json_schema_for::<BatchToolInput>(format)
     }
 

crates/assistant_tools/src/code_action_tool.rs

@@ -2,7 +2,7 @@ use anyhow::{Context as _, Result, anyhow};
 use assistant_tool::{ActionLog, Tool};
 use gpui::{App, Entity, Task};
 use language::{self, Anchor, Buffer, ToPointUtf16};
-use language_model::LanguageModelRequestMessage;
+use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
 use project::{self, LspAction, Project};
 use regex::Regex;
 use schemars::JsonSchema;
@@ -97,10 +97,7 @@ impl Tool for CodeActionTool {
         IconName::Wand
     }
 
-    fn input_schema(
-        &self,
-        format: language_model::LanguageModelToolSchemaFormat,
-    ) -> serde_json::Value {
+    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> Result<serde_json::Value> {
         json_schema_for::<CodeActionToolInput>(format)
     }
 

crates/assistant_tools/src/code_symbols_tool.rs

@@ -91,7 +91,7 @@ impl Tool for CodeSymbolsTool {
         IconName::Code
     }
 
-    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
+    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> Result<serde_json::Value> {
         json_schema_for::<CodeSymbolsInput>(format)
     }
 

crates/assistant_tools/src/copy_path_tool.rs

@@ -55,7 +55,7 @@ impl Tool for CopyPathTool {
         IconName::Clipboard
     }
 
-    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
+    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> Result<serde_json::Value> {
         json_schema_for::<CopyPathToolInput>(format)
     }
 

crates/assistant_tools/src/create_directory_tool.rs

@@ -45,7 +45,7 @@ impl Tool for CreateDirectoryTool {
         IconName::Folder
     }
 
-    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
+    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> Result<serde_json::Value> {
         json_schema_for::<CreateDirectoryToolInput>(format)
     }
 

crates/assistant_tools/src/create_file_tool.rs

@@ -52,7 +52,7 @@ impl Tool for CreateFileTool {
         IconName::FileCreate
     }
 
-    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
+    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> Result<serde_json::Value> {
         json_schema_for::<CreateFileToolInput>(format)
     }
 

crates/assistant_tools/src/delete_path_tool.rs

@@ -45,7 +45,7 @@ impl Tool for DeletePathTool {
         IconName::FileDelete
     }
 
-    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
+    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> Result<serde_json::Value> {
         json_schema_for::<DeletePathToolInput>(format)
     }
 

crates/assistant_tools/src/diagnostics_tool.rs

@@ -58,7 +58,7 @@ impl Tool for DiagnosticsTool {
         IconName::XCircle
     }
 
-    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
+    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> Result<serde_json::Value> {
         json_schema_for::<DiagnosticsToolInput>(format)
     }
 

crates/assistant_tools/src/fetch_tool.rs

@@ -128,7 +128,7 @@ impl Tool for FetchTool {
         IconName::Globe
     }
 
-    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
+    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> Result<serde_json::Value> {
         json_schema_for::<FetchToolInput>(format)
     }
 

crates/assistant_tools/src/find_replace_file_tool.rs

@@ -151,7 +151,7 @@ impl Tool for FindReplaceFileTool {
         IconName::Pencil
     }
 
-    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
+    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> Result<serde_json::Value> {
         json_schema_for::<FindReplaceFileToolInput>(format)
     }
 

crates/assistant_tools/src/list_directory_tool.rs

@@ -56,7 +56,7 @@ impl Tool for ListDirectoryTool {
         IconName::Folder
     }
 
-    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
+    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> Result<serde_json::Value> {
         json_schema_for::<ListDirectoryToolInput>(format)
     }
 

crates/assistant_tools/src/move_path_tool.rs

@@ -54,7 +54,7 @@ impl Tool for MovePathTool {
         IconName::ArrowRightLeft
     }
 
-    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
+    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> Result<serde_json::Value> {
         json_schema_for::<MovePathToolInput>(format)
     }
 

crates/assistant_tools/src/now_tool.rs

@@ -45,7 +45,7 @@ impl Tool for NowTool {
         IconName::Info
     }
 
-    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
+    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> Result<serde_json::Value> {
         json_schema_for::<NowToolInput>(format)
     }
 

crates/assistant_tools/src/open_tool.rs

@@ -35,7 +35,7 @@ impl Tool for OpenTool {
         IconName::ArrowUpRight
     }
 
-    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
+    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> Result<serde_json::Value> {
         json_schema_for::<OpenToolInput>(format)
     }
 

crates/assistant_tools/src/path_search_tool.rs

@@ -53,7 +53,7 @@ impl Tool for PathSearchTool {
         IconName::SearchCode
     }
 
-    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
+    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> Result<serde_json::Value> {
         json_schema_for::<PathSearchToolInput>(format)
     }
 

crates/assistant_tools/src/read_file_tool.rs

@@ -63,7 +63,7 @@ impl Tool for ReadFileTool {
         IconName::FileSearch
     }
 
-    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
+    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> Result<serde_json::Value> {
         json_schema_for::<ReadFileToolInput>(format)
     }
 

crates/assistant_tools/src/regex_search_tool.rs

@@ -60,7 +60,7 @@ impl Tool for RegexSearchTool {
         IconName::Regex
     }
 
-    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
+    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> Result<serde_json::Value> {
         json_schema_for::<RegexSearchToolInput>(format)
     }
 

crates/assistant_tools/src/rename_tool.rs

@@ -2,7 +2,7 @@ use anyhow::{Context as _, Result, anyhow};
 use assistant_tool::{ActionLog, Tool};
 use gpui::{App, Entity, Task};
 use language::{self, Buffer, ToPointUtf16};
-use language_model::LanguageModelRequestMessage;
+use language_model::{LanguageModelRequestMessage, LanguageModelToolSchemaFormat};
 use project::Project;
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
@@ -68,10 +68,7 @@ impl Tool for RenameTool {
         IconName::Pencil
     }
 
-    fn input_schema(
-        &self,
-        format: language_model::LanguageModelToolSchemaFormat,
-    ) -> serde_json::Value {
+    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> Result<serde_json::Value> {
         json_schema_for::<RenameToolInput>(format)
     }
 

crates/assistant_tools/src/schema.rs

@@ -5,23 +5,20 @@ use schemars::{
     schema::{RootSchema, Schema, SchemaObject},
 };
 
-pub fn json_schema_for<T: JsonSchema>(format: LanguageModelToolSchemaFormat) -> serde_json::Value {
+pub fn json_schema_for<T: JsonSchema>(
+    format: LanguageModelToolSchemaFormat,
+) -> Result<serde_json::Value> {
     let schema = root_schema_for::<T>(format);
-    schema_to_json(&schema, format).expect("Failed to convert tool calling schema to JSON")
+    schema_to_json(&schema, format)
 }
 
-pub fn schema_to_json(
+fn schema_to_json(
     schema: &RootSchema,
     format: LanguageModelToolSchemaFormat,
 ) -> Result<serde_json::Value> {
     let mut value = serde_json::to_value(schema)?;
-    match format {
-        LanguageModelToolSchemaFormat::JsonSchema => Ok(value),
-        LanguageModelToolSchemaFormat::JsonSchemaSubset => {
-            transform_fields_to_json_schema_subset(&mut value);
-            Ok(value)
-        }
-    }
+    assistant_tool::adapt_schema_to_format(&mut value, format)?;
+    Ok(value)
 }
 
 fn root_schema_for<T: JsonSchema>(format: LanguageModelToolSchemaFormat) -> RootSchema {
@@ -79,42 +76,3 @@ impl schemars::visit::Visitor for TransformToJsonSchemaSubsetVisitor {
         schemars::visit::visit_schema_object(self, schema)
     }
 }
-
-fn transform_fields_to_json_schema_subset(json: &mut serde_json::Value) {
-    if let serde_json::Value::Object(obj) = json {
-        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(), serde_json::Value::Bool(true));
-            }
-        }
-
-        // If a type is not specified for an input parameter we need to add it.
-        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(),
-                serde_json::Value::String("string".to_string()),
-            );
-        }
-
-        //Format field is only partially supported (e.g. not uint compatibility)
-        obj.remove("format");
-
-        for (_, value) in obj.iter_mut() {
-            if let serde_json::Value::Object(_) | serde_json::Value::Array(_) = value {
-                transform_fields_to_json_schema_subset(value);
-            }
-        }
-    } else if let serde_json::Value::Array(arr) = json {
-        for item in arr.iter_mut() {
-            transform_fields_to_json_schema_subset(item);
-        }
-    }
-}

crates/assistant_tools/src/symbol_info_tool.rs

@@ -84,7 +84,7 @@ impl Tool for SymbolInfoTool {
         IconName::Code
     }
 
-    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
+    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> Result<serde_json::Value> {
         json_schema_for::<SymbolInfoToolInput>(format)
     }
 

crates/assistant_tools/src/terminal_tool.rs

@@ -44,7 +44,7 @@ impl Tool for TerminalTool {
         IconName::Terminal
     }
 
-    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
+    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> Result<serde_json::Value> {
         json_schema_for::<TerminalToolInput>(format)
     }
 

crates/assistant_tools/src/thinking_tool.rs

@@ -36,7 +36,7 @@ impl Tool for ThinkingTool {
         IconName::LightBulb
     }
 
-    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> serde_json::Value {
+    fn input_schema(&self, format: LanguageModelToolSchemaFormat) -> Result<serde_json::Value> {
         json_schema_for::<ThinkingToolInput>(format)
     }